fix(logging): track whether headers are written properly

1 files changed, 20 insertions, 3 deletions

diff --git a/internal/server/logging.go b/internal/server/logging.go
index 6a16f42..ac8468f 100644
--- a/internal/server/logging.go
+++ b/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 {