refactor(handler/forward/local): split handler into 6 focused files with 40 tests at 100% coverage

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