diff --git a/internal/amneziawgnet/manager.go b/internal/amneziawgnet/manager.go index f9bae32c2..898741b0d 100644 --- a/internal/amneziawgnet/manager.go +++ b/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 } diff --git a/internal/amneziawgnet/manager_test.go b/internal/amneziawgnet/manager_test.go index 6baaeb751..497b4360d 100644 --- a/internal/amneziawgnet/manager_test.go +++ b/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 diff --git a/internal/amneziawgnet/netstack.go b/internal/amneziawgnet/netstack.go index a9a08c41f..b05467f53 100644 --- a/internal/amneziawgnet/netstack.go +++ b/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 { diff --git a/internal/amneziawgnet/portfwd.go b/internal/amneziawgnet/portfwd.go index 80cbef491..0a7aa5a4d 100644 --- a/internal/amneziawgnet/portfwd.go +++ b/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 - closing chan struct{} + ln net.Listener + 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 -// left to finish on their own -- there's no shared state to tear down early -// for, and an abrupt cut would just look like a network error to whichever -// external client was mid-transfer. +// Close stops accepting and cuts the connections already relaying, as +// udpForwardListener.Close does: a withdrawn forward must not stay usable. func (l *tcpForwardListener) Close() { - close(l.closing) + l.cancel() l.ln.Close() } diff --git a/internal/amneziawgnet/portfwd_test.go b/internal/amneziawgnet/portfwd_test.go index d33405e46..43eeef3e1 100644 --- a/internal/amneziawgnet/portfwd_test.go +++ b/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) + } } diff --git a/internal/amneziawgnet/relay.go b/internal/amneziawgnet/relay.go index be8cb915b..45b110c8b 100644 --- a/internal/amneziawgnet/relay.go +++ b/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() {