refactor(handler/forward): extract 5 helpers from Handle, add 39 unit tests

handler/forward/local/handler.go

@@ -11,6 +11,7 @@ import (
 	"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"
@@ -87,25 +88,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 	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,
@@ -165,54 +149,19 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 		}
 
 		conn = xnet.NewReadWriteConn(br, conn, conn)
-		switch proto {
-		case sniffing.ProtoHTTP, sniffing.ProtoTLS:
-			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()
-				return proxyproto.WrapClientConn(
-					h.md.proxyProtocol,
-					xctx.SrcAddrFromContext(ctx),
-					xctx.DstAddrFromContext(ctx),
-					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,
-			}
-			if proto == 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),
-				)
-			}
-			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
 		}
 	}
 
-	// h.hop may be nil when no forwarder is configured — this is valid.
-	// The sniffer and Select() both guard against nil hop internally.
+	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,
@@ -287,6 +236,89 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 	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