Handle only IP

proxy/socks/protocol.go

@@ -210,11 +210,16 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer net.Co
 		responsePort = net.Port(udpHub.LocalAddr().(*net.UDPAddr).Port)
 		tempUDPConn = NewTempUDPConn(udpHub, writer)
 		if !(request.Address.IP().IsUnspecified() && request.Port == 0) {
-			var udpRemote gonet.Addr = &gonet.UDPAddr{
-				IP:   request.Address.IP(),
-				Port: int(request.Port),
+			// only specified an IP without port
+			if request.Port == 0 {
+				tempUDPConn.predefinedRemoteIP = request.Address.String()
+			} else { // specified both IP and port
+				var udpRemote gonet.Addr = &gonet.UDPAddr{
+					IP:   request.Address.IP(),
+					Port: int(request.Port),
+				}
+				tempUDPConn.remote.Store(&udpRemote)
 			}
-			tempUDPConn.remote.Store(&udpRemote)
 		}
 	}
 	if err := writeSocks5Response(writer, statusSuccess, responseAddress, responsePort); err != nil {

proxy/socks/temp_udp_listen.go

@@ -23,6 +23,8 @@ type TempUDPConn struct {
 	net.PacketConn
 	AssociateTCPConn net.Conn
 
+	predefinedRemoteIP string
+
 	timer  *signal.ActivityTimer
 	remote atomic.Pointer[net.Addr]
 }
@@ -35,9 +37,14 @@ func (c *TempUDPConn) Read(b []byte) (n int, err error) {
 			return n, err
 		}
 		if c.remote.Load() == nil {
-			tcpRemote, _, _ := net.SplitHostPort(c.AssociateTCPConn.RemoteAddr().String())
+			var expectedRemote string
+			if c.predefinedRemoteIP != "" {
+				expectedRemote = c.predefinedRemoteIP
+			} else {
+				expectedRemote, _, _ = net.SplitHostPort(c.AssociateTCPConn.RemoteAddr().String())
+			}
 			udpRemote, _, _ := net.SplitHostPort(remote.String())
-			if tcpRemote != udpRemote {
+			if expectedRemote != udpRemote {
 				continue
 			} else {
 				c.remote.CompareAndSwap(nil, &remote)