From 5075b57370377e22958271df12d8703057c78a46 Mon Sep 17 00:00:00 2001
From: Fangliding <Fangliding.fshxy@outlook.com>
Date: Thu, 18 Dec 2025 02:47:00 +0800
Subject: [PATCH] Dokodemo: Fix stats conn unwrap

---
 proxy/dokodemo/dokodemo.go            |  3 ++-
 proxy/http/client.go                  |  5 +----
 proxy/proxy.go                        |  5 +----
 proxy/trojan/server.go                |  6 +-----
 proxy/vless/inbound/inbound.go        |  5 +----
 proxy/vless/outbound/outbound.go      |  5 +----
 proxy/vmess/inbound/inbound.go        |  5 +----
 transport/internet/stat/connection.go | 10 ++++++++++
 8 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index 90fc53e8..e56363db 100644
--- 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
index 9544ad37..1f804d0b 100644
--- 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
index d6a797cb..29548d9f 100644
--- 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
index 8ed3b0e6..d66219c4 100644
--- 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
index 89ed0e72..eeb1a25f 100644
--- 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
index 17e82210..a4fe1e93 100644
--- 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
index 7975551b..6a8591ad 100644
--- 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
index 6921943d..b039b29c 100644
--- 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
+}