refactor(handler/sni): extract helpers, add SnifferBuilder, comprehensive tests
- 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/.
This commit is contained in:
+79
-89
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user