diff --git a/internal/net/gphttp/middleware/middleware.go b/internal/net/gphttp/middleware/middleware.go index 3c86ced6..91803f2a 100644 --- a/internal/net/gphttp/middleware/middleware.go +++ b/internal/net/gphttp/middleware/middleware.go @@ -213,14 +213,13 @@ func (m *Middleware) ServeHTTP(next http.HandlerFunc, w http.ResponseWriter, r * return } isBodyModifier := isBodyResponseModifier(exec) - - shouldBuffer := canBufferAndModifyResponseBody if !isBodyModifier { - // Header-only response modifiers do not need body rewrite capability checks. - // We still respect max buffer limits and may fall back to passthrough for large bodies. - shouldBuffer = func(http.Header) bool { return true } + mrw := httputils.NewModifyResponseWriter(w, r, exec.modifyResponse) + next(mrw, r) + return } - lrm := httputils.NewLazyResponseModifier(w, shouldBuffer) + + lrm := httputils.NewLazyResponseModifier(w, canBufferAndModifyResponseBody) lrm.SetMaxBufferedBytes(maxModifiableBody) defer func() { _, err := lrm.FlushRelease() diff --git a/internal/net/gphttp/middleware/middleware_test.go b/internal/net/gphttp/middleware/middleware_test.go index 6ae2ea2d..5d961ade 100644 --- a/internal/net/gphttp/middleware/middleware_test.go +++ b/internal/net/gphttp/middleware/middleware_test.go @@ -260,6 +260,36 @@ func TestMiddlewareResponseRewriteGateServeHTTP(t *testing.T) { } } +func TestMiddlewareHeaderRewriteDoesNotBufferLargeBody(t *testing.T) { + headerMid, err := responseHeaderRewrite.New(OptionsRaw{ + "status_code": http.StatusAccepted, + "header_key": "X-Rewrite", + "header_val": "1", + }) + expect.NoError(t, err) + + req := httptest.NewRequest(http.MethodGet, "http://example.com", nil) + rw := httptest.NewRecorder() + + next := func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "video/mp4") + w.Header().Set("Content-Length", strconv.Itoa(64*1024*1024)) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("video")) + } + + headerMid.ServeHTTP(next, rw, req) + + resp := rw.Result() + defer resp.Body.Close() + data, readErr := io.ReadAll(resp.Body) + expect.NoError(t, readErr) + + expect.Equal(t, resp.StatusCode, http.StatusAccepted) + expect.Equal(t, resp.Header.Get("X-Rewrite"), "1") + expect.Equal(t, string(data), "video") +} + func TestThemedSkipsBodyRewriteWhenRewriteBlocked(t *testing.T) { result, err := newMiddlewareTest(Themed, &testArgs{ middlewareOpt: OptionsRaw{