pr feedback

.vscode/settings.json

@@ -3,7 +3,7 @@
   // to develop this project.
   "files.eol": "\n",
   "editor.formatOnSave": true,
-  "go.buildTags": "linux",
+  "go.buildTags": "linux,integration",
   "go.toolsEnvVars": {
     "CGO_ENABLED": "0"
   },

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

@@ -13,6 +13,7 @@ import (
 type udpAssociation struct {
 	id               uint64
 	clientAddrPort   netip.AddrPort
+	expectedAddrPort netip.AddrPort
 	controlConnAddr  netip.Addr
 	packetCh         chan *bytes.Buffer
 }
@@ -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)
@@ -81,10 +82,18 @@ func (r *udpRouter) registerAssociation(controlConn net.Conn) (udpAssociation, e
 
 	association := udpAssociation{
 		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]
 	pendingAssociations = append(pendingAssociations, association)
 	r.clientIPToPendingAssociations[controlConnAddr] = pendingAssociations
@@ -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,
 ) {
@@ -209,7 +239,7 @@ func (r *udpRouter) clientAddrPortForAssociation(associationID uint64) (
 
 func (r *udpRouter) runAssociationHandler(ctx context.Context, association udpAssociation) {
 	config := &net.ListenConfig{}
-	socket, err := config.ListenPacket(ctx, "udp", "127.0.0.1:0")
+	socket, err := config.ListenPacket(ctx, "udp", ":0")
 	if err != nil {
 		r.logger.Warnf("creating per-association UDP socket: %s", err)
 		return

internal/socks5/udp_router_integration_test.go

@@ -0,0 +1,162 @@
+//go:build integration
+
+package socks5
+
+import (
+	"bytes"
+	"context"
+	"math/rand/v2"
+	"net"
+	"net/netip"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_udpRouter_ResolveGithubFromCloudflareDNS(t *testing.T) {
+	ctx := t.Context()
+	var cancel context.CancelFunc
+	deadline, hasDeadline := ctx.Deadline()
+	if hasDeadline {
+		const deadlineBuffer = 500 * time.Millisecond
+		deadline = deadline.Add(-deadlineBuffer)
+	} else {
+		const defaultTimeout = 10 * time.Second
+		deadline = time.Now().Add(defaultTimeout)
+	}
+	ctx, cancel = context.WithDeadline(ctx, deadline)
+
+	ctrl := gomock.NewController(t)
+	logger := NewMockLogger(ctrl)
+
+	router, err := newUDPRouter(ctx, "127.0.0.1:0", logger)
+	require.NoError(t, err)
+
+	routerRunErrCh := make(chan error)
+	go func() {
+		routerRunErrCh <- router.run(ctx)
+	}()
+
+	t.Cleanup(func() {
+		cancel()
+		err := router.close()
+		assert.NoError(t, err, "closing router")
+		runErr := <-routerRunErrCh
+		assert.NoError(t, runErr)
+	})
+
+	controlListener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		err := controlListener.Close()
+		assert.NoError(t, err, "closing control listener")
+	})
+
+	acceptedConnCh := make(chan net.Conn)
+	go func() {
+		acceptedConn, acceptErr := controlListener.Accept()
+		assert.NoError(t, acceptErr, "accepting control connection")
+		if acceptErr != nil {
+			return
+		}
+		acceptedConnCh <- acceptedConn
+	}()
+
+	clientControlConn, err := (&net.Dialer{}).DialContext(ctx, "tcp", controlListener.Addr().String())
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		err = clientControlConn.Close()
+		assert.NoError(t, err, "closing client control connection")
+	})
+
+	serverControlConn := <-acceptedConnCh
+	t.Cleanup(func() {
+		err := serverControlConn.Close()
+		assert.NoError(t, err, "closing server control connection")
+	})
+
+	association, err := router.registerAssociation(serverControlConn, netip.AddrPort{})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		router.unregisterAssociation(association)
+	})
+
+	associationCtx, associationCancel := context.WithCancel(ctx)
+	handlerDoneCh := make(chan struct{})
+	go func() {
+		router.runAssociationHandler(associationCtx, association)
+		close(handlerDoneCh)
+	}()
+	t.Cleanup(func() {
+		associationCancel()
+		<-handlerDoneCh
+	})
+
+	udpRouterAddress, err := net.ResolveUDPAddr("udp", router.localAddress().String())
+	require.NoError(t, err)
+
+	clientUDPConn, err := net.DialUDP("udp", nil, udpRouterAddress)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		err := clientUDPConn.Close()
+		assert.NoError(t, err, "closing client UDP connection")
+	})
+
+	queryID := uint16(rand.Uint())
+	dnsRequest := &dns.Msg{
+		MsgHdr: dns.MsgHdr{
+			Id:               queryID,
+			RecursionDesired: true,
+		},
+		Question: []dns.Question{{
+			Name:   dns.Fqdn("github.com"),
+			Qtype:  dns.TypeA,
+			Qclass: dns.ClassINET,
+		}},
+	}
+	dnsQuery, err := dnsRequest.Pack()
+	require.NoError(t, err)
+
+	targetAddrPort := netip.MustParseAddrPort("1.1.1.1:53")
+	socksDatagramBuffer := bytes.NewBuffer(nil)
+	err = encodeUDPDatagramToBuffer(socksDatagramBuffer, targetAddrPort, dnsQuery)
+	require.NoError(t, err)
+	socksDatagram := socksDatagramBuffer.Bytes()
+
+	err = clientUDPConn.SetDeadline(deadline)
+	require.NoError(t, err)
+
+	_, err = clientUDPConn.Write(socksDatagram)
+	require.NoError(t, err)
+
+	responseBuffer := make([]byte, maxUDPPacketLength)
+	responseLength, err := clientUDPConn.Read(responseBuffer)
+	require.NoError(t, err)
+
+	responseDestination, responsePayload, err := decodeUDPDatagram(responseBuffer[:responseLength])
+	require.NoError(t, err)
+
+	responseHost, responsePortString, err := net.SplitHostPort(responseDestination)
+	require.NoError(t, err)
+	responsePort, err := strconv.ParseUint(responsePortString, 10, 16)
+	require.NoError(t, err)
+	assert.Equal(t, uint64(53), responsePort)
+	assert.NotEmpty(t, responseHost)
+
+	dnsResponse := new(dns.Msg)
+	err = dnsResponse.Unpack(responsePayload)
+	require.NoError(t, err)
+	assert.Equal(t, queryID, dnsResponse.Id)
+	assert.True(t, dnsResponse.Response)
+	assert.Equal(t, dns.RcodeSuccess, dnsResponse.Rcode)
+	require.NotEmpty(t, dnsResponse.Question)
+	assert.Equal(t, dns.Fqdn("github.com"), dnsResponse.Question[0].Name)
+	assert.Equal(t, uint16(dns.TypeA), dnsResponse.Question[0].Qtype)
+	assert.NotEmpty(t, dnsResponse.Answer)
+	require.NoError(t, err)
+}