From f4ae264dfda58727aa08cb8a25771f49f619d581 Mon Sep 17 00:00:00 2001 From: Plastic-time <247457381+Plastic-time@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:33:41 -0400 Subject: [PATCH 1/2] fix: commit response status before flushing --- response.go | 12 +++++++++++ response_test.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/response.go b/response.go index e218ef4d6..7b161c8bf 100644 --- a/response.go +++ b/response.go @@ -80,6 +80,12 @@ func (r *Response) Write(b []byte) (n int, err error) { // buffered data to the client. // See [http.Flusher](https://golang.org/pkg/net/http/#Flusher) func (r *Response) Flush() { + if !r.Committed { + if r.Status == 0 { + r.Status = http.StatusOK + } + r.WriteHeader(r.Status) + } err := http.NewResponseController(r.ResponseWriter).Flush() if err != nil && errors.Is(err, http.ErrNotSupported) { panic(fmt.Errorf("echo: response writer %T does not support flushing (http.Flusher interface)", r.ResponseWriter)) @@ -158,6 +164,12 @@ func (w *delayedStatusWriter) Write(data []byte) (int, error) { } func (w *delayedStatusWriter) Flush() { + if !w.committed { + if w.status == 0 { + w.status = http.StatusOK + } + w.WriteHeader(w.status) + } err := http.NewResponseController(w.ResponseWriter).Flush() if err != nil && errors.Is(err, http.ErrNotSupported) { panic(errors.New("response writer flushing is not supported")) diff --git a/response_test.go b/response_test.go index 4ad849c04..dbc76cb4c 100644 --- a/response_test.go +++ b/response_test.go @@ -92,6 +92,60 @@ func TestResponse_Flush(t *testing.T) { assert.True(t, rec.Flushed) } +func TestResponse_FlushBeforeWrite(t *testing.T) { + for _, status := range []int{0, http.StatusAccepted} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + rec := httptest.NewRecorder() + res := NewResponse(rec, New().Logger) + res.Status = status + beforeCalls := 0 + res.Before(func() { + beforeCalls++ + res.Header().Set("X-Before", "called") + }) + + res.Flush() + wantStatus := status + if wantStatus == 0 { + wantStatus = http.StatusOK + } + assert.Equal(t, wantStatus, rec.Code) + assert.Equal(t, wantStatus, res.Status) + assert.True(t, res.Committed) + assert.True(t, rec.Flushed) + assert.Equal(t, "called", rec.Result().Header.Get("X-Before")) + + res.Flush() + _, err := res.Write([]byte("streamed")) + assert.NoError(t, err) + assert.Equal(t, 1, beforeCalls) + assert.Equal(t, "streamed", rec.Body.String()) + }) + } +} + +func TestDelayedStatusWriter_FlushBeforeWrite(t *testing.T) { + for _, status := range []int{0, http.StatusCreated} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + rec := httptest.NewRecorder() + w := &delayedStatusWriter{ResponseWriter: rec, status: status} + w.Flush() + + wantStatus := status + if wantStatus == 0 { + wantStatus = http.StatusOK + } + assert.Equal(t, wantStatus, rec.Code) + assert.True(t, w.committed) + assert.True(t, rec.Flushed) + + _, err := w.Write([]byte("streamed")) + assert.NoError(t, err) + assert.Equal(t, wantStatus, rec.Code) + }) + } +} + type testResponseWriter struct { } From 0e5cd11328896aae0a84db4c51214e66d558a704 Mon Sep 17 00:00:00 2001 From: Plastic-time <247457381+Plastic-time@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:22:56 -0400 Subject: [PATCH 2/2] fix: defer flushes until response headers are committed --- response.go | 33 +++++++++++++-- response_test.go | 103 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 4 deletions(-) diff --git a/response.go b/response.go index 7b161c8bf..4da210c9c 100644 --- a/response.go +++ b/response.go @@ -22,10 +22,12 @@ type Response struct { // beforeFuncs are functions that are called just before the response (status) is written. Happens only once, during WriteHeader call. beforeFuncs []func() // afterFuncs are functions that are called just after the response is written. During every `Write` method call. - afterFuncs []func() - Status int - Size int64 - Committed bool + afterFuncs []func() + writingHeader bool + flushPending bool + Status int + Size int64 + Committed bool } // NewResponse creates a new instance of Response. @@ -53,11 +55,24 @@ func (r *Response) WriteHeader(code int) { return } r.Status = code + r.writingHeader = true + defer func() { + r.writingHeader = false + if !r.Committed { + // A failed Before callback or header write must not leave a deferred flush. + r.flushPending = false + } + }() for _, fn := range r.beforeFuncs { fn() } r.ResponseWriter.WriteHeader(r.Status) r.Committed = true + r.writingHeader = false + if r.flushPending { + r.flushPending = false + r.Flush() + } } // Write writes the data to the connection as part of an HTTP reply. @@ -78,13 +93,21 @@ func (r *Response) Write(b []byte) (n int, err error) { // Flush implements the http.Flusher interface to allow an HTTP handler to flush // buffered data to the client. +// Calls from Before callbacks are deferred until all callbacks have completed +// and the final response headers have been committed. // See [http.Flusher](https://golang.org/pkg/net/http/#Flusher) func (r *Response) Flush() { + if r.writingHeader { + r.flushPending = true + return + } if !r.Committed { + r.flushPending = true if r.Status == 0 { r.Status = http.StatusOK } r.WriteHeader(r.Status) + return } err := http.NewResponseController(r.ResponseWriter).Flush() if err != nil && errors.Is(err, http.ErrNotSupported) { @@ -120,6 +143,8 @@ func (r *Response) reset(w http.ResponseWriter) { r.Size = 0 r.Status = http.StatusOK r.Committed = false + r.writingHeader = false + r.flushPending = false } // UnwrapResponse unwraps given ResponseWriter to return contexts original Echo Response. rw has to implement diff --git a/response_test.go b/response_test.go index dbc76cb4c..9600b71ff 100644 --- a/response_test.go +++ b/response_test.go @@ -124,6 +124,109 @@ func TestResponse_FlushBeforeWrite(t *testing.T) { } } +func TestResponse_FlushDuringBefore(t *testing.T) { + for _, trigger := range []string{"Flush", "Write", "WriteHeader"} { + t.Run(trigger, func(t *testing.T) { + rec := &flushCountingResponseWriter{ResponseRecorder: httptest.NewRecorder()} + res := NewResponse(rec, New().Logger) + res.Status = http.StatusAccepted + var beforeCalls [2]int + afterCalls := 0 + res.Before(func() { + beforeCalls[0]++ + if beforeCalls[0] > 1 { + t.Fatal("Before callback reentered") + } + res.Flush() + res.Flush() + assert.False(t, res.Committed) + assert.False(t, rec.Flushed) + res.Status = http.StatusCreated + res.Header().Set("X-First", "after flush") + }) + res.Before(func() { + beforeCalls[1]++ + res.Flush() + assert.False(t, rec.Flushed) + res.Status = http.StatusNonAuthoritativeInfo + res.Header().Set("X-Second", "after flush") + }) + res.After(func() { afterCalls++ }) + + switch trigger { + case "Flush": + res.Flush() + case "Write": + _, err := res.Write([]byte("body")) + assert.NoError(t, err) + case "WriteHeader": + res.WriteHeader(http.StatusAccepted) + } + + assert.Equal(t, [2]int{1, 1}, beforeCalls) + assert.True(t, res.Committed) + assert.Equal(t, http.StatusNonAuthoritativeInfo, res.Status) + assert.Equal(t, res.Status, rec.Code) + assert.Equal(t, 1, rec.flushes) + assert.Equal(t, "after flush", rec.Result().Header.Get("X-First")) + assert.Equal(t, "after flush", rec.Result().Header.Get("X-Second")) + if trigger == "Write" { + assert.Equal(t, "body", rec.Body.String()) + assert.Equal(t, 1, afterCalls) + } else { + assert.Equal(t, 0, afterCalls) + } + + res.Flush() + assert.Equal(t, 2, rec.flushes) + assert.Equal(t, [2]int{1, 1}, beforeCalls) + }) + } +} + +func TestResponse_FlushDuringBeforePanic(t *testing.T) { + rec := &flushCountingResponseWriter{ResponseRecorder: httptest.NewRecorder()} + res := NewResponse(rec, New().Logger) + firstCall := true + res.Before(func() { + if firstCall { + firstCall = false + res.Flush() + panic("before failed") + } + }) + + assert.PanicsWithValue(t, "before failed", func() { + res.WriteHeader(http.StatusAccepted) + }) + assert.False(t, res.Committed) + assert.False(t, rec.Flushed) + + res.WriteHeader(http.StatusInternalServerError) + assert.True(t, res.Committed) + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.Equal(t, 0, rec.flushes) + res.Flush() + assert.Equal(t, 1, rec.flushes) + + next := &flushCountingResponseWriter{ResponseRecorder: httptest.NewRecorder()} + res.reset(next) + res.Flush() + assert.True(t, res.Committed) + assert.Equal(t, http.StatusOK, next.Code) + assert.Equal(t, 1, next.flushes) +} + +type flushCountingResponseWriter struct { + *httptest.ResponseRecorder + flushes int +} + +func (w *flushCountingResponseWriter) Flush() { + w.flushes++ + w.ResponseRecorder.Flush() +} + func TestDelayedStatusWriter_FlushBeforeWrite(t *testing.T) { for _, status := range []int{0, http.StatusCreated} { t.Run(fmt.Sprint(status), func(t *testing.T) {