From 14a4b0f922f8fe5d0cfa128da30c5f566f51382e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Tue, 10 Dec 2024 20:36:09 +0800
Subject: [PATCH] Fix socks5 UDP implementation

---
 protocol/http/inbound.go  |  2 +-
 protocol/mixed/inbound.go |  4 ++--
 protocol/socks/inbound.go |  2 +-
 protocol/tor/proxy.go     |  2 +-
 route/route.go            | 11 +++++++----
 5 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/protocol/http/inbound.go b/protocol/http/inbound.go
index 65675bcf..92fd22b0 100644
--- a/protocol/http/inbound.go
+++ b/protocol/http/inbound.go
@@ -91,7 +91,7 @@ func (h *Inbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata a
 			return
 		}
 	}
-	err = http.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), h.authenticator, nil, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
+	err = http.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), h.authenticator, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
 	if err != nil {
 		N.CloseOnHandshakeFailure(conn, onClose, err)
 		h.logger.ErrorContext(ctx, E.Cause(err, "process connection from ", metadata.Source))
diff --git a/protocol/mixed/inbound.go b/protocol/mixed/inbound.go
index d1d343a9..f7ebe043 100644
--- a/protocol/mixed/inbound.go
+++ b/protocol/mixed/inbound.go
@@ -85,9 +85,9 @@ func (h *Inbound) newConnection(ctx context.Context, conn net.Conn, metadata ada
 	}
 	switch headerBytes[0] {
 	case socks4.Version, socks5.Version:
-		return socks.HandleConnectionEx(ctx, conn, reader, h.authenticator, nil, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, metadata.Destination, onClose)
+		return socks.HandleConnectionEx(ctx, conn, reader, h.authenticator, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
 	default:
-		return http.HandleConnectionEx(ctx, conn, reader, h.authenticator, nil, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
+		return http.HandleConnectionEx(ctx, conn, reader, h.authenticator, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
 	}
 }
 
diff --git a/protocol/socks/inbound.go b/protocol/socks/inbound.go
index 1490cee5..36e45e61 100644
--- a/protocol/socks/inbound.go
+++ b/protocol/socks/inbound.go
@@ -62,7 +62,7 @@ func (h *Inbound) Close() error {
 }
 
 func (h *Inbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) {
-	err := socks.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), h.authenticator, nil, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, metadata.Destination, onClose)
+	err := socks.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), h.authenticator, adapter.NewUpstreamHandlerEx(metadata, h.newUserConnection, h.streamUserPacketConnection), metadata.Source, onClose)
 	N.CloseOnHandshakeFailure(conn, onClose, err)
 	if err != nil {
 		if E.IsClosedOrCanceled(err) {
diff --git a/protocol/tor/proxy.go b/protocol/tor/proxy.go
index 1ed30335..feab7971 100644
--- a/protocol/tor/proxy.go
+++ b/protocol/tor/proxy.go
@@ -99,7 +99,7 @@ func (l *ProxyListener) acceptLoop() {
 }
 
 func (l *ProxyListener) accept(ctx context.Context, conn *net.TCPConn) error {
-	return socks.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), l.authenticator, nil, l, M.SocksaddrFromNet(conn.RemoteAddr()), M.Socksaddr{}, nil)
+	return socks.HandleConnectionEx(ctx, conn, std_bufio.NewReader(conn), l.authenticator, l, M.SocksaddrFromNet(conn.RemoteAddr()), nil)
 }
 
 func (l *ProxyListener) NewConnectionEx(ctx context.Context, conn net.Conn, source M.Socksaddr, destination M.Socksaddr, onClose N.CloseHandlerFunc) {
diff --git a/route/route.go b/route/route.go
index 05e22c25..fb2de85d 100644
--- a/route/route.go
+++ b/route/route.go
@@ -461,8 +461,12 @@ match:
 			break match
 		}
 	}
-	if !preMatch && metadata.Destination.Addr.IsUnspecified() {
-		newBuffer, newPacketBuffers, newErr := r.actionSniff(ctx, metadata, &rule.RuleActionSniff{}, inputConn, inputPacketConn)
+	if !preMatch && inputPacketConn != nil && !metadata.Destination.IsFqdn() && !metadata.Destination.Addr.IsGlobalUnicast() {
+		var timeout time.Duration
+		if metadata.InboundType == C.TypeSOCKS {
+			timeout = C.TCPTimeout
+		}
+		newBuffer, newPacketBuffers, newErr := r.actionSniff(ctx, metadata, &rule.RuleActionSniff{Timeout: timeout}, inputConn, inputPacketConn)
 		if newErr != nil {
 			fatalErr = newErr
 			return
@@ -558,8 +562,7 @@ func (r *Router) actionSniff(
 					return
 				}
 			} else {
-				// TODO: maybe always override destination
-				if metadata.Destination.Addr.IsUnspecified() {
+				if !metadata.Destination.Addr.IsGlobalUnicast() {
 					metadata.Destination = destination
 				}
 				if len(packetBuffers) > 0 {