fix: reject logs after response flush failure
This commit is contained in:
+10
-5
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user