refactor(handler/forward): extract 6 helpers from remote Handle, add 44 unit tests

Extract newRecorderObject, selectTarget, sniffingDial, buildSniffer,
handleSniffedProtocol, and handleRawForwarding from Handle (200→87 lines).
Move nil Router guard before sniffing for early failure. Log previously
silent sniff and pipe errors at Debug level. Apply readTimeout to Pipe.
Append :0 when target addr lacks a port. Reorder metadata fields and
add doc comments on all exported symbols.
This commit is contained in:
ginuerzh
2026-05-28 20:40:13 +08:00
parent 9eeb88e6a2
commit 776f045896
3 changed files with 1317 additions and 96 deletions
+143 -91
View File
@@ -11,7 +11,8 @@ import (
"github.com/go-gost/core/chain"
"github.com/go-gost/core/handler"
"github.com/go-gost/core/hop"
mdata "github.com/go-gost/core/metadata"
"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"
@@ -42,6 +43,8 @@ type forwardHandler struct {
certPool tls_util.CertPool
}
// NewHandler creates a remote forwarding handler with the given options.
// The handler registers for "rtcp" and "rudp" protocols.
func NewHandler(opts ...handler.Option) handler.Handler {
options := handler.Options{}
for _, opt := range opts {
@@ -53,7 +56,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
}
}
func (h *forwardHandler) Init(md mdata.Metadata) (err error) {
func (h *forwardHandler) Init(md md.Metadata) (err error) {
if err = h.parseMetadata(md); err != nil {
return
}
@@ -77,29 +80,16 @@ func (h *forwardHandler) Forward(hop hop.Hop) {
h.hop = hop
}
// Handle forwards the accepted connection to a node selected from the hop.
// When sniffing is enabled it inspects the initial bytes to detect HTTP or TLS
// traffic and delegates to the protocol-specific sniffer; otherwise it pipes the
// raw stream through the router.
func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) (err error) {
defer conn.Close()
start := time.Now()
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()
}
network := "tcp"
if _, ok := conn.(net.PacketConn); ok {
network = "udp"
}
ro.Network = network
ro := h.newRecorderObject(ctx, conn, start)
network := ro.Network
log := h.options.Logger.WithFields(map[string]any{
"network": ro.Network,
@@ -135,6 +125,12 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return rate_limiter.ErrRateLimit
}
if h.options.Router == nil {
err := errors.New("router not available")
log.Error(err)
return err
}
var proto string
if network == "tcp" && h.md.sniffing {
if h.md.sniffingTimeout > 0 {
@@ -142,101 +138,61 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
br := bufio.NewReader(conn)
proto, _ = sniffing.Sniff(ctx, br)
proto, err = sniffing.Sniff(ctx, br)
ro.Proto = proto
if err != nil {
log.Debugf("sniff: %v", err)
}
if h.md.sniffingTimeout > 0 {
conn.SetReadDeadline(time.Time{})
}
dial := func(ctx context.Context, network, address string) (net.Conn, error) {
var buf bytes.Buffer
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address)
ro.Route = buf.String()
cc = proxyproto.WrapClientConn(
h.md.proxyProtocol,
xctx.SrcAddrFromContext(ctx),
xctx.DstAddrFromContext(ctx),
cc)
return cc, err
}
sniffer := &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,
}
conn = xnet.NewReadWriteConn(br, conn, conn)
switch proto {
case sniffing.ProtoHTTP:
return 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),
)
case sniffing.ProtoTLS:
return 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),
)
handled, sniffErr := h.handleSniffedProtocol(ctx, conn, ro, log, proto)
if handled {
return sniffErr
}
}
var target *chain.Node
if host := mdutil.GetString(ictx.MetadataFromContext(ctx), "host"); host != "" {
target = &chain.Node{
Addr: host,
}
}
if h.hop != nil {
target = h.hop.Select(ctx,
hop.ProtocolSelectOption(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 := h.selectTarget(ctx, 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 = target.Addr
ro.Host = addr
log = log.WithFields(map[string]any{
"node": target.Name,
"dst": target.Addr,
"dst": addr,
"network": network,
})
log.Debugf("%s >> %s", conn.RemoteAddr(), target.Addr)
log.Debugf("%s >> %s", conn.RemoteAddr(), addr)
var buf bytes.Buffer
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, target.Addr)
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, addr)
ro.Route = buf.String()
if err != nil {
log.Error(err)
@@ -247,14 +203,10 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
return err
}
defer cc.Close()
if marker := target.Marker(); marker != nil {
marker.Reset()
}
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
ro.SrcAddr = cc.LocalAddr().String()
ro.DstAddr = cc.RemoteAddr().String()
defer cc.Close()
cc = proxyproto.WrapClientConn(
h.md.proxyProtocol,
@@ -262,10 +214,15 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
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)
xnet.Pipe(ctx, 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)
@@ -273,6 +230,101 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return nil
}
// selectTarget picks a forwarding target, preferring the host from context
// metadata when set, then falling back to hop selection.
func (h *forwardHandler) selectTarget(ctx context.Context, proto string) *chain.Node {
if host := mdutil.GetString(ictx.MetadataFromContext(ctx), "host"); host != "" {
return &chain.Node{Addr: host}
}
if h.hop != nil {
return h.hop.Select(ctx, hop.ProtocolSelectOption(proto))
}
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
File diff suppressed because it is too large Load Diff
+6 -5
View File
@@ -7,15 +7,15 @@ import (
"time"
"github.com/go-gost/core/bypass"
mdata "github.com/go-gost/core/metadata"
md "github.com/go-gost/core/metadata"
mdutil "github.com/go-gost/x/metadata/util"
"github.com/go-gost/x/registry"
)
type metadata struct {
readTimeout time.Duration
proxyProtocol int
httpKeepalive bool
proxyProtocol int
sniffing bool
sniffingTimeout time.Duration
@@ -28,11 +28,13 @@ type metadata struct {
mitmBypass bypass.Bypass
}
func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
func (h *forwardHandler) parseMetadata(md md.Metadata) (err error) {
h.md.readTimeout = mdutil.GetDuration(md, "readTimeout")
if h.md.readTimeout <= 0 {
h.md.readTimeout = 15 * time.Second
}
h.md.httpKeepalive = mdutil.GetBool(md, "http.keepalive")
h.md.proxyProtocol = mdutil.GetInt(md, "proxyProtocol")
h.md.sniffing = mdutil.GetBool(md, "sniffing")
@@ -40,8 +42,6 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
h.md.httpKeepalive = mdutil.GetBool(md, "http.keepalive")
certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")
if certFile != "" && keyFile != "" {
@@ -57,5 +57,6 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
}
h.md.alpn = mdutil.GetString(md, "mitm.alpn")
h.md.mitmBypass = registry.BypassRegistry().Get(mdutil.GetString(md, "mitm.bypass"))
return
}