Fix missing trailer with custom errors middleware

Co-authored-by: LandryBe <lbenguigui@gmail.com>

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
+// codeCatcher is a response writer that detects as soon as possible
-// response is a code within the ranges of codes it watches for. If it is, it
+// whether the response is a code within the ranges of codes it watches for.
-// simply drops the data from the response. Otherwise, it forwards it directly to
+// If it is, it simply drops the data from the response.
-// the original client (its responseWriter) without any buffering.
+// 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
+// WriteHeader sends the headers, with the enforced code (the code in argument is always ignored),
-// is always ignored), if it hasn't already been done.
+// 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()
+}