fix: reject logs after response flush failure

This commit is contained in:
MiMoCode
2026-07-10 18:39:00 +08:00
parent 3292eb38fc
commit d2cb0107a6
2 changed files with 38 additions and 5 deletions
+10 -5
View File
@@ -66,16 +66,21 @@ func (c *captureWriter) Write(p []byte) (int, error) {
if isStreamResponse(c.Header()) { if isStreamResponse(c.Header()) {
c.sse.Write(p[:n]) c.sse.Write(p[:n])
} }
if f, ok := c.ResponseWriter.(http.Flusher); ok { c.flush()
f.Flush()
}
} }
return n, err return n, err
} }
func (c *captureWriter) Flush() { func (c *captureWriter) Flush() {
if f, ok := c.ResponseWriter.(http.Flusher); ok { c.flush()
f.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
} }
} }
+28
View File
@@ -55,6 +55,18 @@ func (w *failingResponseWriter) Write(p []byte) (int, error) {
return 0, errors.New("write failed") 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 { func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler {
t.Helper() t.Helper()
filter, err := config.NewFilter(config.FilterDisabled, nil) 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) { func TestTrustedProxyClientIP(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "ok") _, _ = io.WriteString(w, "ok")