Copy headers before returning reponse writter

docker/traefik-config/dynamic-configuration.yml

@@ -7,7 +7,7 @@ http:
             - status: 200
               rewrites:
                 - regex: foo
-                  replacement: Bar
+                  replacement: Bara
             - status: 400-499
               rewrites:
                 - regex: .*

main.go

@@ -144,56 +144,26 @@ type responseWriter struct {
 	responses []parsedResponse
 }
 
-// Headers implements the http.ResponseWriter interface.
-func (rw *responseWriter) Header() http.Header {
-	if rw.headersSent {
-		return rw.ResponseWriter.Header()
-	}
-
-	if rw.headerMap == nil {
-		rw.headerMap = make(http.Header)
-	}
-
-	return rw.headerMap
-}
-
 // WriteHeader implements the http.ResponseWriter interface.
 // It intercepts the response status code and stores it in the responseWriter struct.
 func (rw *responseWriter) WriteHeader(statusCode int) {
 	if rw.headersSent {
 		return
 	}
 
-	// Handling informational headers.
-	if statusCode >= 100 && statusCode <= 199 {
-		// Multiple informational status codes can be used,
-		// so here the copy is not appending the values to not repeat them.
-		for k, v := range rw.Header() {
-			rw.ResponseWriter.Header()[k] = v
-		}
-
-		rw.ResponseWriter.WriteHeader(statusCode)
-		return
-	}
-
 	rw.code = statusCode
 
 	// Check if the status code is in the list of status codes to rewrite.
-	var bodyRewrited bool
 	for _, response := range rw.responses {
 		if !response.status.Contains(statusCode) {
 			continue
 		}
-		bodyRewrited = true
-	}
-
-	// Remove Content-Length header if the body is being rewritten.
-	if bodyRewrited {
 		rw.ResponseWriter.Header().Del("Content-Length")
+		break
 	}
-
-	rw.ResponseWriter.WriteHeader(rw.code)
 	rw.headersSent = true
+
+	rw.ResponseWriter.WriteHeader(statusCode)
 }
 
 // Write implements the http.ResponseWriter interface.

main_test.go

@@ -12,13 +12,14 @@ import (
 
 func TestServeHTTP(t *testing.T) {
 	tests := []struct {
-		desc            string
-		httpStatus      int
-		contentEncoding string
-		responses       []Response
-		lastModified    bool
-		resBody         string
-		expResBody      string
+		desc               string
+		httpStatus         int
+		contentEncoding    string
+		responses          []Response
+		lastModified       bool
+		resBody            string
+		expResBody         string
+		contentLengthExist bool
 	}{
 		{
 			desc:       "should replace foo by bar",
@@ -34,8 +35,9 @@ func TestServeHTTP(t *testing.T) {
 					},
 				},
 			},
-			resBody:    "foo is the new bar",
-			expResBody: "bar is the new bar",
+			resBody:            "foo is the new bar",
+			expResBody:         "bar is the new bar",
+			contentLengthExist: true,
 		},
 		{
 			desc:       "should replace nothing",
@@ -72,8 +74,9 @@ func TestServeHTTP(t *testing.T) {
 					},
 				},
 			},
-			resBody:    "foo is the new bar",
-			expResBody: "foo is the new foo",
+			resBody:            "foo is the new bar",
+			expResBody:         "foo is the new foo",
+			contentLengthExist: true,
 		},
 		{
 			desc:       "should replace fo by bar if content encoding is not identity or empty & match response",
@@ -89,9 +92,10 @@ func TestServeHTTP(t *testing.T) {
 					},
 				},
 			},
-			contentEncoding: "gzip",
-			resBody:         "foo is the new bar",
-			expResBody:      "bar is the new bar",
+			contentEncoding:    "gzip",
+			resBody:            "foo is the new bar",
+			expResBody:         "bar is the new bar",
+			contentLengthExist: true,
 		},
 		{
 			desc:       "should replace foo by bar if content encoding is identity",
@@ -107,28 +111,10 @@ func TestServeHTTP(t *testing.T) {
 					},
 				},
 			},
-			contentEncoding: "identity",
-			resBody:         "foo is the new bar",
-			expResBody:      "bar is the new bar",
-		},
-		{
-			desc:       "should not remove the last modified header",
-			httpStatus: http.StatusOK,
-			responses: []Response{
-				{
-					Status: "200-299",
-					Rewrites: []Rewrite{
-						{
-							Regex:       "foo",
-							Replacement: "bar",
-						},
-					},
-				},
-			},
-			contentEncoding: "identity",
-			lastModified:    true,
-			resBody:         "foo is the new bar",
-			expResBody:      "bar is the new bar",
+			contentEncoding:    "identity",
+			resBody:            "foo is the new bar",
+			expResBody:         "bar is the new bar",
+			contentLengthExist: true,
 		},
 	}
 
@@ -157,7 +143,7 @@ func TestServeHTTP(t *testing.T) {
 
 			rewriteBody.ServeHTTP(recorder, req)
 
-			if _, exists := recorder.Result().Header["Content-Length"]; exists {
+			if _, exists := recorder.Result().Header["Content-Length"]; exists == test.contentLengthExist {
 				t.Error("The Content-Length Header must be deleted")
 			}