From 29f16cd054d141270c2fc42aa328dc571d32297a Mon Sep 17 00:00:00 2001
From: database64128 <free122448@hotmail.com>
Date: Wed, 16 Dec 2020 16:39:14 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=B3=20Refine=20socks5=20server=20UdpAs?=
 =?UTF-8?q?sociate=20response=20behavior=20(#523)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Previously, without specifying the server IP, the remote address in the response to a UdpAssoicate command is `127.0.0.1`, which might break UDP for non-localhost clients.
- This commit changes it so that, localhost clients get responses with the corresponding loopback IP, non-localhost clients get responses with the corresponding `net.AnyIP` or `net.AnyIPv6`.
- The new behavior is also consistent with many other implementations. So the compatibility is guaranteed. It also makes specifying server IP optional.
---
 proxy/socks/protocol.go | 25 ++++++++++++++++---------
 proxy/socks/server.go   |  6 ++++--
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/proxy/socks/protocol.go b/proxy/socks/protocol.go
index 9fb268c5d..39d21bc45 100644
--- a/proxy/socks/protocol.go
+++ b/proxy/socks/protocol.go
@@ -41,8 +41,10 @@ var addrParser = protocol.NewAddressParser(
 )
 
 type ServerSession struct {
-	config *ServerConfig
-	port   net.Port
+	config        *ServerConfig
+	address       net.Address
+	port          net.Port
+	clientAddress net.Address
 }
 
 func (s *ServerSession) handshake4(cmd byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
@@ -187,15 +189,20 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 	request.Address = addr
 	request.Port = port
 
-	responseAddress := net.AnyIP
-	responsePort := net.Port(1717)
+	responseAddress := s.address
+	responsePort := s.port
+	//nolint:gocritic // Use if else chain for clarity
 	if request.Command == protocol.RequestCommandUDP {
-		addr := s.config.Address.AsAddress()
-		if addr == nil {
-			addr = net.LocalHostIP
+		if s.config.Address != nil {
+			// Use configured IP as remote address in the response to UdpAssociate
+			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
 		}
-		responseAddress = addr
-		responsePort = s.port
 	}
 	if err := writeSocks5Response(writer, statusSuccess, responseAddress, responsePort); err != nil {
 		return nil, err
diff --git a/proxy/socks/server.go b/proxy/socks/server.go
index 3f421bb78..b398c8457 100644
--- a/proxy/socks/server.go
+++ b/proxy/socks/server.go
@@ -91,8 +91,10 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 	}
 
 	svrSession := &ServerSession{
-		config: s.config,
-		port:   inbound.Gateway.Port,
+		config:        s.config,
+		address:       inbound.Gateway.Address,
+		port:          inbound.Gateway.Port,
+		clientAddress: inbound.Source.Address,
 	}
 
 	reader := &buf.BufferedReader{Reader: buf.NewReader(conn)}