diff --git a/internal/amneziawgnet/bench_test.go b/internal/amneziawgnet/bench_test.go new file mode 100644 index 000000000..037ac6b01 --- /dev/null +++ b/internal/amneziawgnet/bench_test.go @@ -0,0 +1,244 @@ +package amneziawgnet + +import ( + "context" + "fmt" + "io" + "net/netip" + "testing" + "time" + + awgconn "github.com/amnezia-vpn/amneziawg-go/v3/conn" + "github.com/amnezia-vpn/amneziawg-go/v3/device" + "github.com/amnezia-vpn/amneziawg-go/v3/tun/netstack" + "gvisor.dev/gvisor/pkg/buffer" + "gvisor.dev/gvisor/pkg/tcpip/adapters/gonet" + "gvisor.dev/gvisor/pkg/tcpip/link/channel" + + "github.com/mhsanaei/3x-ui/v3/internal/amneziawg" + "github.com/mhsanaei/3x-ui/v3/internal/util/wireguard" +) + +// BenchmarkStackTunWrite measures the upload path's per-packet cost: one +// decrypted packet handed from amneziawg-go into the gVisor stack. +func BenchmarkStackTunWrite(b *testing.B) { + tun := &stackTun{ep: channel.New(tunQueueDepth, 1420, ""), mtu: 1420} + defer tun.ep.Close() + + packet := make([]byte, 1400) + packet[0] = 0x45 + bufs := [][]byte{packet} + + b.SetBytes(int64(len(packet))) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := tun.Write(bufs, 0); err != nil { + b.Fatalf("Write: %v", err) + } + } +} + +// BenchmarkStackTunRead measures the download path's per-packet cost: one +// packet drained out of the stack for amneziawg-go to encrypt. +func BenchmarkStackTunRead(b *testing.B) { + tun := &stackTun{incomingPacket: make(chan *buffer.View, tunQueueDepth)} + packet := make([]byte, 1400) + buf := [][]byte{make([]byte, 2048)} + sizes := make([]int, 1) + + b.SetBytes(int64(len(packet))) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + tun.incomingPacket <- buffer.NewViewWithData(packet) + if _, err := tun.Read(buf, sizes, 0); err != nil { + b.Fatalf("Read: %v", err) + } + } +} + +// BenchmarkUDPDatagramDelivery measures one datagram travelling the whole +// inbound path: stack injection, routing, and the UDP transport handler. +func BenchmarkUDPDatagramDelivery(b *testing.B) { + tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.78.0.1")}, 1420) + if err != nil { + b.Fatalf("createNetTUNWithStack: %v", err) + } + defer tun.Close() + + src := netip.MustParseAddrPort("10.78.0.2:40000") + dst := netip.MustParseAddrPort("10.78.9.9:5353") + payload := make([]byte, 1024) + AttachUDPHandler(gstack, func(netip.AddrPort, netip.AddrPort, []byte) {}) + + bufs := [][]byte{udpDatagram(src, dst, payload)} + st := tun.(*stackTun) + + b.SetBytes(int64(len(payload))) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := st.Write(bufs, 0); err != nil { + b.Fatalf("Write: %v", err) + } + } +} + +// Two destinations the benchmark forwarder tells apart: one drains what the +// client sends, the other streams at the client. Neither is routed anywhere. +const ( + benchDiscardPort = 9001 + benchSourcePort = 9002 +) + +// benchTunnel is a live AmneziaWG pair -- this package's server Device and a +// stock amneziawg-go client -- talking real encrypted UDP over loopback. +type benchTunnel struct { + clientNet *netstack.Net + closeFn func() +} + +// newBenchTunnel brings up both devices and blocks until the handshake has +// actually completed, so no setup cost lands inside the measured loop. +func newBenchTunnel(b *testing.B, listenPort int, serverAddr, clientAddr string) *benchTunnel { + b.Helper() + serverPriv, serverPub, err := wireguard.GenerateWireguardKeypair() + if err != nil { + b.Fatalf("server keypair: %v", err) + } + clientPriv, clientPub, err := wireguard.GenerateWireguardKeypair() + if err != nil { + b.Fatalf("client keypair: %v", err) + } + + inst := amneziawg.Instance{ + Id: 90, + InterfaceName: "awgbench", + ListenPort: listenPort, + PrivateKey: serverPriv, + PublicKey: serverPub, + Address: []string{serverAddr + "/24"}, + MTU: 1420, + Obfuscation: amneziawg.Obfuscation31{ + Jc: 4, Jmin: 40, Jmax: 70, + S1: 20, S2: 30, S3: 20, S4: 20, + }, + Peers: []amneziawg.Peer{{ + Email: "bench@example.com", + PublicKey: clientPub, + AllowedIPs: []string{clientAddr + "/32"}, + }}, + } + + dev, err := newUnconfiguredDevice(inst, DeviceOptions{}) + if err != nil { + b.Fatalf("newUnconfiguredDevice: %v", err) + } + AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) { + defer conn.Close() + switch dest.Port() { + case benchDiscardPort: + _, _ = io.Copy(io.Discard, conn) + case benchSourcePort: + chunk := make([]byte, 64<<10) + for { + if _, err := conn.Write(chunk); err != nil { + return + } + } + } + }) + if err := dev.Configure(inst, DeviceOptions{}); err != nil { + b.Fatalf("Configure: %v", err) + } + + clientTun, clientNet, err := netstack.CreateNetTUN( + []netip.Addr{netip.MustParseAddr(clientAddr)}, + []netip.Addr{netip.MustParseAddr("1.1.1.1")}, 1420) + if err != nil { + b.Fatalf("client CreateNetTUN: %v", err) + } + clientDev := device.NewDevice(clientTun, awgconn.NewDefaultBind(), device.NewLogger(device.LogLevelSilent, "")) + + clientPrivHex, err := wireguard.KeyToHex(clientPriv) + if err != nil { + b.Fatalf("client key to hex: %v", err) + } + serverPubHex, err := wireguard.KeyToHex(serverPub) + if err != nil { + b.Fatalf("server key to hex: %v", err) + } + conf := fmt.Sprintf( + "private_key=%s\njc=4\njmin=40\njmax=70\ns1=20\ns2=30\ns3=20\ns4=20\npublic_key=%s\nendpoint=127.0.0.1:%d\nallowed_ip=0.0.0.0/0\n", + clientPrivHex, serverPubHex, listenPort) + if err := clientDev.IpcSet(conf); err != nil { + b.Fatalf("client IpcSet: %v", err) + } + if err := clientDev.Up(); err != nil { + b.Fatalf("client Up: %v", err) + } + + t := &benchTunnel{clientNet: clientNet, closeFn: func() { + clientDev.Close() + dev.Close() + }} + // Prove the handshake really completed before anything is timed. + probe := t.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchDiscardPort))) + probe.Close() + return t +} + +// dial opens one tunnelled connection, retrying while the handshake settles. +func (t *benchTunnel) dial(b *testing.B, dest netip.AddrPort) *gonet.TCPConn { + b.Helper() + deadline := time.Now().Add(15 * time.Second) + for { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + conn, err := t.clientNet.DialContextTCPAddrPort(ctx, dest) + cancel() + if err == nil { + return conn + } + if time.Now().After(deadline) { + b.Fatalf("dial %v through tunnel: %v", dest, err) + } + } +} + +// BenchmarkTunnelThroughput is the end-to-end number: real bytes through a +// real handshaked AmneziaWG tunnel, in both directions. +func BenchmarkTunnelThroughput(b *testing.B) { + const chunkSize = 64 << 10 + const serverAddr = "10.203.0.1" + tun := newBenchTunnel(b, 58714, serverAddr, "10.203.0.2") + defer tun.closeFn() + + b.Run("upload", func(b *testing.B) { + conn := tun.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchDiscardPort))) + defer conn.Close() + chunk := make([]byte, chunkSize) + b.SetBytes(chunkSize) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := conn.Write(chunk); err != nil { + b.Fatalf("upload write: %v", err) + } + } + }) + + b.Run("download", func(b *testing.B) { + conn := tun.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchSourcePort))) + defer conn.Close() + chunk := make([]byte, chunkSize) + b.SetBytes(chunkSize) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := io.ReadFull(conn, chunk); err != nil { + b.Fatalf("download read: %v", err) + } + } + }) +} diff --git a/internal/amneziawgnet/netstack.go b/internal/amneziawgnet/netstack.go index 5b3045141..f1762a36d 100644 --- a/internal/amneziawgnet/netstack.go +++ b/internal/amneziawgnet/netstack.go @@ -105,7 +105,8 @@ func (t *stackTun) Events() <-chan awgtun.Event { return t.events } func (t *stackTun) MTU() (int, error) { return t.mtu, nil } func (t *stackTun) BatchSize() int { return 1 } -// Read drains incomingPacket into buf, supporting batched reads. +// Read drains incomingPacket into buf, supporting batched reads. Each view is +// released once copied out, so the download path reuses gVisor's pooled chunks. func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) { var view *buffer.View select { @@ -114,6 +115,7 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) { case view = <-t.incomingPacket: } n, err := view.Read(buf[0][offset:]) + view.Release() if err != nil { return 0, err } @@ -123,6 +125,7 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) { select { case view = <-t.incomingPacket: n, err := view.Read(buf[count][offset:]) + view.Release() if err != nil { return count, nil } @@ -135,6 +138,8 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) { return count, nil } +// Write injects each packet into the stack. The injector owns the packet +// buffer -- DecRef returns it and its chunk to gVisor's pools (see loopback.go). func (t *stackTun) Write(buf [][]byte, offset int) (int, error) { for _, b := range buf { packet := b[offset:] @@ -148,8 +153,10 @@ func (t *stackTun) Write(buf [][]byte, offset int) (int, error) { case 6: t.ep.InjectInbound(header.IPv6ProtocolNumber, pkb) default: + pkb.DecRef() return 0, syscall.EAFNOSUPPORT } + pkb.DecRef() } return len(buf), nil } @@ -176,6 +183,7 @@ func (t *stackTun) WriteNotify() { select { case t.incomingPacket <- view: case <-t.done: + view.Release() } } diff --git a/internal/amneziawgnet/netstack_test.go b/internal/amneziawgnet/netstack_test.go index 52fda2d18..f50ae42fc 100644 --- a/internal/amneziawgnet/netstack_test.go +++ b/internal/amneziawgnet/netstack_test.go @@ -4,6 +4,7 @@ import ( "testing" "gvisor.dev/gvisor/pkg/buffer" + "gvisor.dev/gvisor/pkg/tcpip/link/channel" ) // TestStackTunReadDrainsBufferedBatch is a regression test for a real @@ -86,3 +87,44 @@ func TestStackTunReadStopsAtBufCapacity(t *testing.T) { t.Errorf("leftover packet = %v, want [3]", got) } } + +// TestStackTunWriteReturnsPacketBuffersToPool locks in gVisor's ownership rule +// for the upload path: whoever calls InjectInbound must DecRef the packet. +func TestStackTunWriteReturnsPacketBuffersToPool(t *testing.T) { + tun := &stackTun{ep: channel.New(tunQueueDepth, 1420, ""), mtu: 1420} + defer tun.ep.Close() + + packet := make([]byte, 1400) + packet[0] = 0x45 // IPv4, version nibble is all Write inspects + bufs := [][]byte{packet} + + allocs := testing.AllocsPerRun(1000, func() { + if _, err := tun.Write(bufs, 0); err != nil { + t.Fatalf("Write: %v", err) + } + }) + // 0 once pooled, 4 when every packet buffer is stranded; -race adds ~1. + if allocs > 1 { + t.Fatalf("Write allocates %v times per packet, want <=1: injected packet buffers are not being returned to gVisor's pools", allocs) + } +} + +// TestStackTunReadReturnsViewsToPool is the download-path counterpart: a view +// that is copied out but never released strands its pooled chunk. +func TestStackTunReadReturnsViewsToPool(t *testing.T) { + tun := &stackTun{incomingPacket: make(chan *buffer.View, tunQueueDepth)} + packet := make([]byte, 1400) + buf := [][]byte{make([]byte, 2048)} + sizes := make([]int, 1) + + allocs := testing.AllocsPerRun(1000, func() { + tun.incomingPacket <- buffer.NewViewWithData(packet) + if _, err := tun.Read(buf, sizes, 0); err != nil { + t.Fatalf("Read: %v", err) + } + }) + // 0 once the drained view goes back to viewPool, 3 when it does not. + if allocs > 1 { + t.Fatalf("Read allocates %v times per packet, want <=1: drained views are not being released", allocs) + } +} diff --git a/internal/amneziawgnet/udp.go b/internal/amneziawgnet/udp.go index 52814de4e..4c7022419 100644 --- a/internal/amneziawgnet/udp.go +++ b/internal/amneziawgnet/udp.go @@ -32,7 +32,9 @@ func AttachUDPHandler(gstack *stack.Stack, handler UDPHandler) { enablePromiscuousRouting(gstack) gstack.SetTransportProtocolHandler(udp.ProtocolNumber, func(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool { - data := pkt.Clone().Data().AsRange().ToSlice() + // ToSlice already returns an owned copy, so cloning pkt here would only + // strand a pooled packet buffer and its chunks on every datagram. + data := pkt.Data().AsRange().ToSlice() src := netip.AddrPortFrom(addrFromTcpip(id.RemoteAddress), id.RemotePort) dst := netip.AddrPortFrom(addrFromTcpip(id.LocalAddress), id.LocalPort) handler(src, dst, data) diff --git a/internal/amneziawgnet/udp_test.go b/internal/amneziawgnet/udp_test.go index 4649b70d6..d24f3a146 100644 --- a/internal/amneziawgnet/udp_test.go +++ b/internal/amneziawgnet/udp_test.go @@ -9,6 +9,8 @@ import ( awgconn "github.com/amnezia-vpn/amneziawg-go/v3/conn" "github.com/amnezia-vpn/amneziawg-go/v3/device" "github.com/amnezia-vpn/amneziawg-go/v3/tun/netstack" + "gvisor.dev/gvisor/pkg/tcpip" + "gvisor.dev/gvisor/pkg/tcpip/header" "github.com/mhsanaei/3x-ui/v3/internal/amneziawg" "github.com/mhsanaei/3x-ui/v3/internal/util/wireguard" @@ -154,3 +156,64 @@ func TestNewDeviceUDPHandlerAndReply(t *testing.T) { return } } + +// udpDatagram builds a complete IPv4/UDP packet, the shape stackTun.Write +// expects from amneziawg-go after decryption. +func udpDatagram(src, dst netip.AddrPort, payload []byte) []byte { + total := header.IPv4MinimumSize + header.UDPMinimumSize + len(payload) + p := make([]byte, total) + ip := header.IPv4(p) + ip.Encode(&header.IPv4Fields{ + TotalLength: uint16(total), + TTL: 64, + Protocol: uint8(header.UDPProtocolNumber), + SrcAddr: tcpip.AddrFromSlice(src.Addr().AsSlice()), + DstAddr: tcpip.AddrFromSlice(dst.Addr().AsSlice()), + }) + ip.SetChecksum(^ip.CalculateChecksum()) + u := header.UDP(p[header.IPv4MinimumSize:]) + u.Encode(&header.UDPFields{ + SrcPort: src.Port(), + DstPort: dst.Port(), + Length: uint16(header.UDPMinimumSize + len(payload)), + }) + copy(p[header.IPv4MinimumSize+header.UDPMinimumSize:], payload) + return p +} + +// TestAttachUDPHandlerDoesNotStrandPacketBuffers drives a real datagram all the +// way through the stack: Range.ToSlice already copies, so cloning pkt only leaks. +func TestAttachUDPHandlerDoesNotStrandPacketBuffers(t *testing.T) { + tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.77.0.1")}, 1420) + if err != nil { + t.Fatalf("createNetTUNWithStack: %v", err) + } + defer tun.Close() + + src := netip.MustParseAddrPort("10.77.0.2:40000") + dst := netip.MustParseAddrPort("10.77.9.9:5353") + payload := make([]byte, 512) + var delivered int + AttachUDPHandler(gstack, func(gotSrc, gotDst netip.AddrPort, got []byte) { + if gotSrc != src || gotDst != dst || len(got) != len(payload) { + t.Errorf("handler got (%v -> %v, %d bytes), want (%v -> %v, %d bytes)", gotSrc, gotDst, len(got), src, dst, len(payload)) + } + delivered++ + }) + + bufs := [][]byte{udpDatagram(src, dst, payload)} + st := tun.(*stackTun) + allocs := testing.AllocsPerRun(500, func() { + if _, err := st.Write(bufs, 0); err != nil { + t.Fatalf("Write: %v", err) + } + }) + if delivered == 0 { + t.Fatal("handler never ran: the datagram never reached the UDP transport handler") + } + // 2 once nothing is stranded (1 is ToSlice itself), 8 with the leaked + // clone plus the un-released packet buffer; -race adds ~1. + if allocs > 4 { + t.Fatalf("UDP delivery allocates %v times per datagram, want <=4: pooled packet buffers are being stranded", allocs) + } +}