fix: commit response status before flushing

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,8 +93,22 @@ 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) {
 		panic(fmt.Errorf("echo: response writer %T does not support flushing (http.Flusher interface)", r.ResponseWriter))
@@ -114,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
@@ -158,6 +189,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"))

response_test.go

@@ -92,6 +92,163 @@ 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 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) {
+			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 {
 }