diff --git a/proxy/writer.go b/proxy/writer.go index 8a502da..4db1dd8 100644 --- a/proxy/writer.go +++ b/proxy/writer.go @@ -66,16 +66,21 @@ func (c *captureWriter) Write(p []byte) (int, error) { if isStreamResponse(c.Header()) { c.sse.Write(p[:n]) } - if f, ok := c.ResponseWriter.(http.Flusher); ok { - f.Flush() - } + c.flush() } return n, err } func (c *captureWriter) Flush() { - if f, ok := c.ResponseWriter.(http.Flusher); ok { - f.Flush() + c.flush() +} + +func (c *captureWriter) flush() { + if _, ok := c.ResponseWriter.(http.Flusher); !ok { + return + } + if err := http.NewResponseController(c.ResponseWriter).Flush(); err != nil { + c.writeFailed = true } } diff --git a/tests/proxy/robustness_test.go b/tests/proxy/robustness_test.go index 2db21d4..5f91259 100644 --- a/tests/proxy/robustness_test.go +++ b/tests/proxy/robustness_test.go @@ -55,6 +55,18 @@ func (w *failingResponseWriter) Write(p []byte) (int, error) { return 0, errors.New("write failed") } +type failingFlushResponseWriter struct{ header http.Header } + +func (w *failingFlushResponseWriter) Header() http.Header { return w.header } +func (*failingFlushResponseWriter) WriteHeader(int) {} +func (*failingFlushResponseWriter) Write(p []byte) (int, error) { + return len(p), nil +} +func (*failingFlushResponseWriter) Flush() {} +func (*failingFlushResponseWriter) FlushError() error { + return errors.New("flush failed") +} + func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler { t.Helper() filter, err := config.NewFilter(config.FilterDisabled, nil) @@ -168,6 +180,22 @@ func TestResponseWriteFailurePreventsCommit(t *testing.T) { } } +func TestResponseFlushFailurePreventsCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "response") + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{}) + w := &failingFlushResponseWriter{header: make(http.Header)} + h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)) + if sub.Len() != 0 { + t.Fatal("failed response flush was committed") + } +} + func TestTrustedProxyClientIP(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "ok")