diff --git a/proxy/proxy.go b/proxy/proxy.go index 312b63d..1730924 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -221,6 +221,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.rp.ServeHTTP(cw, r) finished := time.Now() + cw.ConfirmDelivery() responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete() if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil { submitOnce.Do(func() { submit(finished) }) diff --git a/proxy/writer.go b/proxy/writer.go index b7388d5..6ee60fd 100644 --- a/proxy/writer.go +++ b/proxy/writer.go @@ -75,6 +75,19 @@ func (c *captureWriter) Flush() { c.flush() } +// ConfirmDelivery establishes an observable delivery boundary for responses +// whose headers were not followed by a body write. +func (c *captureWriter) ConfirmDelivery() { + if c.written != 0 { + return + } + if _, ok := c.ResponseWriter.(http.Flusher); !ok { + c.writeFailed = true + return + } + c.flush() +} + func (c *captureWriter) flush() { if _, ok := c.ResponseWriter.(http.Flusher); !ok { return diff --git a/tests/proxy/robustness_test.go b/tests/proxy/robustness_test.go index 0bbe6f3..2d2e579 100644 --- a/tests/proxy/robustness_test.go +++ b/tests/proxy/robustness_test.go @@ -67,6 +67,14 @@ func (*failingFlushResponseWriter) FlushError() error { return errors.New("flush failed") } +type unflushableResponseWriter struct{ header http.Header } + +func (w *unflushableResponseWriter) Header() http.Header { return w.header } +func (*unflushableResponseWriter) WriteHeader(int) {} +func (*unflushableResponseWriter) Write(p []byte) (int, error) { + return len(p), nil +} + func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler { t.Helper() filter, err := config.NewFilter(config.FilterDisabled, nil) @@ -220,6 +228,70 @@ func TestResponseFlushFailurePreventsCommit(t *testing.T) { } } +func TestHeaderOnlyResponseFlushFailurePreventsCommit(t *testing.T) { + tests := []struct { + name string + method string + status int + }{ + {name: "head", method: http.MethodHead, status: http.StatusOK}, + {name: "no content", method: http.MethodGet, status: http.StatusNoContent}, + {name: "not modified", method: http.MethodGet, status: http.StatusNotModified}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(tc.status) + })) + 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(tc.method, "http://proxy.test/v1/test", nil)) + if sub.Len() != 0 { + t.Fatal("response with undelivered headers was committed") + } + }) + } +} + +func TestHeaderOnlyResponseSuccessfulFlushCommits(t *testing.T) { + for _, status := range []int{http.StatusNoContent, http.StatusNotModified} { + t.Run(http.StatusText(status), func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{}) + h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)) + if sub.Len() != 1 { + t.Fatalf("entries=%d, want successfully delivered response committed", sub.Len()) + } + }) + } +} + +func TestHeaderOnlyResponseWithoutDeliveryBoundaryDoesNotCommit(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + defer upstream.Close() + u, _ := url.Parse(upstream.URL) + + sub := &captureSubmitter{} + h := newTestHandler(t, u, sub, proxy.Options{}) + w := &unflushableResponseWriter{header: make(http.Header)} + h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)) + if sub.Len() != 0 { + t.Fatal("header-only response without a delivery boundary was committed") + } +} + func TestTrustedProxyClientIP(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "ok")