fix(amneziawg): three defects in the embedded relay's connection handling

Half-close. Both TCP relays -- RelayTCP into Xray's SOCKS5 inbound and
relayTCPForward into a peer's tunnel address -- waited on a single `done`
receive and then closed both sides. A client that finished sending and shut
down its write side therefore had the connection torn down before the
response came back. pipeBothWays now runs both directions to completion and
propagates the half-close via CloseWrite (which *net.TCPConn and
*gonet.TCPConn both implement), falling back to a full Close for anything
that does not.

Waiting for both directions reintroduces the risk the old single-receive was
implicitly avoiding: a peer that vanishes mid-transfer would pin the pair
forever. guardedReader bounds that, but as an idle window rather than a
total one -- the deadline is re-armed on every read once armed -- so a slow
transfer is never cut, while a silent peer is. Two minutes matches the idle
window UDPRelay.pump and portForwardUDPIdleTimeout already use.

UDP session retirement. pump's teardown deleted the map entry by key alone,
so a session that lost a create race evicted whichever session currently
held that source, orphaning a live flow. It now retires only its own entry,
and Handle keeps the already-published session when it loses the race. The
map is keyed on netip.AddrPort rather than src.String(), matching
udpForwardListener next door and dropping one allocation per relayed
datagram.

SOCKS5 reply decoding. bytesReader had a value receiver, so each Read
restarted at the head of the slice, and receive never advanced past a
domain-form address because its switch only handled ATYP 0x01 and 0x04 -- a
0x03 reply decoded to a wrong source, port and payload. splitSocks5Addr
replaces it: all three address forms, length-checked at every step, with the
domain form accepting only a literal. Resolving there would have put a
blocking DNS lookup on the receive path, and a datagram's own source is an
address already. Unreachable against Xray's own inbound, which always
answers with an IP, so this is a latent-bug fix rather than an observed one.

internal/amneziawgnet/portfwd.go

@@ -24,7 +24,6 @@ package amneziawgnet
 import (
 	"context"
 	"fmt"
-	"io"
 	"net"
 	"net/netip"
 	"sync"
@@ -327,10 +326,7 @@ func relayTCPForward(gstack *stack.Stack, conn net.Conn, inboundID int, key port
 	}
 	defer tunnelConn.Close()
 
-	done := make(chan struct{}, 2)
-	go func() { _, _ = io.Copy(tunnelConn, conn); done <- struct{}{} }()
-	go func() { _, _ = io.Copy(conn, tunnelConn); done <- struct{}{} }()
-	<-done
+	pipeBothWays(conn, tunnelConn)
 }
 
 // Close stops accepting new connections. Already-relaying connections are

internal/amneziawgnet/relay.go

@@ -15,6 +15,7 @@ import (
 	"net"
 	"net/netip"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/net/proxy"
@@ -60,7 +61,7 @@ func SocksInboundSettings(emails []string, password string) ([]byte, error) {
 }
 
 // RelayTCP dials r.Addr, authenticates as email, issues a SOCKS5 CONNECT to
-// dest, and pipes bytes both ways until either side closes or errors.
+// dest, and pipes bytes both ways until both directions end.
 // Blocks until the relay ends; meant to be called from (or as) an
 // AttachTCPForwarder handler, which already runs each connection on its own
 // goroutine.
@@ -80,10 +81,57 @@ func (r SocksRelay) RelayTCP(conn *gonet.TCPConn, email string, dest netip.AddrP
 	}
 	defer upstream.Close()
 
-	done := make(chan struct{}, 2)
-	go func() { _, _ = io.Copy(upstream, conn); done <- struct{}{} }()
-	go func() { _, _ = io.Copy(conn, upstream); done <- struct{}{} }()
-	<-done
+	pipeBothWays(conn, upstream)
+}
+
+// halfCloseIdle bounds how long the surviving direction of a half-closed pair
+// may sit idle, so a peer that vanished mid-transfer cannot pin it forever.
+const halfCloseIdle = 2 * time.Minute
+
+// closeWriter is the half-close half of *net.TCPConn and *gonet.TCPConn.
+type closeWriter interface{ CloseWrite() error }
+
+// guardedReader reads one side of a relayed pair, re-arming its read deadline
+// on every read once armed, so the bound is an idle window, not a total one.
+type guardedReader struct {
+	conn  net.Conn
+	armed atomic.Bool
+}
+
+func (r *guardedReader) Read(p []byte) (int, error) {
+	if r.armed.Load() {
+		_ = r.conn.SetReadDeadline(time.Now().Add(halfCloseIdle))
+	}
+	return r.conn.Read(p)
+}
+
+// arm bounds this side's remaining reads, including one already in flight.
+func (r *guardedReader) arm() {
+	r.armed.Store(true)
+	_ = r.conn.SetReadDeadline(time.Now().Add(halfCloseIdle))
+}
+
+// pipeBothWays copies a and b into each other until BOTH directions end,
+// half-closing each far side in turn so a half-closed peer still gets its reply.
+func pipeBothWays(a, b net.Conn) {
+	ga, gb := &guardedReader{conn: a}, &guardedReader{conn: b}
+	var wg sync.WaitGroup
+	wg.Add(2)
+	// Arming dst bounds the direction still reading from it -- the one this
+	// copy just signalled EOF to.
+	pipe := func(dst, src *guardedReader) {
+		defer wg.Done()
+		_, _ = io.Copy(dst.conn, src)
+		if cw, ok := dst.conn.(closeWriter); ok {
+			_ = cw.CloseWrite()
+		} else {
+			_ = dst.conn.Close()
+		}
+		dst.arm()
+	}
+	go pipe(gb, ga)
+	go pipe(ga, gb)
+	wg.Wait()
 }
 
 // socks5UDPSession is one established SOCKS5 UDP ASSOCIATE session: udpConn
@@ -278,36 +326,45 @@ func (s *socks5UDPSession) receive(buf []byte) (netip.AddrPort, []byte, error) {
 	if len(data) < 4 {
 		return netip.AddrPort{}, nil, fmt.Errorf("amneziawgnet: short SOCKS5 UDP reply (%d bytes)", n)
 	}
-	atyp := data[3]
-	data = data[4:]
-	addr, err := readSocks5Addr(bytesReader{data}, atyp)
+	addr, rest, err := splitSocks5Addr(data[4:], data[3])
 	if err != nil {
 		return netip.AddrPort{}, nil, err
 	}
-	switch atyp {
-	case 0x01:
-		data = data[4:]
-	case 0x04:
-		data = data[16:]
-	}
-	if len(data) < 2 {
+	if len(rest) < 2 {
 		return netip.AddrPort{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 UDP reply port")
 	}
-	port := binary.BigEndian.Uint16(data[:2])
-	return netip.AddrPortFrom(addr, port), data[2:], nil
+	return netip.AddrPortFrom(addr, binary.BigEndian.Uint16(rest[:2])), rest[2:], nil
 }
 
-// bytesReader is the minimal io.Reader readSocks5Addr needs, over an
-// in-memory slice that's already fully available (a received UDP
-// datagram) -- avoids pulling in bytes.Reader just for this.
-type bytesReader struct{ b []byte }
-
-func (r bytesReader) Read(p []byte) (int, error) {
-	n := copy(p, r.b)
-	if n < len(p) {
-		return n, io.ErrUnexpectedEOF
+// splitSocks5Addr decodes the address at the head of b for address type atyp
+// and returns it with whatever follows, length-checked at every step.
+func splitSocks5Addr(b []byte, atyp byte) (netip.Addr, []byte, error) {
+	switch atyp {
+	case 0x01:
+		if len(b) < 4 {
+			return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 IPv4 reply address")
+		}
+		return netip.AddrFrom4([4]byte(b[:4])), b[4:], nil
+	case 0x04:
+		if len(b) < 16 {
+			return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 IPv6 reply address")
+		}
+		return netip.AddrFrom16([16]byte(b[:16])), b[16:], nil
+	case 0x03:
+		// Resolving here would block the receive loop on DNS, and a datagram's
+		// own source is an address already -- so only a literal is accepted.
+		if len(b) < 1 || len(b) < 1+int(b[0]) {
+			return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 domain reply address")
+		}
+		name := string(b[1 : 1+int(b[0])])
+		addr, err := netip.ParseAddr(name)
+		if err != nil {
+			return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: SOCKS5 UDP reply from non-literal address %q", name)
+		}
+		return addr, b[1+int(b[0]):], nil
+	default:
+		return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: unsupported SOCKS5 address type %d", atyp)
 	}
-	return n, nil
 }
 
 // UDPRelay tracks one SOCKS5 UDP ASSOCIATE session per source (tunnel-
@@ -318,13 +375,15 @@ type UDPRelay struct {
 	relay  SocksRelay
 	gstack *stack.Stack
 
+	// Keyed by the comparable netip.AddrPort, like udpForwardListener's own
+	// session map: src.String() would allocate on every relayed datagram.
 	mu       sync.Mutex
-	sessions map[string]*socks5UDPSession
+	sessions map[netip.AddrPort]*socks5UDPSession
 }
 
 // NewUDPRelay creates a UDPRelay for one embedded AmneziaWG Device's stack.
 func NewUDPRelay(relay SocksRelay, gstack *stack.Stack) *UDPRelay {
-	return &UDPRelay{relay: relay, gstack: gstack, sessions: map[string]*socks5UDPSession{}}
+	return &UDPRelay{relay: relay, gstack: gstack, sessions: map[netip.AddrPort]*socks5UDPSession{}}
 }
 
 // Handle relays one packet from src (the peer's tunnel-internal source) to
@@ -334,20 +393,28 @@ func NewUDPRelay(relay SocksRelay, gstack *stack.Stack) *UDPRelay {
 // reusing it for subsequent packets from the same src.
 func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte) {
 	u.mu.Lock()
-	sess, ok := u.sessions[src.String()]
+	sess, ok := u.sessions[src]
 	u.mu.Unlock()
 
 	if !ok {
-		var err error
-		sess, err = newSocks5UDPSession(u.relay.Addr, email, u.relay.Password)
+		fresh, err := newSocks5UDPSession(u.relay.Addr, email, u.relay.Password)
 		if err != nil {
 			logger.Warningf("amneziawgnet: UDPRelay: SOCKS5 associate for %q: %v", email, err)
 			return
 		}
 		u.mu.Lock()
-		u.sessions[src.String()] = sess
-		u.mu.Unlock()
-		go u.pump(src, sess)
+		// Associating happens off-lock, so a concurrent Handle for the same src
+		// may already have published one; keep it, so the key has a single pump.
+		if existing, dup := u.sessions[src]; dup {
+			u.mu.Unlock()
+			fresh.Close()
+			sess = existing
+		} else {
+			u.sessions[src] = fresh
+			u.mu.Unlock()
+			sess = fresh
+			go u.pump(src, fresh)
+		}
 	}
 	if err := sess.sendTo(dst, payload); err != nil {
 		logger.Warningf("amneziawgnet: UDPRelay: send to %s: %v", dst, err)
@@ -360,7 +427,11 @@ func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte)
 func (u *UDPRelay) pump(src netip.AddrPort, sess *socks5UDPSession) {
 	defer func() {
 		u.mu.Lock()
-		delete(u.sessions, src.String())
+		// Only retire our own entry: a delete by key alone would evict whichever
+		// session currently holds src, orphaning a live one.
+		if u.sessions[src] == sess {
+			delete(u.sessions, src)
+		}
 		u.mu.Unlock()
 		sess.Close()
 	}()

internal/amneziawgnet/relay_test.go

@@ -0,0 +1,257 @@
+package amneziawgnet
+
+import (
+	"io"
+	"net"
+	"net/netip"
+	"testing"
+	"time"
+)
+
+// newDeadUDPSession builds a socks5UDPSession over real but already-closed
+// sockets, so pump's receive fails immediately and its teardown runs at once.
+func newDeadUDPSession(t *testing.T) *socks5UDPSession {
+	t.Helper()
+	peer, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+	t.Cleanup(func() { _ = peer.Close() })
+	udpConn, err := net.DialUDP("udp", nil, peer.LocalAddr().(*net.UDPAddr))
+	if err != nil {
+		t.Fatalf("dial udp: %v", err)
+	}
+	ctrlLn, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen tcp: %v", err)
+	}
+	t.Cleanup(func() { _ = ctrlLn.Close() })
+	ctrl, err := net.Dial("tcp", ctrlLn.Addr().String())
+	if err != nil {
+		t.Fatalf("dial tcp: %v", err)
+	}
+	_ = udpConn.Close()
+	return &socks5UDPSession{ctrl: ctrl, udpConn: udpConn}
+}
+
+// TestUDPRelayPumpOnlyRetiresItsOwnSession pins the flow that survives a
+// duplicate associate: a losing pump must not evict the published session.
+func TestUDPRelayPumpOnlyRetiresItsOwnSession(t *testing.T) {
+	relay := NewUDPRelay(SocksRelay{Addr: "127.0.0.1:1", Password: "x"}, nil)
+	src := netip.MustParseAddrPort("10.8.1.5:51820")
+
+	live := newDeadUDPSession(t)
+	superseded := newDeadUDPSession(t)
+	relay.sessions[src] = live
+
+	// Returns as soon as receive fails on the closed socket, so no wait is needed.
+	relay.pump(src, superseded)
+
+	got, ok := relay.sessions[src]
+	if !ok {
+		t.Fatal("live session was evicted: a retiring pump deleted src's entry regardless of which session held it")
+	}
+	if got != live {
+		t.Fatalf("sessions[%v] = %p, want the live session %p", src, got, live)
+	}
+}
+
+// TestUDPRelayCloseDropsEverySession keeps Close's contract explicit now that
+// pump's teardown is conditional on still owning the key.
+func TestUDPRelayCloseDropsEverySession(t *testing.T) {
+	relay := NewUDPRelay(SocksRelay{Addr: "127.0.0.1:1", Password: "x"}, nil)
+	for _, s := range []string{"10.8.1.5:51820", "10.8.1.6:2000"} {
+		relay.sessions[netip.MustParseAddrPort(s)] = newDeadUDPSession(t)
+	}
+
+	relay.Close()
+
+	if n := len(relay.sessions); n != 0 {
+		t.Fatalf("Close left %d sessions behind, want 0", n)
+	}
+}
+
+// tcpPair returns a connected pair of real loopback TCP conns; net.Pipe would
+// not do, since these tests turn on CloseWrite, which it does not implement.
+func tcpPair(t *testing.T) (client, server net.Conn) {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	defer ln.Close()
+
+	type accepted struct {
+		conn net.Conn
+		err  error
+	}
+	ch := make(chan accepted, 1)
+	go func() {
+		c, err := ln.Accept()
+		ch <- accepted{c, err}
+	}()
+	client, err = net.Dial("tcp", ln.Addr().String())
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	got := <-ch
+	if got.err != nil {
+		t.Fatalf("accept: %v", got.err)
+	}
+	t.Cleanup(func() { _ = client.Close(); _ = got.conn.Close() })
+	return client, got.conn
+}
+
+// TestPipeBothWaysDeliversReplyAfterHalfClose is the half-close regression: a
+// client that shuts down its write side must still receive the full response.
+func TestPipeBothWaysDeliversReplyAfterHalfClose(t *testing.T) {
+	const request = "GET / HTTP/1.0\r\n\r\n"
+	const response = "the reply that arrives only after the request is complete"
+
+	client, a := tcpPair(t)
+	b, server := tcpPair(t)
+	go pipeBothWays(a, b)
+
+	if _, err := client.Write([]byte(request)); err != nil {
+		t.Fatalf("client write: %v", err)
+	}
+	// The half-close the old relay treated as "tear the whole pair down".
+	if err := client.(*net.TCPConn).CloseWrite(); err != nil {
+		t.Fatalf("client CloseWrite: %v", err)
+	}
+
+	_ = server.SetReadDeadline(time.Now().Add(10 * time.Second))
+	gotReq, err := io.ReadAll(server)
+	if err != nil {
+		t.Fatalf("server read: %v", err)
+	}
+	if string(gotReq) != request {
+		t.Fatalf("server got request %q, want %q", gotReq, request)
+	}
+
+	if _, err := server.Write([]byte(response)); err != nil {
+		t.Fatalf("server write: %v", err)
+	}
+	if err := server.(*net.TCPConn).CloseWrite(); err != nil {
+		t.Fatalf("server CloseWrite: %v", err)
+	}
+
+	_ = client.SetReadDeadline(time.Now().Add(10 * time.Second))
+	gotResp, err := io.ReadAll(client)
+	if err != nil {
+		t.Fatalf("client read: %v", err)
+	}
+	if string(gotResp) != response {
+		t.Fatalf("client got response %q, want %q: the reply was cut off by the half-close", gotResp, response)
+	}
+}
+
+// TestPipeBothWaysClosesWhenBothSidesFinish keeps the teardown contract: both
+// directions ending must return, not hang on the idle bound.
+func TestPipeBothWaysClosesWhenBothSidesFinish(t *testing.T) {
+	client, a := tcpPair(t)
+	b, server := tcpPair(t)
+
+	done := make(chan struct{})
+	go func() { defer close(done); pipeBothWays(a, b) }()
+
+	_ = client.(*net.TCPConn).CloseWrite()
+	_, _ = io.ReadAll(server)
+	_ = server.(*net.TCPConn).CloseWrite()
+	_, _ = io.ReadAll(client)
+
+	select {
+	case <-done:
+	case <-time.After(10 * time.Second):
+		t.Fatal("pipeBothWays did not return after both directions ended")
+	}
+}
+
+// liveUDPSession returns a session whose udpConn is connected to the returned
+// peer, so a test can hand receive() one exact reply datagram.
+func liveUDPSession(t *testing.T) (*socks5UDPSession, *net.UDPConn) {
+	t.Helper()
+	peer, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+	t.Cleanup(func() { _ = peer.Close() })
+	udpConn, err := net.DialUDP("udp", nil, peer.LocalAddr().(*net.UDPAddr))
+	if err != nil {
+		t.Fatalf("dial udp: %v", err)
+	}
+	t.Cleanup(func() { _ = udpConn.Close() })
+	return &socks5UDPSession{udpConn: udpConn}, peer
+}
+
+// sendReply delivers one raw datagram to sess's socket.
+func sendReply(t *testing.T, sess *socks5UDPSession, peer *net.UDPConn, datagram []byte) {
+	t.Helper()
+	if _, err := peer.WriteToUDP(datagram, sess.udpConn.LocalAddr().(*net.UDPAddr)); err != nil {
+		t.Fatalf("write reply: %v", err)
+	}
+	_ = sess.udpConn.SetReadDeadline(time.Now().Add(5 * time.Second))
+}
+
+// TestSocks5ReceiveDecodesReplyAddressTypes covers all three ATYP forms; the
+// domain form used to misread its own length byte and never skip the name.
+func TestSocks5ReceiveDecodesReplyAddressTypes(t *testing.T) {
+	tests := []struct {
+		name     string
+		addrPart []byte
+		wantAddr string
+	}{
+		{"IPv4", []byte{0x01, 10, 0, 0, 7}, "10.0.0.7"},
+		{"IPv6", append([]byte{0x04}, netip.MustParseAddr("2001:db8::5").AsSlice()...), "2001:db8::5"},
+		{"domain holding a literal", append([]byte{0x03, 8}, []byte("10.0.0.9")...), "10.0.0.9"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sess, peer := liveUDPSession(t)
+			payload := []byte("the-actual-datagram-payload")
+			datagram := append([]byte{0x00, 0x00, 0x00}, tt.addrPart...)
+			datagram = append(datagram, 0x1f, 0x90) // port 8080
+			datagram = append(datagram, payload...)
+			sendReply(t, sess, peer, datagram)
+
+			buf := make([]byte, 4096)
+			from, got, err := sess.receive(buf)
+			if err != nil {
+				t.Fatalf("receive: %v", err)
+			}
+			want := netip.AddrPortFrom(netip.MustParseAddr(tt.wantAddr), 8080)
+			if from != want {
+				t.Errorf("source = %v, want %v", from, want)
+			}
+			if string(got) != string(payload) {
+				t.Errorf("payload = %q, want %q", got, payload)
+			}
+		})
+	}
+}
+
+// TestSocks5ReceiveRejectsTruncatedReplies pins that a short datagram is an
+// error, not a slice-bounds panic in the relay's own pump goroutine.
+func TestSocks5ReceiveRejectsTruncatedReplies(t *testing.T) {
+	tests := []struct {
+		name     string
+		datagram []byte
+	}{
+		{"header only, IPv4 announced", []byte{0x00, 0x00, 0x00, 0x01}},
+		{"IPv4 address cut short", []byte{0x00, 0x00, 0x00, 0x01, 10, 0}},
+		{"IPv6 address cut short", []byte{0x00, 0x00, 0x00, 0x04, 0x20, 0x01}},
+		{"domain length past the end", []byte{0x00, 0x00, 0x00, 0x03, 40, 'a', 'b'}},
+		{"address complete but port missing", []byte{0x00, 0x00, 0x00, 0x01, 10, 0, 0, 7}},
+		{"unsupported address type", []byte{0x00, 0x00, 0x00, 0x09, 1, 2, 3, 4, 0, 80}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sess, peer := liveUDPSession(t)
+			sendReply(t, sess, peer, tt.datagram)
+			buf := make([]byte, 4096)
+			if _, _, err := sess.receive(buf); err == nil {
+				t.Fatal("receive accepted a malformed datagram instead of returning an error")
+			}
+		})
+	}
+}