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 {
-	return &responseModifier{
+func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*http.Response) error) http.ResponseWriter {
+	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.
-func (r *responseModifier) CloseNotify() <-chan bool {
-	return r.rw.(http.CloseNotifier).CloseNotify()
+type responseModifierWithCloseNotify struct {
+	*responseModifier
+}
+
+// CloseNotify implements http.CloseNotifier.
+func (r *responseModifierWithCloseNotify) CloseNotify() <-chan bool {
+	return r.responseModifier.rw.(http.CloseNotifier).CloseNotify()
 }