fix(amneziawg): drop a removed client's flows before its address is reused

Removing or re-keying a client reconfigured the device in place but left
its TCP and UDP flows standing, and withdrawing a port forward stopped
only new connections. The next client given the same tunnel address
inherited those flows: it could read what outside hosts sent to the
former client and write back into those connections.

After an in-place reconfigure, every stack endpoint and UDP relay
session pinned to an address whose peer was removed or changed is
aborted, and closing a port forward cuts the connections it is relaying.

internal/amneziawgnet/manager.go

@@ -70,6 +70,27 @@ func (m *managed) handleUDP(src, dst netip.AddrPort, payload []byte) {
 	m.udpRelay.Handle(src, dst, peer.Email, payload)
 }
 
+// dropReassignedFlows aborts every stack flow pinned to a tunnel address its
+// peer in before lost (removed or re-keyed): the next client inherits the address.
+func (m *managed) dropReassignedFlows(before *PeerIndex) {
+	reassigned := func(addr netip.Addr) bool {
+		was, ok := before.Lookup(addr)
+		if !ok {
+			return false
+		}
+		now, ok := m.lookupPeer(addr)
+		return !ok || now.PublicKey != was.PublicKey
+	}
+	// Abort, not Close: a closed endpoint keeps retransmitting unacknowledged
+	// data, and that would reach whoever holds the address next.
+	for _, ep := range m.dev.Stack.RegisteredEndpoints() {
+		if remote, ok := endpointRemoteAddr(ep); ok && reassigned(remote) {
+			ep.Abort()
+		}
+	}
+	m.udpRelay.CloseSessions(reassigned)
+}
+
 func (m *managed) close() {
 	m.portForwards.Close()
 	// Stop packet delivery before closing the relay so an in-flight handler
@@ -177,11 +198,12 @@ 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.Store(NewPeerIndex(inst.Peers))
+		oldPeers := cur.peers.Swap(NewPeerIndex(inst.Peers))
 		cur.inst = inst
 		cur.uapiConfig = conf
 		applyV6Aliases(diffV6Aliases(oldInst, inst))
 		cur.portForwards.Reconcile(inst)
+		cur.dropReassignedFlows(oldPeers)
 		return nil
 	}
 

internal/amneziawgnet/manager_test.go

@@ -1,6 +1,7 @@
 package amneziawgnet
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/netip"
@@ -8,6 +9,7 @@ import (
 	"time"
 
 	"github.com/amnezia-vpn/amneziawg-go/v3/device"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 
 	"github.com/mhsanaei/3x-ui/v3/internal/amneziawg"
 	"github.com/mhsanaei/3x-ui/v3/internal/util/wireguard"
@@ -330,6 +332,116 @@ func TestForwardedPortsOnlyChangeStillReconcilesPortForwards(t *testing.T) {
 	conn.Close()
 }
 
+// The stack survives an in-place reconfiguration, so a flow pinned to a removed
+// peer's address would otherwise reach whichever client gets that address next.
+func TestEnsureDropsFlowsOfReassignedTunnelAddress(t *testing.T) {
+	priv, pub, err := wireguard.GenerateWireguardKeypair()
+	if err != nil {
+		t.Fatalf("generate server keypair: %v", err)
+	}
+	var peerPubs [3]string
+	for i := range peerPubs {
+		if _, peerPubs[i], err = wireguard.GenerateWireguardKeypair(); err != nil {
+			t.Fatalf("generate peer keypair: %v", err)
+		}
+	}
+	reused, kept := netip.MustParseAddr("10.206.0.2"), netip.MustParseAddr("10.206.0.3")
+
+	m := &Manager{ifaces: map[int]*managed{}}
+	inst := amneziawg.Instance{
+		Id:            9,
+		InterfaceName: "awgtest9",
+		ListenPort:    58720,
+		PrivateKey:    priv,
+		PublicKey:     pub,
+		Address:       []string{"10.206.0.1/24"},
+		MTU:           1420,
+		Obfuscation: amneziawg.Obfuscation31{
+			Jc: 4, Jmin: 40, Jmax: 70,
+			S1: 20, S2: 30, S3: 20, S4: 20,
+		},
+		Peers: []amneziawg.Peer{
+			{Email: "old@test", PublicKey: peerPubs[0], AllowedIPs: []string{reused.String() + "/32"}},
+			{Email: "other@test", PublicKey: peerPubs[1], AllowedIPs: []string{kept.String() + "/32"}},
+		},
+	}
+	defer m.StopAll()
+	if err := m.Ensure(Desired{Instance: inst}); err != nil {
+		t.Fatalf("Ensure (create): %v", err)
+	}
+	cur := m.ifaces[inst.Id]
+
+	endpointsTo := func(addr netip.Addr) int {
+		n := 0
+		for _, ep := range cur.dev.Stack.RegisteredEndpoints() {
+			if remote, ok := endpointRemoteAddr(ep); ok && remote == addr {
+				n++
+			}
+		}
+		return n
+	}
+	waitEndpointsTo := func(addr netip.Addr, want int) {
+		t.Helper()
+		for deadline := time.Now().Add(2 * time.Second); endpointsTo(addr) != want; time.Sleep(5 * time.Millisecond) {
+			if time.Now().After(deadline) {
+				t.Fatalf("endpoints connected to %s = %d, want %d", addr, endpointsTo(addr), want)
+			}
+		}
+	}
+
+	for _, addr := range []netip.Addr{reused, kept} {
+		raddr := tunnelFullAddress(addr, 8080)
+		udpConn, err := gonet.DialUDP(cur.dev.Stack, nil, &raddr, tunnelNetwork(addr))
+		if err != nil {
+			t.Fatalf("DialUDP %s: %v", addr, err)
+		}
+		defer udpConn.Close()
+
+		ctrl, _ := net.Pipe()
+		relayConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+		if err != nil {
+			t.Fatalf("ListenUDP: %v", err)
+		}
+		cur.udpRelay.sessions[netip.AddrPortFrom(addr, 1234)] = &socks5UDPSession{ctrl: ctrl, udpConn: relayConn}
+	}
+	dialCtx, cancelDial := context.WithCancel(context.Background())
+	dialed := make(chan struct{})
+	go func() {
+		defer close(dialed)
+		if conn, err := gonet.DialContextTCP(dialCtx, cur.dev.Stack, tunnelFullAddress(reused, 8080), tunnelNetwork(reused)); err == nil {
+			conn.Close()
+		}
+	}()
+	defer func() { cancelDial(); <-dialed }()
+	waitEndpointsTo(reused, 2)
+	waitEndpointsTo(kept, 1)
+
+	next := inst
+	next.Peers = []amneziawg.Peer{
+		{Email: "new@test", PublicKey: peerPubs[2], AllowedIPs: []string{reused.String() + "/32"}},
+		inst.Peers[1],
+	}
+	if err := m.Ensure(Desired{Instance: next}); err != nil {
+		t.Fatalf("Ensure (address reassigned): %v", err)
+	}
+	if m.ifaces[inst.Id] != cur {
+		t.Fatal("peer change rebuilt the interface instead of reconfiguring it in place")
+	}
+
+	waitEndpointsTo(reused, 0)
+	waitEndpointsTo(kept, 1)
+	cur.udpRelay.mu.Lock()
+	_, reusedSession := cur.udpRelay.sessions[netip.AddrPortFrom(reused, 1234)]
+	_, keptSession := cur.udpRelay.sessions[netip.AddrPortFrom(kept, 1234)]
+	cur.udpRelay.mu.Unlock()
+	if reusedSession {
+		t.Error("the removed peer's UDP relay session survived its address being reassigned")
+	}
+	if !keptSession {
+		t.Error("an unaffected peer's UDP relay session was torn down")
+	}
+}
+
 // TestEnsureHeaderProtectionKeyChangeReconfiguresInPlace is a regression
 // test for the Phase 3.7 AWG 3.0 wiring: proves that populating
 // Desired.Options with a real HeaderProtectionKey/ContentPaddingAddition

internal/amneziawgnet/netstack.go

@@ -226,6 +226,20 @@ func enablePromiscuousRouting(gstack *stack.Stack) {
 	gstack.SetSpoofing(1, true)
 }
 
+// endpointRemoteAddr returns the address ep is connected to; false for an
+// endpoint that has none (listening, unconnected).
+func endpointRemoteAddr(ep stack.TransportEndpoint) (netip.Addr, bool) {
+	tep, ok := ep.(tcpip.Endpoint)
+	if !ok {
+		return netip.Addr{}, false
+	}
+	info, ok := tep.Info().(*stack.TransportEndpointInfo)
+	if !ok || info.ID.RemoteAddress.Len() == 0 {
+		return netip.Addr{}, false
+	}
+	return addrFromTcpip(info.ID.RemoteAddress), true
+}
+
 // addrFromTcpip converts a gVisor tcpip.Address to netip.Addr.
 func addrFromTcpip(a tcpip.Address) netip.Addr {
 	if a.Len() == 4 {

internal/amneziawgnet/portfwd.go

@@ -305,10 +305,11 @@ func tunnelFullAddress(addr netip.Addr, port int) tcpip.FullAddress {
 }
 
 // tcpForwardListener is one open host-facing TCP listener for a single
-// portForwardKey.
+// portForwardKey; ctx ends with it and takes every relay it started along.
 type tcpForwardListener struct {
-	ln      net.Listener
+	ln     net.Listener
-	closing chan struct{}
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 // listenPortForwardTCP opens a host-facing TCP listener on key.port and
@@ -324,7 +325,8 @@ func listenPortForwardTCP(gstack *stack.Stack, inboundID int, key portForwardKey
 		logger.Warningf("amneziawgnet: port-forward: inbound %d peer %q: listen tcp :%d: %v", inboundID, key.email, key.port, err)
 		return nil
 	}
-	l := &tcpForwardListener{ln: ln, closing: make(chan struct{})}
+	l := &tcpForwardListener{ln: ln}
+	l.ctx, l.cancel = context.WithCancel(context.Background())
 	logger.Infof("amneziawgnet: port-forward: inbound %d peer %q: listening tcp :%d", inboundID, key.email, key.port)
 	go l.acceptLoop(gstack, inboundID, key, target)
 	return l
@@ -335,29 +337,29 @@ func (l *tcpForwardListener) acceptLoop(gstack *stack.Stack, inboundID int, key
 		conn, err := l.ln.Accept()
 		if err != nil {
 			select {
-			case <-l.closing:
+			case <-l.ctx.Done():
 				return // intentional shutdown, not a real accept error
 			default:
 			}
 			logger.Warningf("amneziawgnet: port-forward: inbound %d peer %q: accept tcp :%d: %v", inboundID, key.email, key.port, err)
 			// Keep accepting: returning would strand a bound socket nobody
 			// serves, under a key Reconcile still sees as open and never reopens.
-			if !waitAcceptRetry(l.closing) {
+			if !waitAcceptRetry(l.ctx.Done()) {
 				return
 			}
 			continue
 		}
-		go relayTCPForward(gstack, conn, inboundID, key, target)
+		go l.relayTCPForward(gstack, conn, inboundID, key, target)
 	}
 }
 
-func relayTCPForward(gstack *stack.Stack, conn net.Conn, inboundID int, key portForwardKey, target portForwardTargetFunc) {
+func (l *tcpForwardListener) relayTCPForward(gstack *stack.Stack, conn net.Conn, inboundID int, key portForwardKey, target portForwardTargetFunc) {
 	defer conn.Close()
 	addr, ok := target(key.email)
 	if !ok {
 		return
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), portForwardDialTimeout)
+	ctx, cancel := context.WithTimeout(l.ctx, portForwardDialTimeout)
 	defer cancel()
 	tunnelConn, err := gonet.DialContextTCP(ctx, gstack, tunnelFullAddress(addr, key.port), tunnelNetwork(addr))
 	if err != nil {
@@ -366,14 +368,17 @@ func relayTCPForward(gstack *stack.Stack, conn net.Conn, inboundID int, key port
 	}
 	defer tunnelConn.Close()
 
+	stop := context.AfterFunc(l.ctx, func() {
+		conn.Close()
+		tunnelConn.Close()
+	})
+	defer stop()
 	pipeBothWays(conn, tunnelConn)
 }
 
-// Close stops accepting new connections. Already-relaying connections are
+// Close stops accepting and cuts the connections already relaying, as
-// left to finish on their own -- there's no shared state to tear down early
+// udpForwardListener.Close does: a withdrawn forward must not stay usable.
-// for, and an abrupt cut would just look like a network error to whichever
-// external client was mid-transfer.
 func (l *tcpForwardListener) Close() {
-	close(l.closing)
+	l.cancel()
 	l.ln.Close()
 }

internal/amneziawgnet/portfwd_test.go

@@ -244,7 +244,8 @@ func TestTCPForwardAcceptLoopSurvivesTransientError(t *testing.T) {
 	if err != nil {
 		t.Fatalf("listen: %v", err)
 	}
-	l := &tcpForwardListener{ln: &failFirstAccept{Listener: host}, closing: make(chan struct{})}
+	l := &tcpForwardListener{ln: &failFirstAccept{Listener: host}}
+	l.ctx, l.cancel = context.WithCancel(context.Background())
 	defer l.Close()
 	noTarget := func(string) (netip.Addr, bool) { return netip.Addr{}, false }
 	go l.acceptLoop(newTestStack(t, "10.211.4.1"), 505, portForwardKey{email: "a@x", proto: tcpForward}, noTarget)
@@ -463,4 +464,16 @@ primed:
 	if string(gotUDP) != wantUDP {
 		t.Errorf("UDP round trip = %q, want %q", gotUDP, wantUDP)
 	}
+
+	// Withdrawing the forward must end the connection already relaying over
+	// it, not just refuse new ones.
+	withdrawn := inst
+	withdrawn.Peers = []amneziawg.Peer{
+		{Email: "client@test", PublicKey: clientPub, AllowedIPs: []string{clientAddr + "/32"}},
+	}
+	set.Reconcile(withdrawn)
+	tcpConn.SetReadDeadline(time.Now().Add(5 * time.Second))
+	if _, err := tcpConn.Read(make([]byte, 1)); !errors.Is(err, io.EOF) {
+		t.Errorf("read on the relayed connection after its forward was withdrawn = %v, want io.EOF", err)
+	}
 }

internal/amneziawgnet/relay.go

@@ -455,6 +455,19 @@ func (u *UDPRelay) pump(src netip.AddrPort, sess *socks5UDPSession) {
 	}
 }
 
+// CloseSessions tears down the sessions whose tunnel-internal source address
+// matches; pump retires each one's goroutine on the resulting read error.
+func (u *UDPRelay) CloseSessions(match func(netip.Addr) bool) {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	for k, s := range u.sessions {
+		if match(k.Addr()) {
+			s.Close()
+			delete(u.sessions, k)
+		}
+	}
+}
+
 // Close tears down every open session. Call when the owning Device is
 // closed.
 func (u *UDPRelay) Close() {