From d3e21a8fc8e18c88c6b07ec8aa845499005e01bf Mon Sep 17 00:00:00 2001 From: Ryan Huber Date: Sat, 17 Oct 2026 18:53:36 +0000 Subject: [PATCH] relay: Batch forwarded packets on the relay host --- interface.go | 3 + outside.go | 40 ++- outside_relay_test.go | 291 ++++++++++++++++++++ overlay/batch/tx_batch.go | 4 +- udp/conn.go | 2 +- udp/udp_linux_writebatch.go | 7 +- udp/udp_linux_writebatch_concurrent_test.go | 61 ++++ 7 files changed, 398 insertions(+), 10 deletions(-) create mode 100644 outside_relay_test.go create mode 100644 udp/udp_linux_writebatch_concurrent_test.go diff --git a/interface.go b/interface.go index 96d44446..5adee9b0 100644 --- a/interface.go +++ b/interface.go @@ -384,6 +384,8 @@ type rxContext struct { hostmapCache map[uint32]*HostInfo lhh *LightHouseHandler ctCache *firewall.ConntrackCacheTicker + relayBatch *batch.SendBatch + relayRemote netip.AddrPort } func newRxContext(f *Interface, q int) *rxContext { @@ -414,6 +416,7 @@ func (f *Interface) listenOut(i int) { } flusher := func() { + f.flushRelayBatch(rxc) if err := f.batchers[i].Flush(); err != nil { f.l.Error("Failed to flush tun coalescer", "error", err) } diff --git a/outside.go b/outside.go index 54a1d393..c14abe4b 100644 --- a/outside.go +++ b/outside.go @@ -25,7 +25,7 @@ var ErrOutOfWindow = errors.New("out of window packet") // readOutsidePackets processes one received underlay packet. // Message payloads are decrypted IN PLACE, so packet must stay untouched -// by the caller until the batcher for queue q has been flushed +// by the caller until the TUN and relay batches for queue q have been flushed. func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxContext) { h := rxc.h err := h.Parse(packet) @@ -77,6 +77,12 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont f.messageMetrics.Rx(h.Type, h.Subtype, 1) } + // Earlier forwards must leave before this packet can generate a response. + isMessageRelay := h.Type == header.Message && h.Subtype == header.MessageRelay + if !isMessageRelay { + f.flushRelayBatch(rxc) + } + // Unencrypted packets switch h.Type { case header.Handshake: @@ -88,9 +94,6 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont return } - // Relay packets are special - isMessageRelay := (h.Type == header.Message && h.Subtype == header.MessageRelay) - var hostinfo *HostInfo if isMessageRelay { hostinfo = f.hostMap.QueryRelayIndex(h.RemoteIndex) @@ -102,6 +105,8 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont // recvError if necessary if hostinfo == nil || hostinfo.ConnectionState == nil { if !via.IsRelayed { + // An unknown relay index can also produce a response. + f.flushRelayBatch(rxc) f.maybeSendRecvError(via.UdpAddr, h.RemoteIndex) } return @@ -233,8 +238,10 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender, // Forward this packet through the relay tunnel, rebuilding it in place. // Encode overwrites the old outer header, and the new AEAD tag lands where the old one was fwdBuf := packet[:0] - //todo it would potentially be nice to batch these - f.SendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true, rxc.q) + out, err := f.prepareSendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true) + if err == nil { + f.queueRelayPacket(out, targetHI.GetRemote(), rxc) + } case TerminalType: hostinfo.logger(f.l).Error("Unexpected Relay Type of Terminal") return @@ -259,6 +266,27 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender, } } +// queueRelayPacket borrows a receive buffer until the batch is flushed. Destination changes and the send batch +// limit end a run, preserving arrival order while allowing WriteBatch to use UDP segmentation offload. +func (f *Interface) queueRelayPacket(packet []byte, remote netip.AddrPort, rxc *rxContext) { + if rxc.relayBatch == nil { + rxc.relayBatch = batch.NewSendBatch(f.writers[rxc.q], batch.SendBatchCap, 0) + } else if rxc.relayRemote != remote { + f.flushRelayBatch(rxc) + } + rxc.relayRemote = remote + rxc.relayBatch.Commit(packet, remote) + if rxc.relayBatch.Len() >= batch.SendBatchCap { + f.flushRelayBatch(rxc) + } +} + +func (f *Interface) flushRelayBatch(rxc *rxContext) { + if rxc.relayBatch != nil { + f.flushSendBatch(rxc.relayBatch, rxc.q) + } +} + // closeTunnel closes a tunnel locally, it does not send a closeTunnel packet to the remote func (f *Interface) closeTunnel(hostInfo *HostInfo) { // Hold the hostmap lock across both the delete and the lighthouse eviction, diff --git a/outside_relay_test.go b/outside_relay_test.go new file mode 100644 index 00000000..4ee574aa --- /dev/null +++ b/outside_relay_test.go @@ -0,0 +1,291 @@ +package nebula + +import ( + "bytes" + "net/netip" + "sync" + "testing" + + "github.com/flynn/noise" + "github.com/gaissmai/bart" + "github.com/rcrowley/go-metrics" + "github.com/slackhq/nebula/header" + "github.com/slackhq/nebula/noiseutil" + "github.com/slackhq/nebula/overlay/batch" + "github.com/slackhq/nebula/test" + "github.com/slackhq/nebula/udp" + "github.com/stretchr/testify/require" +) + +type relayBatchWrite struct { + packets [][]byte + remotes []netip.AddrPort + batched bool +} + +type relayBatchConn struct { + udp.NoopConn + read func(udp.EncReader, func()) + writes []relayBatchWrite +} + +func (c *relayBatchConn) ListenOut(r udp.EncReader, flush func()) error { + c.read(r, flush) + return nil +} + +func (c *relayBatchConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) { + w := relayBatchWrite{remotes: append([]netip.AddrPort(nil), addrs...), batched: true} + for _, b := range bufs { + w.packets = append(w.packets, bytes.Clone(b)) + } + c.writes = append(c.writes, w) + return len(bufs), nil +} + +func (c *relayBatchConn) WriteTo(b []byte, addr netip.AddrPort) error { + c.writes = append(c.writes, relayBatchWrite{packets: [][]byte{bytes.Clone(b)}, remotes: []netip.AddrPort{addr}}) + return nil +} + +func (c *relayBatchConn) packets() [][]byte { + var packets [][]byte + for _, w := range c.writes { + packets = append(packets, w.packets...) + } + return packets +} + +type relayBatchPath struct { + f *Interface + conn *relayBatchConn + source, target *HostInfo + incoming, outgoing *Relay + sourceKey noiseutil.CipherState + cipher noise.CipherFunc + counter uint64 +} + +// Separate cipher instances let the reference SendVia path use the same keys and +// counters without rewinding an AEAD instance (including FIPS nonce checks). +func relayBatchKey(cipher noise.CipherFunc, id byte) noiseutil.CipherState { + suite := noise.NewCipherSuite(noise.DH25519, cipher, noise.HashSHA256) + return noiseutil.NewCipherState(noise.UnsafeNewCipherState(suite, [32]byte{id}, 0), cipher) +} + +func newRelayBatchPath(t *testing.T, cipher noise.CipherFunc) *relayBatchPath { + t.Helper() + l := test.NewLogger() + conn := &relayBatchConn{} + f := &Interface{ + ctx: t.Context(), l: l, hostMap: newHostMap(l), outside: conn, + writers: []udp.Conn{conn}, lightHouse: &LightHouse{l: l}, + myVpnNetworksTable: &bart.Lite{}, metricTxDropped: metrics.NewCounter(), + batchers: []*batch.MultiCoalescer{batch.NewMultiCoalescer(&capturingTun{}, l)}, + } + f.connectionManager = &connectionManager{ + intf: f, relayUsed: map[uint32]struct{}{}, relayUsedLock: &sync.RWMutex{}, + } + newHost := func(addr string, remote string, id byte) *HostInfo { + h := &HostInfo{ + vpnAddrs: []netip.Addr{netip.MustParseAddr(addr)}, + ConnectionState: &ConnectionState{ + eKey: relayBatchKey(cipher, id), dKey: relayBatchKey(cipher, id), window: NewBits(ReplayWindow), + }, + relayState: RelayState{relayForByAddr: map[netip.Addr]*Relay{}, relayForByIdx: map[uint32]*Relay{}}, + } + u := netip.MustParseAddrPort(remote) + h.remote.Store(&u) + f.hostMap.Hosts[h.vpnAddrs[0]] = h + return h + } + source := newHost("192.0.2.1", "198.51.100.1:4242", 1) + target := newHost("192.0.2.2", "198.51.100.2:4242", 2) + incoming := &Relay{Type: ForwardingType, State: Established, LocalIndex: 301, RemoteIndex: 302, PeerAddr: target.vpnAddrs[0]} + outgoing := &Relay{Type: ForwardingType, State: Established, LocalIndex: 303, RemoteIndex: 304, PeerAddr: source.vpnAddrs[0]} + source.relayState.InsertRelay(incoming.PeerAddr, incoming.LocalIndex, incoming) + target.relayState.InsertRelay(outgoing.PeerAddr, outgoing.LocalIndex, outgoing) + f.hostMap.Relays[incoming.LocalIndex] = source + f.hostMap.Indexes[100] = source + return &relayBatchPath{f: f, conn: conn, source: source, target: target, incoming: incoming, outgoing: outgoing, sourceKey: relayBatchKey(cipher, 1), cipher: cipher} +} + +func (p *relayBatchPath) wrap(t *testing.T, payload []byte) []byte { + t.Helper() + p.counter++ + b := header.Encode(make([]byte, header.Len), header.Version, header.Message, header.MessageRelay, p.incoming.LocalIndex, p.counter) + b = append(b, payload...) + b, err := p.sourceKey.EncryptDanger(b, b, nil, p.counter, make([]byte, 12)) + require.NoError(t, err) + return b +} + +func (p *relayBatchPath) reference(payloads [][]byte) [][]byte { + conn := &relayBatchConn{} + f := &Interface{writers: []udp.Conn{conn}, connectionManager: p.f.connectionManager, l: p.f.l} + h := &HostInfo{ConnectionState: &ConnectionState{eKey: relayBatchKey(p.cipher, 2)}} + u := p.target.GetRemote() + h.remote.Store(&u) + for _, payload := range payloads { + f.SendVia(h, p.outgoing, payload, make([]byte, 12), make([]byte, 0, len(payload)+header.MaxOverhead), false, 0) + } + return conn.packets() +} + +func TestRelayReceiveBatch(t *testing.T) { + for _, cipher := range []noise.CipherFunc{noiseutil.CipherAESGCM, noise.CipherChaChaPoly} { + t.Run(cipher.CipherName(), func(t *testing.T) { + p := newRelayBatchPath(t, cipher) + const count = 8 + var payloads, wire [][]byte + for i := range count * 2 { + payload := bytes.Repeat([]byte{byte(i)}, 64) + payloads = append(payloads, payload) + wire = append(wire, p.wrap(t, payload)) + } + want := p.reference(payloads) + // Model a listener reusing the same slots for the next receive batch. + buffers := make([][]byte, count) + for i := range buffers { + buffers[i] = make([]byte, len(wire[i])) + } + p.conn.read = func(r udp.EncReader, flush func()) { + for b := range 2 { + for i := range count { + copy(buffers[i], wire[b*count+i]) + r(p.source.GetRemote(), buffers[i]) + require.Len(t, p.conn.writes, b, "forwarding waits for the batch boundary") + } + flush() + require.Len(t, p.conn.writes, b+1, "flush must send before receive slots are reused") + for _, buffer := range buffers { + clear(buffer) + } + } + } + p.f.listenOut(0) + require.Equal(t, want, p.conn.packets()) + for _, w := range p.conn.writes { + require.True(t, w.batched) + require.Len(t, w.packets, count) + for _, remote := range w.remotes { + require.Equal(t, p.target.GetRemote(), remote) + } + } + }) + } +} + +func TestRelayBatchDestinationChange(t *testing.T) { + p := newRelayBatchPath(t, noise.CipherChaChaPoly) + first, second := p.target.GetRemote(), netip.MustParseAddrPort("198.51.100.3:4242") + remotes := []netip.AddrPort{first, first, second, second, first} + var payloads, wire [][]byte + for i := range remotes { + payload := []byte{byte(i)} + payloads = append(payloads, payload) + wire = append(wire, p.wrap(t, payload)) + } + want := p.reference(payloads) + p.conn.read = func(r udp.EncReader, flush func()) { + for i, remote := range remotes { + p.target.remote.Store(&remote) + r(p.source.GetRemote(), wire[i]) + } + flush() + } + p.f.listenOut(0) + require.Equal(t, want, p.conn.packets()) + require.Len(t, p.conn.writes, 3) + for i, wantRemotes := range [][]netip.AddrPort{{first, first}, {second, second}, {first}} { + require.Equal(t, wantRemotes, p.conn.writes[i].remotes) + } +} + +func TestRelayBatchCapacity(t *testing.T) { + p := newRelayBatchPath(t, noise.CipherChaChaPoly) + p.conn.read = func(r udp.EncReader, flush func()) { + for i := range batch.SendBatchCap + 1 { + r(p.source.GetRemote(), p.wrap(t, []byte{byte(i)})) + } + require.Len(t, p.conn.writes, 1) + require.Len(t, p.conn.writes[0].packets, batch.SendBatchCap) + flush() + } + p.f.listenOut(0) + require.Len(t, p.conn.writes, 2) + require.Len(t, p.conn.writes[1].packets, 1) +} + +func TestRelayBatchBeforeResponse(t *testing.T) { + for _, unknownRelay := range []bool{false, true} { + name := "test request" + if unknownRelay { + name = "unknown relay" + } + t.Run(name, func(t *testing.T) { + p := newRelayBatchPath(t, noise.CipherChaChaPoly) + first, last := p.wrap(t, []byte("first")), p.wrap(t, []byte("last")) + p.counter++ + request := header.Encode(make([]byte, header.Len), header.Version, header.Test, header.TestRequest, 100, p.counter) + request, err := p.sourceKey.EncryptDanger(request, request, []byte("ping"), p.counter, make([]byte, 12)) + require.NoError(t, err) + if unknownRelay { + request = header.Encode(make([]byte, header.Len), header.Version, header.Message, header.MessageRelay, 999, p.counter) + } + p.conn.read = func(r udp.EncReader, flush func()) { + r(p.source.GetRemote(), first) + require.Empty(t, p.conn.writes) + r(p.source.GetRemote(), request) + require.Len(t, p.conn.writes, 2) + r(p.source.GetRemote(), last) + flush() + } + p.f.listenOut(0) + require.Len(t, p.conn.writes, 3) + require.True(t, p.conn.writes[0].batched) + require.False(t, p.conn.writes[1].batched) + require.True(t, p.conn.writes[2].batched) + require.Equal(t, p.source.GetRemote(), p.conn.writes[1].remotes[0]) + h := &header.H{} + require.NoError(t, h.Parse(p.conn.writes[1].packets[0])) + if unknownRelay { + require.Equal(t, header.RecvError, h.Type) + require.EqualValues(t, 999, h.RemoteIndex) + } else { + require.Equal(t, header.Test, h.Type) + require.Equal(t, header.TestReply, h.Subtype) + reply := p.conn.writes[1].packets[0] + out, err := p.sourceKey.DecryptDanger(nil, reply[:header.Len], reply[header.Len:], h.MessageCounter, make([]byte, 12)) + require.NoError(t, err) + require.Equal(t, []byte("ping"), out) + } + }) + } +} + +func TestRelayBatchRemovedTarget(t *testing.T) { + for _, remove := range []string{"host", "relay", "disestablished"} { + t.Run(remove, func(t *testing.T) { + p := newRelayBatchPath(t, noise.CipherChaChaPoly) + first, removed := p.wrap(t, []byte("first")), p.wrap(t, []byte("removed")) + want := p.reference([][]byte{[]byte("first")}) + p.conn.read = func(r udp.EncReader, flush func()) { + r(p.source.GetRemote(), first) + switch remove { + case "host": + delete(p.f.hostMap.Hosts, p.target.vpnAddrs[0]) + case "relay": + delete(p.target.relayState.relayForByAddr, p.source.vpnAddrs[0]) + case "disestablished": + p.target.relayState.UpdateRelayForByIpState(p.source.vpnAddrs[0], Disestablished) + } + r(p.source.GetRemote(), removed) + flush() + } + p.f.listenOut(0) + require.Len(t, p.conn.writes, 1) + require.Equal(t, want, p.conn.packets(), "only the packet accepted before removal is sent") + }) + } +} diff --git a/overlay/batch/tx_batch.go b/overlay/batch/tx_batch.go index 4f6f7da2..0cb81ccf 100644 --- a/overlay/batch/tx_batch.go +++ b/overlay/batch/tx_batch.go @@ -10,8 +10,8 @@ type batchWriter interface { } // SendBatch accumulates encrypted UDP packets and flushes them via WriteBatch. -// One SendBatch is owned by each listenIn goroutine; no locking is needed. -// Slots are backed by an Arena (see its docs) +// Each reader goroutine owns its own SendBatch; no locking is needed. +// Slots can borrow receive buffers or be backed by an Arena (see its docs). type SendBatch struct { out batchWriter bufs [][]byte diff --git a/udp/conn.go b/udp/conn.go index 62806569..4375b9fd 100644 --- a/udp/conn.go +++ b/udp/conn.go @@ -41,7 +41,7 @@ type Conn interface { // // Returns the number of packets successfully written. A destination the kernel rejects costs only // its own packet, so a short count means some peers were undeliverable, not that the batch failed. - // Not safe for concurrent use on the same Conn. + // Safe for concurrent use on the same Conn. WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) // ReloadConfig applies the listen settings that can change, r is nil at startup ReloadConfig(c *config.C, r *config.Reload) diff --git a/udp/udp_linux_writebatch.go b/udp/udp_linux_writebatch.go index bd846629..6a051f2b 100644 --- a/udp/udp_linux_writebatch.go +++ b/udp/udp_linux_writebatch.go @@ -12,6 +12,7 @@ import ( "net/netip" "strconv" "strings" + "sync" "unsafe" "golang.org/x/sys/unix" @@ -20,7 +21,7 @@ import ( // batchWriter owns the sendmmsg(2)/UDP-GSO transmit path for a StdConn: the // scratch WriteBatch packs mmsghdr entries into, plus the GSO capability // state probed at socket creation. Each queue has its own StdConn and -// batchWriter, so no locking is needed. +// batchWriter; mu serializes batches from the TUN reader and relay forwarding. // // Terminology, smallest to largest: // @@ -39,6 +40,7 @@ import ( // chunk the entries packed for one sendmmsg call, at most MaxWriteBatch. // batch the caller's whole bufs/addrs pair, processed as one or more chunks. type batchWriter struct { + mu sync.Mutex fd int isV4 bool @@ -192,6 +194,9 @@ func parseRelease(r string) (major, minor int) { // Returns the number of packets sent. An error means the call itself failed. // A short count means some destinations were undeliverable. func (w *batchWriter) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + if len(bufs) != len(addrs) { return 0, fmt.Errorf("WriteBatch: len(bufs)=%d != len(addrs)=%d", len(bufs), len(addrs)) } diff --git a/udp/udp_linux_writebatch_concurrent_test.go b/udp/udp_linux_writebatch_concurrent_test.go new file mode 100644 index 00000000..18ab2fd9 --- /dev/null +++ b/udp/udp_linux_writebatch_concurrent_test.go @@ -0,0 +1,61 @@ +//go:build linux && !android && !e2e_testing + +package udp + +import ( + "net/netip" + "runtime" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +// The TUN reader and relay forwarding share a socket. Exercise their concurrent +// batches against the actual GSO planner and scratch, without relying on a NIC. +func TestWriteBatchConcurrent(t *testing.T) { + const batches = 64 + w := newRewindTestWriter() + sent := make(chan [][]byte, 2*batches) + w.sendFn = func(start, n int) (int, error) { + runtime.Gosched() + sent <- capturePrepared(w, start, n) + return n, nil + } + conn := &StdConn{bw: w} + remote := netip.MustParseAddrPort("192.0.2.1:4242") + var wg sync.WaitGroup + start := make(chan struct{}) + for producer := range 2 { + wg.Go(func() { + <-start + for i := range batches { + var bufs [][]byte + var addrs []netip.AddrPort + for j := range 4 { + bufs = append(bufs, []byte{byte(producer), byte(i), byte(j), 0}) + addrs = append(addrs, remote) + } + n, err := conn.WriteBatch(bufs, addrs) + if err != nil || n != len(bufs) { + t.Errorf("WriteBatch = %d, %v; want %d, nil", n, err, len(bufs)) + return + } + } + }) + } + close(start) + wg.Wait() + close(sent) + var counts [2]int + for packets := range sent { + require.Len(t, packets, 4) + producer := packets[0][0] + require.Less(t, int(producer), len(counts)) + for j, packet := range packets { + require.Equal(t, []byte{producer, byte(counts[producer]), byte(j), 0}, packet) + } + counts[producer]++ + } + require.Equal(t, [2]int{batches, batches}, counts) +}