From 2fab6c3d5032f214220dd9a2e09e356eadbf415a Mon Sep 17 00:00:00 2001 From: Aleksey Fedorov <320380254+Dv2Corn@users.noreply.github.com> Date: Sat, 14 Nov 2026 04:43:03 +1000 Subject: [PATCH] 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 | 4 +- internal/amneziawgnet/race_off_test.go | 5 + internal/amneziawgnet/race_on_test.go | 5 + internal/amneziawgnet/relay.go | 64 +++++- internal/amneziawgnet/relay_halfclose_test.go | 203 ++++++++++++++++++ internal/amneziawgnet/relay_pool_test.go | 198 +++++++++++++++++ 6 files changed, 469 insertions(+), 10 deletions(-) create mode 100644 internal/amneziawgnet/race_off_test.go create mode 100644 internal/amneziawgnet/race_on_test.go create mode 100644 internal/amneziawgnet/relay_halfclose_test.go create mode 100644 internal/amneziawgnet/relay_pool_test.go diff --git a/internal/amneziawgnet/egress.go b/internal/amneziawgnet/egress.go index 0507f447e..ce2942a19 100644 --- a/internal/amneziawgnet/egress.go +++ b/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() { _, _ = io.Copy(upstream, tunnelConn); done <- struct{}{} }() + go func() { copyPooled(tunnelConn, upstream); done <- struct{}{} }() + go func() { copyPooled(upstream, tunnelConn); done <- struct{}{} }() <-done } diff --git a/internal/amneziawgnet/race_off_test.go b/internal/amneziawgnet/race_off_test.go new file mode 100644 index 000000000..e18d7ba7c --- /dev/null +++ b/internal/amneziawgnet/race_off_test.go @@ -0,0 +1,5 @@ +//go:build !race + +package amneziawgnet + +const raceEnabled = false diff --git a/internal/amneziawgnet/race_on_test.go b/internal/amneziawgnet/race_on_test.go new file mode 100644 index 000000000..e580a5808 --- /dev/null +++ b/internal/amneziawgnet/race_on_test.go @@ -0,0 +1,5 @@ +//go:build race + +package amneziawgnet + +const raceEnabled = true diff --git a/internal/amneziawgnet/relay.go b/internal/amneziawgnet/relay.go index f1e76f916..7bfd96839 100644 --- a/internal/amneziawgnet/relay.go +++ b/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} - 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()) + upstream, err := r.dialUpstream(email, dest) 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 { diff --git a/internal/amneziawgnet/relay_halfclose_test.go b/internal/amneziawgnet/relay_halfclose_test.go new file mode 100644 index 000000000..ace28789f --- /dev/null +++ b/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) + } +} diff --git a/internal/amneziawgnet/relay_pool_test.go b/internal/amneziawgnet/relay_pool_test.go new file mode 100644 index 000000000..12aee4bf2 --- /dev/null +++ b/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 +}