Socks5 server: More standard UDP ASSOCIATE (RFC 1928) (#6149)

https://github.com/XTLS/Xray-core/pull/6149#issuecomment-4529069218

Fixes https://github.com/XTLS/Xray-core/issues/6145#issuecomment-4467482623

---------

Co-authored-by: RPRX <63339210+RPRX@users.noreply.github.com>

proxy/socks/protocol.go

@@ -1,14 +1,17 @@
 package socks
 
 import (
+	"context"
 	"encoding/binary"
 	"io"
+	gonet "net"
 
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/protocol"
+	"github.com/xtls/xray-core/transport/internet"
 )
 
 const (
@@ -137,13 +140,13 @@ func (s *ServerSession) auth5(nMethod byte, reader io.Reader, writer io.Writer)
 	return "", nil
 }
 
-func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
+func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer net.Conn) (*protocol.RequestHeader, *TempUDPConn, error) {
 	var (
 		username string
 		err      error
 	)
 	if username, err = s.auth5(nMethod, reader, writer); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	var cmd byte
@@ -151,7 +154,7 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 		buffer := buf.StackNew()
 		if _, err := buffer.ReadFullFrom(reader, 3); err != nil {
 			buffer.Release()
-			return nil, errors.New("failed to read request").Base(err)
+			return nil, nil, errors.New("failed to read request").Base(err)
 		}
 		cmd = buffer.Byte(1)
 		buffer.Release()
@@ -168,28 +171,29 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 	case cmdUDPAssociate:
 		if !s.config.UdpEnabled {
 			writeSocks5Response(writer, statusCmdNotSupport, net.AnyIP, net.Port(0))
-			return nil, errors.New("UDP is not enabled.")
+			return nil, nil, errors.New("UDP is not enabled.")
 		}
 		request.Command = protocol.RequestCommandUDP
 	case cmdTCPBind:
 		writeSocks5Response(writer, statusCmdNotSupport, net.AnyIP, net.Port(0))
-		return nil, errors.New("TCP bind is not supported.")
+		return nil, nil, errors.New("TCP bind is not supported.")
 	default:
 		writeSocks5Response(writer, statusCmdNotSupport, net.AnyIP, net.Port(0))
-		return nil, errors.New("unknown command ", cmd)
+		return nil, nil, errors.New("unknown command ", cmd)
 	}
 
 	request.Version = socks5Version
 
 	addr, port, err := addrParser.ReadAddressPort(nil, reader)
 	if err != nil {
-		return nil, errors.New("failed to read address").Base(err)
+		return nil, nil, errors.New("failed to read address").Base(err)
 	}
 	request.Address = addr
 	request.Port = port
 
 	responseAddress := s.address
 	responsePort := s.port
+	var tempUDPConn *TempUDPConn
 	//nolint:gocritic // Use if else chain for clarity
 	if request.Command == protocol.RequestCommandUDP {
 		if s.config.Address != nil {
@@ -199,20 +203,34 @@ func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Wri
 			// Use conn.LocalAddr() IP as remote address in the response by default
 			responseAddress = s.localAddress
 		}
+		udpHub, err := internet.ListenSystemPacket(context.Background(), &net.UDPAddr{IP: responseAddress.IP(), Port: 0}, nil)
+		if err != nil {
+			return nil, nil, errors.New("failed to create UDP listener").Base(err)
+		}
+		responsePort = net.Port(udpHub.LocalAddr().(*net.UDPAddr).Port)
+		expectedRemote := &gonet.UDPAddr{}
+		if request.Address.IP().IsUnspecified() {
+			expectedRemote.IP = writer.RemoteAddr().(*net.TCPAddr).IP // unix?
+		} else {
+			expectedRemote.IP = request.Address.IP() // panic?
+			expectedRemote.Port = int(request.Port)  // 0 is allowed
+		}
+		tempUDPConn = NewTempUDPConn(udpHub, writer, expectedRemote)
 	}
 	if err := writeSocks5Response(writer, statusSuccess, responseAddress, responsePort); err != nil {
-		return nil, err
+		common.CloseIfExists(tempUDPConn)
+		return nil, nil, err
 	}
 
-	return request, nil
+	return request, tempUDPConn, nil
 }
 
 // Handshake performs a Socks4/4a/5 handshake.
-func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
+func (s *ServerSession) Handshake(reader io.Reader, writer net.Conn) (*protocol.RequestHeader, *TempUDPConn, error) {
 	buffer := buf.StackNew()
 	if _, err := buffer.ReadFullFrom(reader, 2); err != nil {
 		buffer.Release()
-		return nil, errors.New("insufficient header").Base(err)
+		return nil, nil, errors.New("insufficient header").Base(err)
 	}
 
 	version := buffer.Byte(0)
@@ -221,11 +239,12 @@ func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol
 
 	switch version {
 	case socks4Version:
-		return s.handshake4(cmd, reader, writer)
+		header, err := s.handshake4(cmd, reader, writer)
+		return header, nil, err
 	case socks5Version:
 		return s.handshake5(cmd, reader, writer)
 	default:
-		return nil, errors.New("unknown Socks version: ", version)
+		return nil, nil, errors.New("unknown Socks version: ", version)
 	}
 }