fix(amneziawg): take a forwarded TCP connection's source from the request

The TCP forwarder handler read the tunnel client's address from the
accepted connection's RemoteAddr, which asks the live endpoint and is
nil once the peer's RST has been processed. A client sending SYN, ACK
and RST through its own tunnel could hit that window and crash the panel
on a goroutine with no recover. The forwarder now hands the handler the
source from the request's endpoint ID, as it already did for the
destination.

internal/amneziawgnet/bench_test.go

@@ -134,7 +134,7 @@ func newBenchTunnel(b *testing.B, listenPort int, serverAddr, clientAddr string)
 	if err != nil {
 		b.Fatalf("newUnconfiguredDevice: %v", err)
 	}
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		switch dest.Port() {
 		case benchDiscardPort:

internal/amneziawgnet/device_test.go

@@ -80,14 +80,9 @@ func TestNewDeviceHandshakeForwarderAndIdentity(t *testing.T) {
 	// purely from the decapsulated packet, not from any routing table.
 	wantDest := netip.MustParseAddrPort("10.201.9.9:9999")
 
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
 		defer conn.Close()
-		srcAddrPort, parseErr := netip.ParseAddrPort(conn.RemoteAddr().String())
+		peer, ok := idx.Lookup(src.Addr().Unmap())
-		var peer amneziawg.Peer
-		var ok bool
-		if parseErr == nil {
-			peer, ok = idx.Lookup(srcAddrPort.Addr().Unmap())
-		}
 		got <- recovered{email: peer.Email, ok: ok, dest: dest}
 		io.Copy(io.Discard, conn)
 	})
@@ -283,7 +278,7 @@ func TestNewDeviceHeaderProtectionAndContentPaddingRoundTrip(t *testing.T) {
 	const wantRequest = "hello from client"
 	const wantReply = "hello from server"
 	serverDone := make(chan error, 1)
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		buf := make([]byte, len(wantRequest))
 		if _, err := io.ReadFull(conn, buf); err != nil {
@@ -461,7 +456,7 @@ func TestNewDeviceRandomTrailersAndDisableCookiesRoundTrip(t *testing.T) {
 	const wantRequest = "hello from client, with a trailer"
 	const wantReply = "hello from server, with a trailer"
 	serverDone := make(chan error, 1)
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, dest netip.AddrPort) {
 		defer conn.Close()
 		buf := make([]byte, len(wantRequest))
 		if _, err := io.ReadFull(conn, buf); err != nil {

internal/amneziawgnet/diagnostics_test.go

@@ -77,7 +77,7 @@ func TestDiagnoseDeviceReportsListenPortAndPeerState(t *testing.T) {
 	// Registering the forwarder doesn't require any peer to be configured
 	// yet, so this ordering is free; it's Configure's IpcSet that must
 	// never run before the forwarder is registered.
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _ netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, _, _ netip.AddrPort) {
 		defer conn.Close()
 		io.Copy(io.Discard, conn)
 	})

internal/amneziawgnet/forwarder.go

@@ -12,15 +12,18 @@ import (
 // AttachTCPForwarder attaches a TCP forwarder to gstack in promiscuous +
 // spoofing mode, so it accepts connections addressed to any destination --
 // not just the stack's own configured local address -- and hands the
-// handler both the accepted connection and the tunnel client's real,
+// handler the accepted connection, the tunnel client's source and its real,
 // dynamically-arbitrary destination (recovered from the connection's own
 // TransportEndpointID, not from any preconfigured routing table). This is
 // the mechanism the whole embedded-AmneziaWG design depends on: what the
 // handler does with that destination (dial it directly, relay it into
 // Xray's SOCKS5 inbound, ...) is entirely up to the caller.
 //
+// Handlers must identify the peer by src, never conn.RemoteAddr(): that asks
+// the live endpoint and is nil once a RST is processed, possibly before the call.
+//
 // Adapted from xtls/xray-core's proxy/wireguard/tun.go createForwarder (MIT).
-func AttachTCPForwarder(gstack *stack.Stack, handler func(conn *gonet.TCPConn, dest netip.AddrPort)) {
+func AttachTCPForwarder(gstack *stack.Stack, handler func(conn *gonet.TCPConn, src, dest netip.AddrPort)) {
 	enablePromiscuousRouting(gstack)
 
 	fwd := tcp.NewForwarder(gstack, 0, 65535, func(r *tcp.ForwarderRequest) {
@@ -33,8 +36,9 @@ func AttachTCPForwarder(gstack *stack.Stack, handler func(conn *gonet.TCPConn, d
 				r.Complete(true)
 				return
 			}
+			src := netip.AddrPortFrom(addrFromTcpip(id.RemoteAddress), id.RemotePort)
 			dest := netip.AddrPortFrom(addrFromTcpip(id.LocalAddress), id.LocalPort)
-			handler(gonet.NewTCPConn(&wq, ep), dest)
+			handler(gonet.NewTCPConn(&wq, ep), src, dest)
 			ep.Close()
 			r.Complete(false)
 		}(r)

internal/amneziawgnet/forwarder_test.go

@@ -0,0 +1,110 @@
+package amneziawgnet
+
+import (
+	"net"
+	"net/netip"
+	"testing"
+	"time"
+
+	awgtun "github.com/amnezia-vpn/amneziawg-go/v3/tun"
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+	"gvisor.dev/gvisor/pkg/tcpip/checksum"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+)
+
+func rawTCPSegment(src, dst netip.AddrPort, seq, ack uint32, flags header.TCPFlags) []byte {
+	srcIP, dstIP := tcpip.AddrFromSlice(src.Addr().AsSlice()), tcpip.AddrFromSlice(dst.Addr().AsSlice())
+	packet := make([]byte, header.IPv4MinimumSize+header.TCPMinimumSize)
+	ip := header.IPv4(packet)
+	ip.Encode(&header.IPv4Fields{
+		TotalLength: uint16(len(packet)),
+		TTL:         64,
+		Protocol:    uint8(header.TCPProtocolNumber),
+		SrcAddr:     srcIP,
+		DstAddr:     dstIP,
+	})
+	ip.SetChecksum(^ip.CalculateChecksum())
+	tcp := header.TCP(packet[header.IPv4MinimumSize:])
+	tcp.Encode(&header.TCPFields{
+		SrcPort:    src.Port(),
+		DstPort:    dst.Port(),
+		SeqNum:     seq,
+		AckNum:     ack,
+		DataOffset: header.TCPMinimumSize,
+		Flags:      flags,
+		WindowSize: 65535,
+	})
+	xsum := header.PseudoHeaderChecksum(header.TCPProtocolNumber, srcIP, dstIP, header.TCPMinimumSize)
+	tcp.SetChecksum(^tcp.CalculateChecksum(checksum.Checksum(nil, xsum)))
+	return packet
+}
+
+func injectFromPeer(t *testing.T, tun awgtun.Device, packet []byte) {
+	t.Helper()
+	if _, err := tun.Write([][]byte{packet}, 0); err != nil {
+		t.Fatalf("inject packet: %v", err)
+	}
+}
+
+// A peer's RST can beat the handler, and a reset endpoint has no RemoteAddr:
+// deriving the source from it panicked the panel on a goroutine with no recover.
+func TestTCPForwarderHandsOverSourceOfAnAlreadyResetConnection(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.9.0.1")}, 1420)
+	if err != nil {
+		t.Fatalf("createNetTUNWithStack: %v", err)
+	}
+	defer tun.Close()
+
+	peer := netip.MustParseAddrPort("10.9.0.2:40000")
+	target := netip.MustParseAddrPort("203.0.113.7:443")
+
+	type handedOver struct {
+		src, dest  netip.AddrPort
+		remoteAddr net.Addr
+	}
+	accepted, resetSent := make(chan struct{}), make(chan struct{})
+	got := make(chan handedOver, 1)
+	AttachTCPForwarder(gstack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
+		defer conn.Close()
+		close(accepted)
+		<-resetSent
+		// Returns only once the stack has processed the RST.
+		_, _ = conn.Read(make([]byte, 1))
+		got <- handedOver{src: src, dest: dest, remoteAddr: conn.RemoteAddr()}
+	})
+
+	const peerISN = 1000
+	injectFromPeer(t, tun, rawTCPSegment(peer, target, peerISN, 0, header.TCPFlagSyn))
+	bufs, sizes := [][]byte{make([]byte, 2048)}, make([]int, 1)
+	if _, err := tun.Read(bufs, sizes, 0); err != nil {
+		t.Fatalf("read SYN-ACK: %v", err)
+	}
+	synAck := header.TCP(bufs[0][header.IPv4MinimumSize:sizes[0]])
+	if synAck.Flags() != header.TCPFlagSyn|header.TCPFlagAck {
+		t.Fatalf("first reply has flags %s, want SYN-ACK", synAck.Flags())
+	}
+	injectFromPeer(t, tun, rawTCPSegment(peer, target, peerISN+1, synAck.SequenceNumber()+1, header.TCPFlagAck))
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("handler never ran")
+	}
+	injectFromPeer(t, tun, rawTCPSegment(peer, target, peerISN+1, 0, header.TCPFlagRst))
+	close(resetSent)
+
+	select {
+	case h := <-got:
+		if h.remoteAddr != nil {
+			t.Fatalf("RemoteAddr() = %v after the RST, want nil: the test no longer reaches the reset state", h.remoteAddr)
+		}
+		if h.src != peer {
+			t.Errorf("src = %s, want %s", h.src, peer)
+		}
+		if h.dest != target {
+			t.Errorf("dest = %s, want %s", h.dest, target)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("handler never saw the reset")
+	}
+}

internal/amneziawgnet/manager.go

@@ -205,15 +205,10 @@ func (m *Manager) ensureLocked(d Desired) error {
 		structFP:     structFP,
 	}
 	next.peers.Store(NewPeerIndex(inst.Peers))
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
-		srcAddrPort, err := netip.ParseAddrPort(conn.RemoteAddr().String())
-		if err != nil {
-			conn.Close()
-			return
-		}
 		// 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())
+		peer, ok := next.lookupPeer(src.Addr().Unmap())
 		if !ok {
 			conn.Close()
 			return

internal/amneziawgnet/relay_e2e_test.go

@@ -161,13 +161,8 @@ func TestSocksRelayAgainstRealXray(t *testing.T) {
 	udpRelay := NewUDPRelay(relay, dev.Stack)
 	defer udpRelay.Close()
 
-	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, src, dest netip.AddrPort) {
-		srcAddrPort, err := netip.ParseAddrPort(conn.RemoteAddr().String())
+		peer, ok := idx.Lookup(src.Addr().Unmap())
-		if err != nil {
-			conn.Close()
-			return
-		}
-		peer, ok := idx.Lookup(srcAddrPort.Addr().Unmap())
 		if !ok {
 			conn.Close()
 			return