fix(plugins): Traefik custom writer (#161)

* fix(plugins): Traefik custom writer

* Update and prepare v1.5.11-beta6

.github/workflows/plugins-master.yml

@@ -32,4 +32,4 @@ jobs:
           sudo apt install xcaddy
       -
         name: Build current Souin as caddy module with referenced Souin core version when merge on master
-        run: cd plugins/caddy && xcaddy build --with github.com/darkweak/souin/plugins/caddy@$(git rev-parse --short "$GITHUB_SHA") --with github.com/darkweak/souin@$(git rev-parse --short "$GITHUB_SHA")
+        run: cd plugins/caddy && xcaddy build --with github.com/darkweak/souin/plugins/caddy@$(git rev-parse --short "$GITHUB_SHA")

README.md

@@ -354,7 +354,7 @@ experimental:
   plugins:
     souin:
       moduleName: github.com/darkweak/souin
-      version: v1.5.11-beta5
+      version: v1.5.11-beta6
 ```
 After that you can declare either the whole configuration at once in the middleware block or by service. See the examples below.
 ```yaml

plugins/caddy/go.mod

@@ -4,8 +4,8 @@ go 1.16
 
 require (
 	github.com/caddyserver/caddy/v2 v2.4.5
-	github.com/darkweak/souin v1.5.11
+	github.com/darkweak/souin v1.5.11-beta6
 	go.uber.org/zap v1.19.1
 )
 
-replace github.com/darkweak/souin v1.5.11 => ../..
+replace github.com/darkweak/souin v1.5.11-beta6 => ../..

plugins/traefik/base.go

@@ -31,43 +31,40 @@ type getterContext struct {
 // CustomWriter is a custom writer
 type CustomWriter struct {
 	Response *http.Response
-	BufPool  *sync.Pool
-	http.ResponseWriter
+	Buf      *bytes.Buffer
+	Rw       http.ResponseWriter
 }
 
 // Header will write the response headers
 func (r *CustomWriter) Header() http.Header {
-	return r.ResponseWriter.Header()
+	return r.Rw.Header()
 }
 
 // WriteHeader will write the response headers
 func (r *CustomWriter) WriteHeader(code int) {
-	if code == 0 {
-		return
+	if code != 0 {
+		r.Response.StatusCode = code
 	}
-	r.Response.StatusCode = code
 }
 
 // Write will write the response body
 func (r *CustomWriter) Write(b []byte) (int, error) {
-	r.Response.Header = r.ResponseWriter.Header()
-	if r.Response.StatusCode == 0 {
-		r.Response.StatusCode = http.StatusOK
-	}
-	r.Response.Body = ioutil.NopCloser(bytes.NewBuffer(b))
+	r.Response.Header = r.Header()
+	r.Buf.Write(b)
+	r.Response.Body = ioutil.NopCloser(r.Buf)
 	return len(b), nil
 }
 
 // Send delays the response to handle Cache-Status
 func (r *CustomWriter) Send() (int, error) {
-	b, _ := ioutil.ReadAll(r.Response.Body)
 	for h, v := range r.Response.Header {
 		if len(v) > 0 {
-			r.Header().Set(h, strings.Join(v, ", "))
+			r.Rw.Header().Set(h, strings.Join(v, ", "))
 		}
 	}
-	r.ResponseWriter.WriteHeader(r.Response.StatusCode)
-	return r.ResponseWriter.Write(b)
+	r.Rw.WriteHeader(r.Response.StatusCode)
+	b, _ := ioutil.ReadAll(r.Response.Body)
+	return r.Rw.Write(b)
 }
 
 func canHandle(r *http.Request, re types.RetrieverResponsePropertiesInterface) bool {
@@ -86,8 +83,8 @@ func InitializeRequest(rw http.ResponseWriter, req *http.Request, next http.Hand
 	req = req.WithContext(ctx)
 	req.Header.Set("Date", time.Now().UTC().Format(time.RFC1123))
 	return &CustomWriter{
-		ResponseWriter: rw,
-		Response:       &http.Response{},
+		Rw:       rw,
+		Response: &http.Response{},
 	}
 }
 

plugins/traefik/main.go

@@ -174,6 +174,7 @@ func (s *SouinTraefikPlugin) ServeHTTP(rw http.ResponseWriter, req *http.Request
 	buf.Reset()
 	defer bufPool.Put(buf)
 	customRW := InitializeRequest(rw, req, s.next)
+	customRW.Buf = buf
 	regexpURL := s.Retriever.GetRegexpUrls().FindString(req.Host + req.URL.Path)
 	url := configurationtypes.URL{
 		TTL:     configurationtypes.Duration{Duration: s.Retriever.GetConfiguration().GetDefaultCache().GetTTL()},