fix(middleware): retain partially written BodyDump responses

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
 }

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())
+		})
+	}
+}