Conditional CloseNotify in header middleware

pkg/middlewares/headers/responsewriter.go

@@ -22,13 +22,18 @@ type responseModifier struct {
 }
 
 // modifier can be nil.
-func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*http.Response) error) *responseModifier {
+func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*http.Response) error) http.ResponseWriter {
-	return &responseModifier{
+	rm := &responseModifier{
 		req:      r,
 		rw:       w,
 		modifier: modifier,
 		code:     http.StatusOK,
 	}
+
+	if _, ok := w.(http.CloseNotifier); ok {
+		return responseModifierWithCloseNotify{responseModifier: rm}
+	}
+	return rm
 }
 
 func (r *responseModifier) WriteHeader(code int) {
@@ -93,7 +98,11 @@ func (r *responseModifier) Flush() {
 	}
 }
 
-// CloseNotify implements http.CloseNotifier.
+type responseModifierWithCloseNotify struct {
-func (r *responseModifier) CloseNotify() <-chan bool {
+	*responseModifier
-	return r.rw.(http.CloseNotifier).CloseNotify()
+}
+
+// CloseNotify implements http.CloseNotifier.
+func (r *responseModifierWithCloseNotify) CloseNotify() <-chan bool {
+	return r.responseModifier.rw.(http.CloseNotifier).CloseNotify()
 }

pkg/middlewares/tracing/status_code.go

@@ -11,6 +11,15 @@ type statusCodeRecoder interface {
 	Status() int
 }
 
+// newStatusCodeRecoder returns an initialized statusCodeRecoder.
+func newStatusCodeRecoder(rw http.ResponseWriter, status int) statusCodeRecoder {
+	recorder := &statusCodeWithoutCloseNotify{rw, status}
+	if _, ok := rw.(http.CloseNotifier); ok {
+		return &statusCodeWithCloseNotify{recorder}
+	}
+	return recorder
+}
+
 type statusCodeWithoutCloseNotify struct {
 	http.ResponseWriter
 	status int
@@ -46,12 +55,3 @@ type statusCodeWithCloseNotify struct {
 func (s *statusCodeWithCloseNotify) CloseNotify() <-chan bool {
 	return s.ResponseWriter.(http.CloseNotifier).CloseNotify()
 }
-
-// newStatusCodeRecoder returns an initialized statusCodeRecoder.
-func newStatusCodeRecoder(rw http.ResponseWriter, status int) statusCodeRecoder {
-	recorder := &statusCodeWithoutCloseNotify{rw, status}
-	if _, ok := rw.(http.CloseNotifier); ok {
-		return &statusCodeWithCloseNotify{recorder}
-	}
-	return recorder
-}