Panic on aborted requests to properly close the connection

pkg/middlewares/recovery/recovery.go

@@ -1,7 +1,10 @@
 package recovery
 
 import (
+	"bufio"
 	"context"
+	"fmt"
+	"net"
 	"net/http"
 	"runtime"
 
@@ -28,12 +31,16 @@ func New(ctx context.Context, next http.Handler) (http.Handler, error) {
 }
 
 func (re *recovery) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
-	defer recoverFunc(rw, req)
-	re.next.ServeHTTP(rw, req)
+	recoveryRW := newRecoveryResponseWriter(rw)
+	defer recoverFunc(recoveryRW, req)
+
+	re.next.ServeHTTP(recoveryRW, req)
 }
 
-func recoverFunc(rw http.ResponseWriter, r *http.Request) {
+func recoverFunc(rw recoveryResponseWriter, r *http.Request) {
 	if err := recover(); err != nil {
+		defer rw.finalizeResponse()
+
 		logger := log.FromContext(middlewares.GetLoggerCtx(r.Context(), middlewareName, typeName))
 		if !shouldLogPanic(err) {
 			logger.Debugf("Request has been aborted [%s - %s]: %v", r.RemoteAddr, r.URL, err)
@@ -45,8 +52,6 @@ func recoverFunc(rw http.ResponseWriter, r *http.Request) {
 		buf := make([]byte, size)
 		buf = buf[:runtime.Stack(buf, false)]
 		logger.Errorf("Stack: %s", buf)
-
-		http.Error(rw, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 	}
 }
 
@@ -56,3 +61,81 @@ func shouldLogPanic(panicValue interface{}) bool {
 	//nolint:errorlint // false-positive because panicValue is an interface.
 	return panicValue != nil && panicValue != http.ErrAbortHandler
 }
+
+type recoveryResponseWriter interface {
+	http.ResponseWriter
+
+	finalizeResponse()
+}
+
+func newRecoveryResponseWriter(rw http.ResponseWriter) recoveryResponseWriter {
+	wrapper := &responseWriterWrapper{rw: rw}
+	if _, ok := rw.(http.CloseNotifier); !ok {
+		return wrapper
+	}
+
+	return &responseWriterWrapperWithCloseNotify{wrapper}
+}
+
+type responseWriterWrapper struct {
+	rw          http.ResponseWriter
+	headersSent bool
+}
+
+func (r *responseWriterWrapper) Header() http.Header {
+	return r.rw.Header()
+}
+
+func (r *responseWriterWrapper) Write(bytes []byte) (int, error) {
+	r.headersSent = true
+	return r.rw.Write(bytes)
+}
+
+func (r *responseWriterWrapper) WriteHeader(code int) {
+	if r.headersSent {
+		return
+	}
+
+	// Handling informational headers.
+	if code >= 100 && code <= 199 {
+		r.rw.WriteHeader(code)
+		return
+	}
+
+	r.headersSent = true
+	r.rw.WriteHeader(code)
+}
+
+func (r *responseWriterWrapper) Flush() {
+	if f, ok := r.rw.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (r *responseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h, ok := r.rw.(http.Hijacker); ok {
+		return h.Hijack()
+	}
+
+	return nil, nil, fmt.Errorf("not a hijacker: %T", r.rw)
+}
+
+func (r *responseWriterWrapper) finalizeResponse() {
+	// If headers have been sent this is not possible to respond with an HTTP error,
+	// and we let the server abort the response silently thanks to the http.ErrAbortHandler sentinel panic value.
+	if r.headersSent {
+		panic(http.ErrAbortHandler)
+	}
+
+	// The response has not yet started to be written,
+	// we can safely return a fresh new error response.
+	http.Error(r.rw, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+}
+
+type responseWriterWrapperWithCloseNotify struct {
+	*responseWriterWrapper
+}
+
+func (r *responseWriterWrapperWithCloseNotify) CloseNotify() <-chan bool {
+	return r.rw.(http.CloseNotifier).CloseNotify()
+}