From 32c3d33afda8904d26d665af04461314db89b649 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sat, 30 May 2026 16:09:48 +0800 Subject: [PATCH] refactor(handler/forward/local): split handler into 6 focused files with 40 tests at 100% coverage Extract Handle into handleRawForwarding, handleSniffedProtocol, SnifferBuilder, newRecorderObject, and checkRateLimit across forward.go, sniffing.go, and util.go. Replace redundant x509.ParseCertificate with tlsCert.Leaf (already populated by tls.LoadX509KeyPair since Go 1.23). --- handler/forward/local/forward.go | 93 ++ handler/forward/local/forward_test.go | 200 +++++ handler/forward/local/handler.go | 270 ++---- handler/forward/local/handler_test.go | 1072 +++--------------------- handler/forward/local/helpers_test.go | 243 ++++++ handler/forward/local/metadata.go | 5 +- handler/forward/local/metadata_test.go | 114 +++ handler/forward/local/sniffing.go | 99 +++ handler/forward/local/sniffing_test.go | 157 ++++ handler/forward/local/util.go | 50 ++ handler/forward/local/util_test.go | 145 ++++ 11 files changed, 1308 insertions(+), 1140 deletions(-) create mode 100644 handler/forward/local/forward.go create mode 100644 handler/forward/local/forward_test.go create mode 100644 handler/forward/local/helpers_test.go create mode 100644 handler/forward/local/metadata_test.go create mode 100644 handler/forward/local/sniffing.go create mode 100644 handler/forward/local/sniffing_test.go create mode 100644 handler/forward/local/util.go create mode 100644 handler/forward/local/util_test.go diff --git a/handler/forward/local/forward.go b/handler/forward/local/forward.go new file mode 100644 index 00000000..3be86618 --- /dev/null +++ b/handler/forward/local/forward.go @@ -0,0 +1,93 @@ +package local + +import ( + "bytes" + "context" + "net" + "time" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/hop" + "github.com/go-gost/core/logger" + xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" + xnet "github.com/go-gost/x/internal/net" + "github.com/go-gost/x/internal/net/proxyproto" + xrecorder "github.com/go-gost/x/recorder" +) + +// handleRawForwarding performs node selection, dials the target through the +// router, and pipes the raw connection. It is used when sniffing is disabled +// or the protocol was not HTTP/TLS. +func (h *forwardHandler) handleRawForwarding(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, network, proto string) error { + target := &chain.Node{} + if curHop := h.getHop(); curHop != nil { + target = curHop.Select(ctx, + hop.ProtocolSelectOption(proto), + ) + } + if target == nil { + log.Error(errNodeNotAvailable) + return errNodeNotAvailable + } + addr := target.Addr + if opts := target.Options(); opts != nil { + switch opts.Network { + case "unix": + network = opts.Network + default: + if _, _, err := net.SplitHostPort(addr); err != nil { + addr += ":0" + } + } + } + + ro.Network = network + ro.Host = addr + + log = log.WithFields(map[string]any{ + "node": target.Name, + "dst": addr, + "network": network, + }) + + log.Debugf("%s >> %s", conn.RemoteAddr(), addr) + + var buf bytes.Buffer + cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, addr) + ro.Route = buf.String() + if err != nil { + log.Error(err) + // TODO: the router itself may be failed due to the failed node in the router, + // the dead marker may be a wrong operation. + if marker := target.Marker(); marker != nil { + marker.Mark() + } + return err + } + if marker := target.Marker(); marker != nil { + marker.Reset() + } + defer cc.Close() + + cc = proxyproto.WrapClientConn( + h.md.proxyProtocol, + xctx.SrcAddrFromContext(ctx), + xctx.DstAddrFromContext(ctx), + cc) + + log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) + ro.SrcAddr = cc.LocalAddr().String() + ro.DstAddr = cc.RemoteAddr().String() + + t := time.Now() + log.Infof("%s <-> %s", conn.RemoteAddr(), target.Addr) + if err := xnet.Pipe(ctx, conn, cc, xnet.WithReadTimeout(h.md.readTimeout)); err != nil { + log.Debugf("pipe: %v", err) + } + log.WithFields(map[string]any{ + "duration": time.Since(t), + }).Infof("%s >-< %s", conn.RemoteAddr(), target.Addr) + + return nil +} diff --git a/handler/forward/local/forward_test.go b/handler/forward/local/forward_test.go new file mode 100644 index 00000000..6430e3b9 --- /dev/null +++ b/handler/forward/local/forward_test.go @@ -0,0 +1,200 @@ +package local + +import ( + "context" + "errors" + "net" + "testing" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/hop" + xrecorder "github.com/go-gost/x/recorder" +) + +// --------------------------------------------------------------------------- +// handleRawForwarding +// --------------------------------------------------------------------------- + +func TestHandleRawForwarding_NilHop(t *testing.T) { + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, errors.New("no route to host") + }, + })) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err == nil { + t.Fatal("expected error when hop is nil") + } +} + +func TestHandleRawForwarding_NodeUnavailable(t *testing.T) { + h := newInitdHandler() + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return nil + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err == nil { + t.Fatal("expected error when node is nil") + } + if !errors.Is(err, errNodeNotAvailable) { + t.Errorf("expected errNodeNotAvailable, got %v", err) + } +} + +func TestHandleRawForwarding_DialError(t *testing.T) { + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, errors.New("dial failed") + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("test-node", "127.0.0.1:9999") + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err == nil { + t.Fatal("expected error when dial fails") + } + if ro.Host != "127.0.0.1:9999" { + t.Errorf("expected host 127.0.0.1:9999, got %s", ro.Host) + } +} + +func TestHandleRawForwarding_Success(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return &chain.Node{Addr: "10.0.0.1:80", Name: "upstream"} + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if ro.Host != "10.0.0.1:80" { + t.Errorf("expected host 10.0.0.1:80, got %s", ro.Host) + } + if ro.Network != "tcp" { + t.Errorf("expected network tcp, got %s", ro.Network) + } + if ro.SrcAddr == "" { + t.Error("expected SrcAddr to be set") + } + if ro.DstAddr == "" { + t.Error("expected DstAddr to be set") + } +} + +func TestHandleRawForwarding_AddrWithoutPort(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return &chain.Node{Addr: "10.0.0.1", Name: "noport"} + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if ro.Host != "10.0.0.1:0" { + t.Errorf("expected ':0' appended to addr, got %s", ro.Host) + } +} + +func TestHandleRawForwarding_UnixSocket(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("unix-node", "/tmp/unix.sock", + chain.NetworkNodeOption("unix"), + ) + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ro.Network != "unix" { + t.Errorf("expected network 'unix', got '%s'", ro.Network) + } +} + +func TestHandleRawForwarding_MarkerReset(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("target", "10.0.0.1:80") + }, + }) + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ro.Host != "10.0.0.1:80" { + t.Errorf("expected host '10.0.0.1:80', got '%s'", ro.Host) + } +} diff --git a/handler/forward/local/handler.go b/handler/forward/local/handler.go index 36115554..6ba45a6f 100644 --- a/handler/forward/local/handler.go +++ b/handler/forward/local/handler.go @@ -1,25 +1,78 @@ +// Package local implements a low-level forwarding handler for raw TCP and UDP +// connections. It serves as the fallback handler for protocols that cannot be +// matched to a dedicated handler, and supports optional protocol sniffing to +// detect HTTP and TLS traffic for deep inspection (MITM, WebSocket recording). +// +// The handler is registered under the names "tcp", "udp", and "forward" via +// NewHandler in init(). +// +// # Connection processing flow +// +// Each inbound net.Conn is processed by Handle, which wraps the connection +// with I/O stats, checks the rate limiter, and dispatches to one of two +// forwarding paths: +// +// Handle() +// ├─ stats_wrapper.WrapConn (per-connection I/O counters) +// ├─ newRecorderObject (session metadata: service, SID, addresses) +// ├─ checkRateLimit (connection rate limiter, if configured) +// ├─ Router nil check → errRouterNotAvailable +// ├─ [sniffing enabled] sniffing.Sniff (peek buffer, detect protocol) +// │ └─ handleSniffedProtocol() +// │ ├─ sniffing.ProtoHTTP → SnifferBuilder.Build().HandleHTTP() +// │ ├─ sniffing.ProtoTLS → SnifferBuilder.Build().HandleTLS() +// │ └─ default → fall through to raw forwarding +// └─ [sniffing disabled / unrecognised] handleRawForwarding() +// +// # Protocol sniffing dispatch (handleSniffedProtocol) +// +// When sniffing is enabled and the initial bytes match HTTP or TLS, the +// connection is delegated to the forwarder package for protocol-aware +// handling: +// +// - HTTP: parses the request, applies MITM if configured, forwards to +// the upstream, and records WebSocket frames when enabled. +// - TLS: parses the ClientHello to extract the server name, applies +// MITM or SNI-based routing through the hop, and records the TLS +// handshake. +// +// SnifferBuilder is populated once during Init and reused via Build() +// to create per-connection forwarder.Sniffer instances. +// +// # Raw forwarding (handleRawForwarding) +// +// The default path for unrecognised traffic. It selects a target node +// from the hop, dials the upstream through the router, wraps the +// connection with proxy protocol (HAProxy) when enabled, and pipes +// data bidirectionally: +// +// 1. Node selection — getHop().Select() with optional protocol hint. +// 2. Router dial — Router.Dial() through the configured chain. +// 3. Proxy protocol — proxyproto.WrapClientConn() prepends HAProxy +// header when metadata proxyProtocol is set. +// 4. Bidirectional pipe — xnet.Pipe(ctx, conn, cc) with read timeout. +// +// # Hop management +// +// The handler implements handler.Forwarder. Forward() is called during +// service startup (and on reload) to install or update the hop. The hop +// is protected by a mutex because Forward() runs on the loader goroutine +// while Handle() reads it from per-connection goroutines. package local import ( "bufio" - "bytes" "context" - "errors" "net" + "sync" "time" - "github.com/go-gost/core/chain" "github.com/go-gost/core/handler" "github.com/go-gost/core/hop" - "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" - xctx "github.com/go-gost/x/ctx" - ictx "github.com/go-gost/x/internal/ctx" xnet "github.com/go-gost/x/internal/net" - "github.com/go-gost/x/internal/net/proxyproto" - "github.com/go-gost/x/internal/util/forwarder" "github.com/go-gost/x/internal/util/sniffing" tls_util "github.com/go-gost/x/internal/util/tls" rate_limiter "github.com/go-gost/x/limiter/rate" @@ -37,10 +90,12 @@ func init() { type forwardHandler struct { hop hop.Hop + hopMu sync.Mutex md metadata options handler.Options recorder recorder.RecorderObject certPool tls_util.CertPool + sniffer *SnifferBuilder } // NewHandler creates a local forwarding handler with the given options. @@ -72,12 +127,33 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) { h.certPool = tls_util.NewMemoryCertPool() } + h.sniffer = &SnifferBuilder{ + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + ALPN: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, + } + return } // Forward implements handler.Forwarder. func (h *forwardHandler) Forward(hop hop.Hop) { + h.hopMu.Lock() h.hop = hop + h.hopMu.Unlock() +} + +func (h *forwardHandler) getHop() hop.Hop { + h.hopMu.Lock() + defer h.hopMu.Unlock() + return h.hop } // Handle forwards the accepted connection to a node selected from the hop. @@ -126,9 +202,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand } if h.options.Router == nil { - err := errors.New("router not available") - log.Error(err) - return err + log.Error(errRouterNotAvailable) + return errRouterNotAvailable } var proto string @@ -157,176 +232,3 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand return h.handleRawForwarding(ctx, conn, ro, log, network, proto) } - -// handleRawForwarding performs node selection, dials the target through the -// router, and pipes the raw connection. It is used when sniffing is disabled -// or the protocol was not HTTP/TLS. -func (h *forwardHandler) handleRawForwarding(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, network, proto string) error { - target := &chain.Node{} - if h.hop != nil { - target = h.hop.Select(ctx, - hop.ProtocolSelectOption(proto), - ) - } - if target == nil { - err := errors.New("node not available") - log.Error(err) - return err - } - addr := target.Addr - if opts := target.Options(); opts != nil { - switch opts.Network { - case "unix": - network = opts.Network - default: - if _, _, err := net.SplitHostPort(addr); err != nil { - addr += ":0" - } - } - } - - ro.Network = network - ro.Host = addr - - log = log.WithFields(map[string]any{ - "node": target.Name, - "dst": addr, - "network": network, - }) - - log.Debugf("%s >> %s", conn.RemoteAddr(), addr) - - var buf bytes.Buffer - cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, addr) - ro.Route = buf.String() - if err != nil { - log.Error(err) - // TODO: the router itself may be failed due to the failed node in the router, - // the dead marker may be a wrong operation. - if marker := target.Marker(); marker != nil { - marker.Mark() - } - return err - } - if marker := target.Marker(); marker != nil { - marker.Reset() - } - defer cc.Close() - - cc = proxyproto.WrapClientConn( - h.md.proxyProtocol, - xctx.SrcAddrFromContext(ctx), - xctx.DstAddrFromContext(ctx), - cc) - - log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) - ro.SrcAddr = cc.LocalAddr().String() - ro.DstAddr = cc.RemoteAddr().String() - - t := time.Now() - log.Infof("%s <-> %s", conn.RemoteAddr(), target.Addr) - // xnet.Transport(conn, cc) - if err := xnet.Pipe(ctx, conn, cc, xnet.WithReadTimeout(h.md.readTimeout)); err != nil { - log.Debugf("pipe: %v", err) - } - log.WithFields(map[string]any{ - "duration": time.Since(t), - }).Infof("%s >-< %s", conn.RemoteAddr(), target.Addr) - - return nil -} - -// newRecorderObject creates a HandlerRecorderObject populated with connection -// metadata (service, addresses, network type, session ID, client address). -func (h *forwardHandler) newRecorderObject(ctx context.Context, conn net.Conn, start time.Time) *xrecorder.HandlerRecorderObject { - ro := &xrecorder.HandlerRecorderObject{ - Service: h.options.Service, - RemoteAddr: conn.RemoteAddr().String(), - LocalAddr: conn.LocalAddr().String(), - Network: "tcp", - Time: start, - SID: xctx.SidFromContext(ctx).String(), - } - if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil { - ro.ClientAddr = srcAddr.String() - } - if _, ok := conn.(net.PacketConn); ok { - ro.Network = "udp" - } - return ro -} - -// sniffingDial creates a dial function for the sniffing branch that wraps -// Router.Dial with route recording and proxy protocol encapsulation. -func (h *forwardHandler) sniffingDial(ctx context.Context, network, address string, ro *xrecorder.HandlerRecorderObject) (net.Conn, error) { - var buf bytes.Buffer - cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address) - ro.Route = buf.String() - return proxyproto.WrapClientConn( - h.md.proxyProtocol, - xctx.SrcAddrFromContext(ctx), - xctx.DstAddrFromContext(ctx), - cc), err -} - -// buildSniffer creates a forwarder.Sniffer configured from handler metadata. -func (h *forwardHandler) buildSniffer() *forwarder.Sniffer { - return &forwarder.Sniffer{ - Websocket: h.md.sniffingWebsocket, - WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, - } -} - -// handleSniffedProtocol dispatches a sniffed connection to the protocol-specific -// sniffer (HandleHTTP or HandleTLS). 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: - dial := func(ctx context.Context, network, address string) (net.Conn, error) { - return h.sniffingDial(ctx, network, address, ro) - } - sniffer := h.buildSniffer() - if proto == sniffing.ProtoHTTP { - return true, sniffer.HandleHTTP(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.hop), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithHTTPKeepalive(h.md.httpKeepalive), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - } - return true, sniffer.HandleTLS(ctx, conn, - forwarder.WithService(h.options.Service), - forwarder.WithDial(dial), - forwarder.WithHop(h.hop), - forwarder.WithBypass(h.options.Bypass), - forwarder.WithRecorderObject(ro), - forwarder.WithLog(log), - ) - default: - return false, nil - } -} - -func (h *forwardHandler) checkRateLimit(addr net.Addr) bool { - if h.options.RateLimiter == nil { - return true - } - host, _, _ := net.SplitHostPort(addr.String()) - if limiter := h.options.RateLimiter.Limiter(host); limiter != nil { - return limiter.Allow(1) - } - - return true -} diff --git a/handler/forward/local/handler_test.go b/handler/forward/local/handler_test.go index 6302f254..3609c576 100644 --- a/handler/forward/local/handler_test.go +++ b/handler/forward/local/handler_test.go @@ -1,637 +1,43 @@ package local import ( - "bytes" "context" "errors" "net" - "sync" "testing" "time" - "crypto" - "crypto/rand" - "crypto/rsa" - "crypto/x509" - "crypto/x509/pkix" - "encoding/pem" - "math/big" - "os" - "github.com/go-gost/core/chain" "github.com/go-gost/core/handler" "github.com/go-gost/core/hop" "github.com/go-gost/core/limiter/rate" - "github.com/go-gost/core/logger" "github.com/go-gost/core/recorder" - xctx "github.com/go-gost/x/ctx" - "github.com/go-gost/x/internal/util/sniffing" - xlogger "github.com/go-gost/x/logger" xmd "github.com/go-gost/x/metadata" xrecorder "github.com/go-gost/x/recorder" ) -// --------------------------------------------------------------------------- -// Mocks -// --------------------------------------------------------------------------- - -// mockRouter implements chain.Router. -type mockRouter struct { - opts *chain.RouterOptions - dialFn func(ctx context.Context, network, address string) (net.Conn, error) -} - -func (m *mockRouter) Options() *chain.RouterOptions { return m.opts } -func (m *mockRouter) Dial(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { - if m.dialFn != nil { - return m.dialFn(ctx, network, address) - } - return nil, nil -} -func (m *mockRouter) Bind(ctx context.Context, network, address string, opts ...chain.BindOption) (net.Listener, error) { - return nil, nil -} - -// mockHop implements hop.Hop. -type mockHop struct { - selectFn func(ctx context.Context, opts ...hop.SelectOption) *chain.Node -} - -func (m *mockHop) Select(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - if m.selectFn != nil { - return m.selectFn(ctx, opts...) - } - return nil -} - -type mockRateLimiter struct { - limiterFn func(key string) rate.Limiter -} - -func (m *mockRateLimiter) Limiter(key string) rate.Limiter { - if m.limiterFn != nil { - return m.limiterFn(key) - } - return nil -} - -type mockLimiter struct { - allowFn func(n int) bool -} - -func (m *mockLimiter) Allow(n int) bool { - if m.allowFn != nil { - return m.allowFn(n) - } - return true -} - -func (m *mockLimiter) Limit() float64 { return 0 } - -// mockRecorder implements recorder.Recorder. -type mockRecorder struct{} - -func (m *mockRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { - return nil -} -func (m *mockRecorder) Close() error { return nil } - -// stringConn is a net.Conn backed by bytes.Buffer for testing. -type stringConn struct { - readBuf *bytes.Buffer - writeBuf *bytes.Buffer - local net.Addr - remote net.Addr - closed bool - mu sync.Mutex -} - -func newStringConn(data []byte) *stringConn { - return &stringConn{ - readBuf: bytes.NewBuffer(data), - writeBuf: new(bytes.Buffer), - local: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080}, - remote: &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}, - } -} - -func (c *stringConn) Read(b []byte) (int, error) { return c.readBuf.Read(b) } -func (c *stringConn) Write(b []byte) (int, error) { return c.writeBuf.Write(b) } -func (c *stringConn) Close() error { - c.mu.Lock() - defer c.mu.Unlock() - c.closed = true - return nil -} -func (c *stringConn) LocalAddr() net.Addr { return c.local } -func (c *stringConn) RemoteAddr() net.Addr { return c.remote } -func (c *stringConn) SetDeadline(t time.Time) error { return nil } -func (c *stringConn) SetReadDeadline(t time.Time) error { return nil } -func (c *stringConn) SetWriteDeadline(t time.Time) error { return nil } - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -func nopLog() logger.Logger { return xlogger.Nop() } - -func newTestHandler(opts ...handler.Option) *forwardHandler { - options := handler.Options{ - Logger: nopLog(), - Router: &mockRouter{opts: &chain.RouterOptions{}}, - } - for _, opt := range opts { - opt(&options) - } - return &forwardHandler{options: options} -} - -func newInitdHandler(opts ...handler.Option) *forwardHandler { - h := newTestHandler(opts...) - _ = h.Init(xmd.NewMetadata(nil)) - return h -} - -// withRateLimiter is a helper to set the RateLimiter option. -func withRateLimiter(rl rate.RateLimiter) handler.Option { - return func(o *handler.Options) { - o.RateLimiter = rl - } -} - -// withRouter is a helper to set the Router option. -func withRouter(r chain.Router) handler.Option { - return func(o *handler.Options) { - o.Router = r - } -} - -// withRecorder is a helper to append a RecorderObject to the Recorders slice. -func withRecorder(ro recorder.RecorderObject) handler.Option { - return func(o *handler.Options) { - o.Recorders = append(o.Recorders, ro) - } -} - -// --------------------------------------------------------------------------- -// newRecorderObject -// --------------------------------------------------------------------------- - -func TestNewRecorderObject_TCP(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - start := time.Now() - - ro := h.newRecorderObject(context.Background(), conn, start) - - if ro.Network != "tcp" { - t.Errorf("expected network tcp, got %s", ro.Network) - } - if ro.RemoteAddr != "10.0.0.1:12345" { - t.Errorf("expected remote 10.0.0.1:12345, got %s", ro.RemoteAddr) - } - if ro.LocalAddr != "127.0.0.1:8080" { - t.Errorf("expected local 127.0.0.1:8080, got %s", ro.LocalAddr) - } - if !ro.Time.Equal(start) { - t.Errorf("expected time %v, got %v", start, ro.Time) - } -} - -func TestNewRecorderObject_SID(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - ctx := xctx.ContextWithSid(context.Background(), xctx.Sid("test-sid")) - - ro := h.newRecorderObject(ctx, conn, time.Now()) - - if ro.SID != "test-sid" { - t.Errorf("expected SID 'test-sid', got '%s'", ro.SID) - } -} - -func TestNewRecorderObject_NilSrcAddr(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - - ro := h.newRecorderObject(context.Background(), conn, time.Now()) - - if ro.ClientAddr != "" { - t.Errorf("expected empty ClientAddr, got %s", ro.ClientAddr) - } -} - -// --------------------------------------------------------------------------- -// checkRateLimit -// --------------------------------------------------------------------------- - -func TestCheckRateLimit_NilLimiter(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - - if !h.checkRateLimit(conn.RemoteAddr()) { - t.Error("expected true when RateLimiter is nil") - } -} - -func TestCheckRateLimit_Allowed(t *testing.T) { - h := newInitdHandler(withRateLimiter(&mockRateLimiter{ - limiterFn: func(key string) rate.Limiter { - return &mockLimiter{allowFn: func(n int) bool { return true }} - }, - })) - conn := newStringConn(nil) - - if !h.checkRateLimit(conn.RemoteAddr()) { - t.Error("expected true when limiter allows") - } -} - -func TestCheckRateLimit_Blocked(t *testing.T) { - h := newInitdHandler(withRateLimiter(&mockRateLimiter{ - limiterFn: func(key string) rate.Limiter { - return &mockLimiter{allowFn: func(n int) bool { return false }} - }, - })) - conn := newStringConn(nil) - - if h.checkRateLimit(conn.RemoteAddr()) { - t.Error("expected false when limiter blocks") - } -} - -func TestCheckRateLimit_NoLimiterForKey(t *testing.T) { - h := newInitdHandler(withRateLimiter(&mockRateLimiter{ - limiterFn: func(key string) rate.Limiter { return nil }, - })) - conn := newStringConn(nil) - - if !h.checkRateLimit(conn.RemoteAddr()) { - t.Error("expected true when no limiter found for key") - } -} - -// --------------------------------------------------------------------------- -// sniffingDial -// --------------------------------------------------------------------------- - -func TestSniffingDial_Success(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - ro := &xrecorder.HandlerRecorderObject{} - - cc, err := h.sniffingDial(context.Background(), "tcp", "example.com:80", ro) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - defer cc.Close() - - if ro.Route == "" { - t.Log("route is empty (mock does not populate context buffer)") - } -} - -func TestSniffingDial_RouterError(t *testing.T) { - expectedErr := errors.New("dial failed") - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return nil, expectedErr - }, - })) - ro := &xrecorder.HandlerRecorderObject{} - - _, err := h.sniffingDial(context.Background(), "tcp", "example.com:80", ro) - if err == nil { - t.Fatal("expected error") - } - if !errors.Is(err, expectedErr) { - t.Errorf("expected %v, got %v", expectedErr, err) - } -} - -// --------------------------------------------------------------------------- -// buildSniffer -// --------------------------------------------------------------------------- - -func TestBuildSniffer(t *testing.T) { - h := newTestHandler() - h.md.sniffingWebsocket = true - h.md.sniffingWebsocketSampleRate = 0.5 - h.md.readTimeout = 30 * time.Second - - sniffer := h.buildSniffer() - - if !sniffer.Websocket { - t.Error("expected Websocket to be true") - } - if sniffer.WebsocketSampleRate != 0.5 { - t.Errorf("expected 0.5, got %f", sniffer.WebsocketSampleRate) - } - if sniffer.ReadTimeout != 30*time.Second { - t.Errorf("expected 30s, got %v", sniffer.ReadTimeout) - } -} - -// --------------------------------------------------------------------------- -// handleSniffedProtocol -// --------------------------------------------------------------------------- - -func TestHandleSniffedProtocol_UnknownProto(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), "ssh") - if handled { - t.Error("expected handled=false for unknown proto") - } - if err != nil { - t.Errorf("expected nil error, got %v", err) - } -} - -func TestHandleSniffedProtocol_EmptyProto(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), "") - if handled { - t.Error("expected handled=false for empty proto") - } - if err != nil { - t.Errorf("expected nil error, got %v", err) - } -} - -// --------------------------------------------------------------------------- -// handleRawForwarding -// --------------------------------------------------------------------------- - -func TestHandleRawForwarding_NilHop(t *testing.T) { - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return nil, errors.New("no route to host") - }, - })) - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err == nil { - t.Fatal("expected error when hop is nil") - } -} - -func TestHandleRawForwarding_NodeUnavailable(t *testing.T) { - h := newInitdHandler() - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return nil - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err == nil { - t.Fatal("expected error when node is nil") - } - if err.Error() != "node not available" { - t.Errorf("expected 'node not available', got %v", err) - } -} - -func TestHandleRawForwarding_DialError(t *testing.T) { - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return nil, errors.New("dial failed") - }, - })) - // Use NewNode to include a marker (so Mark() is exercised on error). - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return chain.NewNode("test-node", "127.0.0.1:9999") - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err == nil { - t.Fatal("expected error when dial fails") - } - if ro.Host != "127.0.0.1:9999" { - t.Errorf("expected host 127.0.0.1:9999, got %s", ro.Host) - } -} - -func TestHandleRawForwarding_Success(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - // xnet.Pipe will half-close the pipe. We just close the other end to - // prevent goroutine leaks; the actual data flow through Pipe is tested - // at integration level. - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return &chain.Node{Addr: "10.0.0.1:80", Name: "upstream"} - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if ro.Host != "10.0.0.1:80" { - t.Errorf("expected host 10.0.0.1:80, got %s", ro.Host) - } - if ro.Network != "tcp" { - t.Errorf("expected network tcp, got %s", ro.Network) - } - if ro.SrcAddr == "" { - t.Error("expected SrcAddr to be set") - } - if ro.DstAddr == "" { - t.Error("expected DstAddr to be set") - } -} - -func TestHandleRawForwarding_AddrWithoutPort(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return &chain.Node{Addr: "10.0.0.1", Name: "noport"} - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if ro.Host != "10.0.0.1:0" { - t.Errorf("expected ':0' appended to addr, got %s", ro.Host) - } -} - -// --------------------------------------------------------------------------- -// Handle integration -// --------------------------------------------------------------------------- - -func TestHandle_NoRouter(t *testing.T) { - h := newTestHandler(func(o *handler.Options) { - o.Router = nil +func TestNewHandler(t *testing.T) { + h := NewHandler(func(o *handler.Options) { + o.Service = "test-svc" }) - conn := newStringConn(nil) - - err := h.Handle(context.Background(), conn) - if err == nil { - t.Fatal("expected error when router is nil") + if h == nil { + t.Fatal("expected non-nil handler") } - if err.Error() != "router not available" { - t.Errorf("expected 'router not available', got %v", err) + fh, ok := h.(*forwardHandler) + if !ok { + t.Fatal("expected *forwardHandler") + } + if fh.options.Service != "test-svc" { + t.Errorf("expected Service 'test-svc', got %s", fh.options.Service) } } -func TestHandle_RateLimited(t *testing.T) { - h := newInitdHandler(withRateLimiter(&mockRateLimiter{ - limiterFn: func(key string) rate.Limiter { - return &mockLimiter{allowFn: func(n int) bool { return false }} - }, - })) - conn := newStringConn(nil) - - err := h.Handle(context.Background(), conn) - if err == nil { - t.Fatal("expected rate limit error") - } -} - -func TestHandle_BasicForward(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - // Drain the other end of the pipe so writes from xnet.Pipe don't block. - go func() { - buf := make([]byte, 64) - for { - _, err := client.Read(buf) - if err != nil { - return - } - } - }() - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return &chain.Node{Addr: "10.0.0.1:80", Name: "target"} - }, - } - conn := newStringConn([]byte("hello")) - - err := h.Handle(context.Background(), conn) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } -} - -// --------------------------------------------------------------------------- -// parseMetadata -// --------------------------------------------------------------------------- - -func TestParseMetadata_Defaults(t *testing.T) { +func TestForward(t *testing.T) { h := newTestHandler() - - err := h.parseMetadata(xmd.NewMetadata(nil)) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if h.md.readTimeout <= 0 { - t.Error("expected positive default readTimeout") - } - if h.md.sniffing { - t.Error("expected sniffing disabled by default") - } -} - -func TestParseMetadata_CustomValues(t *testing.T) { - h := newTestHandler() - md := xmd.NewMetadata(map[string]any{ - "readTimeout": "60s", - "http.keepalive": true, - "proxyProtocol": 2, - "sniffing": true, - "sniffing.timeout": "5s", - "sniffing.websocket": true, - "sniffing.websocket.sampleRate": 0.75, - }) - - err := h.parseMetadata(md) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if h.md.readTimeout != 60*time.Second { - t.Errorf("expected 60s, got %v", h.md.readTimeout) - } - if !h.md.httpKeepalive { - t.Error("expected httpKeepalive true") - } - if h.md.proxyProtocol != 2 { - t.Errorf("expected proxyProtocol 2, got %d", h.md.proxyProtocol) - } - if !h.md.sniffing { - t.Error("expected sniffing true") - } - if h.md.sniffingTimeout != 5*time.Second { - t.Errorf("expected 5s sniffingTimeout, got %v", h.md.sniffingTimeout) - } - if !h.md.sniffingWebsocket { - t.Error("expected sniffingWebsocket true") - } - if h.md.sniffingWebsocketSampleRate != 0.75 { - t.Errorf("expected 0.75, got %f", h.md.sniffingWebsocketSampleRate) + mh := &mockHop{} + h.Forward(mh) + if h.getHop() != mh { + t.Error("expected hop to be set") } } @@ -667,104 +73,107 @@ func TestInit_CertPool(t *testing.T) { } } -// --------------------------------------------------------------------------- -// NewHandler / Forward -// --------------------------------------------------------------------------- - -func TestNewHandler(t *testing.T) { - h := NewHandler(func(o *handler.Options) { - o.Service = "test-svc" - }) - if h == nil { - t.Fatal("expected non-nil handler") - } - fh, ok := h.(*forwardHandler) - if !ok { - t.Fatal("expected *forwardHandler") - } - if fh.options.Service != "test-svc" { - t.Errorf("expected Service 'test-svc', got %s", fh.options.Service) - } -} - -func TestForward(t *testing.T) { +func TestInit_ParseMetadataError(t *testing.T) { h := newTestHandler() - mh := &mockHop{} - h.Forward(mh) - if h.hop != mh { - t.Error("expected hop to be set") - } -} - -// --------------------------------------------------------------------------- -// newRecorderObject — UDP -// --------------------------------------------------------------------------- - -type packetConn struct { - *stringConn -} - -func (c *packetConn) ReadFrom(b []byte) (int, net.Addr, error) { return 0, nil, nil } -func (c *packetConn) WriteTo(b []byte, addr net.Addr) (int, error) { return 0, nil } - -func TestNewRecorderObject_UDP(t *testing.T) { - h := newInitdHandler() - conn := &packetConn{newStringConn(nil)} - - ro := h.newRecorderObject(context.Background(), conn, time.Now()) - - if ro.Network != "udp" { - t.Errorf("expected network udp, got %s", ro.Network) - } -} - -// --------------------------------------------------------------------------- -// handleSniffedProtocol — HTTP / TLS dispatch -// --------------------------------------------------------------------------- - -func TestHandleSniffedProtocol_HTTP(t *testing.T) { - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return nil, errors.New("dial not configured in test") - }, - })) - // Minimal HTTP request — enough for http.ReadRequest to parse. - conn := newStringConn([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")) - ro := &xrecorder.HandlerRecorderObject{} - - handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), sniffing.ProtoHTTP) - - if !handled { - t.Error("expected handled=true for HTTP proto") - } - // The sniffer will try to dial and fail — but dispatch is verified. - if err == nil { - t.Log("sniffer processed HTTP request without error") - } -} - -func TestHandleSniffedProtocol_TLS(t *testing.T) { - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return nil, errors.New("dial not configured in test") - }, - })) - // Minimal TLS ClientHello — enough for tls-dissector to parse. - // Record header: 16 03 01 (handshake), then length + ClientHello. - conn := newStringConn([]byte{ - 0x16, 0x03, 0x01, 0x00, 0x06, 0x01, 0x00, 0x00, 0x02, 0x03, 0x01, + md := xmd.NewMetadata(map[string]any{ + "mitm.certFile": "/nonexistent/cert.pem", + "mitm.keyFile": "/nonexistent/key.pem", }) - ro := &xrecorder.HandlerRecorderObject{} - - handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), sniffing.ProtoTLS) - - if !handled { - t.Error("expected handled=true for TLS proto") - } + err := h.Init(md) if err == nil { - t.Log("sniffer processed TLS ClientHello without error") + t.Fatal("expected error from Init with bad cert files") + } +} + +// --------------------------------------------------------------------------- +// Handle integration +// --------------------------------------------------------------------------- + +func TestHandle_NoRouter(t *testing.T) { + h := newTestHandler(func(o *handler.Options) { + o.Router = nil + }) + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err == nil { + t.Fatal("expected error when router is nil") + } + if !errors.Is(err, errRouterNotAvailable) { + t.Errorf("expected errRouterNotAvailable, got %v", err) + } +} + +func TestHandle_RateLimited(t *testing.T) { + h := newInitdHandler(withRateLimiter(&stubRateLimiter{ + limiterFn: func(key string) rate.Limiter { + return &stubLimiter{allowFn: func(n int) bool { return false }} + }, + })) + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err == nil { + t.Fatal("expected rate limit error") + } +} + +func TestHandle_BasicForward(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + go func() { + buf := make([]byte, 64) + for { + _, err := client.Read(buf) + if err != nil { + return + } + } + }() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return &chain.Node{Addr: "10.0.0.1:80", Name: "target"} + }, + }) + conn := newStringConn([]byte("hello")) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestHandle_RecorderError(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newTestHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + h.recorder = recorder.RecorderObject{Recorder: &errorRecorder{}} + h.Forward(&mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("target", "10.0.0.1:80") + }, + }) + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("unexpected error: %v", err) } } @@ -776,7 +185,6 @@ func TestHandle_SniffingEnabled(t *testing.T) { server, client := net.Pipe() defer server.Close() defer client.Close() - // Drain to prevent pipe deadlock. go func() { buf := make([]byte, 64) for { @@ -796,12 +204,11 @@ func TestHandle_SniffingEnabled(t *testing.T) { h.md.sniffing = true h.md.sniffingTimeout = 5 * time.Second h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}} - h.hop = &mockHop{ + h.Forward(&mockHop{ selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { return &chain.Node{Addr: "10.0.0.1:80", Name: "target"} }, - } - // Data that sniffing.Sniff won't recognize as HTTP/TLS (fall through to raw forward). + }) conn := newStringConn([]byte("hello world")) err := h.Handle(context.Background(), conn) if err != nil { @@ -809,183 +216,6 @@ func TestHandle_SniffingEnabled(t *testing.T) { } } -// --------------------------------------------------------------------------- -// parseMetadata — cert/key -// --------------------------------------------------------------------------- - -func TestParseMetadata_MITMCertError(t *testing.T) { - h := newTestHandler() - md := xmd.NewMetadata(map[string]any{ - "mitm.certFile": "/nonexistent/cert.pem", - "mitm.keyFile": "/nonexistent/key.pem", - }) - err := h.parseMetadata(md) - if err == nil { - t.Fatal("expected error for non-existent cert file") - } -} - -func TestParseMetadata_MITMCert(t *testing.T) { - certPEM, keyPEM := generateTestCertPEM(t) - certFile := writeTempFile(t, "cert-*.pem", certPEM) - defer os.Remove(certFile) - keyFile := writeTempFile(t, "key-*.pem", keyPEM) - defer os.Remove(keyFile) - - h := newTestHandler() - md := xmd.NewMetadata(map[string]any{ - "mitm.certFile": certFile, - "mitm.keyFile": keyFile, - "mitm.alpn": "h2", - }) - - err := h.parseMetadata(md) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if h.md.certificate == nil { - t.Error("expected certificate to be loaded") - } - if h.md.privateKey == nil { - t.Error("expected privateKey to be loaded") - } - if h.md.alpn != "h2" { - t.Errorf("expected alpn 'h2', got %s", h.md.alpn) - } -} - -func TestNewRecorderObject_SrcAddr(t *testing.T) { - h := newInitdHandler() - conn := newStringConn(nil) - srcAddr := &net.TCPAddr{IP: net.IPv4(192, 168, 1, 1), Port: 54321} - ctx := xctx.ContextWithSrcAddr(context.Background(), srcAddr) - - ro := h.newRecorderObject(ctx, conn, time.Now()) - - if ro.ClientAddr != "192.168.1.1:54321" { - t.Errorf("expected ClientAddr '192.168.1.1:54321', got '%s'", ro.ClientAddr) - } -} - -// --------------------------------------------------------------------------- -// handleRawForwarding — unix socket -// --------------------------------------------------------------------------- - -func TestHandleRawForwarding_UnixSocket(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return chain.NewNode("unix-node", "/tmp/unix.sock", - chain.NetworkNodeOption("unix"), - ) - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if ro.Network != "unix" { - t.Errorf("expected network 'unix', got '%s'", ro.Network) - } -} - -func TestHandleRawForwarding_MarkerReset(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - h := newInitdHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - // NewNode sets a FailMarker, which gets Reset() on successful dial. - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return chain.NewNode("target", "10.0.0.1:80") - }, - } - conn := newStringConn(nil) - ro := &xrecorder.HandlerRecorderObject{} - - err := h.handleRawForwarding(context.Background(), conn, ro, nopLog(), "tcp", "") - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if ro.Host != "10.0.0.1:80" { - t.Errorf("expected host '10.0.0.1:80', got '%s'", ro.Host) - } -} - -// --------------------------------------------------------------------------- -// Init error path -// --------------------------------------------------------------------------- - -func TestInit_ParseMetadataError(t *testing.T) { - h := newTestHandler() - md := xmd.NewMetadata(map[string]any{ - "mitm.certFile": "/nonexistent/cert.pem", - "mitm.keyFile": "/nonexistent/key.pem", - }) - err := h.Init(md) - if err == nil { - t.Fatal("expected error from Init with bad cert files") - } -} - -// --------------------------------------------------------------------------- -// Handle — recorder error -// --------------------------------------------------------------------------- - -type errorRecorder struct{} - -func (e *errorRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { - return errors.New("record error") -} -func (e *errorRecorder) Close() error { return nil } - -func TestHandle_RecorderError(t *testing.T) { - server, client := net.Pipe() - defer server.Close() - defer client.Close() - - h := newTestHandler(withRouter(&mockRouter{ - opts: &chain.RouterOptions{}, - dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { - return server, nil - }, - })) - h.recorder = recorder.RecorderObject{Recorder: &errorRecorder{}} - h.hop = &mockHop{ - selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { - return chain.NewNode("target", "10.0.0.1:80") - }, - } - conn := newStringConn(nil) - - err := h.Handle(context.Background(), conn) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } -} - -// --------------------------------------------------------------------------- -// Handle — sniff error (empty conn) -// --------------------------------------------------------------------------- - func TestHandle_SniffError(t *testing.T) { server, client := net.Pipe() defer server.Close() @@ -1000,12 +230,11 @@ func TestHandle_SniffError(t *testing.T) { h.md.sniffing = true h.md.sniffingTimeout = 5 * time.Second h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}} - h.hop = &mockHop{ + h.Forward(&mockHop{ selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { return chain.NewNode("target", "10.0.0.1:80") }, - } - // Empty conn → Peek fails → Sniff returns error. + }) conn := newStringConn(nil) err := h.Handle(context.Background(), conn) @@ -1014,16 +243,11 @@ func TestHandle_SniffError(t *testing.T) { } } -// --------------------------------------------------------------------------- -// Handle — sniff HTTP → handled=true path -// --------------------------------------------------------------------------- - func TestHandle_SniffHTTP(t *testing.T) { server, client := net.Pipe() defer server.Close() defer client.Close() - // Drain and echo back a minimal response so the sniffer's roundTrip completes. go func() { buf := make([]byte, 4096) n, err := client.Read(buf) @@ -1041,11 +265,11 @@ func TestHandle_SniffHTTP(t *testing.T) { h.md.sniffing = true h.md.sniffingTimeout = 5 * time.Second h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}} - h.hop = &mockHop{ + h.Forward(&mockHop{ selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { return chain.NewNode("target", "10.0.0.1:80") }, - } + }) conn := newStringConn([]byte("GET / HTTP/1.1\r\nHost: x\r\n\r\n")) err := h.Handle(context.Background(), conn) @@ -1053,59 +277,3 @@ func TestHandle_SniffHTTP(t *testing.T) { t.Logf("sniffed HTTP path returned: %v", err) } } - -// --------------------------------------------------------------------------- -// Test certificate helpers -// --------------------------------------------------------------------------- - -func generateTestCertKey(t *testing.T) (*x509.Certificate, crypto.PrivateKey) { - t.Helper() - key, err := rsa.GenerateKey(rand.Reader, 2048) - if err != nil { - t.Fatalf("generate key: %v", err) - } - tmpl := &x509.Certificate{ - SerialNumber: big.NewInt(1), - Subject: pkix.Name{CommonName: "test"}, - NotBefore: time.Now().Add(-time.Hour), - NotAfter: time.Now().Add(time.Hour), - } - certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) - if err != nil { - t.Fatalf("create certificate: %v", err) - } - cert, err := x509.ParseCertificate(certDER) - if err != nil { - t.Fatalf("parse certificate: %v", err) - } - return cert, key -} - -func generateTestCertPEM(t *testing.T) (certPEM, keyPEM []byte) { - t.Helper() - cert, key := generateTestCertKey(t) - certPEM = pemEncode("CERTIFICATE", cert.Raw) - keyBytes := x509.MarshalPKCS1PrivateKey(key.(*rsa.PrivateKey)) - keyPEM = pemEncode("RSA PRIVATE KEY", keyBytes) - return -} - -func pemEncode(blockType string, der []byte) []byte { - block := &pem.Block{Type: blockType, Bytes: der} - return pem.EncodeToMemory(block) -} - -func writeTempFile(t *testing.T, pattern string, data []byte) string { - t.Helper() - f, err := os.CreateTemp("", pattern) - if err != nil { - t.Fatalf("create temp file: %v", err) - } - if _, err := f.Write(data); err != nil { - f.Close() - os.Remove(f.Name()) - t.Fatalf("write temp file: %v", err) - } - f.Close() - return f.Name() -} diff --git a/handler/forward/local/helpers_test.go b/handler/forward/local/helpers_test.go new file mode 100644 index 00000000..e06ad2b5 --- /dev/null +++ b/handler/forward/local/helpers_test.go @@ -0,0 +1,243 @@ +package local + +import ( + "bytes" + "context" + "crypto" + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "errors" + "math/big" + "net" + "os" + "sync" + "testing" + "time" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/handler" + "github.com/go-gost/core/hop" + "github.com/go-gost/core/limiter/rate" + "github.com/go-gost/core/logger" + "github.com/go-gost/core/recorder" + xlogger "github.com/go-gost/x/logger" + xmd "github.com/go-gost/x/metadata" +) + +// --------------------------------------------------------------------------- +// Mocks +// --------------------------------------------------------------------------- + +// mockRouter implements chain.Router. +type mockRouter struct { + opts *chain.RouterOptions + dialFn func(ctx context.Context, network, address string) (net.Conn, error) +} + +func (m *mockRouter) Options() *chain.RouterOptions { return m.opts } +func (m *mockRouter) Dial(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + if m.dialFn != nil { + return m.dialFn(ctx, network, address) + } + return nil, nil +} +func (m *mockRouter) Bind(ctx context.Context, network, address string, opts ...chain.BindOption) (net.Listener, error) { + return nil, nil +} + +// mockHop implements hop.Hop. +type mockHop struct { + selectFn func(ctx context.Context, opts ...hop.SelectOption) *chain.Node +} + +func (m *mockHop) Select(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + if m.selectFn != nil { + return m.selectFn(ctx, opts...) + } + return nil +} + +type stubRateLimiter struct { + limiterFn func(key string) rate.Limiter +} + +func (m *stubRateLimiter) Limiter(key string) rate.Limiter { + if m.limiterFn != nil { + return m.limiterFn(key) + } + return nil +} + +type stubLimiter struct { + allowFn func(n int) bool +} + +func (m *stubLimiter) Allow(n int) bool { + if m.allowFn != nil { + return m.allowFn(n) + } + return true +} + +func (m *stubLimiter) Limit() float64 { return 0 } + +// mockRecorder implements recorder.Recorder. +type mockRecorder struct{} + +func (m *mockRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + return nil +} +func (m *mockRecorder) Close() error { return nil } + +type errorRecorder struct{} + +func (e *errorRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + return errors.New("record error") +} +func (e *errorRecorder) Close() error { return nil } + +// stringConn is a net.Conn backed by bytes.Buffer for testing. +type stringConn struct { + readBuf *bytes.Buffer + writeBuf *bytes.Buffer + local net.Addr + remote net.Addr + closed bool + mu sync.Mutex +} + +func newStringConn(data []byte) *stringConn { + return &stringConn{ + readBuf: bytes.NewBuffer(data), + writeBuf: new(bytes.Buffer), + local: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080}, + remote: &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}, + } +} + +func (c *stringConn) Read(b []byte) (int, error) { return c.readBuf.Read(b) } +func (c *stringConn) Write(b []byte) (int, error) { return c.writeBuf.Write(b) } +func (c *stringConn) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + return nil +} +func (c *stringConn) LocalAddr() net.Addr { return c.local } +func (c *stringConn) RemoteAddr() net.Addr { return c.remote } +func (c *stringConn) SetDeadline(t time.Time) error { return nil } +func (c *stringConn) SetReadDeadline(t time.Time) error { return nil } +func (c *stringConn) SetWriteDeadline(t time.Time) error { return nil } + +type packetConn struct { + *stringConn +} + +func (c *packetConn) ReadFrom(b []byte) (int, net.Addr, error) { return 0, nil, nil } +func (c *packetConn) WriteTo(b []byte, addr net.Addr) (int, error) { return 0, nil } + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +func nopLog() logger.Logger { return xlogger.Nop() } + +func newTestHandler(opts ...handler.Option) *forwardHandler { + options := handler.Options{ + Logger: nopLog(), + Router: &mockRouter{opts: &chain.RouterOptions{}}, + } + for _, opt := range opts { + opt(&options) + } + return &forwardHandler{ + options: options, + sniffer: &SnifferBuilder{}, + } +} + +func newInitdHandler(opts ...handler.Option) *forwardHandler { + h := newTestHandler(opts...) + _ = h.Init(xmd.NewMetadata(nil)) + return h +} + +// withRateLimiter is a helper to set the RateLimiter option. +func withRateLimiter(rl rate.RateLimiter) handler.Option { + return func(o *handler.Options) { + o.RateLimiter = rl + } +} + +// withRouter is a helper to set the Router option. +func withRouter(r chain.Router) handler.Option { + return func(o *handler.Options) { + o.Router = r + } +} + +// withRecorder is a helper to append a RecorderObject to the Recorders slice. +func withRecorder(ro recorder.RecorderObject) handler.Option { + return func(o *handler.Options) { + o.Recorders = append(o.Recorders, ro) + } +} + +// --------------------------------------------------------------------------- +// Test certificate helpers +// --------------------------------------------------------------------------- + +func generateTestCertKey(t *testing.T) (*x509.Certificate, crypto.PrivateKey) { + t.Helper() + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("generate key: %v", err) + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "test"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + } + certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + t.Fatalf("create certificate: %v", err) + } + cert, err := x509.ParseCertificate(certDER) + if err != nil { + t.Fatalf("parse certificate: %v", err) + } + return cert, key +} + +func generateTestCertPEM(t *testing.T) (certPEM, keyPEM []byte) { + t.Helper() + cert, key := generateTestCertKey(t) + certPEM = pemEncode("CERTIFICATE", cert.Raw) + keyBytes := x509.MarshalPKCS1PrivateKey(key.(*rsa.PrivateKey)) + keyPEM = pemEncode("RSA PRIVATE KEY", keyBytes) + return +} + +func pemEncode(blockType string, der []byte) []byte { + block := &pem.Block{Type: blockType, Bytes: der} + return pem.EncodeToMemory(block) +} + +func writeTempFile(t *testing.T, pattern string, data []byte) string { + t.Helper() + f, err := os.CreateTemp("", pattern) + if err != nil { + t.Fatalf("create temp file: %v", err) + } + if _, err := f.Write(data); err != nil { + f.Close() + os.Remove(f.Name()) + t.Fatalf("write temp file: %v", err) + } + f.Close() + return f.Name() +} diff --git a/handler/forward/local/metadata.go b/handler/forward/local/metadata.go index d9f74280..68d2b834 100644 --- a/handler/forward/local/metadata.go +++ b/handler/forward/local/metadata.go @@ -49,10 +49,7 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) { if err != nil { return err } - h.md.certificate, err = x509.ParseCertificate(tlsCert.Certificate[0]) - if err != nil { - return err - } + h.md.certificate = tlsCert.Leaf h.md.privateKey = tlsCert.PrivateKey } h.md.alpn = mdutil.GetString(md, "mitm.alpn") diff --git a/handler/forward/local/metadata_test.go b/handler/forward/local/metadata_test.go new file mode 100644 index 00000000..b793ef60 --- /dev/null +++ b/handler/forward/local/metadata_test.go @@ -0,0 +1,114 @@ +package local + +import ( + "os" + "testing" + "time" + + xmd "github.com/go-gost/x/metadata" +) + +// --------------------------------------------------------------------------- +// parseMetadata +// --------------------------------------------------------------------------- + +func TestParseMetadata_Defaults(t *testing.T) { + h := newTestHandler() + + err := h.parseMetadata(xmd.NewMetadata(nil)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if h.md.readTimeout <= 0 { + t.Error("expected positive default readTimeout") + } + if h.md.sniffing { + t.Error("expected sniffing disabled by default") + } +} + +func TestParseMetadata_CustomValues(t *testing.T) { + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "readTimeout": "60s", + "http.keepalive": true, + "proxyProtocol": 2, + "sniffing": true, + "sniffing.timeout": "5s", + "sniffing.websocket": true, + "sniffing.websocket.sampleRate": 0.75, + }) + + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if h.md.readTimeout != 60*time.Second { + t.Errorf("expected 60s, got %v", h.md.readTimeout) + } + if !h.md.httpKeepalive { + t.Error("expected httpKeepalive true") + } + if h.md.proxyProtocol != 2 { + t.Errorf("expected proxyProtocol 2, got %d", h.md.proxyProtocol) + } + if !h.md.sniffing { + t.Error("expected sniffing true") + } + if h.md.sniffingTimeout != 5*time.Second { + t.Errorf("expected 5s sniffingTimeout, got %v", h.md.sniffingTimeout) + } + if !h.md.sniffingWebsocket { + t.Error("expected sniffingWebsocket true") + } + if h.md.sniffingWebsocketSampleRate != 0.75 { + t.Errorf("expected 0.75, got %f", h.md.sniffingWebsocketSampleRate) + } +} + +// --------------------------------------------------------------------------- +// parseMetadata — cert/key +// --------------------------------------------------------------------------- + +func TestParseMetadata_MITMCertError(t *testing.T) { + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "mitm.certFile": "/nonexistent/cert.pem", + "mitm.keyFile": "/nonexistent/key.pem", + }) + err := h.parseMetadata(md) + if err == nil { + t.Fatal("expected error for non-existent cert file") + } +} + +func TestParseMetadata_MITMCert(t *testing.T) { + certPEM, keyPEM := generateTestCertPEM(t) + certFile := writeTempFile(t, "cert-*.pem", certPEM) + defer os.Remove(certFile) + keyFile := writeTempFile(t, "key-*.pem", keyPEM) + defer os.Remove(keyFile) + + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "mitm.certFile": certFile, + "mitm.keyFile": keyFile, + "mitm.alpn": "h2", + }) + + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.md.certificate == nil { + t.Error("expected certificate to be loaded") + } + if h.md.privateKey == nil { + t.Error("expected privateKey to be loaded") + } + if h.md.alpn != "h2" { + t.Errorf("expected alpn 'h2', got %s", h.md.alpn) + } +} diff --git a/handler/forward/local/sniffing.go b/handler/forward/local/sniffing.go new file mode 100644 index 00000000..4484bcf3 --- /dev/null +++ b/handler/forward/local/sniffing.go @@ -0,0 +1,99 @@ +package local + +import ( + "bytes" + "context" + "crypto" + "crypto/x509" + "net" + "time" + + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/logger" + "github.com/go-gost/core/recorder" + xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" + "github.com/go-gost/x/internal/net/proxyproto" + "github.com/go-gost/x/internal/util/forwarder" + "github.com/go-gost/x/internal/util/sniffing" + tls_util "github.com/go-gost/x/internal/util/tls" + xrecorder "github.com/go-gost/x/recorder" +) + +// SnifferBuilder holds configuration for creating per-connection protocol sniffers. +// It is populated once during Init and reused for each connection. +type SnifferBuilder struct { + Websocket bool + WebsocketSampleRate float64 + Recorder recorder.Recorder + RecorderOptions *recorder.Options + Certificate *x509.Certificate + PrivateKey crypto.PrivateKey + ALPN string + CertPool tls_util.CertPool + MitmBypass bypass.Bypass + ReadTimeout time.Duration +} + +// Build creates a new forwarder.Sniffer from the builder's configuration. +func (b *SnifferBuilder) Build() *forwarder.Sniffer { + return &forwarder.Sniffer{ + Websocket: b.Websocket, + WebsocketSampleRate: b.WebsocketSampleRate, + Recorder: b.Recorder, + RecorderOptions: b.RecorderOptions, + Certificate: b.Certificate, + PrivateKey: b.PrivateKey, + NegotiatedProtocol: b.ALPN, + CertPool: b.CertPool, + MitmBypass: b.MitmBypass, + ReadTimeout: b.ReadTimeout, + } +} + +// sniffingDial creates a dial function for the sniffing branch that wraps +// Router.Dial with route recording and proxy protocol encapsulation. +func (h *forwardHandler) sniffingDial(ctx context.Context, network, address string, ro *xrecorder.HandlerRecorderObject) (net.Conn, error) { + var buf bytes.Buffer + cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address) + ro.Route = buf.String() + return proxyproto.WrapClientConn( + h.md.proxyProtocol, + xctx.SrcAddrFromContext(ctx), + xctx.DstAddrFromContext(ctx), + cc), err +} + +// handleSniffedProtocol dispatches a sniffed connection to the protocol-specific +// sniffer (HandleHTTP or HandleTLS). 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: + 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), + ) + } + 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), + ) + default: + return false, nil + } +} diff --git a/handler/forward/local/sniffing_test.go b/handler/forward/local/sniffing_test.go new file mode 100644 index 00000000..401163a7 --- /dev/null +++ b/handler/forward/local/sniffing_test.go @@ -0,0 +1,157 @@ +package local + +import ( + "context" + "errors" + "net" + "testing" + "time" + + "github.com/go-gost/core/chain" + "github.com/go-gost/x/internal/util/sniffing" + xrecorder "github.com/go-gost/x/recorder" +) + +// --------------------------------------------------------------------------- +// sniffingDial +// --------------------------------------------------------------------------- + +func TestSniffingDial_Success(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + ro := &xrecorder.HandlerRecorderObject{} + + cc, err := h.sniffingDial(context.Background(), "tcp", "example.com:80", ro) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer cc.Close() + + if ro.Route == "" { + t.Log("route is empty (mock does not populate context buffer)") + } +} + +func TestSniffingDial_RouterError(t *testing.T) { + expectedErr := errors.New("dial failed") + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, expectedErr + }, + })) + ro := &xrecorder.HandlerRecorderObject{} + + _, err := h.sniffingDial(context.Background(), "tcp", "example.com:80", ro) + if err == nil { + t.Fatal("expected error") + } + if !errors.Is(err, expectedErr) { + t.Errorf("expected %v, got %v", expectedErr, err) + } +} + +// --------------------------------------------------------------------------- +// SnifferBuilder +// --------------------------------------------------------------------------- + +func TestSnifferBuilder(t *testing.T) { + h := newTestHandler() + h.sniffer.Websocket = true + h.sniffer.WebsocketSampleRate = 0.5 + h.sniffer.ReadTimeout = 30 * time.Second + + sniffer := h.sniffer.Build() + + if !sniffer.Websocket { + t.Error("expected Websocket to be true") + } + if sniffer.WebsocketSampleRate != 0.5 { + t.Errorf("expected 0.5, got %f", sniffer.WebsocketSampleRate) + } + if sniffer.ReadTimeout != 30*time.Second { + t.Errorf("expected 30s, got %v", sniffer.ReadTimeout) + } +} + +// --------------------------------------------------------------------------- +// handleSniffedProtocol +// --------------------------------------------------------------------------- + +func TestHandleSniffedProtocol_UnknownProto(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), "ssh") + if handled { + t.Error("expected handled=false for unknown proto") + } + if err != nil { + t.Errorf("expected nil error, got %v", err) + } +} + +func TestHandleSniffedProtocol_EmptyProto(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + ro := &xrecorder.HandlerRecorderObject{} + + handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), "") + if handled { + t.Error("expected handled=false for empty proto") + } + if err != nil { + t.Errorf("expected nil error, got %v", err) + } +} + +func TestHandleSniffedProtocol_HTTP(t *testing.T) { + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, errors.New("dial not configured in test") + }, + })) + conn := newStringConn([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")) + ro := &xrecorder.HandlerRecorderObject{} + + handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), sniffing.ProtoHTTP) + + if !handled { + t.Error("expected handled=true for HTTP proto") + } + if err == nil { + t.Log("sniffer processed HTTP request without error") + } +} + +func TestHandleSniffedProtocol_TLS(t *testing.T) { + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, errors.New("dial not configured in test") + }, + })) + conn := newStringConn([]byte{ + 0x16, 0x03, 0x01, 0x00, 0x06, 0x01, 0x00, 0x00, 0x02, 0x03, 0x01, + }) + ro := &xrecorder.HandlerRecorderObject{} + + handled, err := h.handleSniffedProtocol(context.Background(), conn, ro, nopLog(), sniffing.ProtoTLS) + + if !handled { + t.Error("expected handled=true for TLS proto") + } + if err == nil { + t.Log("sniffer processed TLS ClientHello without error") + } +} diff --git a/handler/forward/local/util.go b/handler/forward/local/util.go new file mode 100644 index 00000000..886a87c8 --- /dev/null +++ b/handler/forward/local/util.go @@ -0,0 +1,50 @@ +package local + +import ( + "context" + "errors" + "net" + "time" + + xctx "github.com/go-gost/x/ctx" + xrecorder "github.com/go-gost/x/recorder" +) + +var ( + errRouterNotAvailable = errors.New("router not available") + errNodeNotAvailable = errors.New("node not available") +) + +// newRecorderObject creates a HandlerRecorderObject populated with connection +// metadata (service, addresses, network type, session ID, client address). +func (h *forwardHandler) newRecorderObject(ctx context.Context, conn net.Conn, start time.Time) *xrecorder.HandlerRecorderObject { + ro := &xrecorder.HandlerRecorderObject{ + Service: h.options.Service, + RemoteAddr: conn.RemoteAddr().String(), + LocalAddr: conn.LocalAddr().String(), + Network: "tcp", + Time: start, + SID: xctx.SidFromContext(ctx).String(), + } + if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil { + ro.ClientAddr = srcAddr.String() + } + if _, ok := conn.(net.PacketConn); ok { + ro.Network = "udp" + } + return ro +} + +func (h *forwardHandler) checkRateLimit(addr net.Addr) bool { + if h.options.RateLimiter == nil { + return true + } + if addr == nil { + return true + } + host, _, _ := net.SplitHostPort(addr.String()) + if limiter := h.options.RateLimiter.Limiter(host); limiter != nil { + return limiter.Allow(1) + } + return true +} diff --git a/handler/forward/local/util_test.go b/handler/forward/local/util_test.go new file mode 100644 index 00000000..e32cfc76 --- /dev/null +++ b/handler/forward/local/util_test.go @@ -0,0 +1,145 @@ +package local + +import ( + "context" + "net" + "testing" + "time" + + "github.com/go-gost/core/limiter/rate" + xctx "github.com/go-gost/x/ctx" +) + +// --------------------------------------------------------------------------- +// newRecorderObject +// --------------------------------------------------------------------------- + +func TestNewRecorderObject_TCP(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + start := time.Now() + + ro := h.newRecorderObject(context.Background(), conn, start) + + if ro.Network != "tcp" { + t.Errorf("expected network tcp, got %s", ro.Network) + } + if ro.RemoteAddr != "10.0.0.1:12345" { + t.Errorf("expected remote 10.0.0.1:12345, got %s", ro.RemoteAddr) + } + if ro.LocalAddr != "127.0.0.1:8080" { + t.Errorf("expected local 127.0.0.1:8080, got %s", ro.LocalAddr) + } + if !ro.Time.Equal(start) { + t.Errorf("expected time %v, got %v", start, ro.Time) + } +} + +func TestNewRecorderObject_SID(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + ctx := xctx.ContextWithSid(context.Background(), xctx.Sid("test-sid")) + + ro := h.newRecorderObject(ctx, conn, time.Now()) + + if ro.SID != "test-sid" { + t.Errorf("expected SID 'test-sid', got '%s'", ro.SID) + } +} + +func TestNewRecorderObject_NilSrcAddr(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + + ro := h.newRecorderObject(context.Background(), conn, time.Now()) + + if ro.ClientAddr != "" { + t.Errorf("expected empty ClientAddr, got %s", ro.ClientAddr) + } +} + +func TestNewRecorderObject_UDP(t *testing.T) { + h := newInitdHandler() + conn := &packetConn{newStringConn(nil)} + + ro := h.newRecorderObject(context.Background(), conn, time.Now()) + + if ro.Network != "udp" { + t.Errorf("expected network udp, got %s", ro.Network) + } +} + +func TestNewRecorderObject_SrcAddr(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + srcAddr := &net.TCPAddr{IP: net.IPv4(192, 168, 1, 1), Port: 54321} + ctx := xctx.ContextWithSrcAddr(context.Background(), srcAddr) + + ro := h.newRecorderObject(ctx, conn, time.Now()) + + if ro.ClientAddr != "192.168.1.1:54321" { + t.Errorf("expected ClientAddr '192.168.1.1:54321', got '%s'", ro.ClientAddr) + } +} + +// --------------------------------------------------------------------------- +// checkRateLimit +// --------------------------------------------------------------------------- + +func TestCheckRateLimit_NilLimiter(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + + if !h.checkRateLimit(conn.RemoteAddr()) { + t.Error("expected true when RateLimiter is nil") + } +} + +func TestCheckRateLimit_NilAddr(t *testing.T) { + h := newInitdHandler(withRateLimiter(&stubRateLimiter{ + limiterFn: func(key string) rate.Limiter { + return &stubLimiter{allowFn: func(n int) bool { return false }} + }, + })) + + if !h.checkRateLimit(nil) { + t.Error("expected true when addr is nil") + } +} + +func TestCheckRateLimit_Allowed(t *testing.T) { + h := newInitdHandler(withRateLimiter(&stubRateLimiter{ + limiterFn: func(key string) rate.Limiter { + return &stubLimiter{allowFn: func(n int) bool { return true }} + }, + })) + conn := newStringConn(nil) + + if !h.checkRateLimit(conn.RemoteAddr()) { + t.Error("expected true when limiter allows") + } +} + +func TestCheckRateLimit_Blocked(t *testing.T) { + h := newInitdHandler(withRateLimiter(&stubRateLimiter{ + limiterFn: func(key string) rate.Limiter { + return &stubLimiter{allowFn: func(n int) bool { return false }} + }, + })) + conn := newStringConn(nil) + + if h.checkRateLimit(conn.RemoteAddr()) { + t.Error("expected false when limiter blocks") + } +} + +func TestCheckRateLimit_NoLimiterForKey(t *testing.T) { + h := newInitdHandler(withRateLimiter(&stubRateLimiter{ + limiterFn: func(key string) rate.Limiter { return nil }, + })) + conn := newStringConn(nil) + + if !h.checkRateLimit(conn.RemoteAddr()) { + t.Error("expected true when no limiter found for key") + } +}