perf(amneziawg): reuse relay buffers and bound the bridge connect time (#7043)

Every TCP connection through a built-in AmneziaWG inbound or outbound
allocated two fresh 32 KiB copy buffers, because io.Copy on a *net.TCPConn
goes through ReadFrom/WriteTo and ignores a caller's buffer. Under a stream
of short connections that is a lot of garbage. The relays now copy with
buffers from a sync.Pool; one side of each pair is always a userspace
netstack connection, so no kernel splice path is given up.

The connect to the local Xray bridge and its SOCKS5 handshake had no time
limit, so a bridge that accepted and never answered pinned the relay's
goroutines and descriptors forever. Both are now bounded at 10 s, and the
deadline is cleared once the handshake succeeds so the relayed connection
keeps its half-close behaviour.

internal/amneziawgnet/egress.go

@@ -469,8 +469,8 @@ func (s *socks5EgressServer) relayTCP(dev *Device, tag string, upstream net.Conn
 		return
 	}
 	done := make(chan struct{}, 2)
-	go func() { _, _ = io.Copy(tunnelConn, upstream); done <- struct{}{} }()
+	go func() { copyPooled(tunnelConn, upstream); done <- struct{}{} }()
-	go func() { _, _ = io.Copy(upstream, tunnelConn); done <- struct{}{} }()
+	go func() { copyPooled(upstream, tunnelConn); done <- struct{}{} }()
 	<-done
 }
 

internal/amneziawgnet/race_off_test.go

@@ -0,0 +1,5 @@
+//go:build !race
+
+package amneziawgnet
+
+const raceEnabled = false

internal/amneziawgnet/race_on_test.go

@@ -0,0 +1,5 @@
+//go:build race
+
+package amneziawgnet
+
+const raceEnabled = true

internal/amneziawgnet/relay.go

@@ -68,13 +68,7 @@ func SocksInboundSettings(emails []string, password string) ([]byte, error) {
 func (r SocksRelay) RelayTCP(conn *gonet.TCPConn, email string, dest netip.AddrPort) {
 	defer conn.Close()
 
-	auth := &proxy.Auth{User: email, Password: r.Password}
+	upstream, err := r.dialUpstream(email, dest)
-	dialer, err := proxy.SOCKS5("tcp", r.Addr, auth, proxy.Direct)
-	if err != nil {
-		logger.Warningf("amneziawgnet: RelayTCP: build SOCKS5 dialer: %v", err)
-		return
-	}
-	upstream, err := dialer.Dial("tcp", dest.String())
 	if err != nil {
 		logger.Warningf("amneziawgnet: RelayTCP: SOCKS5 CONNECT to %s as %q: %v", dest, email, err)
 		return
@@ -84,6 +78,60 @@ func (r SocksRelay) RelayTCP(conn *gonet.TCPConn, email string, dest netip.AddrP
 	pipeBothWays(conn, upstream)
 }
 
+// dialUpstream bounds the bridge connect and SOCKS5 handshake so a silent core cannot pin the relay.
+// It returns the raw *net.TCPConn with no deadline left: pipeBothWays needs its CloseWrite.
+func (r SocksRelay) dialUpstream(email string, dest netip.AddrPort) (net.Conn, error) {
+	bridge := net.Dialer{Timeout: socks5ConnectTimeout}
+	conn, err := bridge.DialContext(context.Background(), "tcp", r.Addr)
+	if err != nil {
+		return nil, err
+	}
+	if err := conn.SetDeadline(time.Now().Add(socks5ConnectTimeout)); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	auth := &proxy.Auth{User: email, Password: r.Password}
+	dialer, err := proxy.SOCKS5("tcp", r.Addr, auth, dialedConn{conn})
+	if err != nil {
+		conn.Close()
+		return nil, fmt.Errorf("build SOCKS5 dialer: %w", err)
+	}
+	upstream, err := dialer.Dial("tcp", dest.String())
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if err := upstream.SetDeadline(time.Time{}); err != nil {
+		upstream.Close()
+		return nil, err
+	}
+	return upstream, nil
+}
+
+// dialedConn hands proxy.SOCKS5 a connection already made under the bound.
+type dialedConn struct{ conn net.Conn }
+
+func (d dialedConn) Dial(string, string) (net.Conn, error) { return d.conn, nil }
+
+var socks5ConnectTimeout = 10 * time.Second
+
+// relayBuffers recycles copy buffers: a fresh 32 KiB pair per relayed connection made a
+// burst of short connections through AmneziaWG allocate and sweep gigabytes.
+var relayBuffers = sync.Pool{New: func() any { b := make([]byte, 32*1024); return &b }}
+
+// Wrappers hide ReadFrom/WriteTo, through which *net.TCPConn would copy with a buffer of its own.
+type (
+	relayWriter struct{ io.Writer }
+	relayReader struct{ io.Reader }
+)
+
+// copyPooled copies src into dst with a pooled buffer.
+func copyPooled(dst io.Writer, src io.Reader) {
+	buf := relayBuffers.Get().(*[]byte)
+	defer relayBuffers.Put(buf)
+	_, _ = io.CopyBuffer(relayWriter{dst}, relayReader{src}, *buf)
+}
+
 // 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
@@ -121,7 +169,7 @@ func pipeBothWays(a, b net.Conn) {
 	// copy just signalled EOF to.
 	pipe := func(dst, src *guardedReader) {
 		defer wg.Done()
-		_, _ = io.Copy(dst.conn, src)
+		copyPooled(dst.conn, src)
 		if cw, ok := dst.conn.(closeWriter); ok {
 			_ = cw.CloseWrite()
 		} else {

internal/amneziawgnet/relay_halfclose_test.go

@@ -0,0 +1,203 @@
+package amneziawgnet
+
+// Half-close through SocksRelay.RelayTCP needs the raw TCP conn to the bridge,
+// not x/net/proxy's *socks.Conn wrapper, which has no CloseWrite.
+
+import (
+	"bufio"
+	"context"
+	"errors"
+	"io"
+	"net"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/amnezia-vpn/amneziawg-go/v3/tun"
+	"github.com/amnezia-vpn/amneziawg-go/v3/tun/netstack"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+)
+
+func hcListen(tb testing.TB) net.Listener {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		tb.Fatalf("listen: %v", err)
+	}
+	return ln
+}
+
+type hcNetstacks struct {
+	client *netstack.Net
+	ln     *gonet.TCPListener
+	addr   netip.AddrPort
+	close  func()
+}
+
+func hcShuttle(src, dst tun.Device) {
+	batch := src.BatchSize()
+	bufs := make([][]byte, batch)
+	for i := range bufs {
+		bufs[i] = make([]byte, 65536)
+	}
+	sizes := make([]int, batch)
+	for {
+		n, err := src.Read(bufs, sizes, 0)
+		if err != nil {
+			return
+		}
+		out := make([][]byte, 0, n)
+		for i := 0; i < n; i++ {
+			out = append(out, bufs[i][:sizes[i]])
+		}
+		if _, err := dst.Write(out, 0); err != nil {
+			return
+		}
+	}
+}
+
+// hcNewNetstacks links a client stack to a server stack whose accepted conns are the
+// *gonet.TCPConn the relay gets in production.
+func hcNewNetstacks(tb testing.TB) *hcNetstacks {
+	ca := netip.MustParseAddr("10.77.0.1")
+	sa := netip.MustParseAddr("10.77.0.2")
+	ctun, cnet, err := netstack.CreateNetTUN([]netip.Addr{ca}, nil, 1420)
+	if err != nil {
+		tb.Fatalf("client netstack: %v", err)
+	}
+	stun, snet, err := netstack.CreateNetTUN([]netip.Addr{sa}, nil, 1420)
+	if err != nil {
+		tb.Fatalf("server netstack: %v", err)
+	}
+	go hcShuttle(ctun, stun)
+	go hcShuttle(stun, ctun)
+	addr := netip.AddrPortFrom(sa, 80)
+	ln, err := snet.ListenTCPAddrPort(addr)
+	if err != nil {
+		tb.Fatalf("server listen: %v", err)
+	}
+	return &hcNetstacks{client: cnet, ln: ln, addr: addr, close: func() {
+		_ = ln.Close()
+		_ = ctun.Close()
+		_ = stun.Close()
+	}}
+}
+
+func (n *hcNetstacks) pair(tb testing.TB) (client, server *gonet.TCPConn) {
+	ch := make(chan net.Conn, 1)
+	go func() {
+		c, err := n.ln.Accept()
+		if err != nil {
+			ch <- nil
+			return
+		}
+		ch <- c
+	}()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	client, err := n.client.DialContextTCPAddrPort(ctx, n.addr)
+	if err != nil {
+		tb.Fatalf("tunnel dial: %v", err)
+	}
+	s := <-ch
+	if s == nil {
+		tb.Fatal("tunnel accept failed")
+	}
+	return client, s.(*gonet.TCPConn)
+}
+
+// hcServeSocks answers one username/password SOCKS5 CONNECT (as Xray's inbound would),
+// reads the 4-byte request, writes payload and half-closes.
+func hcServeSocks(c net.Conn, payload []byte) error {
+	br := bufio.NewReader(c)
+	buf := make([]byte, 512)
+	read := func(n int) ([]byte, error) { _, err := io.ReadFull(br, buf[:n]); return buf[:n], err }
+	hdr, err := read(2)
+	if err != nil {
+		return err
+	}
+	if _, err := read(int(hdr[1])); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{5, 2}); err != nil {
+		return err
+	}
+	v, err := read(2)
+	if err != nil {
+		return err
+	}
+	if _, err := read(int(v[1])); err != nil {
+		return err
+	}
+	pl, err := read(1)
+	if err != nil {
+		return err
+	}
+	if _, err := read(int(pl[0])); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{1, 0}); err != nil {
+		return err
+	}
+	req, err := read(4)
+	if err != nil {
+		return err
+	}
+	alen := map[byte]int{1: 4, 4: 16}[req[3]]
+	if req[3] == 3 {
+		l, err := read(1)
+		if err != nil {
+			return err
+		}
+		alen = int(l[0])
+	}
+	if _, err := read(alen + 2); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{5, 0, 0, 1, 0, 0, 0, 0, 0, 0}); err != nil {
+		return err
+	}
+	if _, err := read(4); err != nil {
+		return err
+	}
+	if _, err := c.Write(payload); err != nil {
+		return err
+	}
+	if tc, ok := c.(*net.TCPConn); ok {
+		return tc.CloseWrite()
+	}
+	return errors.New("not tcp")
+}
+
+// TestHCHalfCloseThroughRelayTCP: a client that sends its request and half-closes
+// must still get the upstream's reply through SocksRelay.RelayTCP.
+func TestHCHalfCloseThroughRelayTCP(t *testing.T) {
+	ns := hcNewNetstacks(t)
+	defer ns.close()
+	bridge := hcListen(t)
+	defer bridge.Close()
+	relay := SocksRelay{Addr: bridge.Addr().String(), Password: "x"}
+	reply := []byte("reply-after-half-close")
+	served := make(chan error, 1)
+	go func() {
+		c, err := bridge.Accept()
+		if err != nil {
+			served <- err
+			return
+		}
+		defer c.Close()
+		served <- hcServeSocks(c, reply)
+	}()
+	client, server := ns.pair(t)
+	defer client.Close()
+	done := make(chan struct{})
+	go func() { relay.RelayTCP(server, "bench@example", netip.MustParseAddrPort("192.0.2.1:443")); close(done) }()
+	_, _ = client.Write([]byte("ping"))
+	_ = client.CloseWrite()
+	_ = client.SetReadDeadline(time.Now().Add(5 * time.Second))
+	got, err := io.ReadAll(client)
+	<-done
+	srvErr := <-served
+	if string(got) != string(reply) {
+		t.Fatalf("client got %q (read err %v, bridge err %v), want %q: the half-close killed the upstream before its reply", got, err, srvErr, reply)
+	}
+}

internal/amneziawgnet/relay_pool_test.go

@@ -0,0 +1,198 @@
+package amneziawgnet
+
+import (
+	"errors"
+	"io"
+	"net"
+	"net/netip"
+	"runtime"
+	"testing"
+	"time"
+)
+
+// A burst of short connections through AmneziaWG used to allocate a fresh 32 KiB copy buffer
+// per direction per connection (gigabytes per hour under churn); the pair must be reused.
+func TestPipeBothWaysReusesCopyBuffersAcrossConnections(t *testing.T) {
+	if raceEnabled {
+		t.Skip("under -race sync.Pool drops a share of Puts at random")
+	}
+	payload := make([]byte, 16*1024)
+	sink := make([]byte, 16*1024)
+	ping := make([]byte, len("ping")) // the server goroutine's own buffer
+	relay := func() {
+		tunnelClient, tunnelSide := tcpPair(t)
+		upstreamSide, upstreamServer := tcpPair(t)
+		done := make(chan struct{})
+		go func() { pipeBothWays(tunnelSide, upstreamSide); close(done) }()
+		go func() {
+			_, _ = io.ReadFull(upstreamServer, ping)
+			_, _ = upstreamServer.Write(payload)
+			_ = upstreamServer.Close()
+		}()
+		_, _ = tunnelClient.Write([]byte("ping"))
+		_ = tunnelClient.(*net.TCPConn).CloseWrite()
+		if _, err := io.ReadFull(tunnelClient, sink); err != nil {
+			t.Fatalf("reply: %v", err)
+		}
+		<-done
+	}
+	for range 5 {
+		relay()
+	}
+	const conns = 100
+	var before, after runtime.MemStats
+	runtime.GC()
+	runtime.ReadMemStats(&before)
+	for range conns {
+		relay()
+	}
+	runtime.ReadMemStats(&after)
+	perConn := (after.TotalAlloc - before.TotalAlloc) / conns
+	if perConn >= 32*1024 {
+		t.Fatalf("each relayed connection allocated %d bytes, want the copy buffers reused (< 32 KiB)", perConn)
+	}
+}
+
+// TestDialUpstreamGivesUpOnSilentBridge: a bridge that accepts but never answers (a hung
+// Xray) must fail the CONNECT within its bound instead of pinning the relay.
+func TestDialUpstreamGivesUpOnSilentBridge(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0") // never accepted: no greeting reply ever comes
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	t.Cleanup(func() { _ = ln.Close() })
+
+	saved := socks5ConnectTimeout
+	socks5ConnectTimeout = 150 * time.Millisecond
+	t.Cleanup(func() { socks5ConnectTimeout = saved })
+
+	relay := SocksRelay{Addr: ln.Addr().String(), Password: "x"}
+	done := make(chan error, 1)
+	go func() {
+		conn, err := relay.dialUpstream("peer@example", netip.MustParseAddrPort("192.0.2.1:443"))
+		if conn != nil {
+			_ = conn.Close()
+		}
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		if err == nil {
+			t.Fatal("CONNECT succeeded through a bridge that never answered")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("dialUpstream still blocked 2s past the connect bound: a silent bridge pins the relay")
+	}
+}
+
+// TestDialUpstreamLeavesNoDeadlineAfterHandshake: the connect bound covers the handshake
+// only, so writes and reads after it has passed still go through.
+func TestDialUpstreamLeavesNoDeadlineAfterHandshake(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	t.Cleanup(func() { _ = ln.Close() })
+
+	saved := socks5ConnectTimeout
+	socks5ConnectTimeout = 150 * time.Millisecond
+	t.Cleanup(func() { socks5ConnectTimeout = saved })
+
+	served := make(chan error, 1)
+	go func() {
+		c, err := ln.Accept()
+		if err != nil {
+			served <- err
+			return
+		}
+		defer c.Close()
+		served <- serveSocks5Connect(c, []byte("ping"), []byte("late"))
+	}()
+
+	relay := SocksRelay{Addr: ln.Addr().String(), Password: "x"}
+	conn, err := relay.dialUpstream("peer@example", netip.MustParseAddrPort("192.0.2.1:443"))
+	if err != nil {
+		t.Fatalf("dialUpstream: %v", err)
+	}
+	defer conn.Close()
+	guard := time.AfterFunc(5*time.Second, func() { _ = conn.Close() })
+	defer guard.Stop()
+	time.Sleep(3 * socks5ConnectTimeout)
+	if _, err := conn.Write([]byte("ping")); err != nil {
+		t.Fatalf("write after the connect bound passed: %v (the handshake deadline leaked onto the relay)", err)
+	}
+	got := make([]byte, 4)
+	if _, err := io.ReadFull(conn, got); err != nil {
+		t.Fatalf("read after the connect bound passed: %v (the handshake deadline leaked onto the relay)", err)
+	}
+	if string(got) != "late" {
+		t.Fatalf("read %q, want %q", got, "late")
+	}
+	if err := <-served; err != nil {
+		t.Fatalf("fake bridge: %v", err)
+	}
+}
+
+// serveSocks5Connect answers one username/password SOCKS5 CONNECT, reads request, then writes reply.
+func serveSocks5Connect(c net.Conn, request, reply []byte) error {
+	buf := make([]byte, 512)
+	read := func(n int) ([]byte, error) { _, err := io.ReadFull(c, buf[:n]); return buf[:n], err }
+	hdr, err := read(2) // VER NMETHODS
+	if err != nil {
+		return err
+	}
+	if _, err := read(int(hdr[1])); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{5, 2}); err != nil { // username/password
+		return err
+	}
+	ver, err := read(2) // VER ULEN
+	if err != nil {
+		return err
+	}
+	if ver[0] != 1 {
+		return errors.New("bad auth version")
+	}
+	if _, err := read(int(ver[1])); err != nil {
+		return err
+	}
+	plen, err := read(1)
+	if err != nil {
+		return err
+	}
+	if _, err := read(int(plen[0])); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{1, 0}); err != nil {
+		return err
+	}
+	req, err := read(4) // VER CMD RSV ATYP
+	if err != nil {
+		return err
+	}
+	addrLen := map[byte]int{1: 4, 4: 16}[req[3]]
+	if req[3] == 3 {
+		l, err := read(1)
+		if err != nil {
+			return err
+		}
+		addrLen = int(l[0])
+	}
+	if _, err := read(addrLen + 2); err != nil {
+		return err
+	}
+	if _, err := c.Write([]byte{5, 0, 0, 1, 0, 0, 0, 0, 0, 0}); err != nil {
+		return err
+	}
+	got, err := read(len(request))
+	if err != nil {
+		return err
+	}
+	if string(got) != string(request) {
+		return errors.New("bridge read an unexpected request")
+	}
+	_, err = c.Write(reply)
+	return err
+}