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..85a0c46734 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 TestInterceptedWrite(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")) +}