From f6a2c6063db6e952a63779fb45e668f1ad009eb3 Mon Sep 17 00:00:00 2001
From: mohamed abdelrhman <mohamed95zoytak@gmail.com>
Date: Tue, 23 Apr 2024 01:07:05 +0200
Subject: [PATCH] feat:remove duplicate headers from rw

---
 plugin.go   | 19 ++++++++++---------
 recorder.go |  8 ++++----
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/plugin.go b/plugin.go
index 6f80b3c..3c73fdb 100644
--- a/plugin.go
+++ b/plugin.go
@@ -90,19 +90,24 @@ func (c *Cache) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	// Cache miss - record the response
-	recorder := &responseRecorder{rw: rw}
+	recorder := &responseRecorder{
+		rw:     rw,
+		header: rw.Header().Clone(), // Initialize with the original headers.
+	}
 	c.next.ServeHTTP(recorder, req)
 
 	// Serialize the response data
 	cachedResponse := CachedResponse{
 		StatusCode: recorder.status,
-		Headers:    recorder.Header().Clone(), // Convert http.Header to a map for serialization
+		Headers:    recorder.Header(), // Convert http.Header to a map for serialization
 		Body:       recorder.body.Bytes(),
 	}
 	var buffer bytes.Buffer
 	enc := gob.NewEncoder(&buffer)
 	if err := enc.Encode(cachedResponse); err != nil {
 		log.Printf("Failed to serialize response for caching: %s", err)
+		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
+		return
 	} else {
 		// Store the serialized response in Redis
 		if err := respClient.SetWithTTL(req.Context(), cacheKey, buffer.String(), c.cacheExpiry); err != nil {
@@ -110,13 +115,9 @@ func (c *Cache) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		}
 	}
 
-	// Write the response to the client
-	for key, values := range recorder.Header() {
-		for _, value := range values {
-			rw.Header().Add(key, value)
-		}
+	if _, err := rw.Write(recorder.body.Bytes()); err != nil {
+		log.Printf("Failed to write response body: %s", err)
+		return
 	}
-	rw.WriteHeader(recorder.status)
-	rw.Write(recorder.body.Bytes())
 	return
 }
diff --git a/recorder.go b/recorder.go
index fb80993..f684bcc 100644
--- a/recorder.go
+++ b/recorder.go
@@ -9,18 +9,18 @@ type responseRecorder struct {
 	rw     http.ResponseWriter
 	status int
 	body   bytes.Buffer
+	header http.Header
 }
 
 func (r *responseRecorder) Header() http.Header {
-	return r.rw.Header()
+	return r.header
 }
 
 func (r *responseRecorder) Write(b []byte) (int, error) {
-	r.body.Write(b)
-	return r.rw.Write(b)
+	return r.body.Write(b) // Just buffer the body, don't write to rw
+
 }
 
 func (r *responseRecorder) WriteHeader(statusCode int) {
 	r.status = statusCode
-	r.rw.WriteHeader(statusCode)
 }