diff --git a/config/config.go b/config/config.go index 3b0e892b..e21802b2 100644 --- a/config/config.go +++ b/config/config.go @@ -402,6 +402,8 @@ type HTTPBodyRewriteConfig struct { Type string Match string Replacement string + // name of the rewriter plugin (via registry) + Rewriter string } type NodeFilterConfig struct { diff --git a/config/parsing/node/parse.go b/config/parsing/node/parse.go index feb09096..ce9b7127 100644 --- a/config/parsing/node/parse.go +++ b/config/parsing/node/parse.go @@ -29,6 +29,29 @@ import ( // metadata-driven options (so_mark, interface, netns, proxy protocol), sets up // bypass rules, node filters, HTTP settings, and TLS node settings. The hop // parameter is used only for logging context. + +func parseBodyRewrites(vs []config.HTTPBodyRewriteConfig, log logger.Logger) []chain.HTTPBodyRewriteSettings { + var out []chain.HTTPBodyRewriteSettings + for _, v := range vs { + pattern, _ := regexp.Compile(v.Match) + rw := chain.HTTPBodyRewriteSettings{ + Type: v.Type, + Pattern: pattern, + Replacement: []byte(v.Replacement), + } + if v.Rewriter != "" { + if !registry.RewriterRegistry().IsRegistered(v.Rewriter) { + log.Warnf("rewriter %q not found in registry for rewrite rule", v.Rewriter) + } + rw.Rewriter = registry.RewriterRegistry().Get(v.Rewriter) + } + if pattern != nil || rw.Rewriter != nil { + out = append(out, rw) + } + } + return out +} + func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.Node, error) { if cfg == nil { return nil, nil @@ -226,33 +249,9 @@ func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.No }) } } - for _, v := range cfg.HTTP.RewriteBody { - if pattern, _ := regexp.Compile(v.Match); pattern != nil { - settings.RewriteResponseBody = append(settings.RewriteResponseBody, chain.HTTPBodyRewriteSettings{ - Type: v.Type, - Pattern: pattern, - Replacement: []byte(v.Replacement), - }) - } - } - for _, v := range cfg.HTTP.RewriteResponseBody { - if pattern, _ := regexp.Compile(v.Match); pattern != nil { - settings.RewriteResponseBody = append(settings.RewriteResponseBody, chain.HTTPBodyRewriteSettings{ - Type: v.Type, - Pattern: pattern, - Replacement: []byte(v.Replacement), - }) - } - } - for _, v := range cfg.HTTP.RewriteRequestBody { - if pattern, _ := regexp.Compile(v.Match); pattern != nil { - settings.RewriteRequestBody = append(settings.RewriteRequestBody, chain.HTTPBodyRewriteSettings{ - Type: v.Type, - Pattern: pattern, - Replacement: []byte(v.Replacement), - }) - } - } + settings.RewriteResponseBody = append(settings.RewriteResponseBody, parseBodyRewrites(cfg.HTTP.RewriteBody, log)...) + settings.RewriteResponseBody = append(settings.RewriteResponseBody, parseBodyRewrites(cfg.HTTP.RewriteResponseBody, log)...) + settings.RewriteRequestBody = append(settings.RewriteRequestBody, parseBodyRewrites(cfg.HTTP.RewriteRequestBody, log)...) opts = append(opts, chain.HTTPNodeOption(settings)) } diff --git a/config/parsing/rewriter/parse.go b/config/parsing/rewriter/parse.go index 3ba5535a..8ad08d6a 100644 --- a/config/parsing/rewriter/parse.go +++ b/config/parsing/rewriter/parse.go @@ -39,6 +39,7 @@ func ParseRewriter(cfg *config.RewriterConfig) rewriter.Rewriter { cfg.Name, cfg.Plugin.Addr, plugin.TokenOption(cfg.Plugin.Token), plugin.TLSConfigOption(tlsCfg), + plugin.TimeoutOption(cfg.Plugin.Timeout), ) } } diff --git a/config/parsing/service/parse.go b/config/parsing/service/parse.go index 2bf7052f..8dded45e 100644 --- a/config/parsing/service/parse.go +++ b/config/parsing/service/parse.go @@ -10,8 +10,8 @@ import ( "github.com/go-gost/core/chain" "github.com/go-gost/core/handler" "github.com/go-gost/core/hop" - "github.com/go-gost/core/rewriter" "github.com/go-gost/core/listener" + "github.com/go-gost/core/rewriter" "github.com/go-gost/core/logger" "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" @@ -356,6 +356,9 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) { var rew rewriter.Rewriter if cfg.Rewriter != "" { + if !registry.RewriterRegistry().IsRegistered(cfg.Rewriter) { + serviceLogger.Warnf("rewriter %q not found in registry", cfg.Rewriter) + } rew = registry.RewriterRegistry().Get(cfg.Rewriter) } diff --git a/go.mod b/go.mod index 7f0eb84a..ddb8e599 100644 --- a/go.mod +++ b/go.mod @@ -7,7 +7,7 @@ require ( github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d github.com/gin-contrib/cors v1.7.2 github.com/gin-gonic/gin v1.10.1 - github.com/go-gost/core v0.5.0 + github.com/go-gost/core v0.5.1 github.com/go-gost/go-shadowsocks2 v0.1.3 github.com/go-gost/gosocks4 v0.1.0 github.com/go-gost/gosocks5 v0.5.0 diff --git a/go.sum b/go.sum index b8d48e92..5149a1cc 100644 --- a/go.sum +++ b/go.sum @@ -53,6 +53,8 @@ github.com/gin-gonic/gin v1.10.1 h1:T0ujvqyCSqRopADpgPgiTT63DUQVSfojyME59Ei63pQ= github.com/gin-gonic/gin v1.10.1/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= github.com/go-gost/core v0.5.0 h1:3SkPojjqw2Av9l5OSyRVlBpRvWhaXgU5BWMXB9aR4SM= github.com/go-gost/core v0.5.0/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A= +github.com/go-gost/core v0.5.1 h1:HbIn3naEOC661Z4SwzaSYUzffWSRbU0P6gkxiHmuG0I= +github.com/go-gost/core v0.5.1/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A= github.com/go-gost/go-shadowsocks2 v0.1.3 h1:6CUZLp+mTWXnKP2aK8/Z9ZP+ERMX9gSbywmPu4kGX/A= github.com/go-gost/go-shadowsocks2 v0.1.3/go.mod h1:866zFNNI3He6Wef1M/IvAjTal74WhcfKfBgRpTlkKys= github.com/go-gost/gosocks4 v0.1.0 h1:eAzev6qw4fzkFQKC9uCHLVNnnPdHyqCggbnfNN80Pmk= diff --git a/internal/util/forwarder/sniffer_http.go b/internal/util/forwarder/sniffer_http.go index 9748b20f..5823d277 100644 --- a/internal/util/forwarder/sniffer_http.go +++ b/internal/util/forwarder/sniffer_http.go @@ -369,7 +369,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, // Rewrite request body before wrapping for recording, // so the recorder sees the rewritten content. - if err = rewriteReqBody(req, reqBodyRewrites...); err != nil { + if err = rewriteReqBody(ctx, req, reqBodyRewrites...); err != nil { log.Errorf("rewrite request body: %v", err) return } @@ -446,7 +446,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, resp.Header.Set("Connection", "close") } - if err = rewriteRespBody(resp, respBodyRewrites...); err != nil { + if err = rewriteRespBody(ctx, resp, respBodyRewrites...); err != nil { log.Errorf("rewrite body: %v", err) return } diff --git a/internal/util/forwarder/sniffer_rewrite.go b/internal/util/forwarder/sniffer_rewrite.go index 71adccd3..7b6f8347 100644 --- a/internal/util/forwarder/sniffer_rewrite.go +++ b/internal/util/forwarder/sniffer_rewrite.go @@ -41,7 +41,7 @@ func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWrite return xnet.Pipe(ctx, rw, cc) } -func rewriteRespBody(resp *http.Response, rewrites ...chain.HTTPBodyRewriteSettings) error { +func rewriteRespBody(ctx context.Context, resp *http.Response, rewrites ...chain.HTTPBodyRewriteSettings) error { if resp == nil || len(rewrites) == 0 || resp.ContentLength <= 0 { return nil } @@ -65,7 +65,12 @@ func rewriteRespBody(resp *http.Response, rewrites ...chain.HTTPBodyRewriteSetti continue } - if rewrite.Pattern != nil { + if rewrite.Rewriter != nil { + body, err = rewrite.Rewriter.Rewrite(ctx, body) + if err != nil { + return err + } + } else if rewrite.Pattern != nil { body = rewrite.Pattern.ReplaceAll(body, rewrite.Replacement) } } @@ -90,7 +95,7 @@ func drainBody(b io.ReadCloser) (body []byte, err error) { return buf.Bytes(), nil } -func rewriteReqBody(req *http.Request, rewrites ...chain.HTTPBodyRewriteSettings) error { +func rewriteReqBody(ctx context.Context, req *http.Request, rewrites ...chain.HTTPBodyRewriteSettings) error { if req == nil || len(rewrites) == 0 || req.Body == nil || req.ContentLength <= 0 { return nil } @@ -114,7 +119,12 @@ func rewriteReqBody(req *http.Request, rewrites ...chain.HTTPBodyRewriteSettings continue } - if rewrite.Pattern != nil { + if rewrite.Rewriter != nil { + body, err = rewrite.Rewriter.Rewrite(ctx, body) + if err != nil { + return err + } + } else if rewrite.Pattern != nil { body = rewrite.Pattern.ReplaceAll(body, rewrite.Replacement) } } diff --git a/internal/util/forwarder/sniffer_test.go b/internal/util/forwarder/sniffer_test.go index dd5adda7..304a46f6 100644 --- a/internal/util/forwarder/sniffer_test.go +++ b/internal/util/forwarder/sniffer_test.go @@ -164,7 +164,7 @@ func TestRewriteReqBody(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.req != nil { - _ = rewriteReqBody(tt.req, tt.rewrites...) + _ = rewriteReqBody(context.Background(), tt.req, tt.rewrites...) if tt.req.Body != nil { body, _ := io.ReadAll(tt.req.Body) tt.req.Body.Close() @@ -178,7 +178,7 @@ func TestRewriteReqBody(t *testing.T) { t.Errorf("ContentLength = %d, want %d", tt.req.ContentLength, tt.wantCL) } } else { - _ = rewriteReqBody(nil) // should not panic + _ = rewriteReqBody(context.Background(), nil) // should not panic } }) } @@ -344,8 +344,8 @@ func TestDrainBody(t *testing.T) { func TestRewriteRespBody(t *testing.T) { t.Run("nil response", func(t *testing.T) { - if err := rewriteRespBody(nil); err != nil { - t.Errorf("rewriteRespBody(nil) = %v, want nil", err) + if err := rewriteRespBody(context.Background(), nil); err != nil { + t.Errorf("rewriteRespBody(context.Background(), nil) = %v, want nil", err) } }) @@ -354,7 +354,7 @@ func TestRewriteRespBody(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("original")), } - _ = rewriteRespBody(resp) + _ = rewriteRespBody(context.Background(), resp) // Body unchanged got, _ := io.ReadAll(resp.Body) if string(got) != "original" { @@ -367,7 +367,7 @@ func TestRewriteRespBody(t *testing.T) { ContentLength: 0, Body: io.NopCloser(strings.NewReader("original")), } - _ = rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{ + _ = rewriteRespBody(context.Background(), resp, chain.HTTPBodyRewriteSettings{ Pattern: regexp.MustCompile("."), Type: "*", Replacement: []byte("replaced"), @@ -385,7 +385,7 @@ func TestRewriteRespBody(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("original")), } - _ = rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{ + _ = rewriteRespBody(context.Background(), resp, chain.HTTPBodyRewriteSettings{ Type: "*", Replacement: []byte("replaced"), }) @@ -413,7 +413,7 @@ func TestRewriteRespBodyContentTypeFilter(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("original")), } - _ = rewriteRespBody(resp, makeRewrite("text/html", "replaced")) + _ = rewriteRespBody(context.Background(), resp, makeRewrite("text/html", "replaced")) got, _ := io.ReadAll(resp.Body) if string(got) != "replaced" { t.Errorf("body = %q, want %q", string(got), "replaced") @@ -426,7 +426,7 @@ func TestRewriteRespBodyContentTypeFilter(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("hello")), } - _ = rewriteRespBody(resp, makeRewrite("*", "world")) + _ = rewriteRespBody(context.Background(), resp, makeRewrite("*", "world")) got, _ := io.ReadAll(resp.Body) if string(got) != "world" { t.Errorf("body = %q, want %q", string(got), "world") @@ -439,7 +439,7 @@ func TestRewriteRespBodyContentTypeFilter(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("original")), } - _ = rewriteRespBody(resp, makeRewrite("text/html", "replaced")) + _ = rewriteRespBody(context.Background(), resp, makeRewrite("text/html", "replaced")) got, _ := io.ReadAll(resp.Body) if string(got) != "original" { t.Errorf("body = %q, want %q (unchanged)", string(got), "original") @@ -454,7 +454,7 @@ func TestRewriteRespBodyContentTypeFilter(t *testing.T) { } // Empty type defaults to "text/html", and response has no Content-Type -> "" containing "text/html"? No. // strings.Contains("text/html", "") is always true, so rewrite should happen. - _ = rewriteRespBody(resp, makeRewrite("", "replaced")) + _ = rewriteRespBody(context.Background(), resp, makeRewrite("", "replaced")) got, _ := io.ReadAll(resp.Body) if string(got) != "replaced" { t.Errorf("body = %q, want %q (should be rewritten)", string(got), "replaced") @@ -1146,7 +1146,7 @@ func TestRewriteRespBody_PatternReplacement(t *testing.T) { ContentLength: 100, Body: io.NopCloser(strings.NewReader("