Block consecutive WriteHeader

handler.go

import (
	"context"
	"sync/atomic"

	"github.com/felixge/httpsnoop"
)

// BlockConsecutiveWriteHeader is a middleware that prevents multiple WriteHeader calls on the http.ResponseWriter.
// This scenario leads to a server-side errors that indicate a wrong handler implementation. Despite this being a
// useful information, often in prod applications we still don't want our logs to be flooded, and it's in these cases
// that this handler comes in handy.
func BlockConsecutiveWriteHeader(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ww := consecutiveWriteHeaderBlocker{ResponseWriter: w}
		hooks := httpsnoop.Hooks{
			Header: func(httpsnoop.HeaderFunc) httpsnoop.HeaderFunc {
				return ww.Header
			},
			WriteHeader: func(httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
				return ww.WriteHeader
			},
			Write: func(httpsnoop.WriteFunc) httpsnoop.WriteFunc {
				return ww.Write
			},
			Flush: func(flushFunc httpsnoop.FlushFunc) httpsnoop.FlushFunc {
				return ww.Flush
			},
		}
		w = httpsnoop.Wrap(w, hooks)

		next.ServeHTTP(w, r)
	})
}


var _ http.ResponseWriter = &consecutiveWriteHeaderBlocker{}

// consecutiveWriteHeaderBlocker wraps a http.ResponseWriter in order prevent multiple WriteHeader calls that would
// result in errors in the net/http layer which prints them as they might indicate wrong handler implementations.
// Nevertheless, some of these implementations are out of our control (e.g. third-party middlewares) and we don't want
// them to flood our error logs. This wrapper helps to avoid this situation.
//
// NOTE: this wrapper doesn't directly implement any of the optional types (e.g. http.Hijacker, http.Pusher,
// http.CloseNotifier) but these are added by httpsnoop under the hood if the original http.ResponseWriter already
// implements them.
type consecutiveWriteHeaderBlocker struct {
	http.ResponseWriter
	wroteHeader atomic.Bool
}

// Write overrides the wrapped http.ResponseWriter's Write to prevent consecutive WriteHeader calls.
func (w *consecutiveWriteHeaderBlocker) Write(p []byte) (int, error) {
	w.writeHeader(http.StatusOK)

	return w.ResponseWriter.Write(p)
}

// WriteHeader overrides the wrapped http.ResponseWriter's WriteHeader to prevent consecutive WriteHeader calls.
func (w *consecutiveWriteHeaderBlocker) WriteHeader(statusCode int) {
	w.writeHeader(statusCode)
}

// Flush implements http.Flusher that also needs to flush headers.
func (w *consecutiveWriteHeaderBlocker) Flush() {
	w.writeHeader(http.StatusOK)

	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// writeHeader propagates the WriteHeader call to the underlying ResponseWriter only the first time it is called.
func (w *consecutiveWriteHeaderBlocker) writeHeader(statusCode int) {
	if w.wroteHeader.CompareAndSwap(false, true) {
		w.ResponseWriter.WriteHeader(statusCode)
	}
}