diff --git a/config/config.go b/config/config.go index 9c10be9..6637cdb 100644 --- a/config/config.go +++ b/config/config.go @@ -464,6 +464,17 @@ type HTTPBodyRewriteConfig struct { MaxChunkSize int `yaml:"maxChunkSize,omitempty" json:"maxChunkSize,omitempty"` } +type HTTPHeaderRewriteConfig struct { + // Name is the header name regex, e.g. "(?i)^(location|set-cookie)$". + Name string + // Match is the header value regex. + Match string + // Replacement is the replacement string. + Replacement string + // Rewriter is the name of the rewriter plugin (via registry). + Rewriter string `yaml:",omitempty" json:"rewriter,omitempty"` +} + type NodeFilterConfig struct { Host string `yaml:",omitempty" json:"host,omitempty"` Protocol string `yaml:",omitempty" json:"protocol,omitempty"` @@ -505,6 +516,9 @@ type ProbeConfig struct { type HTTPNodeConfig struct { // rewrite host header Host string `yaml:",omitempty" json:"host,omitempty"` + // regex matched against the request path to derive Host; Host becomes a + // replacement template ($1, $2, ...) when set. + HostPattern string `yaml:"hostPattern,omitempty" json:"hostPattern,omitempty"` // Deprecated: use requestHeader instead Header map[string]string `yaml:",omitempty" json:"header,omitempty"` // additional request header @@ -522,6 +536,11 @@ type HTTPNodeConfig struct { // rewrite response body RewriteResponseBody []HTTPBodyRewriteConfig `yaml:"rewriteResponseBody,omitempty" json:"rewriteResponseBody,omitempty"` + // rewrite request header values + RewriteRequestHeader []HTTPHeaderRewriteConfig `yaml:"rewriteRequestHeader,omitempty" json:"rewriteRequestHeader,omitempty"` + // rewrite response header values + RewriteResponseHeader []HTTPHeaderRewriteConfig `yaml:"rewriteResponseHeader,omitempty" json:"rewriteResponseHeader,omitempty"` + // comma-separated response status codes marking the node failed, e.g. "429,5xx" FailCodes string `yaml:"failCodes,omitempty" json:"failCodes,omitempty"` diff --git a/config/parsing/node/parse.go b/config/parsing/node/parse.go index 27cb6bb..a82104f 100644 --- a/config/parsing/node/parse.go +++ b/config/parsing/node/parse.go @@ -143,6 +143,38 @@ func parseBodyRewrites(vs []config.HTTPBodyRewriteConfig, log logger.Logger) []c return out } +func parseHeaderRewrites(vs []config.HTTPHeaderRewriteConfig, log logger.Logger) []chain.HTTPHeaderRewriteSettings { + var out []chain.HTTPHeaderRewriteSettings + for _, v := range vs { + var name, pattern *regexp.Regexp + if v.Name != "" { + name, _ = regexp.Compile(v.Name) + } + if v.Match != "" { + pattern, _ = regexp.Compile(v.Match) + } + + rw := chain.HTTPHeaderRewriteSettings{ + Name: name, + 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) + } + // Gate on the config string, not the compiled regexp: regexp.Compile("") + // returns a non-nil match-all regexp, so an empty Name with no rewriter + // would otherwise become a silent no-op rule. + if v.Name != "" || 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 @@ -322,6 +354,10 @@ func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.No settings.RequestHeader = cfg.HTTP.Header } + if v := cfg.HTTP.HostPattern; v != "" { + settings.HostPattern, _ = regexp.Compile(v) + } + if auth := cfg.HTTP.Auth; auth != nil && auth.Username != "" { settings.Auther = xauth.NewAuthenticator( xauth.AuthsOption(map[string]string{auth.Username: auth.Password}), @@ -348,6 +384,8 @@ func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.No 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)...) + settings.RewriteRequestHeader = parseHeaderRewrites(cfg.HTTP.RewriteRequestHeader, log) + settings.RewriteResponseHeader = parseHeaderRewrites(cfg.HTTP.RewriteResponseHeader, log) if v := strings.TrimSpace(cfg.HTTP.FailCodes); v != "" { settings.FailCodes = parseFailCodes(v, nodeLogger) diff --git a/config/parsing/node/parse_test.go b/config/parsing/node/parse_test.go index 6b367b6..26cdb09 100644 --- a/config/parsing/node/parse_test.go +++ b/config/parsing/node/parse_test.go @@ -387,6 +387,44 @@ func TestParseNode_WithHTTPDeprecatedHeader(t *testing.T) { } } +func TestParseNode_WithHeaderRewrite(t *testing.T) { + n, err := ParseNode("test-hop", &config.NodeConfig{ + Name: "hdr-node", + Addr: "example.com:8080", + HTTP: &config.HTTPNodeConfig{ + RewriteRequestHeader: []config.HTTPHeaderRewriteConfig{ + {Name: "(?i)^referer$", Match: ".*", Replacement: ""}, + }, + RewriteResponseHeader: []config.HTTPHeaderRewriteConfig{ + {Name: "(?i)^location$", Match: "https://github.com", Replacement: "https://127.0.0.1:8080"}, + }, + }, + }, testLogger()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n == nil { + t.Fatal("expected non-nil node") + } + settings := n.Options().HTTP + if settings == nil { + t.Fatal("expected HTTP settings") + } + if len(settings.RewriteRequestHeader) != 1 || len(settings.RewriteResponseHeader) != 1 { + t.Fatalf("rewrite header counts = %d/%d, want 1/1", len(settings.RewriteRequestHeader), len(settings.RewriteResponseHeader)) + } +} + +func TestParseHeaderRewrites_EmptyNameNoRewriter(t *testing.T) { + // Empty Name with no Rewriter must not be appended (silent no-op guard). + out := parseHeaderRewrites([]config.HTTPHeaderRewriteConfig{ + {Name: "", Match: ".*", Replacement: ""}, + }, testLogger()) + if len(out) != 0 { + t.Fatalf("out = %d rules, want 0", len(out)) + } +} + func TestParseNode_WithTLSNode(t *testing.T) { n, err := ParseNode("test-hop", &config.NodeConfig{ Name: "tls-node", diff --git a/internal/util/forwarder/rewrite_metadata.go b/internal/util/forwarder/rewrite_metadata.go new file mode 100644 index 0000000..958a687 --- /dev/null +++ b/internal/util/forwarder/rewrite_metadata.go @@ -0,0 +1,39 @@ +package forwarder + +import "maps" + +// Rewriter-plugin metadata common fields (JSON convention, opaque to +// core/plugin — carried in RewriteRequest.Metadata). +const ( + MetaKeySid = "sid" + MetaKeyDirection = "direction" + MetaKeyURI = "uri" + MetaKeyKind = "kind" +) + +// Kind values: what the plugin is receiving (byte payload vs header block). +const ( + KindBody = "body" + KindHeader = "header" +) + +// rewriteMeta builds the standard plugin metadata: sid/direction/uri/kind +// plus dimension-specific extras. Empty common fields are omitted so plugins +// can rely on presence rather than empty strings. +func rewriteMeta(sid, direction, uri, kind string, extras map[string]any) map[string]any { + md := map[string]any{} + if sid != "" { + md[MetaKeySid] = sid + } + if direction != "" { + md[MetaKeyDirection] = direction + } + if uri != "" { + md[MetaKeyURI] = uri + } + if kind != "" { + md[MetaKeyKind] = kind + } + maps.Copy(md, extras) + return md +} diff --git a/internal/util/forwarder/sniffer_header_rewrite_test.go b/internal/util/forwarder/sniffer_header_rewrite_test.go new file mode 100644 index 0000000..6f3b56e --- /dev/null +++ b/internal/util/forwarder/sniffer_header_rewrite_test.go @@ -0,0 +1,235 @@ +package forwarder + +import ( + "bytes" + "context" + "net/http" + "regexp" + "testing" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/rewriter" +) + +// headerRewriter is a rewriter.Rewriter that records the metadata it receives. +type headerRewriter struct { + cb func(b []byte) []byte + called bool + metadata map[string]any +} + +func (m *headerRewriter) Rewrite(_ context.Context, b []byte, opts ...rewriter.RewriteOption) ([]byte, error) { + m.called = true + var ro rewriter.RewriteOptions + for _, o := range opts { + o(&ro) + } + if md, ok := ro.Metadata.(map[string]any); ok { + m.metadata = md + } + if m.cb != nil { + return m.cb(b), nil + } + return b, nil +} + +func compileHeaderRewrite(name, match, replacement string) chain.HTTPHeaderRewriteSettings { + var n, p *regexp.Regexp + if name != "" { + n = regexp.MustCompile(name) + } + if match != "" { + p = regexp.MustCompile(match) + } + return chain.HTTPHeaderRewriteSettings{ + Name: n, + Pattern: p, + Replacement: []byte(replacement), + } +} + +func TestRewriteHeaderBlock_Regex(t *testing.T) { + t.Run("location rewrite", func(t *testing.T) { + h := http.Header{"Location": {"https://github.com/foo"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{compileHeaderRewrite(`(?i)^location$`, `https://(github\.com)`, `https://127.0.0.1:8080/$1`)}, + "response", "/") + if err != nil { + t.Fatal(err) + } + if got := h.Get("Location"); got != "https://127.0.0.1:8080/github.com/foo" { + t.Errorf("Location = %q", got) + } + }) + + t.Run("multi-value set-cookie per-value", func(t *testing.T) { + h := http.Header{ + "Set-Cookie": { + "a=1; Domain=.github.com; Path=/", + "b=2; Domain=github.com", + }, + } + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{compileHeaderRewrite(`(?i)^set-cookie$`, `Domain=\.?github\.com;?\s*`, ``)}, + "response", "/") + if err != nil { + t.Fatal(err) + } + for _, v := range h.Values("Set-Cookie") { + if regexp.MustCompile(`Domain`).MatchString(v) { + t.Errorf("Set-Cookie still has Domain: %q", v) + } + } + }) + + t.Run("delete header on empty value", func(t *testing.T) { + h := http.Header{"Content-Security-Policy": {"default-src 'self'"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{compileHeaderRewrite(`(?i)^content-security-policy$`, `.*`, ``)}, + "response", "/") + if err != nil { + t.Fatal(err) + } + if _, ok := h["Content-Security-Policy"]; ok { + t.Errorf("header should be deleted, got %v", h) + } + }) + + t.Run("mixed-case name matching", func(t *testing.T) { + h := http.Header{"Referer": {"https://127.0.0.1:8080/github.com/x"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{compileHeaderRewrite(`(?i)^(referer|origin)$`, `.*`, ``)}, + "request", "/github.com/x") + if err != nil { + t.Fatal(err) + } + if _, ok := h["Referer"]; ok { + t.Errorf("Referer should be deleted, got %v", h) + } + }) + + t.Run("name no-match leaves header untouched", func(t *testing.T) { + h := http.Header{"X-Forwarded-Host": {"example.com"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{compileHeaderRewrite(`(?i)^location$`, `.*`, ``)}, + "response", "/") + if err != nil { + t.Fatal(err) + } + if got := h.Get("X-Forwarded-Host"); got != "example.com" { + t.Errorf("X-Forwarded-Host = %q, want untouched", got) + } + }) +} + +func TestRewriteHeaderBlock_Plugin(t *testing.T) { + t.Run("round-trip with in-place replacement", func(t *testing.T) { + rw := &headerRewriter{cb: func(b []byte) []byte { + return []byte("X-Added: yes\r\nLocation: https://127.0.0.1:8080/github.com/x\r\n\r\n") + }} + h := http.Header{"Location": {"https://github.com/x"}, "Keep-Me": {"v"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{{Rewriter: rw}}, + "response", "/github.com/x") + if err != nil { + t.Fatal(err) + } + if !rw.called { + t.Fatal("rewriter not called") + } + if got := h.Get("Location"); got != "https://127.0.0.1:8080/github.com/x" { + t.Errorf("Location = %q", got) + } + if got := h.Get("X-Added"); got != "yes" { + t.Errorf("X-Added = %q", got) + } + if _, ok := h["Keep-Me"]; ok { + t.Errorf("Keep-Me should be gone after plugin replaced the block, got %v", h) + } + }) + + t.Run("metadata contains kind header", func(t *testing.T) { + rw := &headerRewriter{} + h := http.Header{"Location": {"https://github.com/x"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{{Rewriter: rw}}, + "response", "/github.com/x") + if err != nil { + t.Fatal(err) + } + if rw.metadata == nil { + t.Fatal("metadata not captured") + } + if got := rw.metadata[MetaKeyKind]; got != KindHeader { + t.Errorf("kind = %v, want %q", got, KindHeader) + } + if got := rw.metadata[MetaKeyDirection]; got != "response" { + t.Errorf("direction = %v, want %q", got, "response") + } + if got := rw.metadata[MetaKeyURI]; got != "/github.com/x" { + t.Errorf("uri = %v, want %q", got, "/github.com/x") + } + }) + + t.Run("plugin output without trailing blank line", func(t *testing.T) { + var got []byte + rw := &headerRewriter{cb: func(b []byte) []byte { + got = append([]byte(nil), b...) + return []byte("Location: https://127.0.0.1:8080/github.com/x\r\n") + }} + h := http.Header{"Location": {"https://github.com/x"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{{Rewriter: rw}}, "response", "/") + if err != nil { + t.Fatalf("ReadMIMEHeader must not fail on a plugin block with no blank line: %v", err) + } + if !bytes.HasSuffix(got, []byte("\r\n\r\n")) { + t.Errorf("plugin input missing trailing blank line: %q", got) + } + if v := h.Get("Location"); v != "https://127.0.0.1:8080/github.com/x" { + t.Errorf("Location = %q", v) + } + }) + + t.Run("name gate skips non-matching", func(t *testing.T) { + rw := &headerRewriter{} + h := http.Header{"X-Other": {"v"}} + err := rewriteHeaderBlock(context.Background(), h, + []chain.HTTPHeaderRewriteSettings{{Name: regexp.MustCompile(`(?i)^location$`), Rewriter: rw}}, + "response", "/") + if err != nil { + t.Fatal(err) + } + if rw.called { + t.Fatal("rewriter should be skipped when name does not match") + } + }) +} + +func TestRewriteMeta(t *testing.T) { + t.Run("omits empty common fields", func(t *testing.T) { + md := rewriteMeta("", "", "", "", nil) + if len(md) != 0 { + t.Errorf("md = %v, want empty", md) + } + }) + + t.Run("includes kind and extras", func(t *testing.T) { + md := rewriteMeta("s", "response", "/u", KindHeader, map[string]any{"sse_phase": "end"}) + if md[MetaKeySid] != "s" || md[MetaKeyDirection] != "response" || md[MetaKeyURI] != "/u" || md[MetaKeyKind] != KindHeader { + t.Errorf("common fields wrong: %v", md) + } + if md["sse_phase"] != "end" { + t.Errorf("extra missing: %v", md) + } + }) +} + +func TestRewriteReqRespHeader_NilGuard(t *testing.T) { + if err := rewriteReqHeader(context.Background(), nil); err != nil { + t.Errorf("rewriteReqHeader(nil) = %v", err) + } + if err := rewriteRespHeader(context.Background(), nil); err != nil { + t.Errorf("rewriteRespHeader(nil) = %v", err) + } +} diff --git a/internal/util/forwarder/sniffer_http.go b/internal/util/forwarder/sniffer_http.go index 480eb71..d5fa835 100644 --- a/internal/util/forwarder/sniffer_http.go +++ b/internal/util/forwarder/sniffer_http.go @@ -130,12 +130,23 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO // --- Dial upstream (miss or stale; fresh cache hit with keep-alive skips) --- if !freshHit { if cc == nil { - node, cc, err = h.dial(ctx, conn, req, &ho) + var dialRes *http.Response + node, cc, dialRes, err = h.dial(ctx, req, &ho) if err != nil { - // h.dial already wrote the error response to conn. + // Prefer serving a stale cached entry over surfacing the + // dial error, when the policy allows it. h.dial no longer + // writes the error response so this stays the sole writer. + var shouldClose bool + if h.serveStale(conn, staleResp, ro, &shouldClose, &ho, log) { + err = nil + return nil + } if staleResp != nil && staleResp.Body != nil { staleResp.Body.Close() } + if dialRes != nil { + dialRes.Write(conn) + } return err } upstreamHost = normalizeHost(ro.HTTP.Host, "80") @@ -297,8 +308,11 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha return node, nil, nil } -// dial selects a node, establishes a connection, and sends the request upstream. -func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) { +// dial selects a node and establishes a connection. It does NOT write the +// error response to conn; instead it returns res (nil on a plain dial failure) +// so the caller can choose between serving a stale cached response and writing +// the error response. +func (h *Sniffer) dial(ctx context.Context, req *http.Request, ho *HandleOptions) (node *chain.Node, cc net.Conn, res *http.Response, err error) { dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext @@ -306,7 +320,7 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho if node = ho.Node; node != nil { cc, err = dial(ctx, "tcp", node.Addr) - return + return node, cc, nil, err } ro := ho.RecorderObject @@ -321,10 +335,7 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho node, res, resolveErr := resolveHTTPNode(ctx, host, req, ho) if resolveErr != nil { ro.HTTP.StatusCode = res.StatusCode - if werr := res.Write(conn); werr != nil { - ho.Log.Warnf("write error response: %v", werr) - } - return nil, nil, resolveErr + return nil, nil, res, resolveErr } // Prepare an error response for potential connection failures. @@ -348,17 +359,14 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho marker.Mark() } ho.Log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) - if werr := res.Write(conn); werr != nil { - ho.Log.Warnf("write error response: %v", werr) - } - return + return nil, nil, res, err } if marker := node.Marker(); marker != nil { marker.Reset() } cc = tlsWrapConn(cc, node.Options().TLS) - return + return node, cc, nil, nil } // httpRoundTrip forwards a single HTTP request/response pair and records @@ -368,6 +376,12 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, node *chain.Node, req *http.Request, pStats stats.Stats, ho *HandleOptions, staleResp *http.Response) (shouldClose bool, err error) { shouldClose = true + // Cache key identity is derived from the client's request (method + Host + + // request-target). Capture it before any rewrite so Lookup (which runs + // before rewrites in HandleHTTP) and Store agree on the same key — the Host + // header is rewritten below for mirror routing and must not change the key. + cacheKey := httpcache.Key(req.Method, req.Host, req.RequestURI) + log := ho.Log ro := &xrecorder.HandlerRecorderObject{} *ro = *ho.RecorderObject @@ -410,11 +424,11 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, // rewrite dimension added there must be added here too, or its originals // won't be captured. httpSettings := node.Options().HTTP - hasReqRewrite := httpSettings != nil && (httpSettings.Host != "" || + hasReqRewrite := httpSettings != nil && (httpSettings.Host != "" || httpSettings.HostPattern != nil || len(httpSettings.RequestHeader) > 0 || len(httpSettings.RewriteURL) > 0 || - len(httpSettings.RewriteRequestBody) > 0) + len(httpSettings.RewriteRequestBody) > 0 || len(httpSettings.RewriteRequestHeader) > 0) hasRespRewrite := httpSettings != nil && (len(httpSettings.ResponseHeader) > 0 || - len(httpSettings.RewriteResponseBody) > 0) + len(httpSettings.RewriteResponseBody) > 0 || len(httpSettings.RewriteResponseHeader) > 0) bodySize := clampBodySize(h.RecorderOptions) res := &http.Response{ @@ -436,6 +450,8 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, var responseHeader map[string]string var respBodyRewrites []chain.HTTPBodyRewriteSettings var reqBodyRewrites []chain.HTTPBodyRewriteSettings + var reqHeaderRewrites []chain.HTTPHeaderRewriteSettings + var respHeaderRewrites []chain.HTTPHeaderRewriteSettings if httpSettings != nil { if auther := httpSettings.Auther; auther != nil { username, password, _ := req.BasicAuth() @@ -461,7 +477,11 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, } } - if httpSettings.Host != "" { + if httpSettings.HostPattern != nil { + if m := httpSettings.HostPattern.FindStringSubmatchIndex(req.URL.Path); m != nil { + req.Host = string(httpSettings.HostPattern.ExpandString(nil, httpSettings.Host, req.URL.Path, m)) + } + } else if httpSettings.Host != "" { req.Host = httpSettings.Host } for k, v := range httpSettings.RequestHeader { @@ -491,9 +511,19 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, } } + reqHeaderRewrites = httpSettings.RewriteRequestHeader + if len(reqHeaderRewrites) > 0 { + if err = rewriteReqHeader(ctx, req, reqHeaderRewrites...); err != nil { + log.Errorf("rewrite request header: %v", err) + return + } + ro.HTTP.Request.Header = req.Header.Clone() + } + responseHeader = httpSettings.ResponseHeader respBodyRewrites = httpSettings.RewriteResponseBody reqBodyRewrites = httpSettings.RewriteRequestBody + respHeaderRewrites = httpSettings.RewriteResponseHeader } // Snapshot the original request body before rewriting, restoring it via @@ -627,6 +657,16 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, resp.Header.Set(k, v) } } + + // Header regex/plugin rewrite runs after static overrides and before the + // recorder clone, so the recorder sees the final (post-rewrite) headers. + if len(respHeaderRewrites) > 0 { + if err = rewriteRespHeader(ctx, resp, respHeaderRewrites...); err != nil { + log.Errorf("rewrite response header: %v", err) + return + } + } + ro.HTTP.Response.Header = resp.Header.Clone() ro.HTTP.Response.ContentLength = resp.ContentLength @@ -656,10 +696,10 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, if tee != nil { if data := tee.Captured(); data != nil { - if serr := h.Cache.Store(ctx, req, data, resp.StatusCode); serr != nil { + if serr := h.Cache.StoreKey(ctx, cacheKey, data, resp.StatusCode); serr != nil { log.Warnf("cache store: %v", serr) } else { - log.Debugf("cache store: %s", httpcache.Key(req.Method, req.Host, req.RequestURI)) + log.Debugf("cache store: %s", cacheKey) } } } diff --git a/internal/util/forwarder/sniffer_rewrite.go b/internal/util/forwarder/sniffer_rewrite.go index 752afd3..da44c68 100644 --- a/internal/util/forwarder/sniffer_rewrite.go +++ b/internal/util/forwarder/sniffer_rewrite.go @@ -1,15 +1,20 @@ package forwarder import ( + "bufio" "bytes" "context" "fmt" "io" "net/http" + "net/textproto" + "regexp" "strings" "github.com/go-gost/core/chain" "github.com/go-gost/core/logger" + "github.com/go-gost/core/rewriter" + xctx "github.com/go-gost/x/ctx" xnet "github.com/go-gost/x/internal/net" xrecorder "github.com/go-gost/x/recorder" "golang.org/x/net/http/httpguts" @@ -105,3 +110,103 @@ func rewriteReqBody(ctx context.Context, req *http.Request, rewrites ...chain.HT return nil } +// rewriteHeaderBlock applies the header rewrite chain to h in place. +func rewriteHeaderBlock(ctx context.Context, h http.Header, rewrites []chain.HTTPHeaderRewriteSettings, direction, uri string) error { + if len(rewrites) == 0 || h == nil { + return nil + } + sid := xctx.SidFromContext(ctx).String() + + for _, rw := range rewrites { + if rw.Rewriter != nil { + // Plugin mode: gate on rw.Name matching any header name (nil Name + // means "always invoke"), serialize the whole header block, and + // parse the plugin's bytes back into the header. + if rw.Name != nil && !anyHeaderNameMatch(h, rw.Name) { + continue + } + var buf bytes.Buffer + if err := h.Write(&buf); err != nil { + return err + } + // http.Header.Write emits "Key: value\r\n" lines with no trailing + // blank line; append one so ReadMIMEHeader can find the end. + buf.WriteString("\r\n") + md := rewriteMeta(sid, direction, uri, KindHeader, nil) + rewritten, err := rw.Rewriter.Rewrite(ctx, buf.Bytes(), rewriter.MetadataRewriteOption(md)) + if err != nil { + return err + } + // Tolerate a plugin that omits the trailing blank line. + if !bytes.HasSuffix(rewritten, []byte("\r\n\r\n")) && !bytes.HasSuffix(rewritten, []byte("\n\n")) { + rewritten = append(rewritten, '\r', '\n') + } + nh, err := textproto.NewReader(bufio.NewReader(bytes.NewReader(rewritten))).ReadMIMEHeader() + if err != nil { + return err + } + // h is a map reference; replace its contents in place, not via + // reassignment (which would only rebind the local variable). + clear(h) + for k, vs := range nh { + h[k] = vs + } + continue + } + + // Regex mode. + if rw.Name == nil { + continue + } + for name, values := range h { + if !rw.Name.MatchString(strings.ToLower(name)) { + continue + } + kept := values[:0] + for _, v := range values { + if rw.Pattern == nil { + kept = append(kept, v) + continue + } + if nv := rw.Pattern.ReplaceAllString(v, string(rw.Replacement)); nv != "" { + kept = append(kept, nv) + } + } + if len(kept) == 0 { + h.Del(name) + } else { + h[name] = kept + } + } + } + return nil +} + +func rewriteReqHeader(ctx context.Context, req *http.Request, rewrites ...chain.HTTPHeaderRewriteSettings) error { + if req == nil { + return nil + } + return rewriteHeaderBlock(ctx, req.Header, rewrites, "request", req.RequestURI) +} + +func rewriteRespHeader(ctx context.Context, resp *http.Response, rewrites ...chain.HTTPHeaderRewriteSettings) error { + if resp == nil { + return nil + } + uri := "" + if resp.Request != nil { + uri = resp.Request.RequestURI + } + return rewriteHeaderBlock(ctx, resp.Header, rewrites, "response", uri) +} + +// anyHeaderNameMatch reports whether any header name matches the regex. +func anyHeaderNameMatch(h http.Header, name *regexp.Regexp) bool { + for k := range h { + if name.MatchString(strings.ToLower(k)) { + return true + } + } + return false +} + diff --git a/internal/util/forwarder/sniffer_sse.go b/internal/util/forwarder/sniffer_sse.go index cf791b0..5c9199b 100644 --- a/internal/util/forwarder/sniffer_sse.go +++ b/internal/util/forwarder/sniffer_sse.go @@ -296,13 +296,10 @@ func (b *rewriteBody) Read(p []byte) (n int, err error) { // This empty-body call lets Rewriter plugins append a // trailing event (e.g. message_stop for LLM conversion). b.ended = true - md := map[string]any{ - "sid": b.sid, - "direction": b.direction, - "uri": b.uri, + md := rewriteMeta(b.sid, b.direction, b.uri, KindBody, map[string]any{ "sse_phase": "end", "event_index": b.eventIndex, - } + }) rewritten, endErr := b.apply(nil, rewriter.MetadataRewriteOption(md)) if endErr != nil { return 0, endErr @@ -320,14 +317,11 @@ func (b *rewriteBody) Read(p []byte) (n int, err error) { if !b.ended { b.ended = true b.scannerErr = err - md := map[string]any{ - "sid": b.sid, - "direction": b.direction, - "uri": b.uri, + md := rewriteMeta(b.sid, b.direction, b.uri, KindBody, map[string]any{ "sse_phase": "end", "event_index": b.eventIndex, "stream_error": err.Error(), - } + }) rewritten, _ := b.apply(nil, rewriter.MetadataRewriteOption(md)) if len(rewritten) > 0 { b.buf.Write(rewritten) @@ -350,13 +344,10 @@ func (b *rewriteBody) Read(p []byte) (n int, err error) { if b.eventIndex == 0 { phase = "start" } - md := map[string]any{ - "sid": b.sid, - "direction": b.direction, - "uri": b.uri, + md := rewriteMeta(b.sid, b.direction, b.uri, KindBody, map[string]any{ "sse_phase": phase, "event_index": b.eventIndex, - } + }) b.eventIndex++ rewritten, err := b.apply(event, rewriter.MetadataRewriteOption(md)) @@ -377,17 +368,7 @@ func (b *rewriteBody) Close() error { // baseMetadata returns rewriter metadata with all available fields. func (b *rewriteBody) baseMetadata() map[string]any { - md := map[string]any{} - if b.sid != "" { - md["sid"] = b.sid - } - if b.direction != "" { - md["direction"] = b.direction - } - if b.uri != "" { - md["uri"] = b.uri - } - return md + return rewriteMeta(b.sid, b.direction, b.uri, KindBody, nil) } // baseMetadataOpt returns a RewriteOption wrapping baseMetadata. diff --git a/internal/util/forwarder/sniffer_test.go b/internal/util/forwarder/sniffer_test.go index 6173155..7d169b6 100644 --- a/internal/util/forwarder/sniffer_test.go +++ b/internal/util/forwarder/sniffer_test.go @@ -691,6 +691,64 @@ func TestWithNode(t *testing.T) { } } +func TestHandleHTTP_HostPattern(t *testing.T) { + // Upstream echoes the Host header it received in the response body. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(r.Host)) + })) + defer upstream.Close() + + rec := &captureRecorder{} + h := &Sniffer{Recorder: rec} + ro := &xrecorder.HandlerRecorderObject{} + + // Node with HostPattern: derive Host from the first path segment. + node := chain.NewNode("githubio", "unused:443", + chain.HTTPNodeOption(&chain.HTTPNodeSettings{ + HostPattern: regexp.MustCompile(`^/([a-z0-9-]+\.github\.io)/`), + Host: "$1", + }), + ) + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- h.HandleHTTP(context.Background(), serverConn, + sniffing.WithService("test"), + sniffing.WithNode(node), + sniffing.WithRecorderObject(ro), + sniffing.WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + return net.Dial("tcp", upstream.Listener.Addr().String()) + }), + sniffing.WithLog(xlogger.Nop()), + ) + }() + + // Request path encodes the upstream host as its first segment. + req, _ := http.NewRequest("GET", "http://gh.home.pi/microsoft.github.io/foo", nil) + req.Close = true + if err := req.Write(clientConn); err != nil { + t.Fatal(err) + } + + resp, err := http.ReadResponse(bufio.NewReader(clientConn), req) + if err != nil { + t.Fatal(err) + } + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + clientConn.Close() + <-errCh // teardown error on net.Pipe keep-alive is expected; ignore it + + if got := string(body); got != "microsoft.github.io" { + t.Errorf("upstream Host = %q, want %q", got, "microsoft.github.io") + } +} + func TestWithHop(t *testing.T) { opts := &HandleOptions{} mh := &mockHop{}