starred/godoxy
1
0
Fork 0
mirror of https://github.com/yusing/godoxy.git synced 2026-04-21 00:11:42 +02:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

fix(middleware): gate only body response modifiers

Browse Source 
Replace the rewrite requirement check with a BodyResponseModifier
marker and treat header and body modifiers separately.

This ensures header/status rewrites still apply when body rewriting is
blocked (for binary, encoded, or chunked responses), while body changes
are skipped safely. It also avoids body reset/close side effects and
returns early on passthrough responses.

Update middleware tests to cover split header/body behavior and themed
middleware body-skip scenarios.
This commit is contained in:
yusing
2026-03-01 04:09:50 +08:00
parent 59238adb5b
commit 41de86de75
7 changed files with 118 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

155
internal/net/gphttp/middleware/middleware_test.go
View File

@@ -1,7 +1,6 @@
package middleware
import (
"errors"
"io"
"net/http"
"net/http/httptest"
@@ -17,50 +16,37 @@ type testPriority struct {
}
var test = NewMiddleware[testPriority]()
var responseRewrite = NewMiddleware[testResponseRewrite]()
var responseHeaderRewrite = NewMiddleware[testHeaderRewrite]()
var responseBodyRewrite = NewMiddleware[testBodyRewrite]()
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 {
type testHeaderRewrite struct {
StatusCode int `json:"status_code"`
HeaderKey string `json:"header_key"`
HeaderVal string `json:"header_val"`
Body string `json:"body"`
}
type closeSensitiveBody struct {
data []byte
offset int
closed bool
}
func (b *closeSensitiveBody) Read(p []byte) (int, error) {
if b.closed {
return 0, errors.New("http: read on closed response body")
}
if b.offset >= len(b.data) {
return 0, io.EOF
}
n := copy(p, b.data[b.offset:])
b.offset += n
return n, nil
}
func (b *closeSensitiveBody) Close() error {
b.closed = true
func (t testHeaderRewrite) modifyResponse(resp *http.Response) error {
resp.StatusCode = t.StatusCode
resp.Header.Set(t.HeaderKey, t.HeaderVal)
return nil
}
func (t testResponseRewrite) modifyResponse(resp *http.Response) error {
resp.StatusCode = t.StatusCode
resp.Header.Set(t.HeaderKey, t.HeaderVal)
type testBodyRewrite struct {
Body string `json:"body"`
}
func (t testBodyRewrite) modifyResponse(resp *http.Response) error {
resp.Body = io.NopCloser(strings.NewReader(t.Body))
return nil
}
func (testBodyRewrite) isBodyResponseModifier() {}
func TestMiddlewarePriority(t *testing.T) {
priorities := []int{4, 7, 9, 0}
chain := make([]*Middleware, len(priorities))
@@ -78,50 +64,66 @@ func TestMiddlewarePriority(t *testing.T) {
}
func TestMiddlewareResponseRewriteGate(t *testing.T) {
opts := OptionsRaw{
headerOpts := OptionsRaw{
"status_code": 418,
"header_key": "X-Rewrite",
"header_val": "1",
"body": "rewritten-body",
}
bodyOpts := OptionsRaw{
"body": "rewritten-body",
}
headerMid, err := responseHeaderRewrite.New(headerOpts)
expect.NoError(t, err)
bodyMid, err := responseBodyRewrite.New(bodyOpts)
expect.NoError(t, err)
tests := []struct {
name string
respHeaders http.Header
respBody []byte
expectBody string
name string
respHeaders http.Header
respBody []byte
expectStatus int
expectHeader string
expectBody string
}{
{
name: "allow_body_rewrite_for_html",
respHeaders: http.Header{
"Content-Type": []string{"text/html; charset=utf-8"},
},
respBody: []byte("<html><body>original</body></html>"),
expectBody: "rewritten-body",
respBody: []byte("<html><body>original</body></html>"),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "rewritten-body",
},
{
name: "allow_body_rewrite_for_json",
respHeaders: http.Header{
"Content-Type": []string{"application/json"},
},
respBody: []byte(`{"message":"original"}`),
expectBody: "rewritten-body",
respBody: []byte(`{"message":"original"}`),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "rewritten-body",
},
{
name: "allow_body_rewrite_for_yaml",
respHeaders: http.Header{
"Content-Type": []string{"application/yaml"},
},
respBody: []byte("k: v"),
expectBody: "rewritten-body",
respBody: []byte("k: v"),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "rewritten-body",
},
{
name: "block_body_rewrite_for_binary_content",
respHeaders: http.Header{
"Content-Type": []string{"application/octet-stream"},
},
respBody: []byte("binary"),
expectBody: "binary",
respBody: []byte("binary"),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "binary",
},
{
name: "block_body_rewrite_for_transfer_encoded_html",
@@ -129,8 +131,10 @@ func TestMiddlewareResponseRewriteGate(t *testing.T) {
"Content-Type": []string{"text/html"},
"Transfer-Encoding": []string{"chunked"},
},
respBody: []byte("<html><body>original</body></html>"),
expectBody: "<html><body>original</body></html>",
respBody: []byte("<html><body>original</body></html>"),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "<html><body>original</body></html>",
},
{
name: "block_body_rewrite_for_content_encoded_html",
@@ -138,34 +142,42 @@ func TestMiddlewareResponseRewriteGate(t *testing.T) {
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"gzip"},
},
respBody: []byte("<html><body>original</body></html>"),
expectBody: "<html><body>original</body></html>",
respBody: []byte("<html><body>original</body></html>"),
expectStatus: http.StatusTeapot,
expectHeader: "1",
expectBody: "<html><body>original</body></html>",
},
}
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,
result, err := newMiddlewaresTest([]*Middleware{headerMid, bodyMid}, &testArgs{
respHeaders: tc.respHeaders,
respBody: tc.respBody,
respStatus: http.StatusOK,
})
expect.NoError(t, err)
expect.Equal(t, result.ResponseStatus, http.StatusTeapot)
expect.Equal(t, result.ResponseHeaders.Get("X-Rewrite"), "1")
expect.Equal(t, result.ResponseStatus, tc.expectStatus)
expect.Equal(t, result.ResponseHeaders.Get("X-Rewrite"), tc.expectHeader)
expect.Equal(t, string(result.Data), tc.expectBody)
})
}
}
func TestMiddlewareResponseRewriteGateServeHTTP(t *testing.T) {
opts := OptionsRaw{
headerOpts := OptionsRaw{
"status_code": 418,
"header_key": "X-Rewrite",
"header_val": "1",
"body": "rewritten-body",
}
bodyOpts := OptionsRaw{
"body": "rewritten-body",
}
headerMid, err := responseHeaderRewrite.New(headerOpts)
expect.NoError(t, err)
bodyMid, err := responseBodyRewrite.New(bodyOpts)
expect.NoError(t, err)
mid := NewMiddlewareChain("test", []*Middleware{headerMid, bodyMid})
tests := []struct {
name string
@@ -221,9 +233,6 @@ func TestMiddlewareResponseRewriteGateServeHTTP(t *testing.T) {
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
mid, err := responseRewrite.New(opts)
expect.NoError(t, err)
req := httptest.NewRequest(http.MethodGet, "http://example.com", nil)
rw := httptest.NewRecorder()
@@ -251,33 +260,17 @@ func TestMiddlewareResponseRewriteGateServeHTTP(t *testing.T) {
}
}
func TestMiddlewareResponseRewriteGateSkipsBodyRewriterWhenRewriteBlocked(t *testing.T) {
originalBody := &closeSensitiveBody{
data: []byte("<html><body>original</body></html>"),
}
req := httptest.NewRequest(http.MethodGet, "http://example.com", nil)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
func TestThemedSkipsBodyRewriteWhenRewriteBlocked(t *testing.T) {
result, err := newMiddlewareTest(Themed, &testArgs{
middlewareOpt: OptionsRaw{
"theme": DarkTheme,
},
respHeaders: http.Header{
"Content-Type": []string{"text/html; charset=utf-8"},
"Transfer-Encoding": []string{"chunked"},
},
Body: originalBody,
ContentLength: -1,
TransferEncoding: []string{"chunked"},
Request: req,
}
themedMid, err := Themed.New(OptionsRaw{
"theme": DarkTheme,
respBody: []byte("<html><body>original</body></html>"),
})
expect.NoError(t, err)
respMod, ok := themedMid.impl.(ResponseModifier)
expect.True(t, ok)
expect.NoError(t, modifyResponseWithBodyRewriteGate(respMod, resp))
data, err := io.ReadAll(resp.Body)
expect.NoError(t, err)
expect.Equal(t, string(data), "<html><body>original</body></html>")
expect.Equal(t, string(result.Data), "<html><body>original</body></html>")
}