From 421b4538d8e1fb734bf70cbd8595830f75d23175 Mon Sep 17 00:00:00 2001 From: Abdullah Alqahtani <61029571+anqorithm@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:59:26 +0100 Subject: [PATCH] fix(middleware): retain partially written BodyDump responses --- middleware/body_dump.go | 30 ++++++++++------------ middleware/body_dump_test.go | 50 ++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 17 deletions(-) diff --git a/middleware/body_dump.go b/middleware/body_dump.go index e92319bc9..4eb207934 100644 --- a/middleware/body_dump.go +++ b/middleware/body_dump.go @@ -118,15 +118,10 @@ func (config BodyDumpConfig) ToMiddleware() (echo.MiddlewareFunc, error) { resBuf.Reset() defer bodyDumpBufferPool.Put(resBuf) - var respWriter io.Writer - if config.MaxResponseBytes > 0 { - respWriter = &limitedWriter{ - response: c.Response(), - dumpBuf: resBuf, - limit: config.MaxResponseBytes, - } - } else { - respWriter = io.MultiWriter(c.Response(), resBuf) + respWriter := &limitedWriter{ + response: c.Response(), + dumpBuf: resBuf, + limit: config.MaxResponseBytes, } origResponse := c.Response() writer := &bodyDumpResponseWriter{ @@ -185,17 +180,18 @@ type limitedWriter struct { func (w *limitedWriter) Write(b []byte) (n int, err error) { // Always write full data to actual response (don't truncate client response) n, err = w.response.Write(b) - if err != nil { - return n, err - } - // Write to dump buffer only up to limit - if w.dumped < w.limit { - remaining := w.limit - w.dumped - toDump := min(int64(n), remaining) + if w.limit < 0 || w.dumped < w.limit { + toDump := int64(n) + if w.limit >= 0 { + toDump = min(toDump, w.limit-w.dumped) + } w.dumpBuf.Write(b[:toDump]) w.dumped += toDump } - return n, nil + if err == nil && n < len(b) { + err = io.ErrShortWrite + } + return n, err } diff --git a/middleware/body_dump_test.go b/middleware/body_dump_test.go index 6a0c49dd8..d53191264 100644 --- a/middleware/body_dump_test.go +++ b/middleware/body_dump_test.go @@ -677,3 +677,53 @@ func TestBodyDump_ErrorHandlerRace(t *testing.T) { } wg.Wait() } + +type partialDumpWriter struct { + *httptest.ResponseRecorder + err error +} + +func (w *partialDumpWriter) Write(b []byte) (int, error) { + n, _ := w.ResponseRecorder.Write(b[:3]) + return n, w.err +} + +func TestBodyDump_PartialWrite(t *testing.T) { + failure := errors.New("write failed") + for _, tc := range []struct { + limit int64 + writeErr error + }{ + {2, failure}, {5, failure}, {-1, failure}, + {2, nil}, {5, nil}, {-1, nil}, + } { + t.Run(fmt.Sprint(tc.limit, tc.writeErr), func(t *testing.T) { + writeErr := tc.writeErr + if writeErr == nil { + writeErr = io.ErrShortWrite + } + rec := &partialDumpWriter{ResponseRecorder: httptest.NewRecorder(), err: tc.writeErr} + e := echo.New() + c := e.NewContext(httptest.NewRequest(http.MethodGet, "/", nil), rec) + var dumped string + h := BodyDumpWithConfig(BodyDumpConfig{ + MaxResponseBytes: tc.limit, + Handler: func(_ *echo.Context, _, response []byte, err error) { + dumped = string(response) + assert.ErrorIs(t, err, writeErr) + }, + })(func(c *echo.Context) error { + n, err := c.Response().Write([]byte("abcdef")) + assert.Equal(t, 3, n) + return err + }) + assert.ErrorIs(t, h(c), writeErr) + want := "abc" + if tc.limit == 2 { + want = "ab" + } + assert.Equal(t, want, dumped) + assert.Equal(t, "abc", rec.Body.String()) + }) + } +}