From a38a701dc4d7c080d91cc64cffc1ba1edf03cf00 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Tue, 25 Aug 2026 21:35:45 +0800 Subject: [PATCH] refactor(sniffing): register protocol handlers via init() registries Replace the 5 duplicated switch statements that dispatch sniffed connections (forward/local, forward/remote, http CONNECT, sni, relay) with per-package registries keyed by proto string. Protocol handlers self-register via init(), following the house registration pattern, so adding a protocol no longer requires touching every dispatch site. Unify the two HandleOptions types (forwarder lacked dialTLS; sniffing lacked hop/node/httpKeepalive) into one exported struct in sniffing; forwarder aliases it. Unknown protocols fall through to raw forwarding exactly as before; SSH remains detection-only. redirect/tcp's special fallback-replay dispatch is intentionally left untouched. Adds Dispatch helpers and 8 registry tests (register/get/duplicate/ unknown/known) in both packages. Net -151 lines. --- handler/forward/local/sniffing.go | 50 +++------- handler/forward/remote/sniffing.go | 50 +++------- handler/http/connect.go | 37 ++----- handler/relay/connect.go | 40 +++----- handler/sni/sniffing.go | 53 +++------- internal/util/forwarder/registry.go | 50 ++++++++++ internal/util/forwarder/registry_test.go | 86 ++++++++++++++++ internal/util/forwarder/sniffer.go | 84 ++-------------- internal/util/forwarder/sniffer_h2.go | 4 +- internal/util/forwarder/sniffer_http.go | 64 ++++++------ internal/util/forwarder/sniffer_redis.go | 28 +++--- internal/util/forwarder/sniffer_test.go | 120 +++++++++++------------ internal/util/forwarder/sniffer_tls.go | 52 +++++----- internal/util/sniffing/registry.go | 48 +++++++++ internal/util/sniffing/registry_test.go | 84 ++++++++++++++++ internal/util/sniffing/sniffer.go | 57 ++++++++--- internal/util/sniffing/sniffer_h2.go | 6 +- internal/util/sniffing/sniffer_http.go | 8 +- internal/util/sniffing/sniffer_redis.go | 8 +- internal/util/sniffing/sniffer_test.go | 32 +++--- internal/util/sniffing/sniffer_tls.go | 16 +-- 21 files changed, 547 insertions(+), 430 deletions(-) create mode 100644 internal/util/forwarder/registry.go create mode 100644 internal/util/forwarder/registry_test.go create mode 100644 internal/util/sniffing/registry.go create mode 100644 internal/util/sniffing/registry_test.go diff --git a/handler/forward/local/sniffing.go b/handler/forward/local/sniffing.go index 7e63a26..1837267 100644 --- a/handler/forward/local/sniffing.go +++ b/handler/forward/local/sniffing.go @@ -72,45 +72,19 @@ func (h *forwardHandler) sniffingDial(ctx context.Context, network, address stri } // handleSniffedProtocol dispatches a sniffed connection to the protocol-specific -// sniffer (HandleHTTP or HandleTLS). It returns (true, err) when the protocol was +// sniffer registered for proto. It returns (true, err) when the protocol was // handled and (false, nil) when the caller should fall through to raw forwarding. func (h *forwardHandler) handleSniffedProtocol(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, proto string) (handled bool, err error) { - switch proto { - case sniffing.ProtoHTTP, sniffing.ProtoTLS, sniffing.ProtoRedis: - dial := func(ctx context.Context, network, address string) (net.Conn, error) { - return h.sniffingDial(ctx, network, address, ro) - } - sniffer := h.sniffer.Build() - if proto == sniffing.ProtoHTTP { - return true, sniffer.HandleHTTP(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithHTTPKeepalive(h.md.httpKeepalive), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - } - if proto == sniffing.ProtoTLS { - return true, sniffer.HandleTLS(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - } - return true, sniffer.HandleRedis(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - default: - return false, nil + dial := func(ctx context.Context, network, address string) (net.Conn, error) { + return h.sniffingDial(ctx, network, address, ro) } + return forwarder.Dispatch(h.sniffer.Build(), ctx, conn, proto, + sniffing.WithService(h.options.Service), + sniffing.WithDial(dial), + sniffing.WithHop(h.getHop()), + sniffing.WithBypass(h.options.Bypass), + sniffing.WithHTTPKeepalive(h.md.httpKeepalive), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), + ) } diff --git a/handler/forward/remote/sniffing.go b/handler/forward/remote/sniffing.go index 59773d1..eedd729 100644 --- a/handler/forward/remote/sniffing.go +++ b/handler/forward/remote/sniffing.go @@ -72,45 +72,19 @@ func (h *forwardHandler) sniffingDial(ctx context.Context, network, address stri } // handleSniffedProtocol dispatches a sniffed connection to the protocol-specific -// sniffer (HandleHTTP or HandleTLS). It returns (true, err) when the protocol was +// sniffer registered for proto. It returns (true, err) when the protocol was // handled and (false, nil) when the caller should fall through to raw forwarding. func (h *forwardHandler) handleSniffedProtocol(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, proto string) (handled bool, err error) { - switch proto { - case sniffing.ProtoHTTP, sniffing.ProtoTLS, sniffing.ProtoRedis: - dial := func(ctx context.Context, network, address string) (net.Conn, error) { - return h.sniffingDial(ctx, network, address, ro) - } - sniffer := h.sniffer.Build() - if proto == sniffing.ProtoHTTP { - return true, sniffer.HandleHTTP(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithHTTPKeepalive(h.md.httpKeepalive), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - } - if proto == sniffing.ProtoTLS { - return true, sniffer.HandleTLS(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - } - return true, sniffer.HandleRedis(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.getHop()), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - default: - return false, nil + dial := func(ctx context.Context, network, address string) (net.Conn, error) { + return h.sniffingDial(ctx, network, address, ro) } + return forwarder.Dispatch(h.sniffer.Build(), ctx, conn, proto, + sniffing.WithService(h.options.Service), + sniffing.WithDial(dial), + sniffing.WithHop(h.getHop()), + sniffing.WithBypass(h.options.Bypass), + sniffing.WithHTTPKeepalive(h.md.httpKeepalive), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), + ) } diff --git a/handler/http/connect.go b/handler/http/connect.go index e177f22..9d29886 100644 --- a/handler/http/connect.go +++ b/handler/http/connect.go @@ -128,35 +128,14 @@ func (h *httpHandler) sniffAndHandle(ctx context.Context, conn net.Conn, cc net. sniffer := h.sniffer.Build() conn = xnet.NewReadWriteConn(br, conn, conn) - switch proto { - case sniffing.ProtoHTTP: - return true, sniffer.HandleHTTP(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - case sniffing.ProtoTLS: - return true, sniffer.HandleTLS(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - case sniffing.ProtoRedis: - return true, sniffer.HandleRedis(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - } - - return false, nil + return sniffing.Dispatch(sniffer, ctx, "tcp", conn, proto, + sniffing.WithService(h.options.Service), + sniffing.WithDial(dial), + sniffing.WithDialTLS(dialTLS), + sniffing.WithBypass(h.options.Bypass), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), + ) } // dial establishes an upstream connection through the proxy chain router. diff --git a/handler/relay/connect.go b/handler/relay/connect.go index b388f50..67e86e1 100644 --- a/handler/relay/connect.go +++ b/handler/relay/connect.go @@ -231,35 +231,19 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network } conn = xnet.NewReadWriteConn(br, conn, conn) - switch proto { - case sniffing.ProtoHTTP: + handled, err := sniffing.Dispatch(sniffer, ctx, "tcp", conn, proto, + sniffing.WithService(h.options.Service), + sniffing.WithDial(dial), + sniffing.WithDialTLS(dialTLS), + sniffing.WithBypass(h.options.Bypass), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), + ) + if handled { + // Reset the recording start time: the sniffer handles the + // connection and manages ro.Time from here on. ro.Time = time.Time{} - return sniffer.HandleHTTP(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - case sniffing.ProtoTLS: - ro.Time = time.Time{} - return sniffer.HandleTLS(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - case sniffing.ProtoRedis: - ro.Time = time.Time{} - return sniffer.HandleRedis(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) + return err } } diff --git a/handler/sni/sniffing.go b/handler/sni/sniffing.go index edc14fd..5a6bfc8 100644 --- a/handler/sni/sniffing.go +++ b/handler/sni/sniffing.go @@ -90,49 +90,24 @@ func (h *sniHandler) sniffingDialTLS(ctx context.Context, network, address strin } // handleSniffedProtocol dispatches a sniffed connection to the -// protocol-specific sniffer ([sniffing.Sniffer.HandleHTTP] or -// [sniffing.Sniffer.HandleTLS]). +// protocol-specific sniffer registered for proto. // // It returns (true, err) when the protocol was handled and (false, nil) // when the protocol is unrecognised and the caller should handle the // connection directly (which for the SNI handler means silently dropping it). func (h *sniHandler) handleSniffedProtocol(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, proto string) (handled bool, err error) { - switch proto { - case sniffing.ProtoHTTP, sniffing.ProtoTLS, sniffing.ProtoRedis: - dial := func(ctx context.Context, network, address string) (net.Conn, error) { - return h.sniffingDial(ctx, network, address, ro) - } - dialTLS := func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { - return h.sniffingDialTLS(ctx, network, address, ro, cfg) - } - sniffer := h.sniffer.Build() - if proto == sniffing.ProtoHTTP { - return true, sniffer.HandleHTTP(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - } - if proto == sniffing.ProtoTLS { - return true, sniffer.HandleTLS(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithDialTLS(dialTLS), - sniffing.WithBypass(h.options.Bypass), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - } - return true, sniffer.HandleRedis(ctx, "tcp", conn, - sniffing.WithService(h.options.Service), - sniffing.WithDial(dial), - sniffing.WithRecorderObject(ro), - sniffing.WithLog(log), - ) - default: - return false, nil + dial := func(ctx context.Context, network, address string) (net.Conn, error) { + return h.sniffingDial(ctx, network, address, ro) } + dialTLS := func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { + return h.sniffingDialTLS(ctx, network, address, ro, cfg) + } + return sniffing.Dispatch(h.sniffer.Build(), ctx, "tcp", conn, proto, + sniffing.WithService(h.options.Service), + sniffing.WithDial(dial), + sniffing.WithDialTLS(dialTLS), + sniffing.WithBypass(h.options.Bypass), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), + ) } diff --git a/internal/util/forwarder/registry.go b/internal/util/forwarder/registry.go new file mode 100644 index 0000000..31dc729 --- /dev/null +++ b/internal/util/forwarder/registry.go @@ -0,0 +1,50 @@ +package forwarder + +import ( + "context" + "net" + "sync" + + "github.com/go-gost/x/internal/util/sniffing" +) + +// HandleFunc handles a connection whose protocol has been detected as proto, +// mirroring sniffing.HandleFunc but without the network parameter. The +// receiver s carries the per-connection Sniffer state. +type HandleFunc func(s *Sniffer, ctx context.Context, conn net.Conn, opts ...sniffing.HandleOption) error + +var protoHandlers sync.Map // map[string]HandleFunc + +// Register associates a protocol with its handler. Duplicate registration +// panics, mirroring the house registry pattern. +func Register(proto string, fn HandleFunc) { + if _, dup := protoHandlers.LoadOrStore(proto, fn); dup { + panic("forwarder: duplicate protocol handler: " + proto) + } +} + +// Get returns the handler registered for proto. +func Get(proto string) (HandleFunc, bool) { + v, ok := protoHandlers.Load(proto) + if !ok { + return nil, false + } + return v.(HandleFunc), true +} + +// Dispatch routes a detected protocol to its registered handler. It returns +// (false, nil) when no handler is registered for proto, signalling the caller +// to fall through to raw forwarding. +func Dispatch(s *Sniffer, ctx context.Context, conn net.Conn, proto string, opts ...sniffing.HandleOption) (handled bool, err error) { + fn, ok := Get(proto) + if !ok { + return false, nil + } + return true, fn(s, ctx, conn, opts...) +} + +func init() { + Register(sniffing.ProtoHTTP, (*Sniffer).HandleHTTP) + Register(sniffing.ProtoTLS, (*Sniffer).HandleTLS) + Register(sniffing.ProtoRedis, (*Sniffer).HandleRedis) +} diff --git a/internal/util/forwarder/registry_test.go b/internal/util/forwarder/registry_test.go new file mode 100644 index 0000000..6a16f8f --- /dev/null +++ b/internal/util/forwarder/registry_test.go @@ -0,0 +1,86 @@ +package forwarder + +import ( + "context" + "net" + "testing" + + "github.com/go-gost/x/internal/util/sniffing" +) + +// TestRegisterGet verifies a handler can be registered and looked up by proto. +func TestRegisterGet(t *testing.T) { + const proto = "testproto-get" + called := false + Register(proto, func(s *Sniffer, _ context.Context, _ net.Conn, _ ...sniffing.HandleOption) error { + called = true + return nil + }) + defer protoHandlers.Delete(proto) + + fn, ok := Get(proto) + if !ok { + t.Fatal("Get: handler not found") + } + if err := fn(&Sniffer{}, context.Background(), nil); err != nil { + t.Fatalf("fn() error = %v", err) + } + if !called { + t.Fatal("registered handler was not called") + } + if _, ok := Get("no-such-proto"); ok { + t.Fatal("Get: unexpected handler for unknown proto") + } +} + +// TestRegisterDuplicate verifies registering the same proto twice panics. +func TestRegisterDuplicate(t *testing.T) { + const proto = "testproto-dup" + Register(proto, func(*Sniffer, context.Context, net.Conn, ...sniffing.HandleOption) error { + return nil + }) + defer protoHandlers.Delete(proto) + + defer func() { + if recover() == nil { + t.Fatal("expected panic on duplicate registration") + } + }() + Register(proto, func(*Sniffer, context.Context, net.Conn, ...sniffing.HandleOption) error { + return nil + }) +} + +// TestDispatch_Unknown verifies Dispatch falls through (false, nil) for an +// unregistered proto. +func TestDispatch_Unknown(t *testing.T) { + handled, err := Dispatch(&Sniffer{}, context.Background(), nil, "no-such-proto") + if handled { + t.Fatal("Dispatch: unknown proto handled = true, want false") + } + if err != nil { + t.Fatalf("Dispatch: err = %v, want nil", err) + } +} + +// TestDispatch_Known verifies Dispatch routes a registered proto to its handler. +func TestDispatch_Known(t *testing.T) { + const proto = "testproto-known" + called := false + Register(proto, func(s *Sniffer, _ context.Context, _ net.Conn, _ ...sniffing.HandleOption) error { + called = true + return nil + }) + defer protoHandlers.Delete(proto) + + handled, err := Dispatch(&Sniffer{}, context.Background(), nil, proto) + if !handled { + t.Fatal("Dispatch: known proto handled = false, want true") + } + if err != nil { + t.Fatalf("Dispatch: err = %v, want nil", err) + } + if !called { + t.Fatal("Dispatch: registered handler was not called") + } +} diff --git a/internal/util/forwarder/sniffer.go b/internal/util/forwarder/sniffer.go index 6f63998..b1fbffe 100644 --- a/internal/util/forwarder/sniffer.go +++ b/internal/util/forwarder/sniffer.go @@ -1,7 +1,6 @@ package forwarder import ( - "context" "crypto" "crypto/tls" "crypto/x509" @@ -11,13 +10,11 @@ import ( "github.com/go-gost/core/bypass" "github.com/go-gost/core/chain" - "github.com/go-gost/core/hop" - "github.com/go-gost/core/logger" "github.com/go-gost/core/recorder" "github.com/go-gost/x/config" - tls_util "github.com/go-gost/x/internal/util/tls" "github.com/go-gost/x/internal/util/httpcache" - xrecorder "github.com/go-gost/x/recorder" + "github.com/go-gost/x/internal/util/sniffing" + tls_util "github.com/go-gost/x/internal/util/tls" ) const ( @@ -33,77 +30,12 @@ const DefaultReadTimeout = 30 * time.Second // DefaultCertPool is the default in-memory certificate pool used for TLS MITM. var DefaultCertPool = tls_util.NewMemoryCertPool() -// HandleOptions holds configuration options for sniffing handlers. -type HandleOptions struct { - service string - dial func(ctx context.Context, network, address string) (net.Conn, error) - httpKeepalive bool - readTimeout time.Duration - node *chain.Node - hop hop.Hop - bypass bypass.Bypass - recorderObject *xrecorder.HandlerRecorderObject - log logger.Logger -} +// HandleOptions aliases sniffing.HandleOptions, the unified option set shared +// by the sniffing and forwarder Sniffer implementations. +type HandleOptions = sniffing.HandleOptions // HandleOption configures HandleOptions for sniffing handlers. -type HandleOption func(opts *HandleOptions) - -// WithService sets the service name for bypass and selection lookups. -func WithService(service string) HandleOption { - return func(opts *HandleOptions) { - opts.service = service - } -} - -// WithDial sets the dial function used to establish upstream connections. -func WithDial(dial func(ctx context.Context, network, address string) (net.Conn, error)) HandleOption { - return func(opts *HandleOptions) { - opts.dial = dial - } -} - -// WithHTTPKeepalive enables or disables HTTP keep-alive on the upstream connection. -func WithHTTPKeepalive(keepalive bool) HandleOption { - return func(opts *HandleOptions) { - opts.httpKeepalive = keepalive - } -} - -// WithNode sets a pre-resolved chain node to connect to, bypassing hop selection. -func WithNode(node *chain.Node) HandleOption { - return func(opts *HandleOptions) { - opts.node = node - } -} - -// WithHop sets the hop used for node selection when routing requests. -func WithHop(h hop.Hop) HandleOption { - return func(opts *HandleOptions) { - opts.hop = h - } -} - -// WithBypass sets the bypass rules for filtering requests by host. -func WithBypass(bypass bypass.Bypass) HandleOption { - return func(opts *HandleOptions) { - opts.bypass = bypass - } -} - -// WithRecorderObject sets the recorder object for capturing traffic metadata. -func WithRecorderObject(ro *xrecorder.HandlerRecorderObject) HandleOption { - return func(opts *HandleOptions) { - opts.recorderObject = ro - } -} - -// WithLog sets the logger for the handler. -func WithLog(log logger.Logger) HandleOption { - return func(opts *HandleOptions) { - opts.log = log - } -} +type HandleOption = sniffing.HandleOption // Sniffer handles HTTP and TLS traffic sniffing, recording, and MITM TLS // termination for protocol-aware forwarding. It can intercept HTTP requests, @@ -171,8 +103,8 @@ func normalizeHost(host, defaultPort string) string { // effectiveReadTimeout returns the read timeout from options, falling back to // the Sniffer's ReadTimeout, then to DefaultReadTimeout. func (h *Sniffer) effectiveReadTimeout(ho *HandleOptions) time.Duration { - if ho.readTimeout > 0 { - return ho.readTimeout + if ho.ReadTimeout > 0 { + return ho.ReadTimeout } if h.ReadTimeout > 0 { return h.ReadTimeout diff --git a/internal/util/forwarder/sniffer_h2.go b/internal/util/forwarder/sniffer_h2.go index 55a0e9c..d93b593 100644 --- a/internal/util/forwarder/sniffer_h2.go +++ b/internal/util/forwarder/sniffer_h2.go @@ -32,8 +32,8 @@ func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions) return errors.New("h2: invalid client preface") } - ro := ho.recorderObject - log := ho.log + ro := ho.RecorderObject + log := ho.Log ro.Time = time.Time{} diff --git a/internal/util/forwarder/sniffer_http.go b/internal/util/forwarder/sniffer_http.go index 3318f01..480eb71 100644 --- a/internal/util/forwarder/sniffer_http.go +++ b/internal/util/forwarder/sniffer_http.go @@ -51,7 +51,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO for _, opt := range opts { opt(&ho) } - ho.readTimeout = h.effectiveReadTimeout(&ho) + ho.ReadTimeout = h.effectiveReadTimeout(&ho) pStats := xstats.Stats{} conn = stats_wrapper.WrapConn(conn, &pStats) @@ -62,13 +62,13 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO return err } - log := ho.log + log := ho.Log if log.IsLevelEnabled(logger.TraceLevel) { dump, _ := httputil.DumpRequest(req, false) log.Trace(string(dump)) } - ro := ho.recorderObject + ro := ho.RecorderObject // Copy ro so that all internal recording (cache-hit serveCachedResponse // and cache-miss httpRoundTrip) happens on a local clone, preventing @@ -76,7 +76,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO ro2 := &xrecorder.HandlerRecorderObject{} *ro2 = *ro ro = ro2 - ho.recorderObject = ro2 + ho.RecorderObject = ro2 if clientIP := xhttp.GetClientIP(req); clientIP != nil { clientAddr := &net.TCPAddr{IP: clientIP} @@ -140,8 +140,8 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO } upstreamHost = normalizeHost(ro.HTTP.Host, "80") - ho.log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) - log = ho.log + ho.Log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) + log = ho.Log log.Debugf("connected to node %s(%s)", node.Name, node.Addr) ro.SrcAddr = cc.LocalAddr().String() @@ -179,7 +179,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO ro.Host = reqHost log = log.WithFields(map[string]any{"host": reqHost}) - ho.log = log + ho.Log = log } } } @@ -227,19 +227,19 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha StatusCode: http.StatusServiceUnavailable, } - if ho.bypass != nil && - ho.bypass.Contains(ctx, "tcp", host, - bypass.WithService(ho.service), + if ho.Bypass != nil && + ho.Bypass.Contains(ctx, "tcp", host, + bypass.WithService(ho.Service), bypass.WithPathOption(req.RequestURI)) { - ho.log.Debugf("bypass: %s %s", host, req.RequestURI) + ho.Log.Debugf("bypass: %s %s", host, req.RequestURI) res.StatusCode = http.StatusForbidden return nil, res, xbypass.ErrBypass } node = &chain.Node{} - if ho.hop != nil { + if ho.Hop != nil { var clientIP net.IP - if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.recorderObject.ClientAddr); clientAddr != nil { + if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.RecorderObject.ClientAddr); clientAddr != nil { clientIP = clientAddr.IP } @@ -248,7 +248,7 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha // without consuming the stream that is still forwarded below. var bodyPrefix []byte var maxBodySize int - if nl, ok := ho.hop.(hop.NodeList); ok { + if nl, ok := ho.Hop.(hop.NodeList); ok { for _, n := range nl.Nodes() { if n == nil { continue @@ -271,7 +271,7 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha } } - node = ho.hop.Select(ctx, + node = ho.Hop.Select(ctx, hop.ClientIPSelectOption(clientIP), hop.ProtocolSelectOption(sniffing.ProtoHTTP), hop.HostSelectOption(host), @@ -283,11 +283,11 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha ) } if node == nil { - ho.log.Warnf("node for %s not found", host) + ho.Log.Warnf("node for %s not found", host) res.StatusCode = http.StatusBadGateway return nil, res, errors.New("node not available") } - ho.recorderObject.Node = node.Name + ho.RecorderObject.Node = node.Name if node.Addr == "" { node = &chain.Node{ Name: node.Name, @@ -299,21 +299,21 @@ func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *Ha // 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 := ho.dial + dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext } - if node = ho.node; node != nil { + if node = ho.Node; node != nil { cc, err = dial(ctx, "tcp", node.Addr) return } - ro := ho.recorderObject + ro := ho.RecorderObject host := normalizeHost(req.Host, "80") if host != "" { ro.Host = host - ho.log = ho.log.WithFields(map[string]any{ + ho.Log = ho.Log.WithFields(map[string]any{ "host": host, }) } @@ -322,7 +322,7 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho if resolveErr != nil { ro.HTTP.StatusCode = res.StatusCode if werr := res.Write(conn); werr != nil { - ho.log.Warnf("write error response: %v", werr) + ho.Log.Warnf("write error response: %v", werr) } return nil, nil, resolveErr } @@ -336,20 +336,20 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho } ro.Host = node.Addr - ho.log = ho.log.WithFields(map[string]any{ + ho.Log = ho.Log.WithFields(map[string]any{ "node": node.Name, "dst": node.Addr, }) - ho.log.Debugf("find node for host %s -> %s(%s)", host, node.Name, node.Addr) + ho.Log.Debugf("find node for host %s -> %s(%s)", host, node.Name, node.Addr) cc, err = dial(ctx, "tcp", node.Addr) if err != nil { if marker := node.Marker(); marker != nil { marker.Mark() } - ho.log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) + 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) + ho.Log.Warnf("write error response: %v", werr) } return } @@ -368,9 +368,9 @@ 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 - log := ho.log + log := ho.Log ro := &xrecorder.HandlerRecorderObject{} - *ro = *ho.recorderObject + *ro = *ho.RecorderObject ro.Time = time.Now() log.Infof("%s <-> %s", ro.RemoteAddr, req.Host) @@ -439,7 +439,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, if httpSettings != nil { if auther := httpSettings.Auther; auther != nil { username, password, _ := req.BasicAuth() - id, ok := auther.Authenticate(ctx, username, password, auth.WithService(ho.service)) + id, ok := auther.Authenticate(ctx, username, password, auth.WithService(ho.Service)) if !ok { res.StatusCode = http.StatusUnauthorized ro.HTTP.StatusCode = res.StatusCode @@ -534,7 +534,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, br := bufio.NewReader(cc) var resp *http.Response for { - xio.SetReadDeadline(cc, time.Now().Add(ho.readTimeout)) + xio.SetReadDeadline(cc, time.Now().Add(ho.ReadTimeout)) resp, err = http.ReadResponse(br, req) if err != nil { log.Errorf("read response: %v", err) @@ -600,7 +600,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, resp.ProtoMinor = req.ProtoMinor } - if !ho.httpKeepalive { + if !ho.HTTPKeepalive { resp.Header.Set("Connection", "close") } @@ -679,7 +679,7 @@ func (h *Sniffer) serveStale(rw io.Writer, staleResp *http.Response, ro *xrecord if h.Cache == nil || staleResp == nil || !h.Cache.ServeStale() { return false } - if !ho.httpKeepalive { + if !ho.HTTPKeepalive { staleResp.Header.Set("Connection", "close") } ro.HTTP.StatusCode = staleResp.StatusCode diff --git a/internal/util/forwarder/sniffer_redis.go b/internal/util/forwarder/sniffer_redis.go index 4a0fa07..732771d 100644 --- a/internal/util/forwarder/sniffer_redis.go +++ b/internal/util/forwarder/sniffer_redis.go @@ -25,20 +25,20 @@ func (h *Sniffer) HandleRedis(ctx context.Context, conn net.Conn, opts ...Handle } buf := new(bytes.Buffer) - if err := sniffing.ParseRedisMetadata(io.TeeReader(conn, buf), ho.recorderObject); err != nil { + if err := sniffing.ParseRedisMetadata(io.TeeReader(conn, buf), ho.RecorderObject); err != nil { return err } - ro := ho.recorderObject + ro := ho.RecorderObject node, cc, err := h.dialRedis(ctx, &ho) if err != nil { return err } defer cc.Close() - ho.node = node + ho.Node = node - log := ho.log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) + log := ho.Log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) log.Debugf("connected to node %s(%s)", node.Name, node.Addr) ro.SrcAddr = cc.LocalAddr().String() @@ -63,12 +63,12 @@ func (h *Sniffer) HandleRedis(ctx context.Context, conn net.Conn, opts ...Handle // and ClientIPSelectOption. func resolveRedisNode(ctx context.Context, ho *HandleOptions) (node *chain.Node, err error) { node = &chain.Node{} - if ho.hop != nil { + if ho.Hop != nil { var clientIP net.IP - if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.recorderObject.ClientAddr); clientAddr != nil { + if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.RecorderObject.ClientAddr); clientAddr != nil { clientIP = clientAddr.IP } - node = ho.hop.Select(ctx, + node = ho.Hop.Select(ctx, hop.ClientIPSelectOption(clientIP), hop.ProtocolSelectOption(sniffing.ProtoRedis), ) @@ -76,19 +76,19 @@ func resolveRedisNode(ctx context.Context, ho *HandleOptions) (node *chain.Node, if node == nil { return nil, fmt.Errorf("node not available") } - ho.recorderObject.Node = node.Name + ho.RecorderObject.Node = node.Name return node, nil } // dialRedis selects a node via resolveRedisNode and establishes a TCP // connection to the node's address. func (h *Sniffer) dialRedis(ctx context.Context, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) { - dial := ho.dial + dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext } - if node = ho.node; node != nil { + if node = ho.Node; node != nil { cc, err = dial(ctx, "tcp", node.Addr) return } @@ -98,7 +98,7 @@ func (h *Sniffer) dialRedis(ctx context.Context, ho *HandleOptions) (node *chain return } - ro := ho.recorderObject + ro := ho.RecorderObject addr := node.Addr network := "tcp" if opts := node.Options(); opts != nil { @@ -117,18 +117,18 @@ func (h *Sniffer) dialRedis(ctx context.Context, ho *HandleOptions) (node *chain } ro.Host = addr - ho.log = ho.log.WithFields(map[string]any{ + ho.Log = ho.Log.WithFields(map[string]any{ "node": node.Name, "dst": fmt.Sprintf("%s/%s", addr, network), }) - ho.log.Debugf("find node for redis -> %s(%s)", node.Name, addr) + ho.Log.Debugf("find node for redis -> %s(%s)", node.Name, addr) cc, err = dial(ctx, network, addr) if err != nil { if marker := node.Marker(); marker != nil { marker.Mark() } - ho.log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) + ho.Log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) return } diff --git a/internal/util/forwarder/sniffer_test.go b/internal/util/forwarder/sniffer_test.go index 056c87b..c7b0aeb 100644 --- a/internal/util/forwarder/sniffer_test.go +++ b/internal/util/forwarder/sniffer_test.go @@ -21,9 +21,11 @@ import ( "github.com/go-gost/core/hop" "github.com/go-gost/core/recorder" "github.com/go-gost/core/rewriter" - "github.com/klauspost/compress/zstd" xlogger "github.com/go-gost/x/logger" xrecorder "github.com/go-gost/x/recorder" + "github.com/klauspost/compress/zstd" + + "github.com/go-gost/x/internal/util/sniffing" ) // ============================================================================= @@ -230,7 +232,6 @@ func TestRewriteReqBody_Gzip(t *testing.T) { }) } - func TestClampBodySize(t *testing.T) { tests := []struct { name string @@ -297,7 +298,7 @@ func TestEffectiveReadTimeout(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { h := &Sniffer{ReadTimeout: tt.snifferTO} - ho := &HandleOptions{readTimeout: tt.hoTO} + ho := &HandleOptions{ReadTimeout: tt.hoTO} got := h.effectiveReadTimeout(ho) if got != tt.want { t.Errorf("effectiveReadTimeout() = %v, want %v", got, tt.want) @@ -555,12 +556,12 @@ func TestHandleHTTP_BasicProxy(t *testing.T) { errCh := make(chan error, 1) go func() { errCh <- h.HandleHTTP(context.Background(), serverConn, - WithService("test-svc"), - WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + sniffing.WithService("test-svc"), + sniffing.WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { return net.Dial("tcp", upstream.Listener.Addr().String()) }), - WithRecorderObject(ro), - WithLog(xlogger.Nop()), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(xlogger.Nop()), ) }() @@ -612,9 +613,9 @@ func TestHandleHTTP_HTTP2Detection(t *testing.T) { errCh := make(chan error, 1) go func() { errCh <- h.HandleHTTP(context.Background(), serverConn, - WithService("test-svc"), - WithRecorderObject(ro), - WithLog(xlogger.Nop()), + sniffing.WithService("test-svc"), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(xlogger.Nop()), ) }() @@ -637,16 +638,16 @@ func TestHandleHTTP_HTTP2Detection(t *testing.T) { func TestWithService(t *testing.T) { opts := &HandleOptions{} - WithService("mysvc")(opts) - if opts.service != "mysvc" { - t.Errorf("service = %q, want %q", opts.service, "mysvc") + sniffing.WithService("mysvc")(opts) + if opts.Service != "mysvc" { + t.Errorf("service = %q, want %q", opts.Service, "mysvc") } } func TestWithHTTPKeepalive(t *testing.T) { opts := &HandleOptions{} - WithHTTPKeepalive(true)(opts) - if !opts.httpKeepalive { + sniffing.WithHTTPKeepalive(true)(opts) + if !opts.HTTPKeepalive { t.Errorf("httpKeepalive = false, want true") } } @@ -654,8 +655,8 @@ func TestWithHTTPKeepalive(t *testing.T) { func TestWithNode(t *testing.T) { opts := &HandleOptions{} node := &chain.Node{Name: "test-node", Addr: "127.0.0.1:8080"} - WithNode(node)(opts) - if opts.node != node { + sniffing.WithNode(node)(opts) + if opts.Node != node { t.Errorf("node not set") } } @@ -663,8 +664,8 @@ func TestWithNode(t *testing.T) { func TestWithHop(t *testing.T) { opts := &HandleOptions{} mh := &mockHop{} - WithHop(mh)(opts) - if opts.hop != mh { + sniffing.WithHop(mh)(opts) + if opts.Hop != mh { t.Errorf("hop not set") } } @@ -672,8 +673,8 @@ func TestWithHop(t *testing.T) { func TestWithBypass(t *testing.T) { opts := &HandleOptions{} bp := &mockBypass{contains: true} - WithBypass(bp)(opts) - if opts.bypass != bp { + sniffing.WithBypass(bp)(opts) + if opts.Bypass != bp { t.Errorf("bypass not set") } } @@ -681,8 +682,8 @@ func TestWithBypass(t *testing.T) { func TestWithRecorderObject(t *testing.T) { opts := &HandleOptions{} ro := &xrecorder.HandlerRecorderObject{} - WithRecorderObject(ro)(opts) - if opts.recorderObject != ro { + sniffing.WithRecorderObject(ro)(opts) + if opts.RecorderObject != ro { t.Errorf("recorderObject not set") } } @@ -690,8 +691,8 @@ func TestWithRecorderObject(t *testing.T) { func TestWithLog(t *testing.T) { opts := &HandleOptions{} log := xlogger.Nop() - WithLog(log)(opts) - if opts.log != log { + sniffing.WithLog(log)(opts) + if opts.Log != log { t.Errorf("log not set") } } @@ -732,11 +733,11 @@ func TestHandleHTTP_DialError(t *testing.T) { errCh := make(chan error, 1) go func() { errCh <- h.HandleHTTP(context.Background(), serverConn, - WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + sniffing.WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { return nil, io.ErrUnexpectedEOF }), - WithRecorderObject(ro), - WithLog(xlogger.Nop()), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(xlogger.Nop()), ) }() @@ -782,11 +783,11 @@ func TestHandleHTTP_HTTP10Request(t *testing.T) { errCh := make(chan error, 1) go func() { errCh <- h.HandleHTTP(context.Background(), serverConn, - WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + sniffing.WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { return net.Dial("tcp", upstream.Listener.Addr().String()) }), - WithRecorderObject(ro), - WithLog(xlogger.Nop()), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(xlogger.Nop()), ) }() @@ -926,10 +927,10 @@ var _ hop.Hop = (*configurableHop)(nil) func TestResolveHTTPNode_Bypass(t *testing.T) { ho := &HandleOptions{ - service: "test-svc", - bypass: &mockBypass{contains: true}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Service: "test-svc", + Bypass: &mockBypass{contains: true}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } req, _ := http.NewRequest("GET", "http://example.com/path", nil) @@ -950,8 +951,8 @@ func TestResolveHTTPNode_Bypass(t *testing.T) { func TestResolveHTTPNode_NoHop(t *testing.T) { ho := &HandleOptions{ - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } req, _ := http.NewRequest("GET", "http://example.com/path", nil) @@ -972,9 +973,9 @@ func TestResolveHTTPNode_NoHop(t *testing.T) { func TestResolveHTTPNode_HopReturnsNil(t *testing.T) { ho := &HandleOptions{ - hop: &configurableHop{node: nil}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: nil}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } req, _ := http.NewRequest("GET", "http://example.com/path", nil) @@ -999,9 +1000,9 @@ func TestResolveHTTPNode_HopReturnsNil(t *testing.T) { func TestResolveHTTPNode_HopReturnsValidNode(t *testing.T) { expectedNode := &chain.Node{Name: "backend", Addr: "10.0.0.1:8080"} ho := &HandleOptions{ - hop: &configurableHop{node: expectedNode}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: expectedNode}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } req, _ := http.NewRequest("GET", "http://example.com/path", nil) @@ -1019,9 +1020,9 @@ func TestResolveHTTPNode_HopReturnsValidNode(t *testing.T) { func TestResolveHTTPNode_HopReturnsNodeWithoutAddr(t *testing.T) { ho := &HandleOptions{ - hop: &configurableHop{node: &chain.Node{Name: "backend"}}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: &chain.Node{Name: "backend"}}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } req, _ := http.NewRequest("GET", "http://example.com/path", nil) @@ -1049,8 +1050,8 @@ func TestResolveHTTPNode_HopReturnsNodeWithoutAddr(t *testing.T) { func TestResolveTLSNode_NoHop(t *testing.T) { ho := &HandleOptions{ - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } node, err := resolveTLSNode(context.Background(), "example.com", ho) @@ -1067,9 +1068,9 @@ func TestResolveTLSNode_NoHop(t *testing.T) { func TestResolveTLSNode_HopReturnsNil(t *testing.T) { ho := &HandleOptions{ - hop: &configurableHop{node: nil}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: nil}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } node, err := resolveTLSNode(context.Background(), "example.com", ho) @@ -1087,9 +1088,9 @@ func TestResolveTLSNode_HopReturnsNil(t *testing.T) { func TestResolveTLSNode_HopReturnsValidNode(t *testing.T) { expectedNode := &chain.Node{Name: "backend", Addr: "10.0.0.1:443"} ho := &HandleOptions{ - hop: &configurableHop{node: expectedNode}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: expectedNode}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } node, err := resolveTLSNode(context.Background(), "example.com", ho) @@ -1103,9 +1104,9 @@ func TestResolveTLSNode_HopReturnsValidNode(t *testing.T) { func TestResolveTLSNode_HopReturnsNodeWithoutAddr(t *testing.T) { ho := &HandleOptions{ - hop: &configurableHop{node: &chain.Node{Name: "backend"}}, - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Hop: &configurableHop{node: &chain.Node{Name: "backend"}}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } node, err := resolveTLSNode(context.Background(), "example.com", ho) @@ -1720,8 +1721,8 @@ func TestServeH2_InvalidPreface(t *testing.T) { Recorder: &noopRecorder{}, } ho := &HandleOptions{ - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } clientConn, serverConn := net.Pipe() @@ -1744,7 +1745,6 @@ func TestServeH2_InvalidPreface(t *testing.T) { // SSE (Server-Sent Events) Tests // ============================================================================= - func TestRewriteRespBody_SSE_ContentType(t *testing.T) { // rewriteRespBody now rewrites SSE per-event via the universal wrapper. resp := &http.Response{ diff --git a/internal/util/forwarder/sniffer_tls.go b/internal/util/forwarder/sniffer_tls.go index c45a616..bc5efe0 100644 --- a/internal/util/forwarder/sniffer_tls.go +++ b/internal/util/forwarder/sniffer_tls.go @@ -32,7 +32,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp for _, opt := range opts { opt(&ho) } - ho.readTimeout = h.effectiveReadTimeout(&ho) + ho.ReadTimeout = h.effectiveReadTimeout(&ho) buf := new(bytes.Buffer) clientHello, err := dissector.ParseClientHello(io.TeeReader(conn, buf)) @@ -40,7 +40,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp return err } - ro := ho.recorderObject + ro := ho.RecorderObject ro.TLS = &xrecorder.TLSRecorderObject{ ServerName: clientHello.ServerName, ClientHello: hex.EncodeToString(buf.Bytes()), @@ -51,14 +51,14 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp host := normalizeHost(clientHello.ServerName, "443") if host == "" { - if ho.log != nil { - ho.log.Debugf("no sni in clienthello from %s", conn.RemoteAddr()) + if ho.Log != nil { + ho.Log.Debugf("no sni in clienthello from %s", conn.RemoteAddr()) } } else { ro.Host = host } - if ho.bypass != nil && ho.bypass.Contains(ctx, "tcp", host, bypass.WithService(ho.service)) { + if ho.Bypass != nil && ho.Bypass.Contains(ctx, "tcp", host, bypass.WithService(ho.Service)) { return xbypass.ErrBypass } @@ -67,9 +67,9 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp return err } defer cc.Close() - ho.node = node + ho.Node = node - log := ho.log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) + log := ho.Log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) log.Debugf("connected to node %s(%s)", node.Name, node.Addr) ro.SrcAddr = cc.LocalAddr().String() @@ -89,7 +89,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp return err } - xio.SetReadDeadline(cc, time.Now().Add(ho.readTimeout)) + xio.SetReadDeadline(cc, time.Now().Add(ho.ReadTimeout)) serverHello, serverHelloErr := dissector.ParseServerHello(io.TeeReader(cc, buf)) xio.SetReadDeadline(cc, time.Time{}) @@ -128,12 +128,12 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp // resolveTLSNode selects a node for a TLS connection by applying hop selection. func resolveTLSNode(ctx context.Context, host string, ho *HandleOptions) (node *chain.Node, err error) { node = &chain.Node{} - if ho.hop != nil { + if ho.Hop != nil { var clientIP net.IP - if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.recorderObject.ClientAddr); clientAddr != nil { + if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.RecorderObject.ClientAddr); clientAddr != nil { clientIP = clientAddr.IP } - node = ho.hop.Select(ctx, + node = ho.Hop.Select(ctx, hop.ClientIPSelectOption(clientIP), hop.HostSelectOption(host), hop.ProtocolSelectOption(sniffing.ProtoTLS), @@ -142,7 +142,7 @@ func resolveTLSNode(ctx context.Context, host string, ho *HandleOptions) (node * if node == nil { return nil, errors.New("node not available") } - ho.recorderObject.Node = node.Name + ho.RecorderObject.Node = node.Name if node.Addr == "" { node = &chain.Node{ Name: node.Name, @@ -154,12 +154,12 @@ func resolveTLSNode(ctx context.Context, host string, ho *HandleOptions) (node * // dialTLS selects a node and establishes a TLS connection. func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) { - dial := ho.dial + dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext } - if node = ho.node; node != nil { + if node = ho.Node; node != nil { cc, err = dial(ctx, "tcp", node.Addr) return } @@ -169,7 +169,7 @@ func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) ( return } - ro := ho.recorderObject + ro := ho.RecorderObject addr := node.Addr network := "tcp" if opts := node.Options(); opts != nil { @@ -188,19 +188,19 @@ func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) ( } ro.Host = addr - ho.log = ho.log.WithFields(map[string]any{ + ho.Log = ho.Log.WithFields(map[string]any{ "host": host, "node": node.Name, "dst": fmt.Sprintf("%s/%s", addr, network), }) - ho.log.Debugf("find node for host %s -> %s(%s)", host, node.Name, addr) + ho.Log.Debugf("find node for host %s -> %s(%s)", host, node.Name, addr) cc, err = dial(ctx, network, addr) if err != nil { if marker := node.Marker(); marker != nil { marker.Mark() } - ho.log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) + ho.Log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err) return } @@ -217,8 +217,8 @@ func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) ( // dynamically generated certificate. The decrypted traffic is then handled // as HTTP. func (h *Sniffer) terminateTLS(ctx context.Context, conn, cc net.Conn, clientHello *dissector.ClientHelloInfo, ho *HandleOptions) error { - ro := ho.recorderObject - log := ho.log + ro := ho.RecorderObject + log := ho.Log nextProtos := clientHello.SupportedProtos if h.NegotiatedProtocol != "" { @@ -317,14 +317,14 @@ func (h *Sniffer) terminateTLS(ctx context.Context, conn, cc net.Conn, clientHel } opts := []HandleOption{ - WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + sniffing.WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { return clientConn, nil }), - WithHTTPKeepalive(true), - WithNode(ho.node), - WithBypass(ho.bypass), - WithRecorderObject(ro), - WithLog(log), + sniffing.WithHTTPKeepalive(true), + sniffing.WithNode(ho.Node), + sniffing.WithBypass(ho.Bypass), + sniffing.WithRecorderObject(ro), + sniffing.WithLog(log), } return h.HandleHTTP(ctx, serverConn, opts...) } diff --git a/internal/util/sniffing/registry.go b/internal/util/sniffing/registry.go new file mode 100644 index 0000000..71c692e --- /dev/null +++ b/internal/util/sniffing/registry.go @@ -0,0 +1,48 @@ +package sniffing + +import ( + "context" + "net" + "sync" +) + +// HandleFunc handles a connection whose protocol has been detected as proto. +// The receiver s carries the per-connection Sniffer state (MITM certificate, +// HTTP cache, recorder, etc.). +type HandleFunc func(s *Sniffer, ctx context.Context, network string, conn net.Conn, opts ...HandleOption) error + +var protoHandlers sync.Map // map[string]HandleFunc + +// Register associates a protocol with its handler. Duplicate registration +// panics, mirroring the house registry pattern. +func Register(proto string, fn HandleFunc) { + if _, dup := protoHandlers.LoadOrStore(proto, fn); dup { + panic("sniffing: duplicate protocol handler: " + proto) + } +} + +// Get returns the handler registered for proto. +func Get(proto string) (HandleFunc, bool) { + v, ok := protoHandlers.Load(proto) + if !ok { + return nil, false + } + return v.(HandleFunc), true +} + +// Dispatch routes a detected protocol to its registered handler. It returns +// (false, nil) when no handler is registered for proto, signalling the caller +// to fall through to raw forwarding. +func Dispatch(s *Sniffer, ctx context.Context, network string, conn net.Conn, proto string, opts ...HandleOption) (handled bool, err error) { + fn, ok := Get(proto) + if !ok { + return false, nil + } + return true, fn(s, ctx, network, conn, opts...) +} + +func init() { + Register(ProtoHTTP, (*Sniffer).HandleHTTP) + Register(ProtoTLS, (*Sniffer).HandleTLS) + Register(ProtoRedis, (*Sniffer).HandleRedis) +} diff --git a/internal/util/sniffing/registry_test.go b/internal/util/sniffing/registry_test.go new file mode 100644 index 0000000..aeffed6 --- /dev/null +++ b/internal/util/sniffing/registry_test.go @@ -0,0 +1,84 @@ +package sniffing + +import ( + "context" + "net" + "testing" +) + +// TestRegisterGet verifies a handler can be registered and looked up by proto. +func TestRegisterGet(t *testing.T) { + const proto = "testproto-get" + called := false + Register(proto, func(s *Sniffer, _ context.Context, _ string, _ net.Conn, _ ...HandleOption) error { + called = true + return nil + }) + defer protoHandlers.Delete(proto) + + fn, ok := Get(proto) + if !ok { + t.Fatal("Get: handler not found") + } + if err := fn(&Sniffer{}, context.Background(), "tcp", nil); err != nil { + t.Fatalf("fn() error = %v", err) + } + if !called { + t.Fatal("registered handler was not called") + } + if _, ok := Get("no-such-proto"); ok { + t.Fatal("Get: unexpected handler for unknown proto") + } +} + +// TestRegisterDuplicate verifies registering the same proto twice panics. +func TestRegisterDuplicate(t *testing.T) { + const proto = "testproto-dup" + Register(proto, func(*Sniffer, context.Context, string, net.Conn, ...HandleOption) error { + return nil + }) + defer protoHandlers.Delete(proto) + + defer func() { + if recover() == nil { + t.Fatal("expected panic on duplicate registration") + } + }() + Register(proto, func(*Sniffer, context.Context, string, net.Conn, ...HandleOption) error { + return nil + }) +} + +// TestDispatch_Unknown verifies Dispatch falls through (false, nil) for an +// unregistered proto. +func TestDispatch_Unknown(t *testing.T) { + handled, err := Dispatch(&Sniffer{}, context.Background(), "tcp", nil, "no-such-proto") + if handled { + t.Fatal("Dispatch: unknown proto handled = true, want false") + } + if err != nil { + t.Fatalf("Dispatch: err = %v, want nil", err) + } +} + +// TestDispatch_Known verifies Dispatch routes a registered proto to its handler. +func TestDispatch_Known(t *testing.T) { + const proto = "testproto-known" + called := false + Register(proto, func(s *Sniffer, _ context.Context, _ string, _ net.Conn, _ ...HandleOption) error { + called = true + return nil + }) + defer protoHandlers.Delete(proto) + + handled, err := Dispatch(&Sniffer{}, context.Background(), "tcp", nil, proto) + if !handled { + t.Fatal("Dispatch: known proto handled = false, want true") + } + if err != nil { + t.Fatalf("Dispatch: err = %v, want nil", err) + } + if !called { + t.Fatal("Dispatch: registered handler was not called") + } +} diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index 1ae6bbc..c57ac5e 100644 --- a/internal/util/sniffing/sniffer.go +++ b/internal/util/sniffing/sniffer.go @@ -10,6 +10,8 @@ import ( "time" "github.com/go-gost/core/bypass" + "github.com/go-gost/core/chain" + "github.com/go-gost/core/hop" "github.com/go-gost/core/logger" "github.com/go-gost/core/recorder" "github.com/go-gost/x/internal/util/httpcache" @@ -33,58 +35,87 @@ const ( var DefaultCertPool = tls_util.NewMemoryCertPool() // HandleOptions holds configuration options for sniffing handlers. +// Fields are exported because forwarder's protocol handlers set them +// across the package boundary. type HandleOptions struct { - service string - dial func(ctx context.Context, network, address string) (net.Conn, error) - dialTLS func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) + Service string + Dial func(ctx context.Context, network, address string) (net.Conn, error) + DialTLS func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) - bypass bypass.Bypass - recorderObject *xrecorder.HandlerRecorderObject - log logger.Logger + HTTPKeepalive bool + ReadTimeout time.Duration + + Node *chain.Node + Hop hop.Hop + + Bypass bypass.Bypass + RecorderObject *xrecorder.HandlerRecorderObject + Log logger.Logger } // HandleOption configures HandleOptions for sniffing handlers. type HandleOption func(opts *HandleOptions) -// WithService sets the service name for bypass lookups. +// WithService sets the service name for bypass and selection lookups. func WithService(service string) HandleOption { return func(opts *HandleOptions) { - opts.service = service + opts.Service = service } } // WithDial sets the dial function used to establish upstream connections. func WithDial(dial func(ctx context.Context, network, address string) (net.Conn, error)) HandleOption { return func(opts *HandleOptions) { - opts.dial = dial + opts.Dial = dial } } // WithDialTLS sets the dial function used for TLS-wrapped upstream connections. func WithDialTLS(dialTLS func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error)) HandleOption { return func(opts *HandleOptions) { - opts.dialTLS = dialTLS + opts.DialTLS = dialTLS + } +} + +// WithHTTPKeepalive enables or disables HTTP keep-alive on the upstream connection. +func WithHTTPKeepalive(keepalive bool) HandleOption { + return func(opts *HandleOptions) { + opts.HTTPKeepalive = keepalive + } +} + +// WithNode sets a pre-resolved chain node to connect to, bypassing hop selection. +func WithNode(node *chain.Node) HandleOption { + return func(opts *HandleOptions) { + opts.Node = node + } +} + +// WithHop sets the hop used for node selection when routing requests. +func WithHop(h hop.Hop) HandleOption { + return func(opts *HandleOptions) { + opts.Hop = h } } // WithBypass sets the bypass rules for filtering requests by host. func WithBypass(bypass bypass.Bypass) HandleOption { return func(opts *HandleOptions) { - opts.bypass = bypass + opts.Bypass = bypass } } // WithRecorderObject sets the recorder object for capturing traffic metadata. func WithRecorderObject(ro *xrecorder.HandlerRecorderObject) HandleOption { return func(opts *HandleOptions) { - opts.recorderObject = ro + opts.RecorderObject = ro } } // WithLog sets the logger for the handler. func WithLog(log logger.Logger) HandleOption { return func(opts *HandleOptions) { - opts.log = log + opts.Log = log } } diff --git a/internal/util/sniffing/sniffer_h2.go b/internal/util/sniffing/sniffer_h2.go index 3e17fdb..ef8e93d 100644 --- a/internal/util/sniffing/sniffer_h2.go +++ b/internal/util/sniffing/sniffer_h2.go @@ -33,14 +33,14 @@ func (h *Sniffer) serveH2(ctx context.Context, network string, conn net.Conn, ho return errors.New("h2: invalid client preface") } - ro := ho.recorderObject - log := ho.log + ro := ho.RecorderObject + log := ho.Log ro.Time = time.Time{} tr := &http2.Transport{ DialTLSContext: func(ctx context.Context, nw, addr string, cfg *tls.Config) (net.Conn, error) { - if dial := ho.dialTLS; dial != nil { + if dial := ho.DialTLS; dial != nil { return dial(ctx, network, addr, cfg) } diff --git a/internal/util/sniffing/sniffer_http.go b/internal/util/sniffing/sniffer_http.go index c1be2f8..f0da963 100644 --- a/internal/util/sniffing/sniffer_http.go +++ b/internal/util/sniffing/sniffer_http.go @@ -44,13 +44,13 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, network string, conn net.Conn, return err } - log := ho.log + log := ho.Log if log.IsLevelEnabled(logger.TraceLevel) { dump, _ := httputil.DumpRequest(req, false) log.Trace(string(dump)) } - ro := ho.recorderObject + ro := ho.RecorderObject // Copy ro so that all internal recording (cache-hit serveCachedResponse // and cache-miss httpRoundTrip) happens on a local clone, preventing @@ -74,12 +74,12 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, network string, conn net.Conn, ro.Host = host log = log.WithFields(map[string]any{"host": host}) - if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + if ho.Bypass != nil && ho.Bypass.Contains(ctx, network, host, bypass.WithService(ho.Service)) { return xbypass.ErrBypass } } - dialFn := ho.dial + dialFn := ho.Dial if dialFn == nil { dialFn = (&net.Dialer{}).DialContext } diff --git a/internal/util/sniffing/sniffer_redis.go b/internal/util/sniffing/sniffer_redis.go index 243e9a7..994012d 100644 --- a/internal/util/sniffing/sniffer_redis.go +++ b/internal/util/sniffing/sniffer_redis.go @@ -153,14 +153,14 @@ func (h *Sniffer) HandleRedis(ctx context.Context, network string, conn net.Conn } buf := new(bytes.Buffer) - if err := ParseRedisMetadata(io.TeeReader(conn, buf), ho.recorderObject); err != nil { + if err := ParseRedisMetadata(io.TeeReader(conn, buf), ho.RecorderObject); err != nil { return err } - log := ho.log - ro := ho.recorderObject + log := ho.Log + ro := ho.RecorderObject - dial := ho.dial + dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext } diff --git a/internal/util/sniffing/sniffer_test.go b/internal/util/sniffing/sniffer_test.go index 858959c..035ef9e 100644 --- a/internal/util/sniffing/sniffer_test.go +++ b/internal/util/sniffing/sniffer_test.go @@ -177,8 +177,8 @@ func TestWrapErr(t *testing.T) { func TestWithService(t *testing.T) { opts := &HandleOptions{} WithService("mysvc")(opts) - if opts.service != "mysvc" { - t.Errorf("service = %q, want %q", opts.service, "mysvc") + if opts.Service != "mysvc" { + t.Errorf("service = %q, want %q", opts.Service, "mysvc") } } @@ -188,7 +188,7 @@ func TestWithDial(t *testing.T) { return nil, nil } WithDial(dial)(opts) - if opts.dial == nil { + if opts.Dial == nil { t.Error("dial should be set") } } @@ -199,7 +199,7 @@ func TestWithDialTLS(t *testing.T) { return nil, nil } WithDialTLS(dialTLS)(opts) - if opts.dialTLS == nil { + if opts.DialTLS == nil { t.Error("dialTLS should be set") } } @@ -208,7 +208,7 @@ func TestWithBypass(t *testing.T) { opts := &HandleOptions{} bp := &mockBypass{contains: true} WithBypass(bp)(opts) - if opts.bypass != bp { + if opts.Bypass != bp { t.Errorf("bypass not set") } } @@ -217,7 +217,7 @@ func TestWithRecorderObject(t *testing.T) { opts := &HandleOptions{} ro := &xrecorder.HandlerRecorderObject{} WithRecorderObject(ro)(opts) - if opts.recorderObject != ro { + if opts.RecorderObject != ro { t.Errorf("recorderObject not set") } } @@ -226,7 +226,7 @@ func TestWithLog(t *testing.T) { opts := &HandleOptions{} log := xlogger.Nop() WithLog(log)(opts) - if opts.log != log { + if opts.Log != log { t.Errorf("log not set") } } @@ -416,8 +416,8 @@ func TestCopyWebsocketFrame_Basic(t *testing.T) { } var frame bytes.Buffer - frame.WriteByte(0x81) // FIN + text opcode - frame.WriteByte(0x82) // MASK + len=2 + frame.WriteByte(0x81) // FIN + text opcode + frame.WriteByte(0x82) // MASK + len=2 frame.Write(mask) frame.Write(maskedPayload) @@ -664,8 +664,8 @@ func TestServeH2_InvalidPreface(t *testing.T) { Recorder: &noopRecorder{}, } ho := &HandleOptions{ - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } clientConn, serverConn := net.Pipe() @@ -692,8 +692,8 @@ func TestServeH2_ShortRead(t *testing.T) { Recorder: &noopRecorder{}, } ho := &HandleOptions{ - log: xlogger.Nop(), - recorderObject: &xrecorder.HandlerRecorderObject{}, + Log: xlogger.Nop(), + RecorderObject: &xrecorder.HandlerRecorderObject{}, } clientConn, serverConn := net.Pipe() @@ -1058,9 +1058,9 @@ func TestHandleTLS_NonTLSData(t *testing.T) { func TestHandleUpgradeResponse_WebsocketEnabled(t *testing.T) { h := &Sniffer{ - Websocket: true, - ReadTimeout: 5 * time.Second, - Recorder: &noopRecorder{}, + Websocket: true, + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, } ro := &xrecorder.HandlerRecorderObject{} diff --git a/internal/util/sniffing/sniffer_tls.go b/internal/util/sniffing/sniffer_tls.go index a05c7ba..9e6dbce 100644 --- a/internal/util/sniffing/sniffer_tls.go +++ b/internal/util/sniffing/sniffer_tls.go @@ -36,9 +36,9 @@ func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, return err } - log := ho.log + log := ho.Log - ro := ho.recorderObject + ro := ho.RecorderObject ro.TLS = &xrecorder.TLSRecorderObject{ ServerName: clientHello.ServerName, ClientHello: hex.EncodeToString(buf.Bytes()), @@ -56,11 +56,11 @@ func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, ro.Host = host } - if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + if ho.Bypass != nil && ho.Bypass.Contains(ctx, network, host, bypass.WithService(ho.Service)) { return xbypass.ErrBypass } - dial := ho.dial + dial := ho.Dial if dial == nil { dial = (&net.Dialer{}).DialContext } @@ -79,7 +79,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, if host == "" { host = ro.Host } - if h.MitmBypass == nil || !h.MitmBypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + if h.MitmBypass == nil || !h.MitmBypass.Contains(ctx, network, host, bypass.WithService(ho.Service)) { return h.terminateTLS(ctx, network, xnet.NewReadWriteConn(io.MultiReader(buf, conn), conn, conn), cc, clientHello, &ho) } } @@ -128,8 +128,8 @@ func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, // dynamically generated certificate. The decrypted traffic is then handled // as HTTP. func (h *Sniffer) terminateTLS(ctx context.Context, network string, conn, cc net.Conn, clientHello *dissector.ClientHelloInfo, ho *HandleOptions) error { - ro := ho.recorderObject - log := ho.log + ro := ho.RecorderObject + log := ho.Log nextProtos := clientHello.SupportedProtos if h.NegotiatedProtocol != "" { @@ -234,7 +234,7 @@ func (h *Sniffer) terminateTLS(ctx context.Context, network string, conn, cc net WithDialTLS(func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { return clientConn, nil }), - WithBypass(ho.bypass), + WithBypass(ho.Bypass), WithRecorderObject(ro), WithLog(log), }