From 1f10aff3e834c6891255452967d4d87259ef8131 Mon Sep 17 00:00:00 2001 From: Quentin McGaw Date: Thu, 28 May 2026 14:43:20 +0000 Subject: [PATCH] pr feedback --- internal/socks5/socks5.go | 52 ++++++++++++- internal/socks5/socks5_test.go | 77 +++++++++++++++++++ internal/socks5/udp_router.go | 50 +++++++++--- .../socks5/udp_router_integration_test.go | 2 +- 4 files changed, 167 insertions(+), 14 deletions(-) diff --git a/internal/socks5/socks5.go b/internal/socks5/socks5.go index 3a4a4e81..7230c7fa 100644 --- a/internal/socks5/socks5.go +++ b/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)}) diff --git a/internal/socks5/socks5_test.go b/internal/socks5/socks5_test.go index dbe2678a..3844805c 100644 --- a/internal/socks5/socks5_test.go +++ b/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) + }) + } +} diff --git a/internal/socks5/udp_router.go b/internal/socks5/udp_router.go index 4d0f3308..83f86c60 100644 --- a/internal/socks5/udp_router.go +++ b/internal/socks5/udp_router.go @@ -11,10 +11,11 @@ import ( ) type udpAssociation struct { - id uint64 - clientAddrPort netip.AddrPort - controlConnAddr netip.Addr - packetCh chan *bytes.Buffer + id uint64 + clientAddrPort netip.AddrPort + expectedAddrPort netip.AddrPort + 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, - controlConnAddr: controlConnAddr, - packetCh: make(chan *bytes.Buffer, udpPacketChannelBuffer), + id: associationID, + expectedAddrPort: expectedAddrPort, + 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] - r.clientIPToPendingAssociations[sourceAddr] = pendingAssociations[1:] + index := -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, ) { diff --git a/internal/socks5/udp_router_integration_test.go b/internal/socks5/udp_router_integration_test.go index c19c155c..cecfa506 100644 --- a/internal/socks5/udp_router_integration_test.go +++ b/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)