From 1c8ac6b7325f690496b006da870a6f717294f192 Mon Sep 17 00:00:00 2001 From: Thibault Jamet 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 --- context.go | 2 +- context_test.go | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/context.go b/context.go index 90d4c6e59e..57232ea6ed 100644 --- a/context.go +++ b/context.go @@ -578,7 +578,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_test.go b/context_test.go index 9024cfc127..6c092d374d 100644 --- a/context_test.go +++ b/context_test.go @@ -1377,3 +1377,42 @@ func TestContextGetRawData(t *testing.T) { assert.Nil(t, err) assert.Equal(t, "Fetch binary post data", string(data)) } + +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")) +}