add sniffer for forwarding handler

handler/forward/local/handler.go

@@ -6,20 +6,23 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net"
+	"strconv"
 	"time"
 
 	"github.com/go-gost/core/chain"
 	"github.com/go-gost/core/handler"
 	"github.com/go-gost/core/hop"
 	md "github.com/go-gost/core/metadata"
+	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/core/recorder"
 	ctxvalue "github.com/go-gost/x/ctx"
-	xio "github.com/go-gost/x/internal/io"
 	xnet "github.com/go-gost/x/internal/net"
+	"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"
+	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 	xrecorder "github.com/go-gost/x/recorder"
 	"github.com/go-gost/x/registry"
 )
@@ -35,6 +38,7 @@ type forwardHandler struct {
 	md       metadata
 	options  handler.Options
 	recorder recorder.RecorderObject
+	certPool tls_util.CertPool
 }
 
 func NewHandler(opts ...handler.Option) handler.Handler {
@@ -60,6 +64,10 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) {
 		}
 	}
 
+	if h.md.certificate != nil && h.md.privateKey != nil {
+		h.certPool = tls_util.NewMemoryCertPool()
+	}
+
 	return
 }
 
@@ -77,6 +85,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 		Service:    h.options.Service,
 		RemoteAddr: conn.RemoteAddr().String(),
 		LocalAddr:  conn.LocalAddr().String(),
+		Network:    "tcp",
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
@@ -87,14 +96,21 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 		"local":  conn.LocalAddr().String(),
 		"sid":    ctxvalue.SidFromContext(ctx),
 	})
-
 	log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
+
+	pStats := stats.Stats{}
+	conn = stats_wrapper.WrapConn(conn, &pStats)
+
 	defer func() {
 		if err != nil {
 			ro.Err = err.Error()
 		}
+		ro.InputBytes = pStats.Get(stats.KindInputBytes)
+		ro.OutputBytes = pStats.Get(stats.KindOutputBytes)
 		ro.Duration = time.Since(start)
-		ro.Record(ctx, h.recorder.Recorder)
+		if err := ro.Record(ctx, h.recorder.Recorder); err != nil {
+			log.Errorf("record: %v", err)
+		}
 
 		log.WithFields(map[string]any{
 			"duration": time.Since(start),
@@ -111,7 +127,6 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 	}
 	ro.Network = network
 
-	var rw io.ReadWriter = conn
 	var proto string
 	if network == "tcp" && h.md.sniffing {
 		if h.md.sniffingTimeout > 0 {
@@ -126,15 +141,42 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 			conn.SetReadDeadline(time.Time{})
 		}
 
-		rw = xio.NewReadWriter(br, conn)
+		dial := func(ctx context.Context, network, address string) (net.Conn, error) {
+			var buf bytes.Buffer
+			cc, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), "tcp", address)
+			ro.Route = buf.String()
+			return cc, err
+		}
+		sniffer := &forwarder.Sniffer{
+			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:
-			ro2 := &xrecorder.HandlerRecorderObject{}
-			*ro2 = *ro
-			ro.Time = time.Time{}
-			return h.handleHTTP(ctx, rw, ro2, log)
+			return sniffer.HandleHTTP(ctx, conn,
+				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 h.handleTLS(ctx, xnet.NewReadWriteConn(rw, rw, conn), ro, log)
+			return sniffer.HandleTLS(ctx, conn,
+				forwarder.WithDial(dial),
+				forwarder.WithHop(h.hop),
+				forwarder.WithBypass(h.options.Bypass),
+				forwarder.WithRecorderObject(ro),
+				forwarder.WithLog(log),
+			)
 		}
 	}
 
@@ -191,7 +233,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 
 	t := time.Now()
 	log.Infof("%s <-> %s", conn.RemoteAddr(), target.Addr)
-	xnet.Transport(rw, cc)
+	xnet.Transport(conn, cc)
 	log.WithFields(map[string]any{
 		"duration": time.Since(t),
 	}).Infof("%s >-< %s", conn.RemoteAddr(), target.Addr)
@@ -210,3 +252,27 @@ func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
 
 	return true
 }
+
+func convertAddr(addr net.Addr) net.Addr {
+	host, sp, _ := net.SplitHostPort(addr.String())
+	ip := net.ParseIP(host)
+	port, _ := strconv.Atoi(sp)
+
+	if ip == nil || ip.Equal(net.IPv6zero) {
+		ip = net.IPv4zero
+	}
+
+	switch addr.Network() {
+	case "tcp", "tcp4", "tcp6":
+		return &net.TCPAddr{
+			IP:   ip,
+			Port: port,
+		}
+
+	default:
+		return &net.UDPAddr{
+			IP:   ip,
+			Port: port,
+		}
+	}
+}