Tunnel/Dokodemo: Fix stats conn unwrap (#5440)

Fangliding · web-flow · commit fa64775f07b5 · 2025-12-23T09:44:54.000Z
Fixes #5439
diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
@@ -111,7 +111,8 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st
 				destinationOverridden = true
 			}
 		}
-		if tlsConn, ok := conn.(tls.Interface); ok && !destinationOverridden {
+		iConn := stat.TryUnwrapStatsConn(conn)
+		if tlsConn, ok := iConn.(tls.Interface); ok && !destinationOverridden {
 			if serverName := tlsConn.HandshakeContextServerName(ctx); serverName != "" {
 				dest.Address = net.DomainAddress(serverName)
 				destinationOverridden = true
diff --git a/proxy/http/client.go b/proxy/http/client.go
@@ -296,10 +296,7 @@ func setUpHTTPTunnel(ctx context.Context, dest net.Destination, target string, u
 		return nil, err
 	}
 
-	iConn := rawConn
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(rawConn)
 
 	nextProto := ""
 	if tlsConn, ok := iConn.(*tls.Conn); ok {
diff --git a/proxy/proxy.go b/proxy/proxy.go
@@ -787,10 +787,7 @@ func readV(ctx context.Context, reader buf.Reader, writer buf.Writer, timer sign
 }
 
 func IsRAWTransportWithoutSecurity(conn stat.Connection) bool {
-	iConn := conn
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(conn)
 	_, ok1 := iConn.(*proxyproto.Conn)
 	_, ok2 := iConn.(*net.TCPConn)
 	_, ok3 := iConn.(*internet.UnixConnWrapper)
diff --git a/proxy/trojan/server.go b/proxy/trojan/server.go
@@ -147,11 +147,7 @@ func (s *Server) Network() []net.Network {
 
 // Process implements proxy.Inbound.Process().
 func (s *Server) Process(ctx context.Context, network net.Network, conn stat.Connection, dispatcher routing.Dispatcher) error {
-	iConn := conn
-	statConn, ok := iConn.(*stat.CounterConnection)
-	if ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(conn)
 
 	sessionPolicy := s.policyManager.ForLevel(0)
 	if err := conn.SetReadDeadline(time.Now().Add(sessionPolicy.Timeouts.Handshake)); err != nil {
diff --git a/proxy/vless/inbound/inbound.go b/proxy/vless/inbound/inbound.go
@@ -265,10 +265,7 @@ func (*Handler) Network() []net.Network {
 
 // Process implements proxy.Inbound.Process().
 func (h *Handler) Process(ctx context.Context, network net.Network, connection stat.Connection, dispatcher routing.Dispatcher) error {
-	iConn := connection
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(connection)
 
 	if h.decryption != nil {
 		var err error
diff --git a/proxy/vless/outbound/outbound.go b/proxy/vless/outbound/outbound.go
@@ -192,10 +192,7 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 
 	ob.Conn = conn // for Vision's pre-connect
 
-	iConn := conn
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(conn)
 	target := ob.Target
 	errors.LogInfo(ctx, "tunneling request to ", target, " via ", rec.Destination.NetAddr())
 
diff --git a/proxy/vmess/inbound/inbound.go b/proxy/vmess/inbound/inbound.go
@@ -229,10 +229,7 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
 		return errors.New("unable to set read deadline").Base(err).AtWarning()
 	}
 
-	iConn := connection
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
+	iConn := stat.TryUnwrapStatsConn(connection)
 	_, isDrain := iConn.(*net.TCPConn)
 	if !isDrain {
 		_, isDrain = iConn.(*net.UnixConn)
diff --git a/transport/internet/stat/connection.go b/transport/internet/stat/connection.go
@@ -32,3 +32,13 @@ func (c *CounterConnection) Write(b []byte) (int, error) {
 	}
 	return nBytes, err
 }
+
+func TryUnwrapStatsConn(conn net.Conn) net.Conn {
+	if conn == nil {
+		return conn
+	}
+	if conn, ok := conn.(*CounterConnection); ok {
+		return conn.Connection
+	}
+	return conn
+}

Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,8 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st`
`111`	`111`	`destinationOverridden = true`
`112`	`112`	`}`
`113`	`113`	`}`
`114`		`- if tlsConn, ok := conn.(tls.Interface); ok && !destinationOverridden {`
	`114`	`+ iConn := stat.TryUnwrapStatsConn(conn)`
	`115`	`+ if tlsConn, ok := iConn.(tls.Interface); ok && !destinationOverridden {`
`115`	`116`	`if serverName := tlsConn.HandshakeContextServerName(ctx); serverName != "" {`
`116`	`117`	`dest.Address = net.DomainAddress(serverName)`
`117`	`118`	`destinationOverridden = true`
Original file line number	Diff line number	Diff line change
`@@ -32,3 +32,13 @@ func (c *CounterConnection) Write(b []byte) (int, error) {`
`32`	`32`	`}`
`33`	`33`	`return nBytes, err`
`34`	`34`	`}`
	`35`	`+`
	`36`	`+func TryUnwrapStatsConn(conn net.Conn) net.Conn {`
	`37`	`+ if conn == nil {`
	`38`	`+ return conn`
	`39`	`+ }`
	`40`	`+ if conn, ok := conn.(*CounterConnection); ok {`
	`41`	`+ return conn.Connection`
	`42`	`+ }`
	`43`	`+ return conn`
	`44`	`+}`