Created
September 18, 2024 09:42
-
-
Save robgee86/1969b44a1e3d365b5a606f70200bc97a to your computer and use it in GitHub Desktop.
Block consecutive WriteHeader
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment