diff --git a/pkg/middlewares/headers/responsewriter.go b/pkg/middlewares/headers/responsewriter.go index 39a171dc2..15201b1fd 100644 --- a/pkg/middlewares/headers/responsewriter.go +++ b/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() } diff --git a/pkg/middlewares/tracing/status_code.go b/pkg/middlewares/tracing/status_code.go index 9cec2e16c..f0c57f81f 100644 --- a/pkg/middlewares/tracing/status_code.go +++ b/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 -}