all repos — searchix @ c5c9f726cae4cf4488312338da9599da53ac8b6c

Search engine for NixOS, nix-darwin, home-manager and NUR users

fix(logging): track whether headers are written properly
Alan Pearce alan@alanpearce.eu
Fri, 17 May 2024 11:00:52 +0200
commit

c5c9f726cae4cf4488312338da9599da53ac8b6c

parent

b8da487f774aa2399b79bbde5f2e6e13be4fce1c

1 files changed, 20 insertions(+), 3 deletions(-)

jump to
M internal/server/logging.gointernal/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 {