From 101b976b1106b66e6ab19e0a5dbe760db760f312 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enzo=20Guti=C3=A9rrez=20Pereira?= Date: Fri, 17 Jul 2026 02:23:23 -0300 Subject: [PATCH] refactor transfer architecture for parallel crypto and chunking for better bandwidth utilization --- cmd/chat.go | 29 +- cmd/connect.go | 156 ++++++-- internal/network/bench_test.go | 200 ++++++++++ internal/network/handshake.go | 54 --- internal/network/helpers.go | 3 +- internal/network/peer.go | 665 ++++++++++++++++++++++++-------- internal/network/transport.go | 220 +++++++++++ internal/transfer/bench_test.go | 49 +++ internal/transfer/transfer.go | 26 +- test/connection_test.go | 2 +- test/hardening_test.go | 2 +- 11 files changed, 1139 insertions(+), 267 deletions(-) create mode 100644 internal/network/bench_test.go delete mode 100644 internal/network/handshake.go create mode 100644 internal/network/transport.go create mode 100644 internal/transfer/bench_test.go diff --git a/cmd/chat.go b/cmd/chat.go index 64f877b..6c7b29d 100644 --- a/cmd/chat.go +++ b/cmd/chat.go @@ -51,34 +51,39 @@ var ChatCmd = &cobra.Command{ publicKey = keyPair.PublicKey } - conn, publicAddr, err := network.OpenUDPConn() + // PSK second factor + Noise prologue, and our published static pubkey. + psk := crypto.DerivePSK(password, sessionId) + prologue := network.Prologue(sessionId) + myPubKeyB64 := base64.StdEncoding.EncodeToString(publicKey) + + tr, err := network.OpenTransport() if err != nil { - fmt.Println("Error opening UDP connection:", err) + fmt.Println("Error opening UDP transport:", err) return } + peerConn := network.NewPeerConn(tr, privateKey, publicKey, psk, prologue) + publicAddr, err := peerConn.DiscoverPublicAddr() + if err != nil { + fmt.Println("Error discovering public address:", err) + peerConn.Close() + return + } fmt.Println("Public addr:", publicAddr) - // PSK second factor + Noise prologue, and our published static pubkey. - psk := crypto.DerivePSK(password, sessionId) - prologue := network.Prologue(sessionId) - myPubKeyB64 := base64.StdEncoding.EncodeToString(publicKey) - peers, err := session.Register(hostname, sessionId, password, publicAddr, myPubKeyB64, new) if err != nil { fmt.Printf("Error registering: %v\n", err) - conn.Close() + peerConn.Close() return } myPublicIP := strings.Split(publicAddr, ":")[0] - peerConn := network.NewPeerConn(conn, privateKey, publicKey, psk, prologue) defer func() { session.Leave(hostname, sessionId, password, publicAddr) peerConn.BroadcastDead() // encrypted "dead" notice so peers tear down promptly - peerConn.Shutdown() // stop any in-flight handshake drivers - conn.Close() + peerConn.Close() // stop handshake drivers/consumers and close the bind }() // knownPeers tracks which resolved peer addresses have been handed to @@ -227,7 +232,7 @@ var ChatCmd = &cobra.Command{ // watch connection go func() { - if err := network.WatchConnection(conn, peerConn.HasPeers); err != nil { + if err := network.WatchConnection(peerConn.HasPeers); err != nil { fmt.Println("Connection lost, exiting chat...") close(quit) } diff --git a/cmd/connect.go b/cmd/connect.go index 454a706..d2e8f68 100644 --- a/cmd/connect.go +++ b/cmd/connect.go @@ -1,6 +1,7 @@ package cmd import ( + "bytes" "encoding/base64" "encoding/json" "fmt" @@ -170,19 +171,27 @@ var ConnectCmd = &cobra.Command{ publicKey = keyPair.PublicKey } - conn, publicAddr, err := network.OpenUDPConn() + // PSK (second factor) from the password via Argon2id, salted with the session + // id; prologue binds the Noise handshake to protocol version + session. Our + // static pubkey is published to the rendezvous so peers can pin it beforehand. + psk := crypto.DerivePSK(password, sessionId) + prologue := network.Prologue(sessionId) + myPubKeyB64 := base64.StdEncoding.EncodeToString(publicKey) + + tr, err := network.OpenTransport() if err != nil { - writeStatus("error: opening UDP connection: " + err.Error()) + writeStatus("error: opening UDP transport: " + err.Error()) os.Remove(pidFile) return } // from here on, defer owns all cleanup var ( - peerConn *network.PeerConn tunDevice bstun.Device registered bool + publicAddr string ) + peerConn := network.NewPeerConn(tr, privateKey, publicKey, psk, prologue) quit := make(chan struct{}) var quitOnce sync.Once closeQuit := func() { quitOnce.Do(func() { close(quit) }) } @@ -195,22 +204,18 @@ var ConnectCmd = &cobra.Command{ if registered { session.Leave(hostname, sessionId, password, publicAddr) } - if peerConn != nil { - peerConn.BroadcastDead() // encrypted "dead" notice so peers tear down promptly - peerConn.Shutdown() // stop any in-flight handshake drivers - } - conn.Close() + peerConn.BroadcastDead() // encrypted "dead" notice so peers tear down promptly + peerConn.Close() // stop handshake drivers/consumers and close the bind os.Remove(pidFile) os.Remove(sessionStopFile()) os.Remove(peersFile()) }() - // PSK (second factor) from the password via Argon2id, salted with the session - // id; prologue binds the Noise handshake to protocol version + session. Our - // static pubkey is published to the rendezvous so peers can pin it beforehand. - psk := crypto.DerivePSK(password, sessionId) - prologue := network.Prologue(sessionId) - myPubKeyB64 := base64.StdEncoding.EncodeToString(publicKey) + publicAddr, err = peerConn.DiscoverPublicAddr() + if err != nil { + writeStatus("error: discovering public address: " + err.Error()) + return + } peers, err := session.Register(hostname, sessionId, password, publicAddr, myPubKeyB64, isNew) if err != nil { @@ -227,7 +232,6 @@ var ConnectCmd = &cobra.Command{ } myPublicIP := strings.Split(publicAddr, ":")[0] - peerConn = network.NewPeerConn(conn, privateKey, publicKey, psk, prologue) // knownPeers tracks which resolved peer addresses have been handed to // AddKnownPeer, so rendezvous announcements are not re-added while a session @@ -329,42 +333,73 @@ var ConnectCmd = &cobra.Command{ } }() - // UDP → TUN: decrypt incoming packets and write into the TUN interface + // tunBufPool recycles packet buffers across both pump directions so the + // steady state allocates nothing per packet. + tunBufPool := sync.Pool{New: func() any { + b := make([]byte, 1600) + return &b + }} + getTunBuf := func() []byte { return *tunBufPool.Get().(*[]byte) } + putTunBuf := func(b []byte) { + if cap(b) >= 1600 { + b = b[:1600] + tunBufPool.Put(&b) + } + } + + // UDP → TUN: drain decrypted tunnel packets in batches (parallel AEAD in + // ReadTunBatch) and hand each surviving batch to the TUN device in a + // single Write call instead of one ring transition per packet. go func() { + batch := peerConn.BatchSize() + bufs := make([][]byte, batch) + for i := range bufs { + bufs[i] = getTunBuf() + } + senders := make([]string, batch) + wr := make([][]byte, 0, batch) for { - pktType, plaintext, addr, err := peerConn.Read() + n, err := peerConn.ReadTunBatch(bufs, senders) if err != nil { - if strings.Contains(err.Error(), "use of closed network connection") { - return - } - continue + return // transport closed } - if pktType != network.PacketTun { - continue + wr = wr[:0] + for i := 0; i < n; i++ { + // Reverse-path filter: the inner IPv4 source address must equal the + // sender's virtual IP. Otherwise a malicious member could inject packets + // spoofing another peer's virtual IP inside the (authenticated) tunnel. + expectedVIP, ok := addrToVIP.Load(senders[i]) + if !ok || !bstun.SrcIPMatchesVirtualIP(bufs[i], expectedVIP.(string)) { + continue + } + wr = append(wr, bufs[i]) } - // Reverse-path filter: the inner IPv4 source address must equal the - // sender's virtual IP. Otherwise a malicious member could inject packets - // spoofing another peer's virtual IP inside the (authenticated) tunnel. - expectedVIP, ok := addrToVIP.Load(addr.String()) - if !ok || !bstun.SrcIPMatchesVirtualIP(plaintext, expectedVIP.(string)) { + if len(wr) == 0 { continue } network.UpdateLastSeen() - tunDevice.Write([][]byte{plaintext}, 0) + tunDevice.Write(wr, 0) } }() - // TUN → UDP: read outbound IP packets and route to the right peer + // TUN → UDP is split into a reader and a sender joined by a channel, so + // packets read one at a time (wintun's batch size is 1) still aggregate + // into batches for the encrypt+send path. + outCh := make(chan []byte, 512) + + // Reader: pull outbound IP packets off the TUN device and hand their + // (pooled) buffers to the sender. go func() { batchSize := tunDevice.BatchSize() bufs := make([][]byte, batchSize) sizes := make([]int, batchSize) for i := range bufs { - bufs[i] = make([]byte, 1500) + bufs[i] = getTunBuf() } for { n, err := tunDevice.Read(bufs, sizes, 0) if err != nil { + close(outCh) return // TUN closed } for i := 0; i < n; i++ { @@ -372,12 +407,65 @@ var ConnectCmd = &cobra.Command{ if len(packet) < 20 || packet[0]>>4 != 4 { continue // not an IPv4 packet } - destIP := net.IP(packet[16:20]).String() + select { + case outCh <- packet: + bufs[i] = getTunBuf() // buffer ownership moved to the sender + case <-quit: + return + } + } + } + }() + + // Sender: aggregate whatever the reader has produced, group consecutive + // packets by destination peer, and push each group through one batched + // encrypt+send (one counter reservation, parallel AEAD, few syscalls). + go func() { + pending := make([][]byte, 0, 128) + flush := func() { + for start := 0; start < len(pending); { + destIP := net.IP(pending[start][16:20]).String() addrVal, ok := virtualIPMap.Load(destIP) if !ok { - continue // no peer with that virtual IP + putTunBuf(pending[start]) // no peer with that virtual IP + start++ + continue + } + end := start + 1 + for end < len(pending) && bytes.Equal(pending[end][16:20], pending[start][16:20]) { + end++ + } + peerConn.SendTunBatch(addrVal.(*net.UDPAddr), pending[start:end]) + for _, b := range pending[start:end] { + putTunBuf(b) + } + start = end + } + pending = pending[:0] + } + for { + select { + case pkt, ok := <-outCh: + if !ok { + return + } + pending = append(pending, pkt) + drain: + for len(pending) < cap(pending) { + select { + case more, ok := <-outCh: + if !ok { + flush() + return + } + pending = append(pending, more) + default: + break drain + } } - peerConn.SendTun(addrVal.(*net.UDPAddr), packet) + flush() + case <-quit: + return } } }() diff --git a/internal/network/bench_test.go b/internal/network/bench_test.go new file mode 100644 index 0000000..3f81578 --- /dev/null +++ b/internal/network/bench_test.go @@ -0,0 +1,200 @@ +package network + +import ( + "fmt" + "net" + "testing" + "time" + + "github.com/neozmmv/blindspot/internal/crypto" +) + +// benchSide is one loopback peer: a real PeerConn over a WrapUDPConn transport. +type benchSide struct { + pc *PeerConn + addr *net.UDPAddr + kp *crypto.KeyPair +} + +func newBenchSide(b *testing.B, psk []byte) *benchSide { + b.Helper() + conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Fatalf("ListenUDP: %v", err) + } + // Large socket buffers so the loopback path measures the pipeline, not the + // default 64 KB kernel queue. + conn.SetReadBuffer(7 << 20) + conn.SetWriteBuffer(7 << 20) + kp, err := crypto.GenerateKeyPair() + if err != nil { + b.Fatalf("GenerateKeyPair: %v", err) + } + pc := NewPeerConn(WrapUDPConn(conn), kp.PrivateKey, kp.PublicKey, psk, Prologue("bench-session")) + return &benchSide{pc: pc, addr: conn.LocalAddr().(*net.UDPAddr), kp: kp} +} + +// benchPair returns two handshaked peers on loopback. +func benchPair(b *testing.B) (*benchSide, *benchSide) { + b.Helper() + psk := crypto.DerivePSK("bench-pass-123", "bench-session") + sa := newBenchSide(b, psk) + sb := newBenchSide(b, psk) + if err := sa.pc.AddKnownPeer(sb.addr, sb.kp.PublicKey); err != nil { + b.Fatalf("AddKnownPeer: %v", err) + } + if err := sb.pc.AddKnownPeer(sa.addr, sa.kp.PublicKey); err != nil { + b.Fatalf("AddKnownPeer: %v", err) + } + for _, s := range []*benchSide{sa, sb} { + select { + case <-s.pc.Connected: + case <-time.After(5 * time.Second): + b.Fatal("handshake did not complete") + } + } + b.Cleanup(func() { + sa.pc.Close() + sb.pc.Close() + }) + return sa, sb +} + +// BenchmarkTunnelSend measures the transmit pipeline — batched counter +// reservation, parallel AES-GCM seal, batched socket writes — through a real +// loopback UDP socket while a receiver actively drains (drops don't stall the +// sender; UDP semantics). Reported MB/s is tunnel payload throughput. +func BenchmarkTunnelSend(b *testing.B) { + sa, sb := benchPair(b) + + // Active receiver: drain and decrypt whatever survives the loopback queue. + go func() { + bufs := make([][]byte, sb.pc.BatchSize()) + for i := range bufs { + bufs[i] = make([]byte, 1600) + } + senders := make([]string, sb.pc.BatchSize()) + for { + if _, err := sb.pc.ReadTunBatch(bufs, senders); err != nil { + return + } + } + }() + + const payloadSize = 1420 + batch := sa.pc.BatchSize() + payloads := make([][]byte, batch) + for i := range payloads { + pkt := make([]byte, payloadSize) + pkt[0] = 0x45 // minimal IPv4-looking header so nothing chokes downstream + payloads[i] = pkt + } + + b.SetBytes(payloadSize) + b.ResetTimer() + for sent := 0; sent < b.N; { + m := batch + if rem := b.N - sent; rem < m { + m = rem + } + if err := sa.pc.SendTunBatch(sb.addr, payloads[:m]); err != nil { + b.Fatalf("SendTunBatch: %v", err) + } + sent += m + } +} + +// BenchmarkTunnelRoundtrip measures delivered end-to-end throughput: the +// receiver must actually authenticate, decrypt, and pass the replay window for +// every counted packet. Loopback UDP can drop under pressure, so the sender +// paces itself with a crude in-flight window driven by the receiver's count. +func BenchmarkTunnelRoundtrip(b *testing.B) { + sa, sb := benchPair(b) + + const payloadSize = 1420 + received := make(chan int, 1024) + go func() { + bufs := make([][]byte, sb.pc.BatchSize()) + for i := range bufs { + bufs[i] = make([]byte, 1600) + } + senders := make([]string, sb.pc.BatchSize()) + for { + n, err := sb.pc.ReadTunBatch(bufs, senders) + if err != nil { + return + } + received <- n + } + }() + + batch := sa.pc.BatchSize() + payloads := make([][]byte, batch) + for i := range payloads { + pkt := make([]byte, payloadSize) + pkt[0] = 0x45 + payloads[i] = pkt + } + + b.SetBytes(payloadSize) + b.ResetTimer() + sent, got := 0, 0 + deadline := time.Now().Add(60 * time.Second) + for got < b.N { + // Keep at most 4 batches in flight to stay under the loopback queue. + for sent-got < 4*batch && sent < b.N+8*batch { + m := batch + if err := sa.pc.SendTunBatch(sb.addr, payloads[:m]); err != nil { + b.Fatalf("SendTunBatch: %v", err) + } + sent += m + } + select { + case n := <-received: + got += n + case <-time.After(100 * time.Millisecond): + // Loss on loopback: top the window back up. + if time.Now().After(deadline) { + b.Fatalf("stalled: sent %d, received %d of %d", sent, got, b.N) + } + sent = got + } + } +} + +var benchSink error + +// BenchmarkSealBatch isolates the parallel encrypt stage (no socket): how fast +// can a batch of tunnel packets be sealed across cores. +func BenchmarkSealBatch(b *testing.B) { + sa, sb := benchPair(b) + _ = sb + s := sa.pc.sessionByAddr(sb.addr) + if s == nil { + b.Fatal("no session") + } + const payloadSize = 1420 + batch := 128 + payloads := make([][]byte, batch) + for i := range payloads { + payloads[i] = make([]byte, payloadSize) + } + s.mu.Lock() + aead := s.txAEAD + s.mu.Unlock() + if aead == nil { + b.Fatal("no tx AEAD") + } + b.SetBytes(payloadSize) + b.ResetTimer() + for done := 0; done < b.N; done += batch { + wire := make([][]byte, batch) + parallelFor(batch, func(i int) { + wire[i] = sealPacket(aead, PacketTun, uint64(done+i), payloads[i]) + }) + for _, w := range wire { + putPacketBuf(w) + } + } + benchSink = fmt.Errorf("%d", b.N) // defeat dead-code elimination +} diff --git a/internal/network/handshake.go b/internal/network/handshake.go deleted file mode 100644 index 0f33f42..0000000 --- a/internal/network/handshake.go +++ /dev/null @@ -1,54 +0,0 @@ -package network - -import ( - "fmt" - "net" - - "github.com/pion/stun" -) - -// OpenUDPConn opens a UDP socket on an ephemeral port and discovers this host's -// public address via a STUN binding request to Google's STUN server. -func OpenUDPConn() (*net.UDPConn, string, error) { - conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: 0}) - if err != nil { - fmt.Printf("Error listening on UDP: %v\n", err) - return nil, "", err - } - - // Large kernel socket buffers (same 7 MB wireguard-go uses). The OS default - // (~64 KB on Windows) holds only ~45 tunnel packets, so a burst from a fast - // TCP sender overflows it while the reader goroutine is busy decrypting; - // every drop looks like congestion to the tunnelled TCP stream and collapses - // throughput. Best effort: the OS may clamp the value (e.g. Linux rmem_max). - conn.SetReadBuffer(7 << 20) - conn.SetWriteBuffer(7 << 20) - - // stun - serverAddr, err := net.ResolveUDPAddr("udp", "stun.l.google.com:19302") - if err != nil { - fmt.Printf("Error resolving STUN server address: %v\n", err) - return nil, "", err - } - - // send bind request to google stun - msg := stun.MustBuild(stun.TransactionID, stun.BindingRequest) - conn.WriteToUDP(msg.Raw, serverAddr) - - buf := make([]byte, 1024) - n, _, err := conn.ReadFromUDP(buf) - if err != nil { - fmt.Printf("Error reading from STUN server: %v\n", err) - return nil, "", err - } - - // decode stun response to get public IP and port - m := &stun.Message{Raw: buf[:n]} - m.Decode() - - var xorAddr stun.XORMappedAddress - xorAddr.GetFrom(m) - - publicAddr := fmt.Sprintf("%s:%d", xorAddr.IP, xorAddr.Port) - return conn, publicAddr, nil -} diff --git a/internal/network/helpers.go b/internal/network/helpers.go index daaac0c..bad6c45 100644 --- a/internal/network/helpers.go +++ b/internal/network/helpers.go @@ -2,7 +2,6 @@ package network import ( "fmt" - "net" "sync" "time" ) @@ -27,7 +26,7 @@ const ( // WatchConnection monitors for activity. hasPeers returns true when there are // currently connected peers — the timeout only fires when peers exist but are silent. -func WatchConnection(conn *net.UDPConn, hasPeers func() bool) error { +func WatchConnection(hasPeers func() bool) error { for { time.Sleep(30 * time.Second) if !hasPeers() { diff --git a/internal/network/peer.go b/internal/network/peer.go index 6e0035a..c60115a 100644 --- a/internal/network/peer.go +++ b/internal/network/peer.go @@ -2,14 +2,20 @@ package network import ( "bytes" + "crypto/aes" + "crypto/cipher" "encoding/binary" "encoding/hex" + "errors" "fmt" "net" + "runtime" "sync" + "sync/atomic" "time" "github.com/flynn/noise" + wgconn "golang.zx2c4.com/wireguard/conn" ) const ( @@ -35,17 +41,17 @@ const ( ) // peerSession holds the per-peer handshake and transport state. All mutable -// fields are guarded by mu; a CipherState is not safe for concurrent use. +// fields are guarded by mu. type peerSession struct { mu sync.Mutex - addr *net.UDPAddr - expected []byte // static key the rendezvous published for this peer (nil until known) - remoteStatic []byte // authenticated static key after the handshake completes + addr *net.UDPAddr // canonical remote address, for events and the public API + key string // canonical addr string; the sessions-map key + ep wgconn.Endpoint // bind endpoint used for all sends to this peer + expected []byte // static key the rendezvous published for this peer (nil until known) + remoteStatic []byte // authenticated static key after the handshake completes initiator bool hs *noise.HandshakeState - tx *noise.CipherState // our send cipher (per-direction, from Noise split) - rx *noise.CipherState // our receive cipher - established bool + established bool // driving is true while a handshake is in flight (armed and not yet // established). It is cleared the moment the session establishes — not when // the retransmit goroutine notices and exits — so a re-arm cannot be blocked @@ -57,19 +63,114 @@ type peerSession struct { initMsg []byte // initiator: cached msg1 packet, retransmitted until established respMsg []byte // responder: cached msg2 packet, retransmitted on duplicate msg1 + // Transport crypto. At establishment the AES-256-GCM keys are extracted from + // the Noise CipherStates (UnsafeKey) and turned into stateless cipher.AEAD + // values: unlike a *noise.CipherState (whose implicit nonce makes every call + // order-dependent), a cipher.AEAD with an explicit nonce is safe for + // concurrent use — this is what lets a whole batch be sealed/opened in + // parallel across cores. The wire format is byte-identical to what the + // CipherStates produced: nonce = 4 zero bytes ‖ big-endian counter, AAD = + // the 10-byte cleartext header. + txAEAD cipher.AEAD + rxAEAD cipher.AEAD + txCtr uint64 // next send counter (the AEAD nonce), guarded by mu + // Anti-replay window for the receive direction, guarded by mu like the rest // of the session state and consulted only after AEAD authentication. replay ReplayWindow } +// aesgcmNonce builds the Noise-spec AESGCM nonce for counter n: 32 zero bits +// followed by the 64-bit big-endian counter (matches flynn/noise exactly). +func aesgcmNonce(n uint64) [12]byte { + var nonce [12]byte + binary.BigEndian.PutUint64(nonce[4:], n) + return nonce +} + +// installTransportKeysLocked derives the per-direction AEADs from the Noise +// split. The initiator sends with cs0 and receives with cs1; the responder the +// reverse. Must be called with s.mu held, before established is set. +func (s *peerSession) installTransportKeysLocked(cs0, cs1 *noise.CipherState) error { + txCS, rxCS := cs0, cs1 + if !s.initiator { + txCS, rxCS = cs1, cs0 + } + txKey := txCS.UnsafeKey() + rxKey := rxCS.UnsafeKey() + txBlock, err := aes.NewCipher(txKey[:]) + if err != nil { + return err + } + txAEAD, err := cipher.NewGCM(txBlock) + if err != nil { + return err + } + rxBlock, err := aes.NewCipher(rxKey[:]) + if err != nil { + return err + } + rxAEAD, err := cipher.NewGCM(rxBlock) + if err != nil { + return err + } + s.txAEAD, s.rxAEAD = txAEAD, rxAEAD + s.txCtr = 0 + return nil +} + +// decrypt authenticates and decrypts a full wire packet +// [version][type][counter][ct], appending the plaintext to dst. The 10-byte +// header is the AAD, so type and counter are authenticated; the anti-replay +// window is consulted only after authentication, so a forged packet can never +// advance it or punch a hole in it. +func (s *peerSession) decrypt(dst []byte, pkt []byte) ([]byte, error) { + if len(pkt) < transportHeaderLen { + return nil, errors.New("transport packet too short") + } + header := pkt[:transportHeaderLen] + counter := binary.BigEndian.Uint64(pkt[2:transportHeaderLen]) + s.mu.Lock() + aead := s.rxAEAD + established := s.established + s.mu.Unlock() + if !established || aead == nil { + return nil, fmt.Errorf("peer %s not established", s.addr) + } + nonce := aesgcmNonce(counter) + plaintext, err := aead.Open(dst, nonce[:], pkt[transportHeaderLen:], header) + if err != nil { + return nil, err + } + s.mu.Lock() + fresh := s.replay.Check(counter) + s.mu.Unlock() + if !fresh { + return nil, fmt.Errorf("replayed or out-of-window packet from %s", s.addr) + } + return plaintext, nil +} + +// rxPacket is a still-encrypted data/tun wire packet queued between the +// receive loops and the consumer (Read / ReadTunBatch). Decryption happens on +// the consumer side so batches can be opened in parallel; buf is pooled and +// owned by the consumer once dequeued. +type rxPacket struct { + typ byte + s *peerSession + buf []byte +} + type PeerConn struct { - conn *net.UDPConn + bind wgconn.Bind + batch int + static noise.DHKey psk []byte prologue []byte mu sync.Mutex - sessions map[string]*peerSession // addr → session + sessions map[string]*peerSession // canonical addr string → session knownStatics map[string]bool // hex(static) → allowed; the session allowlist from the rendezvous missedPings map[string]int // addr → consecutive unanswered pings @@ -78,18 +179,32 @@ type PeerConn struct { stop chan struct{} // closed by Shutdown to stop background loops stopOnce sync.Once - // readBuf is the receive buffer reused across Read calls. Read is only ever - // called from a single reader goroutine; allocating 64 KB per packet would - // generate hundreds of MB/s of garbage at tunnel line rate. - readBuf []byte + // rx carries encrypted data/tun packets from the receive loops to the + // consumer. Sends block when it is full: backpressure ripples to the socket + // instead of silently dropping authenticated-to-be traffic. + rx chan rxPacket + // recvDone is closed when every receive loop has exited (bind closed), so + // consumers blocked in Read/ReadTunBatch wake up at teardown. + recvDone chan struct{} + recvWG sync.WaitGroup + + // stunWaiter, when set, receives copies of non-protocol packets so + // DiscoverPublicAddr can catch the STUN response on the tunnel socket. + stunWaiter atomic.Pointer[chan []byte] + + // tunPkts/tunOK are ReadTunBatch scratch space (single consumer). + tunPkts []rxPacket + tunOK []bool } -// NewPeerConn creates a PeerConn. privateKey/publicKey are this peer's static -// X25519 keypair; psk is the Argon2id-derived pre-shared key (second factor); -// prologue binds the handshake to the protocol version and session id. -func NewPeerConn(conn *net.UDPConn, privateKey, publicKey, psk, prologue []byte) *PeerConn { - return &PeerConn{ - conn: conn, +// NewPeerConn creates a PeerConn on an opened Transport and starts its receive +// loops. privateKey/publicKey are this peer's static X25519 keypair; psk is the +// Argon2id-derived pre-shared key (second factor); prologue binds the handshake +// to the protocol version and session id. +func NewPeerConn(t *Transport, privateKey, publicKey, psk, prologue []byte) *PeerConn { + p := &PeerConn{ + bind: t.bind, + batch: t.batch, static: noise.DHKey{Private: privateKey, Public: publicKey}, psk: psk, prologue: prologue, @@ -99,15 +214,37 @@ func NewPeerConn(conn *net.UDPConn, privateKey, publicKey, psk, prologue []byte) Connected: make(chan *net.UDPAddr, 32), Dead: make(chan *net.UDPAddr, 10), stop: make(chan struct{}), + rx: make(chan rxPacket, 4*t.batch), + recvDone: make(chan struct{}), + } + p.recvWG.Add(len(t.recvFns)) + for _, fn := range t.recvFns { + go p.recvLoop(fn) } + go func() { + p.recvWG.Wait() + close(p.recvDone) + }() + return p } -// Shutdown signals background loops (handshake drivers) to stop. Safe to call -// multiple times; it does not close the underlying UDP connection. +// BatchSize is the number of packets a caller should size ReadTunBatch / +// SendBatch batches to. +func (p *PeerConn) BatchSize() int { return p.batch } + +// Shutdown signals background loops (handshake drivers, blocked consumers) to +// stop. Safe to call multiple times; it does not close the underlying bind. func (p *PeerConn) Shutdown() { p.stopOnce.Do(func() { close(p.stop) }) } +// Close shuts down background loops and closes the underlying bind, waking the +// receive loops. +func (p *PeerConn) Close() error { + p.Shutdown() + return p.bind.Close() +} + func staticKeyHex(pub []byte) string { return hex.EncodeToString(pub) } func (p *PeerConn) isKnownStatic(pub []byte) bool { @@ -125,6 +262,117 @@ func buildPacket(pktType byte, body []byte) []byte { return pkt } +// recvLoop services one of the bind's receive functions: batches of datagrams +// come in, handshake/punch/control packets are handled inline (rare, small), +// and data/tun packets are queued — still encrypted — for the consumer, which +// decrypts them in parallel. +func (p *PeerConn) recvLoop(fn wgconn.ReceiveFunc) { + defer p.recvWG.Done() + bufs := make([][]byte, p.batch) + sizes := make([]int, p.batch) + eps := make([]wgconn.Endpoint, p.batch) + for i := range bufs { + bufs[i] = make([]byte, maxUDPPacket) + } + for { + n, err := fn(bufs, sizes, eps) + if err != nil { + if errors.Is(err, net.ErrClosed) { + return + } + select { + case <-p.stop: + return + default: + continue // transient receive error (e.g. ICMP-triggered reset) + } + } + for i := 0; i < n; i++ { + if sizes[i] == 0 { + continue + } + p.handlePacket(bufs[i][:sizes[i]], eps[i]) + } + } +} + +// handlePacket classifies one received datagram. pkt aliases the receive +// buffer and is only valid for the duration of the call; anything that +// outlives it (queued data/tun packets) is copied into a pooled buffer. +func (p *PeerConn) handlePacket(pkt []byte, ep wgconn.Endpoint) { + // Every packet is [version][type][body...]. Drop anything too short or with a + // mismatched protocol version — no silent downgrade to an older protocol. + if len(pkt) < 2 || pkt[0] != ProtocolVersion { + // Not a blindspot packet. During public-address discovery the STUN + // response arrives on this same socket — hand it to the waiter. + if w := p.stunWaiter.Load(); w != nil { + cp := append([]byte(nil), pkt...) + select { + case *w <- cp: + default: + } + } + return + } + pktType := pkt[1] + body := pkt[2:] + switch pktType { + case PacketPunch: + return + case PacketHandshakeInit: + p.handleHandshakeInit(ep, body) + case PacketHandshakeResp: + p.handleHandshakeResp(ep, body) + case PacketControl: + s := p.sessionForEndpoint(ep) + if s == nil { + return + } + plaintext, err := s.decrypt(nil, pkt) + if err != nil || len(plaintext) < 1 { + return + } + switch plaintext[0] { + case CtrlPing: + UpdateLastSeen() + p.sendControl(s.addr, CtrlPong) + case CtrlPong: + UpdateLastSeen() + p.mu.Lock() + p.missedPings[s.key] = 0 + p.mu.Unlock() + case CtrlDead: + // Authenticated "I'm leaving" from the peer: drop its session and + // surface a Dead event — other peers are unaffected. + p.RemovePeer(s.addr) + } + case PacketData, PacketTun: + s := p.sessionForEndpoint(ep) + if s == nil { + return + } + buf := getPacketBuf(len(pkt)) + copy(buf, pkt) + select { + case p.rx <- rxPacket{typ: pktType, s: s, buf: buf}: + case <-p.stop: + putPacketBuf(buf) + } + } +} + +// sessionForEndpoint resolves the session for a received packet's source. +func (p *PeerConn) sessionForEndpoint(ep wgconn.Endpoint) *peerSession { + key, _, ok := canonEndpointKey(ep) + if !ok { + return nil + } + p.mu.Lock() + s := p.sessions[key] + p.mu.Unlock() + return s +} + // AddKnownPeer registers a peer learned from the trusted rendezvous: its UDP // address and its static public key (the pubkey the rendezvous published). The // key is added to the session allowlist, the local peer's handshake role is @@ -139,12 +387,19 @@ func (p *PeerConn) AddKnownPeer(addr *net.UDPAddr, remoteStatic []byte) error { key := make([]byte, 32) copy(key, remoteStatic) + ap := canonAddrPort(addr.AddrPort()) + sessKey := ap.String() + ep, err := p.bind.ParseEndpoint(sessKey) + if err != nil { + return fmt.Errorf("parsing peer endpoint %s: %w", sessKey, err) + } + p.mu.Lock() p.knownStatics[staticKeyHex(key)] = true - s, ok := p.sessions[addr.String()] + s, ok := p.sessions[sessKey] if !ok { - s = &peerSession{addr: addr} - p.sessions[addr.String()] = s + s = &peerSession{addr: net.UDPAddrFromAddrPort(ap), key: sessKey, ep: ep} + p.sessions[sessKey] = s } p.mu.Unlock() @@ -231,6 +486,7 @@ func (p *PeerConn) driveHandshake(s *peerSession, gen uint64, deadline time.Time established := s.established initiator := s.initiator initMsg := s.initMsg + ep := s.ep s.mu.Unlock() if superseded || established { return @@ -239,9 +495,9 @@ func (p *PeerConn) driveHandshake(s *peerSession, gen uint64, deadline time.Time return } if initiator && initMsg != nil { - p.conn.WriteToUDP(initMsg, s.addr) + p.bind.Send([][]byte{initMsg}, ep) } else { - p.conn.WriteToUDP(punch, s.addr) + p.bind.Send([][]byte{punch}, ep) } attempts++ if attempts >= handshakeAggressiveAttempts && interval < handshakeMaxInterval { @@ -258,75 +514,151 @@ func (p *PeerConn) driveHandshake(s *peerSession, gen uint64, deadline time.Time } } -// Read returns the packet type, decrypted payload, sender address, and any error. +// Read returns one decrypted packet: its type, payload, and sender address. // Callers filter on PacketData (chat) or PacketTun (VPN). Handshake, punch, and -// keepalive packets are handled internally and never returned. +// keepalive packets are handled internally and never returned. High-throughput +// consumers should use ReadTunBatch instead. func (p *PeerConn) Read() (byte, []byte, *net.UDPAddr, error) { - if p.readBuf == nil { - p.readBuf = make([]byte, 65536) - } - buf := p.readBuf for { - n, addr, err := p.conn.ReadFromUDP(buf) + select { + case pkt := <-p.rx: + plaintext, err := pkt.s.decrypt(nil, pkt.buf) + addr := pkt.s.addr + typ := pkt.typ + putPacketBuf(pkt.buf) + if err != nil { + continue + } + return typ, plaintext, addr, nil + case <-p.stop: + return 0, nil, nil, fmt.Errorf("error reading from peer: %w", net.ErrClosed) + case <-p.recvDone: + return 0, nil, nil, fmt.Errorf("error reading from peer: %w", net.ErrClosed) + } + } +} + +// ReadTunBatch blocks for at least one tunnel packet, then drains whatever else +// is immediately available (up to len(bufs)) and decrypts the whole batch in +// parallel. Plaintexts land in bufs[:n] (slices are swapped/replaced as +// needed); senders[:n] receives each packet's canonical remote address string, +// suitable for keying reverse-path maps. Non-tun packets and packets that fail +// authentication or replay checks are dropped. Not safe for concurrent use. +func (p *PeerConn) ReadTunBatch(bufs [][]byte, senders []string) (int, error) { + if len(bufs) == 0 { + return 0, nil + } + pkts := p.tunPkts[:0] + + // Block for the first packet. + select { + case pkt := <-p.rx: + pkts = append(pkts, pkt) + case <-p.stop: + return 0, net.ErrClosed + case <-p.recvDone: + return 0, net.ErrClosed + } + // Opportunistically drain the rest of the burst. +drain: + for len(pkts) < len(bufs) { + select { + case pkt := <-p.rx: + pkts = append(pkts, pkt) + default: + break drain + } + } + + ok := p.tunOK[:0] + for range pkts { + ok = append(ok, false) + } + p.tunPkts, p.tunOK = pkts, ok // keep scratch capacity for the next call + + parallelFor(len(pkts), func(i int) { + if pkts[i].typ != PacketTun { + return + } + plaintext, err := pkts[i].s.decrypt(bufs[i][:0], pkts[i].buf) if err != nil { - return 0, nil, nil, fmt.Errorf("error reading from peer: %w", err) + return } - // Every packet is [version][type][body...]. Drop anything too short or with a - // mismatched protocol version — no silent downgrade to an older protocol. - if n < 2 || buf[0] != ProtocolVersion { + bufs[i] = plaintext // may have been reallocated by Open; keep the header + ok[i] = true + }) + + n := 0 + for i := range pkts { + putPacketBuf(pkts[i].buf) + if !ok[i] { continue } - pktType := buf[1] - body := buf[2:n] - switch pktType { - case PacketPunch: - continue - case PacketHandshakeInit: - p.handleHandshakeInit(addr, body) - continue - case PacketHandshakeResp: - p.handleHandshakeResp(addr, body) - continue - case PacketControl: - plaintext, err := p.openTransport(addr, buf[:n]) - if err != nil || len(plaintext) < 1 { - continue - } - switch plaintext[0] { - case CtrlPing: - UpdateLastSeen() - p.sendControl(addr, CtrlPong) - case CtrlPong: - UpdateLastSeen() - p.mu.Lock() - p.missedPings[addr.String()] = 0 - p.mu.Unlock() - case CtrlDead: - // Authenticated "I'm leaving" from the peer: drop its session, surface - // a Dead event, and keep reading — other peers are unaffected. - p.RemovePeer(addr) - } - continue - case PacketData, PacketTun: - plaintext, err := p.openTransport(addr, buf[:n]) - if err != nil { - continue - } - return pktType, plaintext, addr, nil + if n != i { + bufs[n], bufs[i] = bufs[i], bufs[n] + } + senders[n] = pkts[i].s.key + n++ + } + return n, nil +} + +// parallelFor runs fn(0..n-1) across up to GOMAXPROCS goroutines, falling back +// to a plain loop for small n where fork-join overhead would dominate. +func parallelFor(n int, fn func(i int)) { + const minPerWorker = 8 + workers := runtime.GOMAXPROCS(0) + if w := (n + minPerWorker - 1) / minPerWorker; w < workers { + workers = w + } + if workers <= 1 { + for i := 0; i < n; i++ { + fn(i) + } + return + } + chunk := (n + workers - 1) / workers + var wg sync.WaitGroup + for w := 0; w < workers; w++ { + lo := w * chunk + hi := lo + chunk + if hi > n { + hi = n + } + if lo >= hi { + break } + wg.Add(1) + go func(lo, hi int) { + defer wg.Done() + for i := lo; i < hi; i++ { + fn(i) + } + }(lo, hi) } + wg.Wait() } // handleHandshakeInit processes an inbound Noise msg1 (we are the responder). -func (p *PeerConn) handleHandshakeInit(addr *net.UDPAddr, msg1 []byte) { +func (p *PeerConn) handleHandshakeInit(ep wgconn.Endpoint, msg1 []byte) { + sessKey, ap, okKey := canonEndpointKey(ep) + if !okKey { + return + } p.mu.Lock() - s, ok := p.sessions[addr.String()] + s, ok := p.sessions[sessKey] if !ok { // A msg1 may arrive before the rendezvous stream told us about this peer. // Create a provisional session; the initiator's static is still validated - // against the session allowlist below. - s = &peerSession{addr: addr} - p.sessions[addr.String()] = s + // against the session allowlist below. The endpoint is re-parsed rather + // than retained: received endpoints may alias bind-internal state. + pep, err := p.bind.ParseEndpoint(sessKey) + if err != nil { + p.mu.Unlock() + return + } + s = &peerSession{addr: net.UDPAddrFromAddrPort(ap), key: sessKey, ep: pep} + p.sessions[sessKey] = s } p.mu.Unlock() @@ -334,9 +666,10 @@ func (p *PeerConn) handleHandshakeInit(addr *net.UDPAddr, msg1 []byte) { if s.established { // Retransmit our cached msg2 in case the initiator's copy was lost. resp := s.respMsg + sendEp := s.ep s.mu.Unlock() if resp != nil { - p.conn.WriteToUDP(resp, addr) + p.bind.Send([][]byte{resp}, sendEp) } return } @@ -376,25 +709,27 @@ func (p *PeerConn) handleHandshakeInit(addr *net.UDPAddr, msg1 []byte) { s.mu.Unlock() return } - // Responder: send with cs1 (responder→initiator), receive with cs0 (initiator→responder). - s.tx, s.rx = cs1, cs0 + if err := s.installTransportKeysLocked(cs0, cs1); err != nil { + s.hs = nil + s.mu.Unlock() + return + } s.remoteStatic = append([]byte(nil), remote...) s.respMsg = buildPacket(PacketHandshakeResp, msg2) s.established = true s.driving = false // the retransmit driver will notice and exit; the session may be re-armed before then resp := s.respMsg + sendEp := s.ep s.mu.Unlock() - p.conn.WriteToUDP(resp, addr) - p.fireConnected(addr) + p.bind.Send([][]byte{resp}, sendEp) + p.fireConnected(s) } // handleHandshakeResp processes an inbound Noise msg2 (we are the initiator). -func (p *PeerConn) handleHandshakeResp(addr *net.UDPAddr, msg2 []byte) { - p.mu.Lock() - s, ok := p.sessions[addr.String()] - p.mu.Unlock() - if !ok { +func (p *PeerConn) handleHandshakeResp(ep wgconn.Endpoint, msg2 []byte) { + s := p.sessionForEndpoint(ep) + if s == nil { return } s.mu.Lock() @@ -417,109 +752,116 @@ func (p *PeerConn) handleHandshakeResp(addr *net.UDPAddr, msg2 []byte) { s.mu.Unlock() return } - // Initiator: send with cs0 (initiator→responder), receive with cs1 (responder→initiator). - s.tx, s.rx = cs0, cs1 + if err := s.installTransportKeysLocked(cs0, cs1); err != nil { + s.mu.Unlock() + return + } s.remoteStatic = append([]byte(nil), remote...) s.established = true s.driving = false // the retransmit driver will notice and exit; the session may be re-armed before then s.mu.Unlock() - p.fireConnected(addr) + p.fireConnected(s) } // transportHeaderLen is the cleartext prefix [version][type][counter] that also // serves as the AEAD additional data. const transportHeaderLen = 2 + counterLen -// openTransport authenticates and decrypts a transport packet -// [version][type][counter][ct], enforcing the anti-replay window. The 10-byte -// header is used as AAD, so the type and counter are authenticated. -func (p *PeerConn) openTransport(addr *net.UDPAddr, pkt []byte) ([]byte, error) { - if len(pkt) < transportHeaderLen { - return nil, fmt.Errorf("transport packet too short") - } - header := pkt[:transportHeaderLen] - counter := binary.BigEndian.Uint64(pkt[2:transportHeaderLen]) - ct := pkt[transportHeaderLen:] +// sealPacket encrypts one payload into a pooled wire packet +// [version][type][counter][ct]. The AEAD appends in place: the pooled buffer +// is sized so Seal never reallocates. +func sealPacket(aead cipher.AEAD, pktType byte, counter uint64, payload []byte) []byte { + buf := getPacketBuf(transportHeaderLen + len(payload) + 16) + buf[0] = ProtocolVersion + buf[1] = pktType + binary.BigEndian.PutUint64(buf[2:transportHeaderLen], counter) + nonce := aesgcmNonce(counter) + return aead.Seal(buf[:transportHeaderLen], nonce[:], payload, buf[:transportHeaderLen]) +} +func (p *PeerConn) sessionByAddr(addr *net.UDPAddr) *peerSession { + key := canonAddrPort(addr.AddrPort()).String() p.mu.Lock() - s, ok := p.sessions[addr.String()] + s := p.sessions[key] p.mu.Unlock() - if !ok { - return nil, fmt.Errorf("unknown peer: %s", addr) - } - s.mu.Lock() - defer s.mu.Unlock() - if !s.established || s.rx == nil { - return nil, fmt.Errorf("peer %s not established", addr) - } - // The counter is the AEAD nonce; using it directly lets us decrypt out of order. - s.rx.SetNonce(counter) - plaintext, err := s.rx.Decrypt(nil, header, ct) - if err != nil { - return nil, err - } - // Only after authentication do we consult the replay window, so a forged packet - // can never advance it or punch a hole in it. - if !s.replay.Check(counter) { - return nil, fmt.Errorf("replayed or out-of-window packet from %s", addr) - } - return plaintext, nil + return s } -func (p *PeerConn) send(addr *net.UDPAddr, data []byte, pktType byte) error { - p.mu.Lock() - s, ok := p.sessions[addr.String()] - p.mu.Unlock() - if !ok { +// SendBatch encrypts payloads in parallel and hands them to the bind as +// batches, which coalesces them into far fewer syscalls (UDP GSO/sendmmsg on +// Linux, RIO on Windows). Counters are reserved contiguously up front, so one +// lock acquisition covers the whole batch. +func (p *PeerConn) SendBatch(addr *net.UDPAddr, payloads [][]byte, pktType byte) error { + s := p.sessionByAddr(addr) + if s == nil { return fmt.Errorf("unknown peer: %s", addr) } - // Capacity covers header + ciphertext + AEAD tag so Encrypt's append never - // reallocates mid-packet. - header := make([]byte, transportHeaderLen, transportHeaderLen+len(data)+16) - header[0] = ProtocolVersion - header[1] = pktType + return p.sendBatchSession(s, payloads, pktType) +} +func (p *PeerConn) sendBatchSession(s *peerSession, payloads [][]byte, pktType byte) error { + m := len(payloads) + if m == 0 { + return nil + } s.mu.Lock() - if !s.established || s.tx == nil { + if !s.established || s.txAEAD == nil { s.mu.Unlock() - return fmt.Errorf("peer %s not established", addr) + return fmt.Errorf("peer %s not established", s.addr) } - counter := s.tx.Nonce() - binary.BigEndian.PutUint64(header[2:], counter) - // AAD is the cleartext header (version, type, counter). Append the ciphertext - // after the header in a single buffer so the returned slice is the whole packet. - pkt, err := s.tx.Encrypt(header, header, data) + aead := s.txAEAD + ep := s.ep + base := s.txCtr + s.txCtr += uint64(m) s.mu.Unlock() - if err != nil { - return err + + wire := make([][]byte, m) + parallelFor(m, func(i int) { + wire[i] = sealPacket(aead, pktType, base+uint64(i), payloads[i]) + }) + + var err error + for off := 0; off < m; off += p.batch { + end := off + p.batch + if end > m { + end = m + } + if e := p.bind.Send(wire[off:end], ep); e != nil && err == nil { + err = e + } + } + for _, w := range wire { + putPacketBuf(w) } - _, err = p.conn.WriteToUDP(pkt, addr) return err } // sendControl sends an encrypted control message (ping/pong/dead) over the same // authenticated, anti-replay-protected transport as data. func (p *PeerConn) sendControl(addr *net.UDPAddr, opcode byte) error { - return p.send(addr, []byte{opcode}, PacketControl) + return p.SendBatch(addr, [][]byte{{opcode}}, PacketControl) } func (p *PeerConn) Send(addr *net.UDPAddr, data []byte) error { - return p.send(addr, data, PacketData) + return p.SendBatch(addr, [][]byte{data}, PacketData) } func (p *PeerConn) SendTun(addr *net.UDPAddr, data []byte) error { - return p.send(addr, data, PacketTun) + return p.SendBatch(addr, [][]byte{data}, PacketTun) +} + +// SendTunBatch sends a batch of tunnelled IP packets to one peer. +func (p *PeerConn) SendTunBatch(addr *net.UDPAddr, packets [][]byte) error { + return p.SendBatch(addr, packets, PacketTun) } // PeerPublicKey returns the authenticated static public key of the established // peer at addr. The key is authenticated because it came from the trusted // rendezvous and was verified by the Noise handshake. func (p *PeerConn) PeerPublicKey(addr *net.UDPAddr) ([]byte, bool) { - p.mu.Lock() - s, ok := p.sessions[addr.String()] - p.mu.Unlock() - if !ok { + s := p.sessionByAddr(addr) + if s == nil { return nil, false } s.mu.Lock() @@ -530,24 +872,25 @@ func (p *PeerConn) PeerPublicKey(addr *net.UDPAddr) ([]byte, bool) { return s.remoteStatic, true } -func (p *PeerConn) fireConnected(addr *net.UDPAddr) { +func (p *PeerConn) fireConnected(s *peerSession) { p.mu.Lock() - p.missedPings[addr.String()] = 0 + p.missedPings[s.key] = 0 p.mu.Unlock() // Deliver reliably: block until the consumer takes the event, or until shutdown. // Dropping it would leave the peer without a virtual-IP mapping, so all of its // tunnel traffic would be silently discarded for the rest of the session. This // runs after s.mu is released, so blocking here cannot deadlock a session. select { - case p.Connected <- addr: + case p.Connected <- s.addr: case <-p.stop: } } func (p *PeerConn) dropSession(addr *net.UDPAddr) { + key := canonAddrPort(addr.AddrPort()).String() p.mu.Lock() - delete(p.sessions, addr.String()) - delete(p.missedPings, addr.String()) + delete(p.sessions, key) + delete(p.missedPings, key) p.mu.Unlock() } @@ -575,8 +918,9 @@ func (p *PeerConn) RemovePeer(addr *net.UDPAddr) { // so the session heals on its own if the outage was transient (the peer never // told the rendezvous it left, so nothing else would ever reconnect the two). func (p *PeerConn) TimeoutPeer(addr *net.UDPAddr) { + key := canonAddrPort(addr.AddrPort()).String() p.mu.Lock() - delete(p.missedPings, addr.String()) + delete(p.missedPings, key) p.mu.Unlock() p.retryHandshake(addr) p.fireDead(addr) @@ -587,10 +931,8 @@ func (p *PeerConn) TimeoutPeer(addr *net.UDPAddr) { // as the pinned key for the next one. The retry is bounded by retryWindow; after // that the session is parked until an inbound msg1 or AddKnownPeer revives it. func (p *PeerConn) retryHandshake(addr *net.UDPAddr) { - p.mu.Lock() - s, ok := p.sessions[addr.String()] - p.mu.Unlock() - if !ok { + s := p.sessionByAddr(addr) + if s == nil { return } s.mu.Lock() @@ -606,7 +948,8 @@ func (p *PeerConn) retryHandshake(addr *net.UDPAddr) { return } s.established = false - s.tx, s.rx = nil, nil + s.txAEAD, s.rxAEAD = nil, nil + s.txCtr = 0 s.remoteStatic = nil s.replay = ReplayWindow{} gen, err := p.armHandshakeLocked(s, remote) diff --git a/internal/network/transport.go b/internal/network/transport.go new file mode 100644 index 0000000..7bb0c92 --- /dev/null +++ b/internal/network/transport.go @@ -0,0 +1,220 @@ +package network + +import ( + "fmt" + "net" + "net/netip" + "sync" + "time" + + "github.com/pion/stun" + wgconn "golang.zx2c4.com/wireguard/conn" +) + +// The tunnel transport rides on wireguard-go's conn.Bind instead of a raw +// *net.UDPConn. The Bind batches datagrams through the kernel — UDP GSO/GRO and +// sendmmsg/recvmmsg on Linux, registered I/O (RIO) rings on Windows — so the +// per-packet syscall cost that capped throughput is amortized across batches of +// up to conn.IdealBatchSize packets. The Bind is protocol-agnostic: everything +// above it (Noise handshake, framing, AEAD, replay protection) is unchanged. + +// maxUDPPacket is the receive buffer size per batch slot. It must hold the +// largest possible datagram: with GRO the kernel can hand back a coalesced +// super-packet of up to 64 KB in a single slot before the bind splits it. +const maxUDPPacket = 65536 + +// packetBufSize is the pooled buffer size for individual wire packets. It +// covers the transport header + a full tunnel MTU payload + AEAD tag with room +// to spare; larger (rare, e.g. long chat lines) packets fall back to plain +// allocation. +const packetBufSize = 2048 + +var packetPool = sync.Pool{ + New: func() any { + b := make([]byte, packetBufSize) + return &b + }, +} + +// getPacketBuf returns a buffer of length n, pooled when n fits the standard +// packet size. +func getPacketBuf(n int) []byte { + if n <= packetBufSize { + return (*packetPool.Get().(*[]byte))[:n] + } + return make([]byte, n) +} + +// putPacketBuf recycles a buffer obtained from getPacketBuf. Buffers that were +// reallocated (or oversized ones) are simply dropped for the GC. +func putPacketBuf(b []byte) { + if cap(b) == packetBufSize { + b = b[:packetBufSize] + packetPool.Put(&b) + } +} + +// canonAddrPort normalizes an address to its unmapped form so that +// "::ffff:1.2.3.4" and "1.2.3.4" produce the same session key regardless of +// which socket (v4/v6) or code path (ParseEndpoint vs. receive) produced it. +func canonAddrPort(ap netip.AddrPort) netip.AddrPort { + return netip.AddrPortFrom(ap.Addr().Unmap(), ap.Port()) +} + +// canonEndpointKey derives the canonical session key ("ip:port") and AddrPort +// for a bind endpoint. +func canonEndpointKey(ep wgconn.Endpoint) (string, netip.AddrPort, bool) { + ap, err := netip.ParseAddrPort(ep.DstToString()) + if err != nil { + return "", netip.AddrPort{}, false + } + ap = canonAddrPort(ap) + return ap.String(), ap, true +} + +// Transport is an opened UDP transport: a bind in the listening state plus its +// receive functions, ready to be handed to NewPeerConn. +type Transport struct { + bind wgconn.Bind + recvFns []wgconn.ReceiveFunc + port uint16 + batch int +} + +// OpenTransport opens the platform's default batched bind (RIO on Windows, +// GSO/sendmmsg-capable StdNetBind elsewhere) on an ephemeral port. The bind's +// own control functions already request large (7 MB) kernel socket buffers. +func OpenTransport() (*Transport, error) { + bind := wgconn.NewDefaultBind() + fns, port, err := bind.Open(0) + if err != nil { + return nil, fmt.Errorf("opening UDP bind: %w", err) + } + return &Transport{bind: bind, recvFns: fns, port: port, batch: bind.BatchSize()}, nil +} + +// Port returns the local UDP port the transport is bound to. +func (t *Transport) Port() int { return int(t.port) } + +// WrapUDPConn adapts an existing, already-bound *net.UDPConn into a Transport. +// It performs no batched I/O on receive (one datagram per call) and is meant +// for tests, which bind explicitly to loopback; production code uses +// OpenTransport. +func WrapUDPConn(c *net.UDPConn) *Transport { + b := &udpBind{conn: c} + return &Transport{ + bind: b, + recvFns: []wgconn.ReceiveFunc{b.receive}, + port: uint16(c.LocalAddr().(*net.UDPAddr).Port), + batch: b.BatchSize(), + } +} + +// udpBind is a minimal conn.Bind over a pre-bound *net.UDPConn. Send accepts +// full batches (looping over WriteToUDPAddrPort) so the batched send path is +// exercised; receive returns one datagram per call. +type udpBind struct{ conn *net.UDPConn } + +type udpEndpoint netip.AddrPort + +func (e udpEndpoint) ClearSrc() {} +func (e udpEndpoint) SrcToString() string { return "" } +func (e udpEndpoint) DstToString() string { return netip.AddrPort(e).String() } +func (e udpEndpoint) DstToBytes() []byte { + b, _ := netip.AddrPort(e).MarshalBinary() + return b +} +func (e udpEndpoint) DstIP() netip.Addr { return netip.AddrPort(e).Addr() } +func (e udpEndpoint) SrcIP() netip.Addr { return netip.Addr{} } + +func (b *udpBind) Open(_ uint16) ([]wgconn.ReceiveFunc, uint16, error) { + return []wgconn.ReceiveFunc{b.receive}, uint16(b.conn.LocalAddr().(*net.UDPAddr).Port), nil +} + +func (b *udpBind) receive(packets [][]byte, sizes []int, eps []wgconn.Endpoint) (int, error) { + n, addr, err := b.conn.ReadFromUDPAddrPort(packets[0]) + if err != nil { + return 0, err + } + sizes[0] = n + eps[0] = udpEndpoint(canonAddrPort(addr)) + return 1, nil +} + +func (b *udpBind) Send(bufs [][]byte, ep wgconn.Endpoint) error { + ue, ok := ep.(udpEndpoint) + if !ok { + return wgconn.ErrWrongEndpointType + } + for _, buf := range bufs { + if _, err := b.conn.WriteToUDPAddrPort(buf, netip.AddrPort(ue)); err != nil { + return err + } + } + return nil +} + +func (b *udpBind) ParseEndpoint(s string) (wgconn.Endpoint, error) { + ap, err := netip.ParseAddrPort(s) + if err != nil { + return nil, err + } + return udpEndpoint(canonAddrPort(ap)), nil +} + +func (b *udpBind) Close() error { return b.conn.Close() } +func (b *udpBind) SetMark(_ uint32) error { return nil } +func (b *udpBind) BatchSize() int { return wgconn.IdealBatchSize } + +// stunServer is the public STUN server used to discover this host's public +// address for NAT traversal. +const stunServer = "stun.l.google.com:19302" + +// DiscoverPublicAddr sends a STUN binding request through the tunnel's own +// bind (so the discovered mapping is the one peers will punch to) and returns +// the public "ip:port". Unlike the old one-shot version it retransmits, since +// a single lost STUN datagram used to hang connection setup forever. +func (p *PeerConn) DiscoverPublicAddr() (string, error) { + raddr, err := net.ResolveUDPAddr("udp4", stunServer) + if err != nil { + return "", fmt.Errorf("resolving STUN server address: %w", err) + } + ep, err := p.bind.ParseEndpoint(canonAddrPort(raddr.AddrPort()).String()) + if err != nil { + return "", fmt.Errorf("parsing STUN endpoint: %w", err) + } + + ch := make(chan []byte, 4) + p.stunWaiter.Store(&ch) + defer p.stunWaiter.Store(nil) + + req := stun.MustBuild(stun.TransactionID, stun.BindingRequest) + for attempt := 0; attempt < 4; attempt++ { + if err := p.bind.Send([][]byte{req.Raw}, ep); err != nil { + return "", fmt.Errorf("sending STUN request: %w", err) + } + timer := time.NewTimer(2 * time.Second) + wait: + for { + select { + case pkt := <-ch: + m := &stun.Message{Raw: pkt} + if err := m.Decode(); err != nil { + continue // unrelated non-protocol packet + } + var xorAddr stun.XORMappedAddress + if err := xorAddr.GetFrom(m); err != nil { + continue + } + timer.Stop() + return fmt.Sprintf("%s:%d", xorAddr.IP, xorAddr.Port), nil + case <-timer.C: + break wait + case <-p.stop: + timer.Stop() + return "", net.ErrClosed + } + } + } + return "", fmt.Errorf("no response from STUN server %s", stunServer) +} diff --git a/internal/transfer/bench_test.go b/internal/transfer/bench_test.go new file mode 100644 index 0000000..3533362 --- /dev/null +++ b/internal/transfer/bench_test.go @@ -0,0 +1,49 @@ +package transfer + +import ( + "context" + "math/rand" + "os" + "path/filepath" + "testing" +) + +// BenchmarkLoopbackTransfer measures the whole transfer path (header framing, +// 1 MB copy buffer, progress writer) over loopback TCP — the tunnel-free +// baseline the tunnelled number is compared against. Reported MB/s is file +// payload throughput. +func BenchmarkLoopbackTransfer(b *testing.B) { + const size = 64 << 20 // 64 MB per op + + srcDir := b.TempDir() + destDir := b.TempDir() + src := filepath.Join(srcDir, "payload.bin") + buf := make([]byte, size) + rand.New(rand.NewSource(1)).Read(buf) + if err := os.WriteFile(src, buf, 0644); err != nil { + b.Fatalf("writing payload: %v", err) + } + + ctx := context.Background() + b.SetBytes(size) + b.ResetTimer() + for i := 0; i < b.N; i++ { + recv, err := Listen(ctx, "127.0.0.1") + if err != nil { + b.Fatalf("Listen: %v", err) + } + errCh := make(chan error, 1) + go func() { + _, err := recv.Accept(ctx, destDir, nil) + errCh <- err + }() + if _, err := Send(ctx, "127.0.0.1", src, nil); err != nil { + recv.Close() + b.Fatalf("Send: %v", err) + } + if err := <-errCh; err != nil { + b.Fatalf("Accept: %v", err) + } + recv.Close() + } +} diff --git a/internal/transfer/transfer.go b/internal/transfer/transfer.go index 093764b..54dcc88 100644 --- a/internal/transfer/transfer.go +++ b/internal/transfer/transfer.go @@ -29,6 +29,25 @@ const dialTimeout = 5 * time.Second // sampleInterval is how often live progress snapshots are emitted during a transfer. const sampleInterval = 500 * time.Millisecond +// copyBufSize is the userspace copy buffer for the file body. io.Copy's default +// 32 KB means one write syscall per 32 KB; 1 MB cuts the syscall count ~32x on a +// bulk transfer for a fixed, one-off allocation. +// +// Note: we deliberately do NOT call SetReadBuffer/SetWriteBuffer on the TCP +// connection. A fixed SO_RCVBUF/SO_SNDBUF disables the OS's TCP window +// autotuning, which on both Linux (tcp_rmem, up to ~6 MB) and Windows (receive +// autotuning, up to 16 MB) already grows past anything we could safely pin — +// and on Linux the fixed value would additionally be clamped to rmem_max +// (~416 KB), a hard regression on high-BDP paths. +const copyBufSize = 1 << 20 + +// bodyReader hides any WriteTo/ReadFrom fast paths of the underlying reader so +// io.CopyBuffer actually uses our large buffer instead of falling back to the +// generic 32 KB path inside os.File.WriteTo. +type bodyReader struct{ r io.Reader } + +func (br bodyReader) Read(p []byte) (int, error) { return br.r.Read(p) } + // progressWriter is an io.Writer that counts the bytes passing through it, so a // background sampler can report progress without disturbing the copy. type progressWriter struct { @@ -136,7 +155,7 @@ func Send(ctx context.Context, peer, path string, prog chan<- Progress) (*Result pw := &progressWriter{w: conn} stopSampler := startSampler(prog, pw, base) start := time.Now() - n, err := io.Copy(pw, f) + n, err := io.CopyBuffer(pw, bodyReader{f}, make([]byte, copyBufSize)) stopSampler() if err != nil { if ctx.Err() != nil { @@ -236,7 +255,10 @@ func (r *Receiver) Accept(ctx context.Context, destDir string, prog chan<- Progr pw := &progressWriter{w: f} stopSampler := startSampler(prog, pw, base) start := time.Now() - n, err := io.CopyN(pw, conn, int64(fileSize)) + n, err := io.CopyBuffer(pw, bodyReader{io.LimitReader(conn, int64(fileSize))}, make([]byte, copyBufSize)) + if err == nil && n < int64(fileSize) { + err = io.EOF // peer closed early; matches io.CopyN's short-read error + } stopSampler() if err != nil { if ctx.Err() != nil { diff --git a/test/connection_test.go b/test/connection_test.go index f0d60b0..95605d0 100644 --- a/test/connection_test.go +++ b/test/connection_test.go @@ -209,7 +209,7 @@ func newTestPeerWithKey(t *testing.T, sessionID, password string, kp *crypto.Key } addr := fmt.Sprintf("127.0.0.1:%d", conn.LocalAddr().(*net.UDPAddr).Port) psk := crypto.DerivePSK(password, sessionID) - pc := network.NewPeerConn(conn, kp.PrivateKey, kp.PublicKey, psk, network.Prologue(sessionID)) + pc := network.NewPeerConn(network.WrapUDPConn(conn), kp.PrivateKey, kp.PublicKey, psk, network.Prologue(sessionID)) return &testPeer{ conn: conn, addr: addr, diff --git a/test/hardening_test.go b/test/hardening_test.go index d6850fc..d5ac374 100644 --- a/test/hardening_test.go +++ b/test/hardening_test.go @@ -28,7 +28,7 @@ func deliverAndRead(t *testing.T, payload []byte) (panicVal any) { t.Fatalf("GenerateKeyPair: %v", err) } psk := crypto.DerivePSK("hardening-pass", "hardening-session") - pc := network.NewPeerConn(conn, kp.PrivateKey, kp.PublicKey, psk, network.Prologue("hardening-session")) + pc := network.NewPeerConn(network.WrapUDPConn(conn), kp.PrivateKey, kp.PublicKey, psk, network.Prologue("hardening-session")) var ( mu sync.Mutex