diff --git a/internal/amneziawgnet/egress.go b/internal/amneziawgnet/egress.go index 1f2d7e793..ab7816696 100644 --- a/internal/amneziawgnet/egress.go +++ b/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() { diff --git a/internal/amneziawgnet/egress_port_test.go b/internal/amneziawgnet/egress_port_test.go index b9dfa6cb7..f2e36c666 100644 --- a/internal/amneziawgnet/egress_port_test.go +++ b/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 }) +}