From 4ea7d51deac965727d145929eb0e8f812fb334f8 Mon Sep 17 00:00:00 2001
From: Julio Guerra <julio@sqreen.com>
Date: Wed, 3 Feb 2021 23:35:39 +0100
Subject: [PATCH] Avoid overwriting the written status code

When the response has been written, the status code saved in `responseWriter`
should no longer be changed by subsequent calls to `WriteHeader()`.

This fixes for example the logging possibly showing the wrong status code, or
any other middleware function checking the status code.
---
 response_writer.go      | 1 +
 response_writer_test.go | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/response_writer.go b/response_writer.go
index 26826689..5f1e777d 100644
--- a/response_writer.go
+++ b/response_writer.go
@@ -61,6 +61,7 @@ func (w *responseWriter) WriteHeader(code int) {
 	if code > 0 && w.status != code {
 		if w.Written() {
 			debugPrint("[WARNING] Headers were already written. Wanted to override status code %d with %d", w.status, code)
+			return
 		}
 		w.status = code
 	}
diff --git a/response_writer_test.go b/response_writer_test.go
index 1f113e74..8b154faf 100644
--- a/response_writer_test.go
+++ b/response_writer_test.go
@@ -73,6 +73,10 @@ func TestResponseWriterWriteHeadersNow(t *testing.T) {
 	writer.size = 10
 	w.WriteHeaderNow()
 	assert.Equal(t, 10, w.Size())
+
+	w.WriteHeader(http.StatusOK)
+	assert.Equal(t, http.StatusMultipleChoices, w.Status())
+	assert.Equal(t, http.StatusMultipleChoices, testWriter.Code)
 }
 
 func TestResponseWriterWrite(t *testing.T) {