predefine

proxy/socks/protocol.go

@@ -208,11 +208,12 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer net.Co
 			return nil, nil, errors.New("failed to create UDP listener").Base(err)
 		}
 		responsePort = net.Port(udpHub.LocalAddr().(*net.UDPAddr).Port)
-		tempUDPConn = NewTempUDPConn(udpHub, writer)
+		expectedRemoteIP, _, _ := net.SplitHostPort(writer.RemoteAddr().String())
+		tempUDPConn = NewTempUDPConn(udpHub, writer, expectedRemoteIP)
 		if !(request.Address.IP().IsUnspecified() && request.Port == 0) {
 			// only specified an IP without port
 			if request.Port == 0 {
-				tempUDPConn.predefinedRemoteIP = request.Address.String()
+				tempUDPConn.ExpectedRemoteIP = request.Address.String()
 			} else { // specified both IP and port
 				var udpRemote gonet.Addr = &gonet.UDPAddr{
 					IP:   request.Address.IP(),

proxy/socks/temp_udp_listen.go

@@ -10,10 +10,11 @@ import (
 	"github.com/xtls/xray-core/common/signal"
 )
 
-func NewTempUDPConn(udpConn net.PacketConn, tcpConn net.Conn) *TempUDPConn {
+func NewTempUDPConn(udpConn net.PacketConn, tcpConn net.Conn, remoteIP string) *TempUDPConn {
 	return &TempUDPConn{
 		PacketConn:       udpConn,
 		AssociateTCPConn: tcpConn,
+		ExpectedRemoteIP: remoteIP,
 	}
 }
 
@@ -22,8 +23,7 @@ func NewTempUDPConn(udpConn net.PacketConn, tcpConn net.Conn) *TempUDPConn {
 type TempUDPConn struct {
 	net.PacketConn
 	AssociateTCPConn net.Conn
-
-	predefinedRemoteIP string
+	ExpectedRemoteIP string
 
 	timer  *signal.ActivityTimer
 	remote atomic.Pointer[net.Addr]
@@ -37,14 +37,8 @@ func (c *TempUDPConn) Read(b []byte) (n int, err error) {
 			return n, err
 		}
 		if c.remote.Load() == nil {
-			var expectedRemote string
-			if c.predefinedRemoteIP != "" {
-				expectedRemote = c.predefinedRemoteIP
-			} else {
-				expectedRemote, _, _ = net.SplitHostPort(c.AssociateTCPConn.RemoteAddr().String())
-			}
 			udpRemote, _, _ := net.SplitHostPort(remote.String())
-			if expectedRemote != udpRemote {
+			if c.ExpectedRemoteIP != udpRemote {
 				continue
 			} else {
 				c.remote.CompareAndSwap(nil, &remote)