pr feedback

internal/socks5/socks5.go

@@ -119,7 +119,7 @@ func (c *socksConn) handleRequest(ctx context.Context) error {
 		}
 		return nil
 	case udpAssociate:
-		err = c.handleUDPAssociateRequest(ctx, socksVersion)
+		err = c.handleUDPAssociateRequest(ctx, socksVersion, request)
 		if err != nil {
 			return fmt.Errorf("handling %s request: %w", request.command, err)
 		}
@@ -196,8 +196,14 @@ func (c *socksConn) handleConnectRequest(ctx context.Context,
 }
 
 func (c *socksConn) handleUDPAssociateRequest(ctx context.Context,
-	socksVersion byte,
+	socksVersion byte, request request,
 ) error {
+	expectedAddrPort, err := udpAssociateExpectedClientEndpoint(request)
+	if err != nil {
+		c.encodeFailedResponse(c.clientConn, socksVersion, addressTypeNotSupported)
+		return fmt.Errorf("deriving expected client address and port from request: %w", err)
+	}
+
 	bindAddress, bindPort, bindAddrType, err := c.udpAssociationAddresses()
 	if err != nil {
 		c.encodeFailedResponse(c.clientConn, socksVersion, generalServerFailure)
@@ -211,7 +217,7 @@ func (c *socksConn) handleUDPAssociateRequest(ctx context.Context,
 		return fmt.Errorf("writing successful %s response: %w", udpAssociate, err)
 	}
 
-	association, err := c.udpRouter.registerAssociation(c.clientConn)
+	association, err := c.udpRouter.registerAssociation(c.clientConn, expectedAddrPort)
 	if err != nil {
 		c.encodeFailedResponse(c.clientConn, socksVersion, generalServerFailure)
 		return fmt.Errorf("registering udp association: %w", err)
@@ -236,6 +242,28 @@ func (c *socksConn) handleUDPAssociateRequest(ctx context.Context,
 	return nil
 }
 
+func udpAssociateExpectedClientEndpoint(request request) (expectedAddrPort netip.AddrPort, err error) {
+	switch request.addressType {
+	case ipv4, ipv6:
+		expectedClientAddress, parseErr := netip.ParseAddr(request.destination)
+		if parseErr != nil {
+			return netip.AddrPort{}, fmt.Errorf("parsing destination address: %w", parseErr)
+		}
+		expectedClientAddress = expectedClientAddress.Unmap()
+		if !expectedClientAddress.IsUnspecified() {
+			return netip.AddrPortFrom(expectedClientAddress, request.port), nil
+		}
+		return netip.AddrPortFrom(netip.Addr{}, request.port), nil
+	case domainName:
+		if request.destination != "" || request.port != 0 {
+			return netip.AddrPort{}, fmt.Errorf("domain name is not supported for UDP associate destination")
+		}
+		return netip.AddrPort{}, nil
+	default:
+		return netip.AddrPort{}, fmt.Errorf("address type %d is not supported", request.addressType)
+	}
+}
+
 func (c *socksConn) udpAssociationAddresses() (bindAddress string,
 	bindPort uint16, bindAddrType addrType, err error,
 ) {
@@ -250,6 +278,14 @@ func (c *socksConn) udpAssociationAddresses() (bindAddress string,
 	}
 	bindAddress = host
 	bindPort = uint16(port)
+	if isUnspecifiedIPAddress(bindAddress) {
+		controlLocalAddress := c.clientConn.LocalAddr().String()
+		controlLocalHost, _, splitErr := net.SplitHostPort(controlLocalAddress)
+		if splitErr != nil {
+			return "", 0, 0, fmt.Errorf("splitting control connection local address: %w", splitErr)
+		}
+		bindAddress = controlLocalHost
+	}
 
 	ipAddress := net.ParseIP(bindAddress)
 	if ipAddress == nil {
@@ -266,6 +302,14 @@ func (c *socksConn) udpAssociationAddresses() (bindAddress string,
 	return bindAddress, bindPort, bindAddrType, nil
 }
 
+func isUnspecifiedIPAddress(address string) bool {
+	ipAddress, err := netip.ParseAddr(address)
+	if err != nil {
+		return false
+	}
+	return ipAddress.IsUnspecified()
+}
+
 func decodeUDPDatagram(packet []byte) (destination string, payload []byte, err error) {
 	const minimumPacketLength = 4
 	if len(packet) < minimumPacketLength {
@@ -364,6 +408,8 @@ func writeUDPDatagramSourceAddress(writer io.Writer, address netip.Addr) error {
 		addrType = ipv6
 		array := address.As16()
 		addressBytes = array[:]
+	default:
+		return fmt.Errorf("address type is not supported: %v", address)
 	}
 
 	_, err := writer.Write([]byte{0, 0, 0, byte(addrType)})

internal/socks5/socks5_test.go

@@ -942,3 +942,80 @@ func Test_cmdType_String(t *testing.T) {
 		})
 	}
 }
+
+func Test_socksConn_udpAssociationAddresses(t *testing.T) {
+	t.Parallel()
+
+	testCases := map[string]struct {
+		routerAddress         string
+		expectAddressFromConn bool
+		expectedAddress       string
+	}{
+		"wildcard_router_address_uses_control_connection_local_ip": {
+			routerAddress:         ":0",
+			expectAddressFromConn: true,
+		},
+		"concrete_router_address_is_kept": {
+			routerAddress:   "127.0.0.1:0",
+			expectedAddress: "127.0.0.1",
+		},
+	}
+
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			router, err := newUDPRouter(t.Context(), testCase.routerAddress, noopLogger{})
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				err := router.close()
+				assert.NoError(t, err)
+			})
+
+			controlListener, err := (&net.ListenConfig{}).Listen(t.Context(), "tcp", "127.0.0.1:0")
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				err := controlListener.Close()
+				assert.NoError(t, err)
+			})
+
+			acceptedConnCh := make(chan net.Conn, 1)
+			go func() {
+				acceptedConn, acceptErr := controlListener.Accept()
+				if acceptErr != nil {
+					return
+				}
+				acceptedConnCh <- acceptedConn
+			}()
+
+			clientControlConn, err := (&net.Dialer{}).DialContext(t.Context(), "tcp", controlListener.Addr().String())
+			require.NoError(t, err)
+			defer clientControlConn.Close()
+
+			serverControlConn := <-acceptedConnCh
+			defer serverControlConn.Close()
+
+			socksConnection := &socksConn{
+				clientConn: clientControlConn,
+				udpRouter:  router,
+			}
+			bindAddress, bindPort, bindAddrType, err := socksConnection.udpAssociationAddresses()
+			require.NoError(t, err)
+
+			if testCase.expectAddressFromConn {
+				clientLocalHost, _, err := net.SplitHostPort(clientControlConn.LocalAddr().String())
+				require.NoError(t, err)
+				assert.Equal(t, clientLocalHost, bindAddress)
+			} else {
+				assert.Equal(t, testCase.expectedAddress, bindAddress)
+			}
+
+			_, routerPortString, err := net.SplitHostPort(router.localAddress().String())
+			require.NoError(t, err)
+			routerPort, err := strconv.ParseUint(routerPortString, 10, 16)
+			require.NoError(t, err)
+			assert.Equal(t, uint16(routerPort), bindPort)
+			assert.Equal(t, ipv4, bindAddrType)
+		})
+	}
+}

internal/socks5/udp_router.go

@@ -11,10 +11,11 @@ import (
 )
 
 type udpAssociation struct {
-	id              uint64
+	id               uint64
-	clientAddrPort  netip.AddrPort
+	clientAddrPort   netip.AddrPort
-	controlConnAddr netip.Addr
+	expectedAddrPort netip.AddrPort
-	packetCh        chan *bytes.Buffer
+	controlConnAddr  netip.Addr
+	packetCh         chan *bytes.Buffer
 }
 
 type udpRouter struct {
@@ -65,7 +66,7 @@ func (r *udpRouter) close() error {
 	return r.listener.Close()
 }
 
-func (r *udpRouter) registerAssociation(controlConn net.Conn) (udpAssociation, error) {
+func (r *udpRouter) registerAssociation(controlConn net.Conn, expectedAddrPort netip.AddrPort) (udpAssociation, error) {
 	controlConnAddrPort, err := netip.ParseAddrPort(controlConn.RemoteAddr().String())
 	if err != nil {
 		return udpAssociation{}, fmt.Errorf("parsing control connection address: %w", err)
@@ -80,9 +81,17 @@ func (r *udpRouter) registerAssociation(controlConn net.Conn) (udpAssociation, e
 	r.nextAssociationID++
 
 	association := udpAssociation{
-		id:              associationID,
+		id:               associationID,
-		controlConnAddr: controlConnAddr,
+		expectedAddrPort: expectedAddrPort,
-		packetCh:        make(chan *bytes.Buffer, udpPacketChannelBuffer),
+		controlConnAddr:  controlConnAddr,
+		packetCh:         make(chan *bytes.Buffer, udpPacketChannelBuffer),
+	}
+
+	if expectedAddrPort.Addr().IsValid() && expectedAddrPort.Port() != 0 {
+		association.clientAddrPort = expectedAddrPort
+		r.clientAddrPortToAssociation[association.clientAddrPort] = association
+		r.associationIDToClientAddrPort[association.id] = association.clientAddrPort
+		return association, nil
 	}
 
 	pendingAssociations := r.clientIPToPendingAssociations[controlConnAddr]
@@ -184,8 +193,19 @@ func (r *udpRouter) findClientAssociation(sourceAddrPort netip.AddrPort) (
 		return udpAssociation{}, false
 	}
 
-	association = pendingAssociations[0]
+	index := -1
-	r.clientIPToPendingAssociations[sourceAddr] = pendingAssociations[1:]
+	for i, pendingAssociation := range pendingAssociations {
+		if matchesExpectedClientEndpoint(pendingAssociation, sourceAddrPort) {
+			association = pendingAssociation
+			index = i
+			break
+		}
+	}
+	if index == -1 {
+		return udpAssociation{}, false
+	}
+
+	r.clientIPToPendingAssociations[sourceAddr] = append(pendingAssociations[:index], pendingAssociations[index+1:]...)
 	if len(r.clientIPToPendingAssociations[sourceAddr]) == 0 {
 		delete(r.clientIPToPendingAssociations, sourceAddr)
 	}
@@ -197,6 +217,16 @@ func (r *udpRouter) findClientAssociation(sourceAddrPort netip.AddrPort) (
 	return association, true
 }
 
+func matchesExpectedClientEndpoint(association udpAssociation, sourceAddrPort netip.AddrPort) bool {
+	switch {
+	case association.expectedAddrPort.Addr().IsValid() && sourceAddrPort.Addr() != association.expectedAddrPort.Addr():
+		return false
+	case association.expectedAddrPort.Port() != 0 && sourceAddrPort.Port() != association.expectedAddrPort.Port():
+		return false
+	}
+	return true
+}
+
 func (r *udpRouter) clientAddrPortForAssociation(associationID uint64) (
 	clientAddrPort netip.AddrPort, ok bool,
 ) {

internal/socks5/udp_router_integration_test.go

@@ -80,7 +80,7 @@ func Test_udpRouter_ResolveGithubFromCloudflareDNS(t *testing.T) {
 		assert.NoError(t, err, "closing server control connection")
 	})
 
-	association, err := router.registerAssociation(serverControlConn)
+	association, err := router.registerAssociation(serverControlConn, netip.AddrPort{})
 	require.NoError(t, err)
 	t.Cleanup(func() {
 		router.unregisterAssociation(association)