From 169358659ad76fdda9532aa940e20c765e8340ce Mon Sep 17 00:00:00 2001 From: yusing Date: Mon, 23 Feb 2026 23:46:41 +0800 Subject: [PATCH] refactor(middleware): improve response body modification gating Refactor response body modification to only allow text-like content types (JSON, YAML, XML, etc.) instead of all HTML responses. Body modification is now blocked for binary content and transfer/content encoded responses, while status code and headers can still be modified. This prevents issues with compressed or streaming responses while maintaining the ability to modify text-based API responses. --- goutils | 2 +- internal/net/gphttp/middleware/README.md | 2 + internal/net/gphttp/middleware/middleware.go | 89 +++++++++++++---- .../net/gphttp/middleware/middleware_chain.go | 16 ++- .../net/gphttp/middleware/middleware_test.go | 98 +++++++++++++++++++ .../net/gphttp/middleware/test_utils_test.go | 21 +++- 6 files changed, 206 insertions(+), 22 deletions(-) diff --git a/goutils b/goutils index 482b5bca..3be815cb 160000 --- a/goutils +++ b/goutils @@ -1 +1 @@ -Subproject commit 482b5bca9f2eae9d293c1cc46674635a201af231 +Subproject commit 3be815cb6e3b7b872d71dd041427ee6674683bda diff --git a/internal/net/gphttp/middleware/README.md b/internal/net/gphttp/middleware/README.md index 8c3a4666..6ed57ed7 100644 --- a/internal/net/gphttp/middleware/README.md +++ b/internal/net/gphttp/middleware/README.md @@ -13,6 +13,8 @@ This package implements a flexible HTTP middleware system for GoDoxy. Middleware - **Bypass Rules**: Skip middleware based on request properties - **Dynamic Loading**: Load middleware definitions from files at runtime +Response body rewriting is only applied to unencoded, text-like content types (for example `text/*`, JSON, YAML, XML). Response status and headers can always be modified. + ## Architecture ```mermaid diff --git a/internal/net/gphttp/middleware/middleware.go b/internal/net/gphttp/middleware/middleware.go index 828cb171..54651327 100644 --- a/internal/net/gphttp/middleware/middleware.go +++ b/internal/net/gphttp/middleware/middleware.go @@ -6,6 +6,7 @@ import ( "net/http" "reflect" "sort" + "strings" "github.com/bytedance/sonic" "github.com/rs/zerolog" @@ -195,21 +196,15 @@ func (m *Middleware) ServeHTTP(next http.HandlerFunc, w http.ResponseWriter, r * } if exec, ok := m.impl.(ResponseModifier); ok { - lrm := httputils.NewLazyResponseModifier(w, needsBuffering) + rm := httputils.NewResponseModifier(w) defer func() { - _, err := lrm.FlushRelease() + _, err := rm.FlushRelease() if err != nil { m.LogError(r).Err(err).Msg("failed to flush response") } }() - next(lrm, r) + next(rm, r) - // Skip modification if response wasn't buffered (non-HTML content) - if !lrm.IsBuffered() { - return - } - - rm := lrm.ResponseModifier() currentBody := rm.BodyReader() currentResp := &http.Response{ StatusCode: rm.StatusCode(), @@ -218,20 +213,31 @@ func (m *Middleware) ServeHTTP(next http.HandlerFunc, w http.ResponseWriter, r * Body: currentBody, Request: r, } - if err := exec.modifyResponse(currentResp); err != nil { + allowBodyModification := canModifyResponseBody(currentResp) + respToModify := currentResp + if !allowBodyModification { + shadow := *currentResp + shadow.Body = eofReader{} + respToModify = &shadow + } + if err := exec.modifyResponse(respToModify); err != nil { log.Err(err).Str("middleware", m.Name()).Str("url", fullURL(r)).Msg("failed to modify response") } // override the response status code - rm.WriteHeader(currentResp.StatusCode) + rm.WriteHeader(respToModify.StatusCode) // overriding the response header - maps.Copy(rm.Header(), currentResp.Header) + maps.Copy(rm.Header(), respToModify.Header) // override the content length and body if changed - if currentResp.Body != currentBody { - if err := rm.SetBody(currentResp.Body); err != nil { - m.LogError(r).Err(err).Msg("failed to set response body") + if respToModify.Body != currentBody { + if allowBodyModification { + if err := rm.SetBody(respToModify.Body); err != nil { + m.LogError(r).Err(err).Msg("failed to set response body") + } + } else { + respToModify.Body.Close() } } } else { @@ -239,10 +245,55 @@ func (m *Middleware) ServeHTTP(next http.HandlerFunc, w http.ResponseWriter, r * } } -// needsBuffering determines if a response should be buffered for modification. -// Only HTML responses need buffering; streaming content (video, audio, etc.) should pass through. -func needsBuffering(header http.Header) bool { - return httputils.GetContentType(header).IsHTML() +func canModifyResponseBody(resp *http.Response) bool { + if hasNonIdentityEncoding(resp.TransferEncoding) { + return false + } + if hasNonIdentityEncoding(resp.Header.Values("Transfer-Encoding")) { + return false + } + if hasNonIdentityEncoding(resp.Header.Values("Content-Encoding")) { + return false + } + return isTextLikeMediaType(string(httputils.GetContentType(resp.Header))) +} + +func hasNonIdentityEncoding(values []string) bool { + for _, value := range values { + for _, token := range strings.Split(value, ",") { + if strings.TrimSpace(token) == "" || strings.EqualFold(strings.TrimSpace(token), "identity") { + continue + } + return true + } + } + return false +} + +func isTextLikeMediaType(contentType string) bool { + if contentType == "" { + return false + } + contentType = strings.ToLower(contentType) + if strings.HasPrefix(contentType, "text/") { + return true + } + if contentType == "application/json" || strings.HasSuffix(contentType, "+json") { + return true + } + if contentType == "application/xml" || strings.HasSuffix(contentType, "+xml") { + return true + } + if strings.Contains(contentType, "yaml") || strings.Contains(contentType, "toml") { + return true + } + if strings.Contains(contentType, "javascript") || strings.Contains(contentType, "ecmascript") { + return true + } + if strings.Contains(contentType, "csv") { + return true + } + return contentType == "application/x-www-form-urlencoded" } func (m *Middleware) LogWarn(req *http.Request) *zerolog.Event { diff --git a/internal/net/gphttp/middleware/middleware_chain.go b/internal/net/gphttp/middleware/middleware_chain.go index d9e2be97..e6e4e801 100644 --- a/internal/net/gphttp/middleware/middleware_chain.go +++ b/internal/net/gphttp/middleware/middleware_chain.go @@ -1,6 +1,7 @@ package middleware import ( + "maps" "net/http" "strconv" @@ -46,10 +47,23 @@ func (m *middlewareChain) modifyResponse(resp *http.Response) error { if len(m.modResps) == 0 { return nil } + allowBodyModification := canModifyResponseBody(resp) for i, mr := range m.modResps { - if err := mr.modifyResponse(resp); err != nil { + respToModify := resp + if !allowBodyModification { + shadow := *resp + shadow.Body = eofReader{} + respToModify = &shadow + } + if err := mr.modifyResponse(respToModify); err != nil { return gperr.PrependSubject(err, strconv.Itoa(i)) } + if !allowBodyModification { + resp.StatusCode = respToModify.StatusCode + if respToModify.Header != nil { + maps.Copy(resp.Header, respToModify.Header) + } + } } return nil } diff --git a/internal/net/gphttp/middleware/middleware_test.go b/internal/net/gphttp/middleware/middleware_test.go index 4f289f45..1242a63a 100644 --- a/internal/net/gphttp/middleware/middleware_test.go +++ b/internal/net/gphttp/middleware/middleware_test.go @@ -1,6 +1,7 @@ package middleware import ( + "io" "net/http" "strconv" "strings" @@ -14,12 +15,27 @@ type testPriority struct { } var test = NewMiddleware[testPriority]() +var responseRewrite = NewMiddleware[testResponseRewrite]() func (t testPriority) before(w http.ResponseWriter, r *http.Request) bool { w.Header().Add("Test-Value", strconv.Itoa(t.Value)) return true } +type testResponseRewrite struct { + StatusCode int `json:"status_code"` + HeaderKey string `json:"header_key"` + HeaderVal string `json:"header_val"` + Body string `json:"body"` +} + +func (t testResponseRewrite) modifyResponse(resp *http.Response) error { + resp.StatusCode = t.StatusCode + resp.Header.Set(t.HeaderKey, t.HeaderVal) + resp.Body = io.NopCloser(strings.NewReader(t.Body)) + return nil +} + func TestMiddlewarePriority(t *testing.T) { priorities := []int{4, 7, 9, 0} chain := make([]*Middleware, len(priorities)) @@ -35,3 +51,85 @@ func TestMiddlewarePriority(t *testing.T) { expect.NoError(t, err) expect.Equal(t, strings.Join(res.ResponseHeaders["Test-Value"], ","), "3,0,1,2") } + +func TestMiddlewareResponseRewriteGate(t *testing.T) { + opts := OptionsRaw{ + "status_code": 418, + "header_key": "X-Rewrite", + "header_val": "1", + "body": "rewritten-body", + } + + tests := []struct { + name string + respHeaders http.Header + respBody []byte + expectBody string + }{ + { + name: "allow_body_rewrite_for_html", + respHeaders: http.Header{ + "Content-Type": []string{"text/html; charset=utf-8"}, + }, + respBody: []byte("original"), + expectBody: "rewritten-body", + }, + { + name: "allow_body_rewrite_for_json", + respHeaders: http.Header{ + "Content-Type": []string{"application/json"}, + }, + respBody: []byte(`{"message":"original"}`), + expectBody: "rewritten-body", + }, + { + name: "allow_body_rewrite_for_yaml", + respHeaders: http.Header{ + "Content-Type": []string{"application/yaml"}, + }, + respBody: []byte("k: v"), + expectBody: "rewritten-body", + }, + { + name: "block_body_rewrite_for_binary_content", + respHeaders: http.Header{ + "Content-Type": []string{"application/octet-stream"}, + }, + respBody: []byte("binary"), + expectBody: "binary", + }, + { + name: "block_body_rewrite_for_transfer_encoded_html", + respHeaders: http.Header{ + "Content-Type": []string{"text/html"}, + "Transfer-Encoding": []string{"chunked"}, + }, + respBody: []byte("original"), + expectBody: "original", + }, + { + name: "block_body_rewrite_for_content_encoded_html", + respHeaders: http.Header{ + "Content-Type": []string{"text/html"}, + "Content-Encoding": []string{"gzip"}, + }, + respBody: []byte("original"), + expectBody: "original", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + result, err := newMiddlewareTest(responseRewrite, &testArgs{ + middlewareOpt: opts, + respHeaders: tc.respHeaders, + respBody: tc.respBody, + respStatus: http.StatusOK, + }) + expect.NoError(t, err) + expect.Equal(t, result.ResponseStatus, 418) + expect.Equal(t, result.ResponseHeaders.Get("X-Rewrite"), "1") + expect.Equal(t, string(result.Data), tc.expectBody) + }) + } +} diff --git a/internal/net/gphttp/middleware/test_utils_test.go b/internal/net/gphttp/middleware/test_utils_test.go index f920c69d..dae8fb79 100644 --- a/internal/net/gphttp/middleware/test_utils_test.go +++ b/internal/net/gphttp/middleware/test_utils_test.go @@ -7,6 +7,7 @@ import ( "maps" "net/http" "net/http/httptest" + "strings" "github.com/bytedance/sonic" "github.com/yusing/godoxy/internal/common" @@ -54,7 +55,7 @@ func (rt *requestRecorder) RoundTrip(req *http.Request) (resp *http.Response, er resp = &http.Response{ Status: http.StatusText(rt.args.respStatus), StatusCode: rt.args.respStatus, - Header: testHeaders, + Header: maps.Clone(testHeaders), Body: io.NopCloser(bytes.NewReader(rt.args.respBody)), ContentLength: int64(len(rt.args.respBody)), Request: req, @@ -65,9 +66,27 @@ func (rt *requestRecorder) RoundTrip(req *http.Request) (resp *http.Response, er return nil, err } maps.Copy(resp.Header, rt.args.respHeaders) + if transferEncoding := resp.Header.Values("Transfer-Encoding"); len(transferEncoding) > 0 { + resp.TransferEncoding = parseHeaderTokens(transferEncoding) + resp.ContentLength = -1 + } return resp, nil } +func parseHeaderTokens(values []string) []string { + var tokens []string + for _, value := range values { + for token := range strings.SplitSeq(value, ",") { + token = strings.TrimSpace(token) + if token == "" { + continue + } + tokens = append(tokens, token) + } + } + return tokens +} + type TestResult struct { RequestHeaders http.Header ResponseHeaders http.Header