From 700966508f5c2c5f4a83eb611c8c557f20cd83a0 Mon Sep 17 00:00:00 2001
From: RPRX <63339210+rprx@users.noreply.github.com>
Date: Sat, 9 Jan 2021 16:36:20 +0000
Subject: [PATCH] Improve the response to UDP Associate in Socks5

---
 proxy/socks/protocol.go | 17 +++++++----------
 proxy/socks/server.go   |  8 ++++----
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/proxy/socks/protocol.go b/proxy/socks/protocol.go
index e35f9ad6..9086e519 100644
--- a/proxy/socks/protocol.go
+++ b/proxy/socks/protocol.go
@@ -39,10 +39,10 @@ var addrParser = protocol.NewAddressParser(
 )
 
 type ServerSession struct {
-	config        *ServerConfig
-	address       net.Address
-	port          net.Port
-	clientAddress net.Address
+	config       *ServerConfig
+	address      net.Address
+	port         net.Port
+	localAddress net.Address
 }
 
 func (s *ServerSession) handshake4(cmd byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
@@ -192,14 +192,11 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 	//nolint:gocritic // Use if else chain for clarity
 	if request.Command == protocol.RequestCommandUDP {
 		if s.config.Address != nil {
-			// Use configured IP as remote address in the response to UdpAssociate
+			// Use configured IP as remote address in the response to UDP Associate
 			responseAddress = s.config.Address.AsAddress()
-		} else if s.clientAddress == net.LocalHostIP || s.clientAddress == net.LocalHostIPv6 {
-			// For localhost clients use loopback IP
-			responseAddress = s.clientAddress
 		} else {
-			// For non-localhost clients use inbound listening address
-			responseAddress = s.address
+			// Use conn.LocalAddr() IP as remote address in the response by default
+			responseAddress = s.localAddress
 		}
 	}
 	if err := writeSocks5Response(writer, statusSuccess, responseAddress, responsePort); err != nil {
diff --git a/proxy/socks/server.go b/proxy/socks/server.go
index dd142a72..626a9654 100644
--- a/proxy/socks/server.go
+++ b/proxy/socks/server.go
@@ -89,10 +89,10 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 	}
 
 	svrSession := &ServerSession{
-		config:        s.config,
-		address:       inbound.Gateway.Address,
-		port:          inbound.Gateway.Port,
-		clientAddress: inbound.Source.Address,
+		config:       s.config,
+		address:      inbound.Gateway.Address,
+		port:         inbound.Gateway.Port,
+		localAddress: net.IPAddress(conn.LocalAddr().(*net.TCPAddr).IP),
 	}
 
 	reader := &buf.BufferedReader{Reader: buf.NewReader(conn)}