fix(logging): track whether headers are written properly
Alan Pearce alan@alanpearce.eu
Fri, 17 May 2024 11:00:52 +0200
1 files changed, 20 insertions(+), 3 deletions(-)
jump to
M internal/server/logging.go → internal/server/logging.go
@@ -4,23 +4,40 @@ import ( "fmt" "io" "net/http" + + "github.com/pkg/errors" ) type LoggingResponseWriter struct { + wroteHeader bool http.ResponseWriter statusCode int } func (lrw *LoggingResponseWriter) WriteHeader(code int) { lrw.statusCode = code - // avoids warning: superfluous response.WriteHeader call - if lrw.statusCode != http.StatusOK { + if !lrw.wroteHeader { lrw.ResponseWriter.WriteHeader(code) + lrw.wroteHeader = true } } +func (lrw *LoggingResponseWriter) Write(b []byte) (int, error) { + if !lrw.wroteHeader { + lrw.statusCode = http.StatusOK + lrw.wroteHeader = true + } + + count, err := lrw.ResponseWriter.Write(b) + if err != nil { + return count, errors.Wrap(err, "failed to write response") + } + + return count, nil +} + func NewLoggingResponseWriter(w http.ResponseWriter) *LoggingResponseWriter { - return &LoggingResponseWriter{w, http.StatusOK} + return &LoggingResponseWriter{false, w, http.StatusOK} } type wrappedHandlerOptions struct {