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:
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -1,25 +1,78 @@
|
||||
// Package local implements a low-level forwarding handler for raw TCP and UDP
|
||||
// connections. It serves as the fallback handler for protocols that cannot be
|
||||
// matched to a dedicated handler, and supports optional protocol sniffing to
|
||||
// detect HTTP and TLS traffic for deep inspection (MITM, WebSocket recording).
|
||||
//
|
||||
// The handler is registered under the names "tcp", "udp", and "forward" via
|
||||
// NewHandler in init().
|
||||
//
|
||||
// # Connection processing flow
|
||||
//
|
||||
// Each inbound net.Conn is processed by Handle, which wraps the connection
|
||||
// with I/O stats, checks the rate limiter, and dispatches to one of two
|
||||
// forwarding paths:
|
||||
//
|
||||
// Handle()
|
||||
// ├─ stats_wrapper.WrapConn (per-connection I/O counters)
|
||||
// ├─ newRecorderObject (session metadata: service, SID, addresses)
|
||||
// ├─ checkRateLimit (connection rate limiter, if configured)
|
||||
// ├─ Router nil check → errRouterNotAvailable
|
||||
// ├─ [sniffing enabled] sniffing.Sniff (peek buffer, detect protocol)
|
||||
// │ └─ handleSniffedProtocol()
|
||||
// │ ├─ sniffing.ProtoHTTP → SnifferBuilder.Build().HandleHTTP()
|
||||
// │ ├─ sniffing.ProtoTLS → SnifferBuilder.Build().HandleTLS()
|
||||
// │ └─ default → fall through to raw forwarding
|
||||
// └─ [sniffing disabled / unrecognised] handleRawForwarding()
|
||||
//
|
||||
// # Protocol sniffing dispatch (handleSniffedProtocol)
|
||||
//
|
||||
// When sniffing is enabled and the initial bytes match HTTP or TLS, the
|
||||
// connection is delegated to the forwarder package for protocol-aware
|
||||
// handling:
|
||||
//
|
||||
// - HTTP: parses the request, applies MITM if configured, forwards to
|
||||
// the upstream, and records WebSocket frames when enabled.
|
||||
// - TLS: parses the ClientHello to extract the server name, applies
|
||||
// MITM or SNI-based routing through the hop, and records the TLS
|
||||
// handshake.
|
||||
//
|
||||
// SnifferBuilder is populated once during Init and reused via Build()
|
||||
// to create per-connection forwarder.Sniffer instances.
|
||||
//
|
||||
// # Raw forwarding (handleRawForwarding)
|
||||
//
|
||||
// The default path for unrecognised traffic. It selects a target node
|
||||
// from the hop, dials the upstream through the router, wraps the
|
||||
// connection with proxy protocol (HAProxy) when enabled, and pipes
|
||||
// data bidirectionally:
|
||||
//
|
||||
// 1. Node selection — getHop().Select() with optional protocol hint.
|
||||
// 2. Router dial — Router.Dial() through the configured chain.
|
||||
// 3. Proxy protocol — proxyproto.WrapClientConn() prepends HAProxy
|
||||
// header when metadata proxyProtocol is set.
|
||||
// 4. Bidirectional pipe — xnet.Pipe(ctx, conn, cc) with read timeout.
|
||||
//
|
||||
// # Hop management
|
||||
//
|
||||
// The handler implements handler.Forwarder. Forward() is called during
|
||||
// service startup (and on reload) to install or update the hop. The hop
|
||||
// is protected by a mutex because Forward() runs on the loader goroutine
|
||||
// while Handle() reads it from per-connection goroutines.
|
||||
package local
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/handler"
|
||||
"github.com/go-gost/core/hop"
|
||||
"github.com/go-gost/core/logger"
|
||||
md "github.com/go-gost/core/metadata"
|
||||
"github.com/go-gost/core/observer/stats"
|
||||
"github.com/go-gost/core/recorder"
|
||||
xctx "github.com/go-gost/x/ctx"
|
||||
ictx "github.com/go-gost/x/internal/ctx"
|
||||
xnet "github.com/go-gost/x/internal/net"
|
||||
"github.com/go-gost/x/internal/net/proxyproto"
|
||||
"github.com/go-gost/x/internal/util/forwarder"
|
||||
"github.com/go-gost/x/internal/util/sniffing"
|
||||
tls_util "github.com/go-gost/x/internal/util/tls"
|
||||
rate_limiter "github.com/go-gost/x/limiter/rate"
|
||||
@@ -37,10 +90,12 @@ func init() {
|
||||
|
||||
type forwardHandler struct {
|
||||
hop hop.Hop
|
||||
hopMu sync.Mutex
|
||||
md metadata
|
||||
options handler.Options
|
||||
recorder recorder.RecorderObject
|
||||
certPool tls_util.CertPool
|
||||
sniffer *SnifferBuilder
|
||||
}
|
||||
|
||||
// NewHandler creates a local forwarding handler with the given options.
|
||||
@@ -72,12 +127,33 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) {
|
||||
h.certPool = tls_util.NewMemoryCertPool()
|
||||
}
|
||||
|
||||
h.sniffer = &SnifferBuilder{
|
||||
Websocket: h.md.sniffingWebsocket,
|
||||
WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
|
||||
Recorder: h.recorder.Recorder,
|
||||
RecorderOptions: h.recorder.Options,
|
||||
Certificate: h.md.certificate,
|
||||
PrivateKey: h.md.privateKey,
|
||||
ALPN: h.md.alpn,
|
||||
CertPool: h.certPool,
|
||||
MitmBypass: h.md.mitmBypass,
|
||||
ReadTimeout: h.md.readTimeout,
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Forward implements handler.Forwarder.
|
||||
func (h *forwardHandler) Forward(hop hop.Hop) {
|
||||
h.hopMu.Lock()
|
||||
h.hop = hop
|
||||
h.hopMu.Unlock()
|
||||
}
|
||||
|
||||
func (h *forwardHandler) getHop() hop.Hop {
|
||||
h.hopMu.Lock()
|
||||
defer h.hopMu.Unlock()
|
||||
return h.hop
|
||||
}
|
||||
|
||||
// Handle forwards the accepted connection to a node selected from the hop.
|
||||
@@ -126,9 +202,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
|
||||
}
|
||||
|
||||
if h.options.Router == nil {
|
||||
err := errors.New("router not available")
|
||||
log.Error(err)
|
||||
return err
|
||||
log.Error(errRouterNotAvailable)
|
||||
return errRouterNotAvailable
|
||||
}
|
||||
|
||||
var proto string
|
||||
@@ -157,176 +232,3 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
|
||||
|
||||
return h.handleRawForwarding(ctx, conn, ro, log, network, proto)
|
||||
}
|
||||
|
||||
// handleRawForwarding performs node selection, dials the target through the
|
||||
// router, and pipes the raw connection. It is used when sniffing is disabled
|
||||
// or the protocol was not HTTP/TLS.
|
||||
func (h *forwardHandler) handleRawForwarding(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, network, proto string) error {
|
||||
target := &chain.Node{}
|
||||
if h.hop != nil {
|
||||
target = h.hop.Select(ctx,
|
||||
hop.ProtocolSelectOption(proto),
|
||||
)
|
||||
}
|
||||
if target == nil {
|
||||
err := errors.New("node not available")
|
||||
log.Error(err)
|
||||
return err
|
||||
}
|
||||
addr := target.Addr
|
||||
if opts := target.Options(); opts != nil {
|
||||
switch opts.Network {
|
||||
case "unix":
|
||||
network = opts.Network
|
||||
default:
|
||||
if _, _, err := net.SplitHostPort(addr); err != nil {
|
||||
addr += ":0"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ro.Network = network
|
||||
ro.Host = addr
|
||||
|
||||
log = log.WithFields(map[string]any{
|
||||
"node": target.Name,
|
||||
"dst": addr,
|
||||
"network": network,
|
||||
})
|
||||
|
||||
log.Debugf("%s >> %s", conn.RemoteAddr(), addr)
|
||||
|
||||
var buf bytes.Buffer
|
||||
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, addr)
|
||||
ro.Route = buf.String()
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
// TODO: the router itself may be failed due to the failed node in the router,
|
||||
// the dead marker may be a wrong operation.
|
||||
if marker := target.Marker(); marker != nil {
|
||||
marker.Mark()
|
||||
}
|
||||
return err
|
||||
}
|
||||
if marker := target.Marker(); marker != nil {
|
||||
marker.Reset()
|
||||
}
|
||||
defer cc.Close()
|
||||
|
||||
cc = proxyproto.WrapClientConn(
|
||||
h.md.proxyProtocol,
|
||||
xctx.SrcAddrFromContext(ctx),
|
||||
xctx.DstAddrFromContext(ctx),
|
||||
cc)
|
||||
|
||||
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
||||
ro.SrcAddr = cc.LocalAddr().String()
|
||||
ro.DstAddr = cc.RemoteAddr().String()
|
||||
|
||||
t := time.Now()
|
||||
log.Infof("%s <-> %s", conn.RemoteAddr(), target.Addr)
|
||||
// xnet.Transport(conn, cc)
|
||||
if err := xnet.Pipe(ctx, conn, cc, xnet.WithReadTimeout(h.md.readTimeout)); err != nil {
|
||||
log.Debugf("pipe: %v", err)
|
||||
}
|
||||
log.WithFields(map[string]any{
|
||||
"duration": time.Since(t),
|
||||
}).Infof("%s >-< %s", conn.RemoteAddr(), target.Addr)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// newRecorderObject creates a HandlerRecorderObject populated with connection
|
||||
// metadata (service, addresses, network type, session ID, client address).
|
||||
func (h *forwardHandler) newRecorderObject(ctx context.Context, conn net.Conn, start time.Time) *xrecorder.HandlerRecorderObject {
|
||||
ro := &xrecorder.HandlerRecorderObject{
|
||||
Service: h.options.Service,
|
||||
RemoteAddr: conn.RemoteAddr().String(),
|
||||
LocalAddr: conn.LocalAddr().String(),
|
||||
Network: "tcp",
|
||||
Time: start,
|
||||
SID: xctx.SidFromContext(ctx).String(),
|
||||
}
|
||||
if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil {
|
||||
ro.ClientAddr = srcAddr.String()
|
||||
}
|
||||
if _, ok := conn.(net.PacketConn); ok {
|
||||
ro.Network = "udp"
|
||||
}
|
||||
return ro
|
||||
}
|
||||
|
||||
// sniffingDial creates a dial function for the sniffing branch that wraps
|
||||
// Router.Dial with route recording and proxy protocol encapsulation.
|
||||
func (h *forwardHandler) sniffingDial(ctx context.Context, network, address string, ro *xrecorder.HandlerRecorderObject) (net.Conn, error) {
|
||||
var buf bytes.Buffer
|
||||
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address)
|
||||
ro.Route = buf.String()
|
||||
return proxyproto.WrapClientConn(
|
||||
h.md.proxyProtocol,
|
||||
xctx.SrcAddrFromContext(ctx),
|
||||
xctx.DstAddrFromContext(ctx),
|
||||
cc), err
|
||||
}
|
||||
|
||||
// buildSniffer creates a forwarder.Sniffer configured from handler metadata.
|
||||
func (h *forwardHandler) buildSniffer() *forwarder.Sniffer {
|
||||
return &forwarder.Sniffer{
|
||||
Websocket: h.md.sniffingWebsocket,
|
||||
WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
|
||||
Recorder: h.recorder.Recorder,
|
||||
RecorderOptions: h.recorder.Options,
|
||||
Certificate: h.md.certificate,
|
||||
PrivateKey: h.md.privateKey,
|
||||
NegotiatedProtocol: h.md.alpn,
|
||||
CertPool: h.certPool,
|
||||
MitmBypass: h.md.mitmBypass,
|
||||
ReadTimeout: h.md.readTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
// handleSniffedProtocol dispatches a sniffed connection to the protocol-specific
|
||||
// sniffer (HandleHTTP or HandleTLS). It returns (true, err) when the protocol was
|
||||
// handled and (false, nil) when the caller should fall through to raw forwarding.
|
||||
func (h *forwardHandler) handleSniffedProtocol(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, proto string) (handled bool, err error) {
|
||||
switch proto {
|
||||
case sniffing.ProtoHTTP, sniffing.ProtoTLS:
|
||||
dial := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return h.sniffingDial(ctx, network, address, ro)
|
||||
}
|
||||
sniffer := h.buildSniffer()
|
||||
if proto == sniffing.ProtoHTTP {
|
||||
return true, sniffer.HandleHTTP(ctx, conn,
|
||||
forwarder.WithService(h.options.Service),
|
||||
forwarder.WithDial(dial),
|
||||
forwarder.WithHop(h.hop),
|
||||
forwarder.WithBypass(h.options.Bypass),
|
||||
forwarder.WithHTTPKeepalive(h.md.httpKeepalive),
|
||||
forwarder.WithRecorderObject(ro),
|
||||
forwarder.WithLog(log),
|
||||
)
|
||||
}
|
||||
return true, sniffer.HandleTLS(ctx, conn,
|
||||
forwarder.WithService(h.options.Service),
|
||||
forwarder.WithDial(dial),
|
||||
forwarder.WithHop(h.hop),
|
||||
forwarder.WithBypass(h.options.Bypass),
|
||||
forwarder.WithRecorderObject(ro),
|
||||
forwarder.WithLog(log),
|
||||
)
|
||||
default:
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
|
||||
if h.options.RateLimiter == nil {
|
||||
return true
|
||||
}
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
|
||||
return limiter.Allow(1)
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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()
|
||||
}
|
||||
@@ -49,10 +49,7 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
h.md.certificate, err = x509.ParseCertificate(tlsCert.Certificate[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
h.md.certificate = tlsCert.Leaf
|
||||
h.md.privateKey = tlsCert.PrivateKey
|
||||
}
|
||||
h.md.alpn = mdutil.GetString(md, "mitm.alpn")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user