Adding UDP support to Socks5 dialer

transport/socks5/packet_listener.go

@@ -0,0 +1,183 @@
+// Copyright 2024 Jigsaw Operations LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package socks5
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"time"
+
+	"github.com/Jigsaw-Code/outline-sdk/internal/slicepool"
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+)
+
+// clientUDPBufferSize is the maximum supported UDP packet size in bytes.
+const clientUDPBufferSize = 16 * 1024
+
+// udpPool stores the byte slices used for storing packets.
+var udpPool = slicepool.MakePool(clientUDPBufferSize)
+
+type packetConn struct {
+	pc net.Conn
+	sc io.Closer
+}
+
+var _ net.PacketConn = (*packetConn)(nil)
+
+func (p *packetConn) LocalAddr() net.Addr {
+	return p.pc.LocalAddr()
+}
+
+func (p *packetConn) SetDeadline(t time.Time) error {
+	return p.pc.SetDeadline(t)
+}
+
+func (p *packetConn) SetReadDeadline(t time.Time) error {
+	return p.pc.SetReadDeadline(t)
+}
+
+func (c *packetConn) SetWriteDeadline(t time.Time) error {
+	return c.pc.SetWriteDeadline(t)
+}
+
+// ReadFrom reads the packet from the SOCKS5 server and extract the payload
+// The packet format is specified in https://datatracker.ietf.org/doc/html/rfc1928#section-7
+func (p *packetConn) ReadFrom(b []byte) (int, net.Addr, error) {
+	lazySlice := udpPool.LazySlice()
+	buffer := lazySlice.Acquire()
+	defer lazySlice.Release()
+
+	n, err := p.pc.Read(buffer)
+	if err != nil {
+		return 0, nil, err
+	}
+	// Minimum packet size
+	if n < 10 {
+		return 0, nil, errors.New("invalid SOCKS5 UDP packet: too short")
+	}
+
+	// Using bytes.Buffer to handle data
+	buf := bytes.NewBuffer(buffer[:n])
+
+	// Read and check reserved bytes
+	rsv := make([]byte, 2)
+	if _, err := buf.Read(rsv); err != nil {
+		return 0, nil, err
+	}
+	if rsv[0] != 0x00 || rsv[1] != 0x00 {
+		return 0, nil, fmt.Errorf("invalid reserved bytes: expected 0x0000, got %#x%#x", rsv[0], rsv[1])
+	}
+
+	// Read fragment byte
+	frag, err := buf.ReadByte()
+	if err != nil {
+		return 0, nil, err
+	}
+	if frag != 0 {
+		return 0, nil, errors.New("fragmentation is not supported")
+	}
+
+	// Read address using socks.ReadAddr which must now accept a bytes.Buffer directly
+	address, err := readAddr(buf)
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed to read address: %w", err)
+	}
+
+	// Convert the address to a net.Addr
+	addr, err := transport.MakeNetAddr("udp", addrToString(address))
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed to convert address: %w", err)
+	}
+
+	// Payload handling: remaining bytes in the buffer are the payload
+	payload := buf.Bytes()
+	payloadLength := len(payload)
+	if payloadLength > len(b) {
+		return 0, nil, io.ErrShortBuffer
+	}
+	copy(b, payload)
+
+	return payloadLength, addr, nil
+}
+
+// WriteTo encapsulates the payload in a SOCKS5 UDP packet as specified in
+// https://datatracker.ietf.org/doc/html/rfc1928#section-7
+// and write it to the SOCKS5 server via the underlying connection.
+func (p *packetConn) WriteTo(b []byte, addr net.Addr) (int, error) {
+
+	// The minimum preallocated header size (10 bytes)
+	lazySlice := udpPool.LazySlice()
+	buffer := lazySlice.Acquire()
+	defer lazySlice.Release()
+	buffer = append(buffer[:0],
+		0x00, 0x00, // Reserved
+		0x00, // Fragment number
+		// To be appended below:
+		// ATYP, IPv4, IPv6, Domain Name, Port
+	)
+	buffer, err := appendSOCKS5Address(buffer, addr.String())
+	if err != nil {
+		return 0, fmt.Errorf("failed to append SOCKS5 address: %w", err)
+	}
+	// Combine the header and the payload
+	return p.pc.Write(append(buffer, b...))
+}
+
+// Close closes both the underlying stream and packet connections.
+func (p *packetConn) Close() error {
+	return errors.Join(p.sc.Close(), p.pc.Close())
+}
+
+// ListenPacket creates a [net.PacketConn] for dialing to SOCKS5 server.
+func (c *Client) ListenPacket(ctx context.Context) (net.PacketConn, error) {
+	// Connect to the SOCKS5 server and perform UDP association
+	// Since local address is not known in advance, we use unspecified address
+	// which means the server is going to accept incoming packets from any address
+	// on the bind port on the server. The bind address is determined and returned by
+	// the server.
+	// https://datatracker.ietf.org/doc/html/rfc1928#section-6
+	// Whoile binding address to specific client address has its advantages, it also creates some
+	// challenges such as NAT traveral if client is behind NAT.
+	sc, bindAddr, err := c.connectAndRequest(ctx, CmdUDPAssociate, "0.0.0.0:0")
+	if err != nil {
+		return nil, err
+	}
+
+	// If the returned bind IP address is unspecified (i.e. "0.0.0.0" or "::"),
+	// then use the IP address of the SOCKS5 server
+	if ipAddr := bindAddr.IP; ipAddr.IsValid() && ipAddr.IsUnspecified() {
+		schost, _, err := net.SplitHostPort(sc.RemoteAddr().String())
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse tcp address: %w", err)
+		}
+
+		bindAddr.IP, err = netip.ParseAddr(schost)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse bind address: %w", err)
+		}
+	}
+
+	proxyConn, err := c.pd.DialPacket(ctx, addrToString(bindAddr))
+	if err != nil {
+		sc.Close()
+		return nil, fmt.Errorf("could not connect to packet endpoint: %w", err)
+	}
+	return &packetConn{pc: proxyConn, sc: sc}, nil
+}

transport/socks5/packet_listener_test.go

@@ -0,0 +1,114 @@
+package socks5
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/things-go/go-socks5"
+)
+
+func TestSOCKS5Associate(t *testing.T) {
+	// Create a local listener.
+	// This creates a UDP server that responded to "ping"
+	// message with "pong" response.
+	locIP := net.ParseIP("127.0.0.1")
+	// Create a local listener
+	echoServerAddr := &net.UDPAddr{IP: locIP, Port: 0}
+	echoServer := setupUDPEchoServer(t, echoServerAddr)
+	defer echoServer.Close()
+
+	// Create a socks server to proxy "ping" message.
+	cator := socks5.UserPassAuthenticator{Credentials: socks5.StaticCredentials{
+		"testusername": "testpassword",
+	}}
+	proxySrv := socks5.NewServer(
+		socks5.WithAuthMethods([]socks5.Authenticator{cator}),
+	)
+
+	// Create SOCKS5 proxy on localhost with a random port.
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer listener.Close()
+	proxyServerAddress := listener.Addr().String()
+
+	go func() {
+		err := proxySrv.Serve(listener)
+		if !errors.Is(err, net.ErrClosed) && err != nil {
+			require.NoError(t, err) // Assert no error if it's not the expected close error
+		}
+	}()
+
+	// Connect to local proxy, auth and start the PacketConn.
+	client, err := NewClient(&transport.TCPEndpoint{Address: proxyServerAddress})
+	require.NotNil(t, client)
+	require.NoError(t, err)
+	err = client.SetCredentials([]byte("testusername"), []byte("testpassword"))
+	require.NoError(t, err)
+	client.EnablePacket(&transport.UDPDialer{})
+	conn, err := client.ListenPacket(context.Background())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	// Send "ping" message.
+	_, err = conn.WriteTo([]byte("ping"), echoServer.LocalAddr())
+	require.NoError(t, err)
+	// Max wait time for response.
+	err = conn.SetDeadline(time.Now().Add(time.Second))
+	require.NoError(t, err)
+	response := make([]byte, 1024)
+	n, addr, err := conn.ReadFrom(response)
+	require.Equal(t, echoServer.LocalAddr().String(), addr.String())
+	require.NoError(t, err)
+	require.Equal(t, []byte("pong"), response[:n])
+}
+
+func TestUDPLoopBack(t *testing.T) {
+	// Create a local listener.
+	locIP := net.ParseIP("127.0.0.1")
+	echoServerAddr := &net.UDPAddr{IP: locIP, Port: 0}
+	echoServer := setupUDPEchoServer(t, echoServerAddr)
+	defer echoServer.Close()
+
+	packDialer := transport.UDPDialer{}
+	conn, err := packDialer.DialPacket(context.Background(), echoServer.LocalAddr().String())
+	require.NoError(t, err)
+	_, err = conn.Write([]byte("ping"))
+	require.NoError(t, err)
+	response := make([]byte, 1024)
+	n, err := conn.Read(response)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("pong"), response[:n])
+}
+
+func setupUDPEchoServer(t *testing.T, serverAddr *net.UDPAddr) *net.UDPConn {
+	server, err := net.ListenUDP("udp", serverAddr)
+	require.NoError(t, err)
+	go func() {
+		buf := make([]byte, 2048)
+		for {
+			n, remote, err := server.ReadFrom(buf)
+			if err != nil {
+				return
+			}
+			if bytes.Equal(buf[:n], []byte("ping")) {
+				_, err := server.WriteTo([]byte("pong"), remote)
+				if err != nil {
+					return
+				}
+			}
+		}
+	}()
+
+	t.Cleanup(func() {
+		server.Close()
+	})
+
+	return server
+}

transport/socks5/socks5.go

@@ -18,7 +18,9 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"io"
 	"net"
+	"net/netip"
 	"strconv"
 )
 
@@ -37,6 +39,13 @@ const (
 	ErrAddressTypeNotSupported       = ReplyCode(0x08)
 )
 
+// SOCKS5 commands, from https://datatracker.ietf.org/doc/html/rfc1928#section-4.
+const (
+	CmdConnect      = byte(1)
+	CmdBind         = byte(2)
+	CmdUDPAssociate = byte(3)
+)
+
 // SOCKS5 authentication methods, as specified in https://datatracker.ietf.org/doc/html/rfc1928#section-3
 const (
 	authMethodNoAuth   = 0x00
@@ -79,6 +88,27 @@ const (
 	addrTypeIPv6 = 0x04
 )
 
+// address is a SOCKS-specific address.
+// Either Name or IP is used exclusively.
+type address struct {
+	Name string // fully-qualified domain name
+	IP   netip.Addr
+	Port uint16
+}
+
+// Address returns a string suitable to dial; prefer returning IP-based
+// address, fallback to Name
+func addrToString(a *address) string {
+	if a == nil {
+		return ""
+	}
+	port := strconv.Itoa(int(a.Port))
+	if a.IP.IsValid() {
+		return net.JoinHostPort(a.IP.String(), port)
+	}
+	return net.JoinHostPort(a.Name, port)
+}
+
 // appendSOCKS5Address adds the address to buffer b in SOCKS5 format,
 // as specified in https://datatracker.ietf.org/doc/html/rfc1928#section-4
 func appendSOCKS5Address(b []byte, address string) ([]byte, error) {
@@ -119,3 +149,47 @@ func appendSOCKS5Address(b []byte, address string) ([]byte, error) {
 	b = binary.BigEndian.AppendUint16(b, uint16(portNum))
 	return b, nil
 }
+
+func readAddr(r io.Reader) (*address, error) {
+	address := &address{}
+
+	var addrType [1]byte
+	if _, err := r.Read(addrType[:]); err != nil {
+		return nil, err
+	}
+
+	switch addrType[0] {
+	case addrTypeIPv4:
+		var addr [4]byte
+		if _, err := io.ReadFull(r, addr[:]); err != nil {
+			return nil, err
+		}
+		address.IP = netip.AddrFrom4(addr)
+	case addrTypeIPv6:
+		var addr [16]byte
+		if _, err := io.ReadFull(r, addr[:]); err != nil {
+			return nil, err
+		}
+		address.IP = netip.AddrFrom16(addr)
+	case addrTypeDomainName:
+		if _, err := r.Read(addrType[:]); err != nil {
+			return nil, err
+		}
+		addrLen := addrType[0]
+		// addrLen btye type maximum value is 255 which
+		// prevents passing larger then 255 values for domain names.
+		fqdn := make([]byte, addrLen)
+		if _, err := io.ReadFull(r, fqdn); err != nil {
+			return nil, err
+		}
+		address.Name = string(fqdn)
+	default:
+		return nil, errors.New("unrecognized address type")
+	}
+	var port [2]byte
+	if _, err := io.ReadFull(r, port[:]); err != nil {
+		return nil, err
+	}
+	address.Port = binary.BigEndian.Uint16(port[:])
+	return address, nil
+}