Avoid per-datagram SOCKS5 reply address allocations

native/mobile/lifecycle_regression_test.go

@@ -115,27 +115,42 @@ type allocationPacketSocket struct {
 
 func (c *allocationPacketSocket) Read(b []byte) (int, error) { return copy(b, c.packet), nil }
 func BenchmarkSOCKS5ReadAllocation(b *testing.B) {
-	target := S.ParseAddrString("192.0.2.1:443")
-	packet, err := S.EncodeUDPPacket(target, make([]byte, 1200))
+	benchmarkSOCKS5Read(b, "192.0.2.1:443", "192.0.2.1:443")
+}
+
+func BenchmarkSOCKS5ReadAllocationPeers(b *testing.B) {
+	for _, tc := range []struct{ name, target, reply string }{
+		{"ipv4", "192.0.2.1:443", "192.0.2.1:443"},
+		{"ipv6", "[2001:db8::1]:443", "[2001:db8::1]:443"},
+		{"mapped", "192.0.2.1:443", "[::ffff:192.0.2.1]:443"},
+		{"domain-ipv4", "origin.test:443", "192.0.2.1:443"},
+		{"domain-ipv6", "origin.test:443", "[2001:db8::1]:443"},
+	} {
+		b.Run(tc.name, func(b *testing.B) { benchmarkSOCKS5Read(b, tc.target, tc.reply) })
+	}
+}
+
+func benchmarkSOCKS5Read(b *testing.B, target, reply string) {
+	packet, err := S.EncodeUDPPacket(S.ParseAddrString(reply), make([]byte, 1200))
 	if err != nil {
 		b.Fatal(err)
 	}
-	p := &socks5PacketConn{Conn: &allocationPacketSocket{packet: packet}, target: target, peer: &net.UDPAddr{IP: net.IPv4(192, 0, 2, 1), Port: 443}}
+	p := &socks5PacketConn{Conn: &allocationPacketSocket{packet: packet}, target: S.ParseAddrString(target), peer: &net.UDPAddr{IP: net.IPv4(192, 0, 2, 1), Port: 443}}
 	buffer := make([]byte, 1400)
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		_, _, err := p.ReadFrom(buffer)
-		if err != nil && !errors.Is(err, io.EOF) {
-			b.Fatal(err)
+		n, peer, err := p.ReadFrom(buffer)
+		if err != nil || n != 1200 || peer != p.peer {
+			b.Fatalf("UDP read: n=%d peer=%v err=%v", n, peer, err)
 		}
 	}
 }
 
 func TestSOCKS5ReadAllocationBudget(t *testing.T) {
 	result := testing.Benchmark(BenchmarkSOCKS5ReadAllocation)
-	if bytes := result.AllocedBytesPerOp(); bytes > 2048 {
-		t.Fatalf("UDP read allocates %d bytes per datagram", bytes)
+	if result.AllocedBytesPerOp() != 0 || result.AllocsPerOp() != 0 {
+		t.Fatalf("UDP read allocates %d bytes / %d objects per datagram", result.AllocedBytesPerOp(), result.AllocsPerOp())
 	}
 }
 

native/mobile/socks5.go

@@ -1,7 +1,9 @@
 package mobile
 
 import (
+	"bytes"
 	"context"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -367,10 +369,20 @@ func (p *socks5PacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
 		}
 		// An IP target is pinned to its original IP and port. Domain targets are
 		// resolved at the proxy and use a virtual peer for the diagnostic QUIC client.
-		if p.target[0] != S.AtypDomainName && addr.String() != p.target.String() {
-			continue
+		// Equal wire addresses need no string allocations. Keep the textual
+		// fallback for equivalent IPv4 and IPv4-mapped IPv6 encodings.
+		if p.target[0] != S.AtypDomainName {
+			if addr[0] == p.target[0] {
+				if !bytes.Equal(addr, p.target) {
+					continue
+				}
+			} else if addr.String() != p.target.String() {
+				continue
+			}
 		}
-		if p.target[0] == S.AtypDomainName && (addr.UDPAddr() == nil || addr.UDPAddr().Port != p.peer.Port) {
+		// DecodeUDPPacket has already validated and sliced the whole address.
+		// Domain replies remain unsupported; numeric replies must keep the port.
+		if p.target[0] == S.AtypDomainName && (addr[0] == S.AtypDomainName || int(binary.BigEndian.Uint16(addr[len(addr)-2:])) != p.peer.Port) {
 			continue
 		}
 		copied := copy(b, payload)

native/mobile/socks5_test.go

@@ -492,3 +492,72 @@ func TestSOCKS5BypassedTargetResetIsNotProxyFailure(t *testing.T) {
 		t.Fatalf("local target must not trigger VPN failover: %s", logs.text())
 	}
 }
+
+// A finite packet source makes a dropped reply observable as EOF, without a network timeout.
+type socks5ReplySocket struct {
+	net.Conn
+	packet []byte
+}
+
+func (c *socks5ReplySocket) Read(p []byte) (int, error) {
+	if c.packet == nil {
+		return 0, io.EOF
+	}
+	n := copy(p, c.packet)
+	c.packet = nil
+	return n, nil
+}
+
+func TestSOCKS5UDPReplyAddressPolicy(t *testing.T) {
+	for _, tc := range []struct {
+		name, target, reply string
+		allowed             bool
+	}{
+		{"ipv4", "192.0.2.1:443", "192.0.2.1:443", true},
+		{"ipv6", "[2001:db8::1]:443", "[2001:db8::1]:443", true},
+		{"ipv4-mapped-reply", "192.0.2.1:443", "[::ffff:192.0.2.1]:443", true},
+		{"ipv4-mapped-target", "[::ffff:192.0.2.1]:443", "192.0.2.1:443", true},
+		{"wrong-ipv4", "192.0.2.1:443", "192.0.2.2:443", false},
+		{"wrong-ipv6", "[2001:db8::1]:443", "[2001:db8::2]:443", false},
+		{"wrong-port", "192.0.2.1:443", "192.0.2.1:444", false},
+		{"wrong-mapped-port", "192.0.2.1:443", "[::ffff:192.0.2.1]:444", false},
+		{"different-family", "192.0.2.1:443", "[2001:db8::1]:443", false},
+		{"domain-ipv4", "origin.test:443", "192.0.2.1:443", true},
+		{"domain-ipv6", "origin.test:443", "[2001:db8::1]:443", true},
+		{"domain-mapped", "origin.test:443", "[::ffff:192.0.2.1]:443", true},
+		{"domain-wrong-port", "origin.test:443", "192.0.2.1:444", false},
+		{"domain-response", "origin.test:443", "origin.test:443", false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			packet, err := S.EncodeUDPPacket(S.ParseAddrString(tc.reply), []byte("ok"))
+			if err != nil {
+				t.Fatal(err)
+			}
+			peer := &net.UDPAddr{IP: net.IPv4(192, 0, 2, 1), Port: 443}
+			conn := &socks5PacketConn{Conn: &socks5ReplySocket{packet: packet}, target: S.ParseAddrString(tc.target), peer: peer}
+			payload := make([]byte, 16)
+			n, from, err := conn.ReadFrom(payload)
+			if tc.allowed {
+				if err != nil || string(payload[:n]) != "ok" || from != peer {
+					t.Fatalf("allowed reply: n=%d from=%v err=%v", n, from, err)
+				}
+			} else if err != io.EOF || n != 0 || from != nil {
+				t.Fatalf("untrusted reply escaped: n=%d from=%v err=%v", n, from, err)
+			}
+		})
+	}
+}
+
+func TestSOCKS5DomainUDPDiscardsMalformedReplies(t *testing.T) {
+	for _, packet := range [][]byte{
+		{}, {0, 0, 0, S.AtypIPv4, 192}, {0, 0, 0, S.AtypIPv6, 0},
+		{0, 0, 0, S.AtypDomainName, 255, 'a'}, {0, 0, 0, 255, 0, 0},
+		{0, 1, 0, S.AtypIPv4, 192, 0, 2, 1, 1, 187},
+		{0, 0, 1, S.AtypIPv4, 192, 0, 2, 1, 1, 187},
+	} {
+		conn := &socks5PacketConn{Conn: &socks5ReplySocket{packet: packet}, target: S.ParseAddrString("origin.test:443"), peer: &net.UDPAddr{Port: 443}}
+		if n, from, err := conn.ReadFrom(make([]byte, 16)); n != 0 || from != nil || err != io.EOF {
+			t.Fatalf("malformed packet escaped: n=%d from=%v err=%v", n, from, err)
+		}
+	}
+}