From ff2d0da7544f0a30e483b8959370f3a68ddd6d3f Mon Sep 17 00:00:00 2001
From: Thibault Jamet <tjamet@users.noreply.github.com>
Date: Wed, 7 Mar 2018 01:11:14 +0100
Subject: [PATCH] Fix intercepting headers in middlewares

As explained in the TestInterceptedHeader test, in case a middleware
filters out the headers, this middleware can be done inefficient in case
one following handler is using c.String or other methods writing to the
response body directly.

This commit fixes the issue by using c.Writer when writing the Status as
done in other c.Header, c.SetCookie and other response writers.

The bug has been originally discovered using
https://github.com/gin-contrib/gzip where a failing test has been added
here: https://github.com/tjamet/gzip/blob/header/gzip_test.go#L71

Signed-off-by: Thibault Jamet <tjamet@users.noreply.github.com>
---
 context.go           |  2 +-
 context_go17_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 context_go17_test.go

diff --git a/context.go b/context.go
index 462357e104..856e5d4e82 100755
--- a/context.go
+++ b/context.go
@@ -603,7 +603,7 @@ func bodyAllowedForStatus(status int) bool {
 
 // Status sets the HTTP response code.
 func (c *Context) Status(code int) {
-	c.writermem.WriteHeader(code)
+	c.Writer.WriteHeader(code)
 }
 
 // Header is a intelligent shortcut for c.Writer.Header().Set(key, value).
diff --git a/context_go17_test.go b/context_go17_test.go
new file mode 100644
index 0000000000..eca089ce09
--- /dev/null
+++ b/context_go17_test.go
@@ -0,0 +1,50 @@
+// +build go1.7
+
+package gin
+
+import (
+	"bytes"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type interceptedWriter struct {
+	ResponseWriter
+	b *bytes.Buffer
+}
+
+func (i interceptedWriter) WriteHeader(code int) {
+	i.Header().Del("X-Test")
+	i.ResponseWriter.WriteHeader(code)
+}
+func TestInterceptedHeader(t *testing.T) {
+	w := httptest.NewRecorder()
+	c, r := CreateTestContext(w)
+
+	r.Use(func(c *Context) {
+		i := interceptedWriter{
+			ResponseWriter: c.Writer,
+			b:              bytes.NewBuffer(nil),
+		}
+		c.Writer = i
+		c.Next()
+		c.Header("X-Test", "overridden")
+		c.Writer = i.ResponseWriter
+	})
+	r.GET("/", func(c *Context) {
+		c.Header("X-Test", "original")
+		c.Header("X-Test-2", "present")
+		c.String(http.StatusOK, "hello world")
+	})
+	c.Request = httptest.NewRequest("GET", "/", nil)
+	r.HandleContext(c)
+	// Result() has headers frozen when WriteHeaderNow() has been called
+	// Compared to this time, this is when the response headers will be flushed
+	// As response is flushed on c.String, the Header cannot be set by the first
+	// middleware. Assert this
+	assert.Equal(t, "", w.Result().Header.Get("X-Test"))
+	assert.Equal(t, "present", w.Result().Header.Get("X-Test-2"))
+}