fix(amneziawg): bound the egress relay's UDP sessions as a whole

Every UDP destination a proxy user reached through the AmneziaWG egress
opened a session with a goroutine, a port and a buffer, per association,
and a Mux frame opens a new association, so one user could hold sessions
without limit. All egress UDP sessions now share one budget of 4096;
when it is spent the longest-idle session is closed to make room.

internal/amneziawgnet/egress.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"net/netip"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -482,6 +483,50 @@ type udpControl struct{ conn net.Conn }
 type egressUDPSession struct {
 	dst  netip.AddrPort
 	conn *gonet.UDPConn
+
+	// Unix nanoseconds of the last datagram either way; the budget evicts by it.
+	lastActive atomic.Int64
+}
+
+func (s *egressUDPSession) touch() { s.lastActive.Store(time.Now().UnixNano()) }
+
+// egressMaxUDPSessions bounds the sessions of all associations together: each
+// pins a goroutine, a port and 64 KiB, and a Mux frame opens an association.
+var egressMaxUDPSessions = 4096
+
+// egressUDPBudget is process-wide, like the egress server: the socks bridge
+// carries no user identity, so every user draws on the same budget.
+type egressUDPBudget struct {
+	mu   sync.Mutex
+	live map[*egressUDPSession]struct{}
+}
+
+var egressUDPSessions = egressUDPBudget{live: map[*egressUDPSession]struct{}{}}
+
+// admit counts sess in. A spent budget closes the longest-idle session rather
+// than refusing sess, so a flood's unanswered flows go before live ones.
+func (b *egressUDPBudget) admit(sess *egressUDPSession) {
+	var idlest *egressUDPSession
+	b.mu.Lock()
+	if len(b.live) >= egressMaxUDPSessions {
+		for s := range b.live {
+			if idlest == nil || s.lastActive.Load() < idlest.lastActive.Load() {
+				idlest = s
+			}
+		}
+		delete(b.live, idlest)
+	}
+	b.live[sess] = struct{}{}
+	b.mu.Unlock()
+	if idlest != nil {
+		idlest.conn.Close()
+	}
+}
+
+func (b *egressUDPBudget) release(sess *egressUDPSession) {
+	b.mu.Lock()
+	delete(b.live, sess)
+	b.mu.Unlock()
 }
 
 // relayUDP answers the associate request and relays datagrams to
@@ -604,6 +649,8 @@ func (s *udpEgressSessions) getOrDial(dev *Device, tag string, udpConn *net.UDPC
 		return nil
 	}
 	sess := &egressUDPSession{dst: dst, conn: conn}
+	sess.touch() // before admit: a zero stamp is the next one evicted
+	egressUDPSessions.admit(sess)
 	s.m[dst] = sess
 	go pumpUDPEgress(udpConn, client, sess, s)
 	return sess
@@ -627,6 +674,7 @@ func (s *socks5EgressServer) deliverUDPDatagram(dev *Device, tag string, udpConn
 	if sess == nil {
 		return
 	}
+	sess.touch()
 	if _, werr := sess.conn.Write(payload); werr != nil {
 		logger.Warningf("amneziawgnet: egress %q: send udp to %s: %v", tag, dst, werr)
 	}
@@ -639,6 +687,7 @@ func pumpUDPEgress(udpConn *net.UDPConn, client netip.AddrPort, sess *egressUDPS
 		sessions.mu.Lock()
 		delete(sessions.m, sess.dst)
 		sessions.mu.Unlock()
+		egressUDPSessions.release(sess)
 		sess.conn.Close()
 	}()
 	buf := make([]byte, 65536)
@@ -649,6 +698,7 @@ func pumpUDPEgress(udpConn *net.UDPConn, client netip.AddrPort, sess *egressUDPS
 		if err != nil {
 			return
 		}
+		sess.touch()
 		hdr := make([]byte, 0, 3+1+16+2+n)
 		hdr = append(hdr, 0x00, 0x00, 0x00) // RSV RSV FRAG(=0)
 		if sess.dst.Addr().Is4() {

internal/amneziawgnet/egress_port_test.go

@@ -5,8 +5,10 @@ import (
 	"encoding/json"
 	"errors"
 	"net"
+	"net/netip"
 	"strconv"
 	"testing"
+	"time"
 )
 
 // holdEgressBasePort occupies EgressBasePort the way another service would; a
@@ -112,3 +114,86 @@ func TestEgressFallbackStaysBelowThePanelRanges(t *testing.T) {
 		t.Fatalf("Port() = %d, inside the panel's reserved range", port)
 	}
 }
+
+// TestEgressUDPSessionsShareOneBudget: a proxy user opens an association per
+// Mux frame, so the session ceiling must hold across all of them together.
+func TestEgressUDPSessionsShareOneBudget(t *testing.T) {
+	saved := egressMaxUDPSessions
+	egressMaxUDPSessions = 64
+	t.Cleanup(func() { egressMaxUDPSessions = saved })
+
+	local := netip.MustParseAddr("10.211.4.1")
+	dev := &Device{Stack: newTestStack(t, local.String()), localAddrs: []netip.Addr{local}}
+	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+	defer udpConn.Close()
+	client := netip.MustParseAddrPort("127.0.0.1:9")
+	dst := func(port int) netip.AddrPort {
+		return netip.AddrPortFrom(netip.MustParseAddr("192.0.2.1"), uint16(port))
+	}
+	var associations []*udpEgressSessions
+	associate := func() *udpEgressSessions {
+		sessions := &udpEgressSessions{m: map[netip.AddrPort]*egressUDPSession{}}
+		associations = append(associations, sessions)
+		return sessions
+	}
+	session := func(sessions *udpEgressSessions, port int) *egressUDPSession {
+		sessions.mu.Lock()
+		defer sessions.mu.Unlock()
+		return sessions.m[dst(port)]
+	}
+	liveSessions := func() int {
+		n := 0
+		for _, sessions := range associations {
+			sessions.mu.Lock()
+			n += len(sessions.m)
+			sessions.mu.Unlock()
+		}
+		return n
+	}
+	slotsTaken := func() int {
+		egressUDPSessions.mu.Lock()
+		defer egressUDPSessions.mu.Unlock()
+		return len(egressUDPSessions.live)
+	}
+	waitFor := func(what string, cond func() bool) {
+		t.Helper()
+		for deadline := time.Now().Add(5 * time.Second); !cond(); time.Sleep(time.Millisecond) {
+			if time.Now().After(deadline) {
+				t.Fatalf("timed out waiting for %s (%d live sessions)", what, liveSessions())
+			}
+		}
+	}
+
+	before := slotsTaken()
+	others := associate()
+	busy := others.getOrDial(dev, "t", udpConn, client, dst(1))
+	if busy == nil || others.getOrDial(dev, "t", udpConn, client, dst(2)) == nil {
+		t.Fatal("no session with the budget unspent")
+	}
+	// Three times the budget, spread over associations that stay small.
+	for range 3 * egressMaxUDPSessions / 16 {
+		// The oldest session of all, but never the longest idle one.
+		GetEgressServer().deliverUDPDatagram(dev, "t", udpConn, client, others, dst(1), []byte("x"))
+		flood := associate()
+		for port := 1; port <= 16; port++ {
+			if flood.getOrDial(dev, "t", udpConn, client, dst(port)) == nil {
+				t.Fatalf("destination %d refused: the budget must evict, not drop", port)
+			}
+		}
+	}
+	waitFor("evicted sessions to retire", func() bool { return liveSessions() <= egressMaxUDPSessions })
+	if session(others, 2) != nil {
+		t.Error("the longest-idle session outlived a spent budget")
+	}
+	if session(others, 1) != busy {
+		t.Error("a session that kept relaying was evicted ahead of idle ones")
+	}
+
+	for _, sessions := range associations {
+		sessions.closeAll()
+	}
+	waitFor("closeAll to hand every session back", func() bool { return liveSessions() == 0 && slotsTaken() <= before })
+}