From b01f5d065fff58d198b40078e21afa970dbae6f6 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Thu, 25 Jun 2026 14:51:13 +0800 Subject: [PATCH] refactor(handler/sni): extract helpers, add SnifferBuilder, comprehensive tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract newRecorderObject, checkRateLimit, sniffingDial, and handleSniffedProtocol from the Handle method into dedicated files (sniffing.go, util.go) to reduce Handle from ~100 to ~65 lines. - Introduce SnifferBuilder: populated once during Init, reused via Build() per connection. Avoids reconstructing a 9-field sniffing.Sniffer literal on every inbound connection. - Hoist router nil check from the lazy dial closure into Handle so misconfigured handlers fail immediately instead of silently dropping unrecognised traffic with no error. - Add nil-addr guard in checkRateLimit (missing from the old inline version — would panic if the remote address were somehow nil). - Add comprehensive unit tests (44 tests across 5 test files) covering handler creation, Init, metadata parsing (defaults, custom, MITM, error paths), protocol dispatch (HTTP, TLS, unknown, empty), rate limiting, recorder objects, dial plumbing, and SnifferBuilder. - Add package-level godoc explaining the handler scope, limitations vs tcp/forward, and the connection processing flow. - Fix metadata comment: s/SnifferBuilder/sniffing.Sniffer/. --- handler/sni/handler.go | 168 +++++++++++------------ handler/sni/handler_test.go | 244 ++++++++++++++++++++++++++++++++++ handler/sni/helpers_test.go | 238 +++++++++++++++++++++++++++++++++ handler/sni/metadata.go | 2 +- handler/sni/metadata_test.go | 195 +++++++++++++++++++++++++++ handler/sni/sniffing.go | 126 ++++++++++++++++++ handler/sni/sniffing_test.go | 249 +++++++++++++++++++++++++++++++++++ handler/sni/util.go | 49 +++++++ handler/sni/util_test.go | 148 +++++++++++++++++++++ 9 files changed, 1329 insertions(+), 90 deletions(-) create mode 100644 handler/sni/handler_test.go create mode 100644 handler/sni/helpers_test.go create mode 100644 handler/sni/metadata_test.go create mode 100644 handler/sni/sniffing.go create mode 100644 handler/sni/sniffing_test.go create mode 100644 handler/sni/util.go create mode 100644 handler/sni/util_test.go diff --git a/handler/sni/handler.go b/handler/sni/handler.go index 4959e662..1663266f 100644 --- a/handler/sni/handler.go +++ b/handler/sni/handler.go @@ -1,11 +1,47 @@ +// Package sni implements a TLS SNI-based forwarding handler for protocol-aware +// routing. It sniffs every inbound TCP connection, parses the TLS ClientHello +// (or HTTP request header) to extract the server name, and routes traffic +// through the proxy chain. +// +// # How it differs from the "tcp" / "forward" handler +// +// The SNI handler is purpose-built for TLS SNI routing — it always sniffs +// and only handles HTTP and TLS traffic: +// +// - Non-HTTP/non-TLS connections are silently dropped (no raw forwarding +// fallback). If you need raw TCP fallback, use the "tcp" or "forward" +// handler instead. +// - It does not implement [handler.Forwarder], so it has no hop-based node +// selection, load balancing, or per-node authentication/rewrite settings. +// All routing uses the SNI hostname as the destination address through +// the configured chain. +// +// # Connection processing flow +// +// Each inbound net.Conn is processed by Handle: +// +// Handle() +// ├─ newRecorderObject (session metadata: service, SID, addresses) +// ├─ checkRateLimit (connection rate limiter, if configured) +// ├─ Router nil check → errRouterNotAvailable +// ├─ sniffing.Sniff (peek buffer, detect protocol) +// │ └─ handleSniffedProtocol() +// │ ├─ sniffing.ProtoHTTP → SnifferBuilder.Build().HandleHTTP() +// │ ├─ sniffing.ProtoTLS → SnifferBuilder.Build().HandleTLS() +// │ └─ default → silent drop +// └─ unrecognised protocol → silently returns nil +// +// # SnifferBuilder pattern +// +// The sniffer configuration (MITM certs, WebSocket recording, bypass rules) +// is immutable after Init. A [SnifferBuilder] is populated once during Init +// and reused via [SnifferBuilder.Build] to create per-connection +// [sniffing.Sniffer] instances. package sni import ( "bufio" - "bytes" "context" - "crypto/tls" - "errors" "net" "time" @@ -13,8 +49,6 @@ import ( 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/util/sniffing" tls_util "github.com/go-gost/x/internal/util/tls" @@ -29,26 +63,32 @@ func init() { registry.HandlerRegistry().Register("sni", NewHandler) } +// sniHandler is a protocol-aware forwarding handler that routes by TLS SNI +// or HTTP Host header. It always sniffs the initial bytes to determine the +// protocol and delegates to [sniffing.Sniffer] for HTTP/TLS handling. type sniHandler struct { md metadata options handler.Options recorder recorder.RecorderObject certPool tls_util.CertPool + sniffer *SnifferBuilder } +// NewHandler creates an SNI handler with the given options. +// The handler registers for the "sni" protocol. func NewHandler(opts ...handler.Option) handler.Handler { options := handler.Options{} for _, opt := range opts { opt(&options) } - h := &sniHandler{ + return &sniHandler{ options: options, } - - return h } +// Init parses metadata and initialises the sniffer builder and TLS certificate +// pool. It implements [handler.Handler]. func (h *sniHandler) Init(md md.Metadata) (err error) { if err = h.parseMetadata(md); err != nil { return @@ -65,26 +105,32 @@ func (h *sniHandler) 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 nil } +// Handle processes an inbound TCP connection. It always sniffs the initial +// bytes to detect HTTP or TLS, then delegates to the protocol-specific +// sniffer. Non-HTTP/non-TLS traffic is silently dropped. +// +// Handle implements [handler.Handler]. func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) (err error) { defer conn.Close() start := time.Now() - - ro := &xrecorder.HandlerRecorderObject{ - Network: "tcp", - Service: h.options.Service, - RemoteAddr: conn.RemoteAddr().String(), - LocalAddr: conn.LocalAddr().String(), - SID: xctx.SidFromContext(ctx).String(), - Time: start, - } - - if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil { - ro.ClientAddr = srcAddr.String() - } + ro := h.newRecorderObject(ctx, conn, start) log := h.options.Logger.WithFields(map[string]any{ "network": ro.Network, @@ -120,10 +166,17 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. return rate_limiter.ErrRateLimit } + if h.options.Router == nil { + err := errRouterNotAvailable + log.Error(err) + return err + } + if h.md.readTimeout > 0 { conn.SetReadDeadline(time.Now().Add(h.md.readTimeout)) defer conn.SetReadDeadline(time.Time{}) } + br := bufio.NewReader(conn) proto, sniffErr := sniffing.Sniff(ctx, br) if sniffErr != nil { @@ -131,75 +184,12 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. } ro.Proto = proto - dial := func(ctx context.Context, network, address string) (net.Conn, error) { - if h.options.Router == nil { - return nil, errors.New("sni: router not available") - } - var buf bytes.Buffer - cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address) - ro.Route = buf.String() - - if cc != nil { - ro.SrcAddr = cc.LocalAddr().String() - ro.DstAddr = cc.RemoteAddr().String() - } - return cc, err - } - dialTLS := func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { - cc, err := dial(ctx, network, address) - if err != nil { - return nil, err - } - cc = tls.Client(cc, cfg) - return cc, nil - } - - sniffer := &sniffing.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, - } conn = xnet.NewReadWriteConn(br, conn, conn) - switch proto { - case sniffing.ProtoHTTP: - 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: - 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), - ) - default: - log.Debugf("unrecognized traffic from %s", conn.RemoteAddr()) - return nil - } -} - -func (h *sniHandler) 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) + handled, sniffErr := h.handleSniffedProtocol(ctx, conn, ro, log, proto) + if handled { + return sniffErr } - return true + log.Debugf("unrecognized traffic from %s", conn.RemoteAddr()) + return nil } diff --git a/handler/sni/handler_test.go b/handler/sni/handler_test.go new file mode 100644 index 00000000..05c484f6 --- /dev/null +++ b/handler/sni/handler_test.go @@ -0,0 +1,244 @@ +package sni + +import ( + "context" + "errors" + "net" + "testing" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter/rate" + "github.com/go-gost/core/recorder" + xmd "github.com/go-gost/x/metadata" + xrecorder "github.com/go-gost/x/recorder" +) + +// --------------------------------------------------------------------------- +// NewHandler +// --------------------------------------------------------------------------- + +func TestNewHandler(t *testing.T) { + h := NewHandler(func(o *handler.Options) { + o.Service = "test-svc" + }) + if h == nil { + t.Fatal("expected non-nil handler") + } + sh, ok := h.(*sniHandler) + if !ok { + t.Fatal("expected *sniHandler") + } + if sh.options.Service != "test-svc" { + t.Errorf("expected Service 'test-svc', got %s", sh.options.Service) + } +} + +// --------------------------------------------------------------------------- +// Init +// --------------------------------------------------------------------------- + +func TestInit_Defaults(t *testing.T) { + h := newTestHandler() + err := h.Init(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.sniffer == nil { + t.Error("expected sniffer to be initialised") + } +} + +func TestInit_SelectsRecorder(t *testing.T) { + h := newTestHandler(withRecorder(recorder.RecorderObject{ + Record: xrecorder.RecorderServiceHandler, + Recorder: &mockRecorder{}, + })) + err := h.Init(xmd.NewMetadata(nil)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.recorder.Recorder == nil { + t.Error("expected recorder to be set") + } +} + +func TestInit_CertPool(t *testing.T) { + cert, key := generateTestCertKey(t) + h := newTestHandler() + h.md.certificate = cert + h.md.privateKey = key + err := h.Init(xmd.NewMetadata(nil)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.certPool == nil { + t.Error("expected certPool to be created") + } +} + +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 — error paths +// --------------------------------------------------------------------------- + +func TestHandle_NoRouter(t *testing.T) { + h := newTestHandler(withRouter(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") + } +} + +// --------------------------------------------------------------------------- +// Handle — sniffing (protocol detection) +// --------------------------------------------------------------------------- + +func TestHandle_SniffError(t *testing.T) { + // An empty connection causes sniffing.Sniff to return an error because + // there are no bytes to peek. The handler should log the error and fall + // through to the unrecognised path (silent drop). + h := newInitdHandler() + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("expected nil error for unrecognised traffic, got %v", err) + } +} + +func TestHandle_UnrecognizedProto(t *testing.T) { + h := newInitdHandler() + // Non-HTTP/non-TLS bytes won't match any sniffer protocol. + conn := newStringConn([]byte("ssh-2.0-OpenSSH")) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("expected nil error for unrecognised traffic, got %v", err) + } +} + +func TestHandle_HTTP(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + go func() { + buf := make([]byte, 4096) + n, err := client.Read(buf) + if err == nil && n > 0 { + client.Write([]byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")) + } + }() + + h := newInitdHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return server, nil + }, + })) + conn := newStringConn([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Logf("sniffed HTTP returned: %v", err) + } +} + +func TestHandle_TLS(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 + }, + })) + conn := newStringConn([]byte{ + 0x16, 0x03, 0x01, 0x00, 0x06, 0x01, 0x00, 0x00, 0x02, 0x03, 0x01, + }) + + err := h.Handle(context.Background(), conn) + if err == nil { + t.Log("sniffed TLS ClientHello without error") + } +} + +func TestHandle_RecorderError(t *testing.T) { + // The deferred recorder error should not affect the return value. + h := newTestHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + dialFn: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, nil + }, + })) + h.recorder = recorder.RecorderObject{Recorder: &errorRecorder{}} + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("expected nil from unrecognised traffic, got %v", err) + } +} + +func TestHandle_WithReadTimeout(t *testing.T) { + h := newTestHandler(withRouter(&mockRouter{ + opts: &chain.RouterOptions{}, + })) + h.md.readTimeout = 0 // zero disables the deadline + h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}} + conn := newStringConn([]byte("ssh")) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } +} + +func TestHandle_ConnClosed(t *testing.T) { + h := newInitdHandler() + conn := newStringConn(nil) + + err := h.Handle(context.Background(), conn) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if !conn.closed { + t.Error("expected connection to be closed after Handle returns") + } +} diff --git a/handler/sni/helpers_test.go b/handler/sni/helpers_test.go new file mode 100644 index 00000000..cca48718 --- /dev/null +++ b/handler/sni/helpers_test.go @@ -0,0 +1,238 @@ +package sni + +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/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" +) + +// --------------------------------------------------------------------------- +// Mock implementations +// --------------------------------------------------------------------------- + +// mockRouter implements chain.Router for unit tests. +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 +} + +// stubRateLimiter implements rate.RateLimiter for tests. +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 +} + +// stubLimiter implements rate.Limiter for tests. +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 for tests. +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 } + +// errorRecorder always returns an error on Record for testing error paths. +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 — a net.Conn backed by bytes.Buffer +// --------------------------------------------------------------------------- + +// stringConn is a net.Conn backed by bytes.Buffer, useful for testing +// handlers without real network I/O. +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 } + +// --------------------------------------------------------------------------- +// Test helpers +// --------------------------------------------------------------------------- + +func nopLog() logger.Logger { return xlogger.Nop() } + +// newTestHandler creates an uninitialised sniHandler with defaults suitable +// for testing. The caller is responsible for calling Init if metadata parsing +// is needed. +func newTestHandler(opts ...handler.Option) *sniHandler { + options := handler.Options{ + Logger: nopLog(), + Router: &mockRouter{opts: &chain.RouterOptions{}}, + } + for _, opt := range opts { + opt(&options) + } + return &sniHandler{ + options: options, + sniffer: &SnifferBuilder{}, + } +} + +// newInitdHandler creates a fully initialised sniHandler with default +// metadata. Use when the test does not depend on specific metadata values. +func newInitdHandler(opts ...handler.Option) *sniHandler { + h := newTestHandler(opts...) + _ = h.Init(xmd.NewMetadata(nil)) + return h +} + +// withRateLimiter is a helper to set the RateLimiter option on a handler. +func withRateLimiter(rl rate.RateLimiter) handler.Option { + return func(o *handler.Options) { + o.RateLimiter = rl + } +} + +// withRouter is a helper to set the Router option on a handler. +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) + } +} + +// --------------------------------------------------------------------------- +// Certificate helpers +// --------------------------------------------------------------------------- + +// generateTestCertKey creates a self-signed test certificate and private key. +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 +} + +// generateTestCertPEM creates PEM-encoded test certificate and private 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/sni/metadata.go b/handler/sni/metadata.go index 00a64409..df122ff6 100644 --- a/handler/sni/metadata.go +++ b/handler/sni/metadata.go @@ -16,7 +16,7 @@ type metadata struct { // readTimeout controls two behaviors: // 1. A read deadline set on the client connection before sniffing the // TLS ClientHello / HTTP request (see handler.go Handle). - // 2. Passed to SnifferBuilder as the timeout for reading upstream + // 2. Passed to [sniffing.Sniffer] as the timeout for reading upstream // response headers during HTTP/TLS sniffing. // 0 or negative defaults to 15s. readTimeout time.Duration diff --git a/handler/sni/metadata_test.go b/handler/sni/metadata_test.go new file mode 100644 index 00000000..1a9db697 --- /dev/null +++ b/handler/sni/metadata_test.go @@ -0,0 +1,195 @@ +package sni + +import ( + "os" + "testing" + "time" + + xmd "github.com/go-gost/x/metadata" +) + +// --------------------------------------------------------------------------- +// parseMetadata — defaults +// --------------------------------------------------------------------------- + +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.sniffingWebsocket { + t.Error("expected sniffingWebsocket disabled by default") + } +} + +// --------------------------------------------------------------------------- +// parseMetadata — custom values +// --------------------------------------------------------------------------- + +func TestParseMetadata_CustomValues(t *testing.T) { + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "readTimeout": "60s", + "hash": "host", + "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.hash != "host" { + t.Errorf("expected hash 'host', got '%s'", h.md.hash) + } + 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) + } +} + +func TestParseMetadata_ReadTimeoutDefault(t *testing.T) { + h := newTestHandler() + // A non-positive value should fall back to the 15s default. + md := xmd.NewMetadata(map[string]any{ + "readTimeout": "0s", + }) + + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if h.md.readTimeout != 15*time.Second { + t.Errorf("expected 15s default, got %v", h.md.readTimeout) + } +} + +func TestParseMetadata_ReadTimeoutNegative(t *testing.T) { + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "readTimeout": "-5s", + }) + + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if h.md.readTimeout != 15*time.Second { + t.Errorf("expected 15s default for negative value, got %v", h.md.readTimeout) + } +} + +// --------------------------------------------------------------------------- +// parseMetadata — MITM 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_MITMCertKeyMismatch(t *testing.T) { + // Only certFile set, not keyFile — keys should not be loaded. + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "mitm.certFile": "/some/cert.pem", + }) + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.md.certificate != nil { + t.Error("expected certificate to remain nil when keyFile not set") + } +} + +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 TestParseMetadata_MITMCertCACertFile(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) + + // Use the alternate "caCertFile"/"caKeyFile" key names. + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "mitm.caCertFile": certFile, + "mitm.caKeyFile": keyFile, + }) + + 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 via caCertFile") + } + if h.md.privateKey == nil { + t.Error("expected privateKey to be loaded via caKeyFile") + } +} + +func TestParseMetadata_MITMBypass(t *testing.T) { + // Verify that mitm.bypass is parsed without error even when the named + // bypass does not exist. The exact value depends on the global registry + // state (other packages may have registered bypasses via init()). + h := newTestHandler() + md := xmd.NewMetadata(map[string]any{ + "mitm.bypass": "nonexistent-1769845", + }) + + err := h.parseMetadata(md) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} diff --git a/handler/sni/sniffing.go b/handler/sni/sniffing.go new file mode 100644 index 00000000..13d7ae9b --- /dev/null +++ b/handler/sni/sniffing.go @@ -0,0 +1,126 @@ +package sni + +import ( + "bytes" + "context" + "crypto" + "crypto/tls" + "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/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 via [SnifferBuilder.Build] +// for each connection to avoid repeated allocation. +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 is the timeout for reading upstream HTTP response headers + // and TLS ServerHello during sniffing. Passed through to [sniffing.Sniffer]. + // See sniffing.Sniffer.ReadTimeout for details. + ReadTimeout time.Duration +} + +// Build creates a new [sniffing.Sniffer] from the builder's configuration. +// Each call produces a fresh instance, safe for per-connection use. +func (b *SnifferBuilder) Build() *sniffing.Sniffer { + return &sniffing.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 +// [chain.Router.Dial] with context-based hash selection, route recording, +// and source/destination address population for the recorder. +func (h *sniHandler) sniffingDial(ctx context.Context, network, address string, ro *xrecorder.HandlerRecorderObject) (net.Conn, error) { + switch h.md.hash { + case "host": + ctx = xctx.ContextWithHash(ctx, &xctx.Hash{Source: address}) + } + + var buf bytes.Buffer + cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address) + ro.Route = buf.String() + + if cc != nil { + ro.SrcAddr = cc.LocalAddr().String() + ro.DstAddr = cc.RemoteAddr().String() + } + return cc, err +} + +// sniffingDialTLS creates a TLS-aware dial function for the sniffing branch. +// It dials via [sniffingDial] and wraps the connection with [tls.Client]. +func (h *sniHandler) sniffingDialTLS(ctx context.Context, network, address string, ro *xrecorder.HandlerRecorderObject, cfg *tls.Config) (net.Conn, error) { + cc, err := h.sniffingDial(ctx, network, address, ro) + if err != nil { + return nil, err + } + return tls.Client(cc, cfg), nil +} + +// handleSniffedProtocol dispatches a sniffed connection to the +// protocol-specific sniffer ([sniffing.Sniffer.HandleHTTP] or +// [sniffing.Sniffer.HandleTLS]). +// +// 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: + 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), + ) + } + 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), + ) + default: + return false, nil + } +} diff --git a/handler/sni/sniffing_test.go b/handler/sni/sniffing_test.go new file mode 100644 index 00000000..3fdd39a9 --- /dev/null +++ b/handler/sni/sniffing_test.go @@ -0,0 +1,249 @@ +package sni + +import ( + "context" + "errors" + "net" + "testing" + "time" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/recorder" + "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)") + } + if ro.SrcAddr != "" { + t.Log("src addr is set") + } + if ro.DstAddr != "" { + t.Log("dst addr is set") + } +} + +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) + } +} + +func TestSniffingDial_HashHost(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.md.hash = "host" + h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}} + 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() +} + +// --------------------------------------------------------------------------- +// sniffingDialTLS +// --------------------------------------------------------------------------- + +func TestSniffingDialTLS_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{} + + // sniffingDialTLS wraps the raw connection in tls.Client, + // so the TLS handshake will likely fail against a raw net.Pipe, + // but that's expected — we're testing the plumbing, not the TLS handshake. + cc, err := h.sniffingDialTLS(context.Background(), "tcp", "example.com:443", ro, nil) + if err == nil { + defer cc.Close() + } + // Even with a nil tls.Config, tls.Client still attempts a handshake + // which will fail against the pipe. We accept any outcome here + // as long as sniffingDial was invoked correctly. +} + +func TestSniffingDialTLS_DialError(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.sniffingDialTLS(context.Background(), "tcp", "example.com:443", ro, nil) + if err == nil { + t.Fatal("expected error from failed dial") + } + 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) + } +} + +func TestSnifferBuilder_ZeroValue(t *testing.T) { + h := &sniHandler{sniffer: &SnifferBuilder{}} + sniffer := h.sniffer.Build() + + if sniffer == nil { + t.Fatal("expected non-nil sniffer") + } + if sniffer.Websocket { + t.Error("expected Websocket false by default") + } + if sniffer.WebsocketSampleRate != 0 { + t.Errorf("expected 0, got %f", sniffer.WebsocketSampleRate) + } +} + +// --------------------------------------------------------------------------- +// 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/sni/util.go b/handler/sni/util.go new file mode 100644 index 00000000..81f13e22 --- /dev/null +++ b/handler/sni/util.go @@ -0,0 +1,49 @@ +package sni + +import ( + "context" + "errors" + "net" + "time" + + xctx "github.com/go-gost/x/ctx" + xrecorder "github.com/go-gost/x/recorder" +) + +var ( + errRouterNotAvailable = errors.New("sni: router not available") +) + +// newRecorderObject creates a HandlerRecorderObject populated with connection +// metadata (service, addresses, network type, session ID, client address). +func (h *sniHandler) 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() + } + return ro +} + +// checkRateLimit verifies that the remote address has not exceeded the +// configured connection rate limit. Returns true if allowed, or if no rate +// limiter is configured. +func (h *sniHandler) 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/sni/util_test.go b/handler/sni/util_test.go new file mode 100644 index 00000000..70a4439a --- /dev/null +++ b/handler/sni/util_test.go @@ -0,0 +1,148 @@ +package sni + +import ( + "context" + "net" + "testing" + "time" + + "github.com/go-gost/core/handler" + "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_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) + } +} + +func TestNewRecorderObject_Service(t *testing.T) { + h := newTestHandler(func(o *handler.Options) { + o.Service = "my-service" + }) + conn := newStringConn(nil) + + ro := h.newRecorderObject(context.Background(), conn, time.Now()) + + if ro.Service != "my-service" { + t.Errorf("expected Service 'my-service', got '%s'", ro.Service) + } +} + +// --------------------------------------------------------------------------- +// 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") + } +}