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

Fixes https://github.com/XTLS/Xray-core/issues/5439

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

proxy/http/client.go

@@ -296,10 +296,7 @@ func setUpHTTPTunnel(ctx context.Context, dest net.Destination, target string, u
 		return nil, err
 	}
 
-	iConn := rawConn
+	iConn := stat.TryUnwrapStatsConn(rawConn)
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
 
 	nextProto := ""
 	if tlsConn, ok := iConn.(*tls.Conn); ok {

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
+	iConn := stat.TryUnwrapStatsConn(conn)
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
 	_, ok1 := iConn.(*proxyproto.Conn)
 	_, ok2 := iConn.(*net.TCPConn)
 	_, ok3 := iConn.(*internet.UnixConnWrapper)

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
+	iConn := stat.TryUnwrapStatsConn(conn)
-	statConn, ok := iConn.(*stat.CounterConnection)
-	if ok {
-		iConn = statConn.Connection
-	}
 
 	sessionPolicy := s.policyManager.ForLevel(0)
 	if err := conn.SetReadDeadline(time.Now().Add(sessionPolicy.Timeouts.Handshake)); err != nil {

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
+	iConn := stat.TryUnwrapStatsConn(connection)
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
 
 	if h.decryption != nil {
 		var err error

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
+	iConn := stat.TryUnwrapStatsConn(conn)
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
 	target := ob.Target
 	errors.LogInfo(ctx, "tunneling request to ", target, " via ", rec.Destination.NetAddr())
 

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
+	iConn := stat.TryUnwrapStatsConn(connection)
-	if statConn, ok := iConn.(*stat.CounterConnection); ok {
-		iConn = statConn.Connection
-	}
 	_, isDrain := iConn.(*net.TCPConn)
 	if !isDrain {
 		_, isDrain = iConn.(*net.UnixConn)

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
+}