fix(amneziawg): give each peer a share of the forwarder and UDP relay

The embedded interface's TCP forwarder admits 65535 connections in
flight for all peers together, and a forwarded request holds its slot
until the handshake ends or gVisor gives up on it. One client sending
bare SYNs from fresh source ports therefore took every slot and left
the other clients' connections unanswered. The UDP relay likewise
opened a SOCKS5 association and a host UDP socket for every fresh
tunnel source port with no per-peer bound.

The forwarder now names the peer on the SYN, refuses a source no peer
owns and caps one peer at 4096 connections; the UDP relay caps one peer
at 1024 sessions and drops the datagrams that would need more.

internal/amneziawgnet/bench_test.go

@@ -134,7 +134,7 @@ func newBenchTunnel(b *testing.B, listenPort int, serverAddr, clientAddr string)
 	if err != nil {
 		b.Fatalf("newUnconfiguredDevice: %v", err)
 	}
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		switch dest.Port() {
 		case benchDiscardPort:

internal/amneziawgnet/device_test.go

@@ -80,7 +80,7 @@ func TestNewDeviceHandshakeForwarderAndIdentity(t *testing.T) {
 	// purely from the decapsulated packet, not from any routing table.
 	wantDest := netip.MustParseAddrPort("10.201.9.9:9999")
 
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
 		defer conn.Close()
 		peer, ok := idx.Lookup(src.Addr().Unmap())
 		got <- recovered{email: peer.Email, ok: ok, dest: dest}
@@ -278,7 +278,7 @@ func TestNewDeviceHeaderProtectionAndContentPaddingRoundTrip(t *testing.T) {
 	const wantRequest = "hello from client"
 	const wantReply = "hello from server"
 	serverDone := make(chan error, 1)
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		buf := make([]byte, len(wantRequest))
 		if _, err := io.ReadFull(conn, buf); err != nil {
@@ -456,7 +456,7 @@ func TestNewDeviceRandomTrailersAndDisableCookiesRoundTrip(t *testing.T) {
 	const wantRequest = "hello from client, with a trailer"
 	const wantReply = "hello from server, with a trailer"
 	serverDone := make(chan error, 1)
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		buf := make([]byte, len(wantRequest))
 		if _, err := io.ReadFull(conn, buf); err != nil {

internal/amneziawgnet/diagnostics_test.go

@@ -77,7 +77,7 @@ func TestDiagnoseDeviceReportsListenPortAndPeerState(t *testing.T) {
 	// Registering the forwarder doesn't require any peer to be configured
 	// yet, so this ordering is free; it's Configure's IpcSet that must
 	// never run before the forwarder is registered.
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, _ netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, _, _ netip.AddrPort) {
 		defer conn.Close()
 		io.Copy(io.Discard, conn)
 	})

internal/amneziawgnet/forwarder.go

@@ -2,6 +2,7 @@ package amneziawgnet
 
 import (
 	"net/netip"
+	"sync"
 
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
@@ -9,6 +10,42 @@ import (
 	"gvisor.dev/gvisor/pkg/waiter"
 )
 
+// One datagram or bare SYN from a fresh source port buys a flow its sockets, goroutine and
+// buffers, so each peer gets a share of an interface instead of the whole of it.
+const (
+	maxTCPFlowsPerPeer    = 4096
+	maxUDPSessionsPerPeer = 1024
+)
+
+// peerSlots counts the flows each peer holds; every successful acquire needs one release.
+type peerSlots struct {
+	mu    sync.Mutex
+	limit int
+	held  map[string]int
+}
+
+func newPeerSlots(limit int) *peerSlots {
+	return &peerSlots{limit: limit, held: map[string]int{}}
+}
+
+func (s *peerSlots) acquire(peer string) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.held[peer] >= s.limit {
+		return false
+	}
+	s.held[peer]++
+	return true
+}
+
+func (s *peerSlots) release(peer string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.held[peer]--; s.held[peer] <= 0 {
+		delete(s.held, peer)
+	}
+}
+
 // AttachTCPForwarder attaches a TCP forwarder to gstack in promiscuous +
 // spoofing mode, so it accepts connections addressed to any destination --
 // not just the stack's own configured local address -- and hands the
@@ -22,12 +59,27 @@ import (
 // Handlers must identify the peer by src, never conn.RemoteAddr(): that asks
 // the live endpoint and is nil once a RST is processed, possibly before the call.
 //
+// owner names the peer a tunnel source belongs to; it is asked on the SYN, before anything
+// is allocated, and refuses a source nobody owns or one whose owner holds its full share.
+//
 // Adapted from xtls/xray-core's proxy/wireguard/tun.go createForwarder (MIT).
-func AttachTCPForwarder(gstack *stack.Stack, handler func(conn *gonet.TCPConn, src, dest netip.AddrPort)) {
+func AttachTCPForwarder(gstack *stack.Stack, owner func(src netip.Addr) (string, bool), handler func(conn *gonet.TCPConn, src, dest netip.AddrPort)) {
 	enablePromiscuousRouting(gstack)
 
+	slots := newPeerSlots(maxTCPFlowsPerPeer)
 	fwd := tcp.NewForwarder(gstack, 0, 65535, func(r *tcp.ForwarderRequest) {
+		peer, ok := owner(addrFromTcpip(r.ID().RemoteAddress).Unmap())
+		if !ok {
+			r.Complete(true)
+			return
+		}
+		// CreateEndpoint holds the request until the handshake ends, so a bare SYN costs a slot.
+		if !slots.acquire(peer) {
+			r.Complete(false)
+			return
+		}
 		go func(r *tcp.ForwarderRequest) {
+			defer slots.release(peer)
 			var wq waiter.Queue
 			id := r.ID()
 

internal/amneziawgnet/forwarder_test.go

@@ -65,7 +65,7 @@ func TestTCPForwarderHandsOverSourceOfAnAlreadyResetConnection(t *testing.T) {
 	}
 	accepted, resetSent := make(chan struct{}), make(chan struct{})
 	got := make(chan handedOver, 1)
-	AttachTCPForwarder(gstack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
+	AttachTCPForwarder(gstack, ownerBySource, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
 		defer conn.Close()
 		close(accepted)
 		<-resetSent
@@ -108,3 +108,55 @@ func TestTCPForwarderHandsOverSourceOfAnAlreadyResetConnection(t *testing.T) {
 		t.Fatal("handler never saw the reset")
 	}
 }
+
+// ownerBySource counts every tunnel source address as a peer of its own.
+func ownerBySource(src netip.Addr) (string, bool) {
+	return src.String(), true
+}
+
+// The forwarder's in-flight limit is shared by the interface, so one peer's bare SYNs used
+// to take every slot and leave the other clients' connections unanswered.
+func TestTCPForwarderPeerCannotStarveOthers(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.8.1.1")}, 1420)
+	if err != nil {
+		t.Fatalf("createNetTUNWithStack: %v", err)
+	}
+	defer tun.Close()
+
+	flooderRange := netip.MustParsePrefix("10.8.2.0/24")
+	AttachTCPForwarder(gstack, func(src netip.Addr) (string, bool) {
+		if flooderRange.Contains(src) {
+			return "flooder", true
+		}
+		return "other", true
+	}, func(conn *gonet.TCPConn, _, _ netip.AddrPort) { conn.Close() })
+
+	other := netip.MustParseAddrPort("10.8.1.3:40000")
+	answered := make(chan struct{})
+	go func() {
+		bufs, sizes := [][]byte{make([]byte, 2048)}, []int{0}
+		for {
+			if _, err := tun.Read(bufs, sizes, 0); err != nil {
+				return
+			}
+			reply := header.IPv4(bufs[0][:sizes[0]])
+			seg := header.TCP(reply.Payload())
+			if reply.DestinationAddress() == tcpip.AddrFrom4(other.Addr().As4()) && seg.Flags() == header.TCPFlagSyn|header.TCPFlagAck {
+				close(answered)
+				return
+			}
+		}
+	}()
+
+	dst := netip.MustParseAddrPort("10.8.1.1:80")
+	for n := range 70000 {
+		src := netip.AddrPortFrom(netip.AddrFrom4([4]byte{10, 8, 2, byte(n)}), uint16(1024+n/256))
+		injectFromPeer(t, tun, rawTCPSegment(src, dst, 1000, 0, header.TCPFlagSyn))
+	}
+	injectFromPeer(t, tun, rawTCPSegment(other, dst, 1000, 0, header.TCPFlagSyn))
+	select {
+	case <-answered:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("no SYN-ACK for %s; dropped at the in-flight limit: %d", other, gstack.Stats().TCP.ForwardMaxInFlightDrop.Value())
+	}
+}

internal/amneziawgnet/manager.go

@@ -62,6 +62,11 @@ func (m *managed) lookupPeer(addr netip.Addr) (amneziawg.Peer, bool) {
 	return peers.Lookup(addr)
 }
 
+func (m *managed) peerEmail(addr netip.Addr) (string, bool) {
+	peer, ok := m.lookupPeer(addr)
+	return peer.Email, ok
+}
+
 func (m *managed) handleUDP(src, dst netip.AddrPort, payload []byte) {
 	peer, ok := m.lookupPeer(src.Addr())
 	if !ok {
@@ -227,7 +232,7 @@ func (m *Manager) ensureLocked(d Desired) error {
 		structFP:     structFP,
 	}
 	next.peers.Store(NewPeerIndex(inst.Peers))
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, next.peerEmail, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
 		// Reload for every connection: in-place reconfiguration swaps the peer
 		// index without reattaching handlers and may hold the lifecycle lock.
 		peer, ok := next.lookupPeer(src.Addr().Unmap())

internal/amneziawgnet/relay.go

@@ -386,11 +386,12 @@ type UDPRelay struct {
 	// session map: src.String() would allocate on every relayed datagram.
 	mu       sync.Mutex
 	sessions map[netip.AddrPort]*socks5UDPSession
+	slots    *peerSlots
 }
 
 // NewUDPRelay creates a UDPRelay for one embedded AmneziaWG Device's stack.
 func NewUDPRelay(relay SocksRelay, gstack *stack.Stack) *UDPRelay {
-	return &UDPRelay{relay: relay, gstack: gstack, sessions: map[netip.AddrPort]*socks5UDPSession{}}
+	return &UDPRelay{relay: relay, gstack: gstack, sessions: map[netip.AddrPort]*socks5UDPSession{}, slots: newPeerSlots(maxUDPSessionsPerPeer)}
 }
 
 // Handle relays one packet from src (the peer's tunnel-internal source) to
@@ -404,8 +405,12 @@ func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte)
 	u.mu.Unlock()
 
 	if !ok {
+		if !u.slots.acquire(email) {
+			return
+		}
 		fresh, err := newSocks5UDPSession(u.relay.Addr, email, u.relay.Password)
 		if err != nil {
+			u.slots.release(email)
 			logger.Warningf("amneziawgnet: UDPRelay: SOCKS5 associate for %q: %v", email, err)
 			return
 		}
@@ -415,12 +420,16 @@ func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte)
 		if existing, dup := u.sessions[src]; dup {
 			u.mu.Unlock()
 			fresh.Close()
+			u.slots.release(email)
 			sess = existing
 		} else {
 			u.sessions[src] = fresh
 			u.mu.Unlock()
 			sess = fresh
-			go u.pump(src, fresh)
+			go func() {
+				defer u.slots.release(email)
+				u.pump(src, fresh)
+			}()
 		}
 	}
 	if err := sess.sendTo(dst, payload); err != nil {

internal/amneziawgnet/relay_e2e_test.go

@@ -161,7 +161,7 @@ func TestSocksRelayAgainstRealXray(t *testing.T) {
 	udpRelay := NewUDPRelay(relay, dev.Stack)
 	defer udpRelay.Close()
 
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, ownerBySource, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
 		peer, ok := idx.Lookup(src.Addr().Unmap())
 		if !ok {
 			conn.Close()

internal/amneziawgnet/relay_test.go

@@ -297,3 +297,76 @@ func TestNewSocks5UDPSessionGivesUpOnSilentServer(t *testing.T) {
 		t.Fatal("newSocks5UDPSession still blocked 2s past the associate deadline: a silent SOCKS5 server parks the tunnel's delivery path")
 	}
 }
+
+// startFakeSocks5UDP serves just enough no-auth SOCKS5 for UDP ASSOCIATE to
+// succeed, pointing every association at one UDP socket nobody reads.
+func startFakeSocks5UDP(t *testing.T) string {
+	t.Helper()
+	sink, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+	t.Cleanup(func() { _ = sink.Close() })
+	port := sink.LocalAddr().(*net.UDPAddr).Port
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen tcp: %v", err)
+	}
+	t.Cleanup(func() { _ = ln.Close() })
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func() {
+				defer c.Close()
+				var greeting [4]byte
+				var associate [10]byte
+				if _, err := io.ReadFull(c, greeting[:]); err != nil {
+					return
+				}
+				_, _ = c.Write([]byte{0x05, 0x00})
+				if _, err := io.ReadFull(c, associate[:]); err != nil {
+					return
+				}
+				_, _ = c.Write([]byte{0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, byte(port >> 8), byte(port)})
+				_, _ = io.Copy(io.Discard, c)
+			}()
+		}
+	}()
+	return ln.Addr().String()
+}
+
+// Each fresh tunnel source port bought a SOCKS5 association and a UDP socket with no
+// per-peer bound, so one client could take the host's UDP ports from everyone.
+func TestUDPRelayCapsSessionsPerPeer(t *testing.T) {
+	relay := NewUDPRelay(SocksRelay{Addr: startFakeSocks5UDP(t), Password: "x"}, nil)
+	defer relay.Close()
+	dst := netip.MustParseAddrPort("192.0.2.1:53")
+
+	for n := 1; n < maxUDPSessionsPerPeer; n++ {
+		relay.slots.acquire("greedy")
+	}
+	relay.Handle(netip.MustParseAddrPort("10.8.1.5:1000"), dst, "greedy", []byte("x"))
+	relay.Handle(netip.MustParseAddrPort("10.8.1.5:1001"), dst, "greedy", []byte("x"))
+	relay.Handle(netip.MustParseAddrPort("10.8.1.6:1000"), dst, "modest", []byte("x"))
+
+	relay.mu.Lock()
+	_, last := relay.sessions[netip.MustParseAddrPort("10.8.1.5:1000")]
+	_, excess := relay.sessions[netip.MustParseAddrPort("10.8.1.5:1001")]
+	_, other := relay.sessions[netip.MustParseAddrPort("10.8.1.6:1000")]
+	relay.mu.Unlock()
+	if !last || excess || !other {
+		t.Fatalf("sessions: greedy's last allowed = %v, greedy's excess = %v, modest's = %v; want true, false, true", last, excess, other)
+	}
+
+	relay.Close()
+	deadline := time.Now().Add(2 * time.Second)
+	for !relay.slots.acquire("greedy") {
+		if time.Now().After(deadline) {
+			t.Fatal("closing greedy's session never released its slot")
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}