fix(amneziawg): avoid manager lock inversion (#6397)

* fix(amneziawg): avoid manager lock inversion

Packet handlers re-entered the manager mutex while device reconfiguration and teardown held it and waited for receiver goroutines. Publish immutable peer indexes atomically so the data path can finish without participating in lifecycle locking.

* test(amneziawg): exercise UDP relay hit path

internal/amneziawgnet/manager.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/amnezia-vpn/amneziawg-go/v3/device"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
@@ -47,12 +48,36 @@ type managed struct {
 	dev          *Device
 	udpRelay     *UDPRelay
 	portForwards *PortForwardSet
-	peers        *PeerIndex
+	peers        atomic.Pointer[PeerIndex]
 	inst         amneziawg.Instance
 	structFP     string
 	uapiConfig   string
 }
 
+func (m *managed) lookupPeer(addr netip.Addr) (amneziawg.Peer, bool) {
+	peers := m.peers.Load()
+	if peers == nil {
+		return amneziawg.Peer{}, false
+	}
+	return peers.Lookup(addr)
+}
+
+func (m *managed) handleUDP(src, dst netip.AddrPort, payload []byte) {
+	peer, ok := m.lookupPeer(src.Addr())
+	if !ok {
+		return
+	}
+	m.udpRelay.Handle(src, dst, peer.Email, payload)
+}
+
+func (m *managed) close() {
+	m.portForwards.Close()
+	// Stop packet delivery before closing the relay so an in-flight handler
+	// cannot publish a new session after the relay has already been swept.
+	m.dev.Close()
+	m.udpRelay.Close()
+}
+
 // Manager owns the set of running embedded AmneziaWG interfaces, keyed by
 // inbound id -- the same shape as internal/mtproto.Manager (GetManager()
 // + sync.Once, mu-guarded map, Ensure/Reconcile/StopAll/HasRunning), so a
@@ -135,7 +160,7 @@ func (m *Manager) ensureLocked(d Desired) error {
 		// buildUAPIConfig actually reads, the way a hand-maintained field
 		// list could.
 		if conf == cur.uapiConfig {
-			cur.peers = NewPeerIndex(inst.Peers)
+			cur.peers.Store(NewPeerIndex(inst.Peers))
 			cur.inst = inst
 			applyV6Aliases(diffV6Aliases(oldInst, inst))
 			// buildUAPIConfig never reads ForwardedPorts (it's a panel-level
@@ -151,7 +176,7 @@ func (m *Manager) ensureLocked(d Desired) error {
 		if err := cur.dev.IpcSet(conf); err != nil {
 			return fmt.Errorf("amneziawgnet: reconfigure inbound %d: %w", inst.Id, err)
 		}
-		cur.peers = NewPeerIndex(inst.Peers)
+		cur.peers.Store(NewPeerIndex(inst.Peers))
 		cur.inst = inst
 		cur.uapiConfig = conf
 		applyV6Aliases(diffV6Aliases(oldInst, inst))
@@ -160,9 +185,7 @@ func (m *Manager) ensureLocked(d Desired) error {
 	}
 
 	if exists {
-		cur.udpRelay.Close()
-		cur.portForwards.Close()
-		cur.dev.Close()
+		cur.close()
 		delete(m.ifaces, inst.Id)
 	}
 	dev, err := newUnconfiguredDevice(inst, opts)
@@ -173,40 +196,30 @@ func (m *Manager) ensureLocked(d Desired) error {
 	relay := socksRelayForInstance(inst)
 	udpRelay := NewUDPRelay(relay, dev.Stack)
 	portForwards := NewPortForwardSet(dev.Stack, inst.Id)
-	inboundID := inst.Id // captured for the closures below, which outlive this call
+	next := &managed{
+		dev:          dev,
+		udpRelay:     udpRelay,
+		portForwards: portForwards,
+		inst:         inst,
+		structFP:     structFP,
+	}
+	next.peers.Store(NewPeerIndex(inst.Peers))
 	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
 		srcAddrPort, err := netip.ParseAddrPort(conn.RemoteAddr().String())
 		if err != nil {
 			conn.Close()
 			return
 		}
-		// Re-fetched on every connection, not captured once at attach time:
-		// a reconfigure-in-place (peers added/removed, no rebuild) replaces
-		// cur.peers without ever re-attaching the forwarder, so a stale
-		// captured index would silently miss newly-added peers.
-		_, peers, ok := m.Lookup(inboundID)
-		if !ok {
-			conn.Close()
-			return
-		}
-		peer, ok := peers.Lookup(srcAddrPort.Addr().Unmap())
+		// Reload for every connection: in-place reconfiguration swaps the peer
+		// index without reattaching handlers and may hold the lifecycle lock.
+		peer, ok := next.lookupPeer(srcAddrPort.Addr().Unmap())
 		if !ok {
 			conn.Close()
 			return
 		}
 		relay.RelayTCP(conn, peer.Email, dest)
 	})
-	AttachUDPHandler(dev.Stack, func(src, dst netip.AddrPort, payload []byte) {
-		_, peers, ok := m.Lookup(inboundID)
-		if !ok {
-			return
-		}
-		peer, ok := peers.Lookup(src.Addr())
-		if !ok {
-			return
-		}
-		udpRelay.Handle(src, dst, peer.Email, payload)
-	})
+	AttachUDPHandler(dev.Stack, next.handleUDP)
 
 	// Handlers are registered on dev.Stack above, BEFORE Configure's IpcSet
 	// can start any peer's receive goroutine -- see newUnconfiguredDevice's
@@ -224,16 +237,8 @@ func (m *Manager) ensureLocked(d Desired) error {
 	// the no-op check above a correct baseline to compare the next tick
 	// against instead of an empty string.
 	conf, _ := buildUAPIConfig(inst, opts)
-
-	m.ifaces[inst.Id] = &managed{
-		dev:          dev,
-		udpRelay:     udpRelay,
-		portForwards: portForwards,
-		peers:        NewPeerIndex(inst.Peers),
-		inst:         inst,
-		structFP:     structFP,
-		uapiConfig:   conf,
-	}
+	next.uapiConfig = conf
+	m.ifaces[inst.Id] = next
 	applyV6Aliases(diffV6Aliases(oldInst, inst))
 	portForwards.Reconcile(inst)
 	logger.Infof("amneziawgnet: started embedded interface %s for inbound %d", inst.InterfaceName, inst.Id)
@@ -275,9 +280,7 @@ func (m *Manager) Reconcile(desired []Desired) {
 			continue
 		}
 		applyV6Aliases(diffV6Aliases(cur.inst, amneziawg.Instance{}))
-		cur.udpRelay.Close()
-		cur.portForwards.Close()
-		cur.dev.Close()
+		cur.close()
 		delete(m.ifaces, id)
 		logger.Infof("amneziawgnet: stopped embedded interface for removed inbound %d", id)
 	}
@@ -300,9 +303,7 @@ func (m *Manager) Remove(id int) {
 		return
 	}
 	applyV6Aliases(diffV6Aliases(cur.inst, amneziawg.Instance{}))
-	cur.udpRelay.Close()
-	cur.portForwards.Close()
-	cur.dev.Close()
+	cur.close()
 	delete(m.ifaces, id)
 	logger.Infof("amneziawgnet: stopped embedded interface for removed inbound %d", id)
 }
@@ -313,9 +314,7 @@ func (m *Manager) StopAll() {
 	defer m.mu.Unlock()
 	for id, cur := range m.ifaces {
 		applyV6Aliases(diffV6Aliases(cur.inst, amneziawg.Instance{}))
-		cur.udpRelay.Close()
-		cur.portForwards.Close()
-		cur.dev.Close()
+		cur.close()
 		delete(m.ifaces, id)
 	}
 }
@@ -327,11 +326,8 @@ func (m *Manager) HasRunning() bool {
 	return len(m.ifaces) > 0
 }
 
-// Lookup returns the running Device and PeerIndex for inbound id, if any --
-// the forwarder/UDP-handler closures ensureLocked attaches use this to
-// re-fetch the current peer index on every connection (see ensureLocked's
-// comment on why), and it's equally available to a test harness or any
-// other caller that wants read access to a managed interface's state.
+// Lookup returns the running device and current peer snapshot for diagnostics,
+// tests, and other callers outside the packet-delivery path.
 func (m *Manager) Lookup(id int) (dev *Device, peers *PeerIndex, ok bool) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -339,5 +335,5 @@ func (m *Manager) Lookup(id int) (dev *Device, peers *PeerIndex, ok bool) {
 	if !exists {
 		return nil, nil, false
 	}
-	return cur.dev, cur.peers, true
+	return cur.dev, cur.peers.Load(), true
 }

internal/amneziawgnet/manager_test.go

@@ -3,6 +3,7 @@ package amneziawgnet
 import (
 	"fmt"
 	"net"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -89,6 +90,34 @@ func TestManagerLifecycle(t *testing.T) {
 	}
 }
 
+func TestManagedUDPHandlerDoesNotWaitForManagerLock(t *testing.T) {
+	cur := &managed{udpRelay: NewUDPRelay(SocksRelay{Addr: "invalid"}, nil)}
+	cur.peers.Store(NewPeerIndex([]amneziawg.Peer{{
+		Email:      "peer@test",
+		AllowedIPs: []string{"10.210.0.2/32"},
+	}}))
+	m := &Manager{}
+
+	done := make(chan struct{})
+	m.mu.Lock()
+	go func() {
+		cur.handleUDP(
+			netip.MustParseAddrPort("10.210.0.2:1234"),
+			netip.MustParseAddrPort("10.210.0.3:53"),
+			[]byte("query"),
+		)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		m.mu.Unlock()
+	case <-time.After(time.Second):
+		m.mu.Unlock()
+		t.Fatal("UDP handler blocked on the manager lifecycle lock")
+	}
+}
+
 // TestEnsureUnchangedInstanceDoesNotResetLivePeers is a regression test for a
 // real production bug: an unchanged Ensure call (the common case on every
 // 10s AmneziaWGJob reconcile tick when no admin edit happened) was calling