diff --git a/pkg/middlewares/customerrors/custom_errors.go b/pkg/middlewares/customerrors/custom_errors.go index b24a6b04c..c764b2cf8 100644 --- a/pkg/middlewares/customerrors/custom_errors.go +++ b/pkg/middlewares/customerrors/custom_errors.go @@ -32,7 +32,7 @@ type serviceBuilder interface { BuildHTTP(ctx context.Context, serviceName string) (http.Handler, error) } -// customErrors is a middleware that provides the custom error pages.. +// customErrors is a middleware that provides the custom error pages. type customErrors struct { name string next http.Handler @@ -131,10 +131,10 @@ type responseInterceptor interface { isFilteredCode() bool } -// codeCatcher is a response writer that detects as soon as possible whether the -// response is a code within the ranges of codes it watches for. If it is, it -// simply drops the data from the response. Otherwise, it forwards it directly to -// the original client (its responseWriter) without any buffering. +// codeCatcher is a response writer that detects as soon as possible +// whether the response is a code within the ranges of codes it watches for. +// If it is, it simply drops the data from the response. +// Otherwise, it forwards it directly to the original client (its responseWriter) without any buffering. type codeCatcher struct { headerMap http.Header code int @@ -144,16 +144,6 @@ type codeCatcher struct { headersSent bool } -type codeCatcherWithCloseNotify struct { - *codeCatcher -} - -// CloseNotify returns a channel that receives at most a -// single value (true) when the client connection has gone away. -func (cc *codeCatcherWithCloseNotify) CloseNotify() <-chan bool { - return cc.responseWriter.(http.CloseNotifier).CloseNotify() -} - func newCodeCatcher(rw http.ResponseWriter, httpCodeRanges types.HTTPCodeRanges) responseInterceptor { catcher := &codeCatcher{ headerMap: make(http.Header), @@ -168,6 +158,10 @@ func newCodeCatcher(rw http.ResponseWriter, httpCodeRanges types.HTTPCodeRanges) } func (cc *codeCatcher) Header() http.Header { + if cc.headersSent { + return cc.responseWriter.Header() + } + if cc.headerMap == nil { cc.headerMap = make(http.Header) } @@ -247,6 +241,16 @@ func (cc *codeCatcher) Flush() { } } +type codeCatcherWithCloseNotify struct { + *codeCatcher +} + +// CloseNotify returns a channel that receives at most a single value (true) +// when the client connection has gone away. +func (cc *codeCatcherWithCloseNotify) CloseNotify() <-chan bool { + return cc.responseWriter.(http.CloseNotifier).CloseNotify() +} + // codeModifier forwards a response back to the client, // while enforcing a given response code. type codeModifier interface { @@ -277,18 +281,12 @@ type codeModifierWithoutCloseNotify struct { responseWriter http.ResponseWriter } -type codeModifierWithCloseNotify struct { - *codeModifierWithoutCloseNotify -} - -// CloseNotify returns a channel that receives at most a -// single value (true) when the client connection has gone away. -func (r *codeModifierWithCloseNotify) CloseNotify() <-chan bool { - return r.responseWriter.(http.CloseNotifier).CloseNotify() -} - // Header returns the response headers. func (r *codeModifierWithoutCloseNotify) Header() http.Header { + if r.headerSent { + return r.responseWriter.Header() + } + if r.headerMap == nil { r.headerMap = make(http.Header) } @@ -303,8 +301,8 @@ func (r *codeModifierWithoutCloseNotify) Write(buf []byte) (int, error) { return r.responseWriter.Write(buf) } -// WriteHeader sends the headers, with the enforced code (the code in argument -// is always ignored), if it hasn't already been done. +// WriteHeader sends the headers, with the enforced code (the code in argument is always ignored), +// if it hasn't already been done. func (r *codeModifierWithoutCloseNotify) WriteHeader(_ int) { if r.headerSent { return @@ -332,3 +330,13 @@ func (r *codeModifierWithoutCloseNotify) Flush() { flusher.Flush() } } + +type codeModifierWithCloseNotify struct { + *codeModifierWithoutCloseNotify +} + +// CloseNotify returns a channel that receives at most a single value (true) +// when the client connection has gone away. +func (r *codeModifierWithCloseNotify) CloseNotify() <-chan bool { + return r.responseWriter.(http.CloseNotifier).CloseNotify() +}