udp: Add opt-in ordered staged receive processing

connection_state.go

@@ -104,7 +104,7 @@ func (cs *ConnectionState) Decrypt(l *slog.Logger, messageCounter uint64, packet
 		return nil, ErrAlreadySeen
 	}
 
-	out, err := cs.dKey.DecryptDanger(packet[header.Len:header.Len], packet[:header.Len], packet[header.Len:], messageCounter, nb)
+	out, err := cs.decryptPacket(messageCounter, packet, nb)
 	if err != nil {
 		return nil, err
 	}
@@ -118,6 +118,11 @@ func (cs *ConnectionState) Decrypt(l *slog.Logger, messageCounter uint64, packet
 	return out, nil
 }
 
+// decryptPacket authenticates and decrypts packet in place, leaving the replay window to the caller.
+func (cs *ConnectionState) decryptPacket(messageCounter uint64, packet, nb []byte) ([]byte, error) {
+	return cs.dKey.DecryptDanger(packet[header.Len:header.Len], packet[:header.Len], packet[header.Len:], messageCounter, nb)
+}
+
 func (cs *ConnectionState) VerifyRelay(l *slog.Logger, messageCounter uint64, packet []byte, nb []byte) error {
 	cs.decryptLock.Lock()
 	result := cs.window.Check(l, messageCounter)

connection_state_test.go

@@ -23,11 +23,20 @@ import (
 // runTestHandshake runs a complete IX handshake between two freshly-built
 // peers and returns the initiator and responder Results. Used to produce
 // real cipher states for tests that need to exercise post-handshake glue.
-func runTestHandshake(t *testing.T) (initR, respR *handshake.Result) {
+func runTestHandshake(t testing.TB) (initR, respR *handshake.Result) {
+	t.Helper()
+	return runTestHandshakeWith(t, noise.CipherChaChaPoly)
+}
+
+// runTestHandshakeWith is runTestHandshake with the data plane cipher picked by the caller.
+func runTestHandshakeWith(t testing.TB, cipher noise.CipherFunc) (initR, respR *handshake.Result) {
 	t.Helper()
 
+	// Valid around now rather than the cert_test default, which is fixed at process start and lapses a minute
+	// later, inside a long -count run
+	now := time.Now()
 	ca, _, caKey, _ := ct.NewTestCaCert(
-		cert.Version2, cert.Curve_CURVE25519, time.Time{}, time.Time{}, nil, nil, nil,
+		cert.Version2, cert.Curve_CURVE25519, now.Add(-time.Minute), now.Add(time.Hour), nil, nil, nil,
 	)
 	caPool := ct.NewTestCAPool(ca)
 
@@ -40,7 +49,7 @@ func runTestHandshake(t *testing.T) (initR, respR *handshake.Result) {
 		require.NoError(t, err)
 		hsBytes, err := c.MarshalForHandshakes()
 		require.NoError(t, err)
-		ncs := noise.NewCipherSuite(noise.DH25519, noise.CipherChaChaPoly, noise.HashSHA256)
+		ncs := noise.NewCipherSuite(noise.DH25519, cipher, noise.HashSHA256)
 		cred := handshake.NewCredential(c, hsBytes, priv, ncs)
 		return func(v cert.Version) *handshake.Credential {
 			if v == cert.Version2 {

e2e/tunnels_test.go

@@ -5,9 +5,11 @@ package e2e
 import (
 	"fmt"
 	"net/netip"
+	"runtime"
 	"testing"
 	"time"
 
+	"github.com/rcrowley/go-metrics"
 	"github.com/slackhq/nebula/cert"
 	"github.com/slackhq/nebula/cert_test"
 	"github.com/slackhq/nebula/e2e/router"
@@ -570,3 +572,29 @@ func TestCloseTunnelAuthenticated(t *testing.T) {
 	myControl.Stop()
 	theirControl.Stop()
 }
+
+// TestStagedTunnel completes a handshake and traffic both ways with staged receive. Every received packet,
+// handshakes included, uses the receive pipeline on Linux. Other platforms must do the same over the inline fallback.
+func TestStagedTunnel(t *testing.T) {
+	t.Parallel()
+	received := metrics.GetOrRegisterCounter("flow.rx.inline", nil)
+	before := received.Count()
+	staged := m{"tun": m{"staged": m{"workers": 2}}}
+	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version1, cert.Curve_CURVE25519, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
+	myControl, myVpnIpNet, myUdpAddr, _ := newSimpleServer(cert.Version1, ca, caKey, "me", "10.128.0.1/24", staged)
+	theirControl, theirVpnIpNet, theirUdpAddr, _ := newSimpleServer(cert.Version1, ca, caKey, "them", "10.128.0.2/24", staged)
+	myControl.InjectLightHouseAddr(theirVpnIpNet[0].Addr(), theirUdpAddr)
+	theirControl.InjectLightHouseAddr(myVpnIpNet[0].Addr(), myUdpAddr)
+	myControl.Start()
+	theirControl.Start()
+	r := router.NewR(t, myControl, theirControl)
+	defer r.RenderFlow()
+	assertTunnel(t, myVpnIpNet[0].Addr(), theirVpnIpNet[0].Addr(), myControl, theirControl, r)
+	myControl.Stop()
+	theirControl.Stop()
+	if runtime.GOOS != "linux" {
+		assert.Zero(t, received.Count()-before)
+		return
+	}
+	assert.GreaterOrEqual(t, received.Count()-before, int64(2))
+}

examples/config.yml

@@ -271,6 +271,14 @@ tun:
   # the use_offloads setting controls if Nebula will attempt to enable GSO and GRO for the tun device. Linux only, not reloadable.
   #use_offloads: false
 
+  # Linux only. staged.workers enables staged receive: each routine hands its received batches to this many decrypt
+  # workers, then an ordered writer commits them in arrival order through the same replay, roaming and firewall
+  # handling as inline receive. 0, the default, keeps receive inline. Batches below 256 packets on an idle listener
+  # stay on the listener. Needs no support from peers. Costs about 24 MiB of buffers per routine at listen.batch 64
+  # with GRO. Not reloadable.
+  #staged:
+  #  workers: 0
+
   # Linux only. pin_threads pins each tun reader/encrypt OS thread to a single CPU. This keeps every goroutine's
   # batched sends flowing through one XPS-selected NIC TX ring, so packets within a flow stay ordered on the wire
   # instead of being sprayed across multiple TX rings and reordered. Not reloadable. Coerced to false if routines <= 1.

flow.go

@@ -0,0 +1,80 @@
+package nebula
+
+import (
+	"fmt"
+	"log/slog"
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/config"
+)
+
+// Staged receive (tun.staged.workers) moves decryption off each udp listener onto worker goroutines, behind which an
+// ordered writer commits every batch in arrival order. It reuses the inline packet handlers: staging changes ownership
+// and scheduling, not policy. Each routine has its own workers and writer, so order is per routine. Full queues block
+// the listener. A read error drains the batches in flight, while shutdown waits for the workers and then drops what
+// is left. Each routine keeps six batches; at listen.batch=64 with GRO their buffers take about 24 MiB per routine.
+//
+// Metrics are registered only for active pipelines. flow.rx.batches counts batches handed to the writer and
+// flow.rx.inline those an idle listener handled itself; their rates show whether staging is used.
+// flow.stage.<stage>.q<queue>.busy_ns locates the limiting stage: rx_read, rx_decrypt or rx_write. Take rates; 1e9
+// ns/s fills a serial stage. These are elapsed times that exclude read syscalls and downstream waits; descheduling
+// inflates them and worker times add up across the pool. Small batches sent to the writer include decryption in
+// writer time.
+
+// stagedConfig is tun.staged. workers is per routine; zero keeps receive inline.
+type stagedConfig struct {
+	workers int
+	// inlineBelow is the packet count below which an idle listener skips the worker handoff
+	inlineBelow int
+}
+
+// parseStagedConfig enables staging on Linux only; conns that cannot hand off batches also stay inline.
+func parseStagedConfig(c *config.C, l *slog.Logger) stagedConfig {
+	sc := stagedConfig{workers: c.GetInt("tun.staged.workers", 0)}
+	if sc.workers <= 0 {
+		return stagedConfig{}
+	}
+	if runtime.GOOS != "linux" {
+		l.Warn("tun.staged is only supported on Linux; packet processing stays inline")
+		return stagedConfig{}
+	}
+	// 256 is the measured crossover
+	sc.inlineBelow = 256
+	l.Info("Staged receive processing enabled", "workers", sc.workers)
+	return sc
+}
+
+// flowJob is one batch of crypto work. It owns its packet buffers until run returns; jobs and queues are reused.
+type flowJob interface{ run(nb []byte) }
+
+// flowPool runs one listener's crypto jobs. The listener is the only producer and closes the pool once its ordered
+// stage has drained.
+type flowPool struct {
+	jobs chan flowJob
+	wg   sync.WaitGroup
+}
+
+func newFlowPool(workers, queueLen int, busy metrics.Counter) *flowPool {
+	p := &flowPool{jobs: make(chan flowJob, queueLen)}
+	for range workers {
+		p.wg.Go(func() {
+			nb := make([]byte, 12)
+			for job := range p.jobs {
+				t := time.Now()
+				job.run(nb)
+				busy.Inc(int64(time.Since(t)))
+			}
+		})
+	}
+	return p
+}
+
+func (p *flowPool) close() { close(p.jobs); p.wg.Wait() }
+
+// stageBusy is the busy time counter of one stage of routine q.
+func stageBusy(stage string, q int) metrics.Counter {
+	return metrics.GetOrRegisterCounter(fmt.Sprintf("flow.stage.%s.q%d.busy_ns", stage, q), nil)
+}

flow_rx.go

@@ -0,0 +1,398 @@
+package nebula
+
+import (
+	"encoding/binary"
+	"net/netip"
+	"time"
+
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/udp"
+)
+
+// handoffConn is a udp conn that can hand each received batch over. The callback returns replacement buffers with
+// the same count and at least the same lengths.
+type handoffConn interface {
+	ListenOutHandoff(udp.EncReader, func([][]byte) [][]byte) error
+}
+
+const stagedRxSlots = 6
+const rxChunkMin = 16
+
+// rxItem is one received packet. A data packet for an established tunnel has ci set; out and err are written by
+// whoever decrypts it.
+type rxItem struct {
+	packet   []byte
+	hostinfo *HostInfo
+	ci       *ConnectionState
+	counter  uint64
+	via      netip.AddrPort
+	out      []byte
+	err      error
+}
+
+// rxChunk is a contiguous run of a batch's items for one worker to decrypt.
+type rxChunk struct {
+	items []rxItem
+	done  chan struct{}
+}
+
+func (c *rxChunk) run(nb []byte) {
+	decryptRxItems(c.items, nb)
+	c.done <- struct{}{}
+}
+
+// decryptRxItems decrypts the data packets among items in place. The replay window is left to the ordered stage, as
+// are packets without a tunnel and the known replays precheck marked.
+func decryptRxItems(items []rxItem, nb []byte) {
+	for i := range items {
+		it := &items[i]
+		if it.ci == nil || it.err != nil {
+			continue
+		}
+		it.out, it.err = it.ci.decryptPacket(it.counter, it.packet, nb)
+	}
+}
+
+// rxBatch is one recvmmsg batch moving through the staged receive pipeline.
+type rxBatch struct {
+	// items is every packet of the batch in arrival order. Data packets are decrypted by a worker; anything else is
+	// handled whole by the writer, in its place.
+	items []rxItem
+	data  int
+	// bufs is the set of recvmmsg buffers the items point into, owned by the batch from handoff until it is reused.
+	bufs [][]byte
+	// chunks are the decrypt jobs, nChunks of them in use. With none the writer decrypts the batch itself.
+	chunks  []rxChunk
+	nChunks int
+}
+
+func (b *rxBatch) reset() {
+	clear(b.items)
+	b.items = b.items[:0]
+	b.data = 0
+	for i := range b.nChunks {
+		b.chunks[i].items = nil
+	}
+	b.nChunks = 0
+}
+
+// stagedRx is the staged receive pipeline of one udp listener. deliver and handoff run on the listener goroutine.
+type stagedRx struct {
+	f    *Interface
+	q    int
+	pool *flowPool
+	// h and hostmapCache are the listener's own. wrxc is the commit context: the writer's, or the listener's while
+	// nothing is in flight.
+	h            *header.H
+	hostmapCache map[uint32]*HostInfo
+	wrxc         *rxContext
+
+	started    time.Time
+	cur        *rxBatch
+	free       []*rxBatch
+	ordered    chan *rxBatch
+	returned   chan *rxBatch
+	inFlight   int
+	writerDone chan struct{}
+
+	readBusy, decryptBusy, writeBusy metrics.Counter
+	metricBatches                    metrics.Counter
+	metricInline                     metrics.Counter
+}
+
+func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
+	p := &stagedRx{
+		f:             f,
+		q:             q,
+		h:             &header.H{},
+		hostmapCache:  map[uint32]*HostInfo{},
+		wrxc:          wrxc,
+		ordered:       make(chan *rxBatch, stagedRxSlots),
+		returned:      make(chan *rxBatch, stagedRxSlots),
+		writerDone:    make(chan struct{}),
+		readBusy:      stageBusy("rx_read", q),
+		decryptBusy:   stageBusy("rx_decrypt", q),
+		writeBusy:     stageBusy("rx_write", q),
+		metricBatches: metrics.GetOrRegisterCounter("flow.rx.batches", nil),
+		metricInline:  metrics.GetOrRegisterCounter("flow.rx.inline", nil),
+	}
+	workers := f.staged.workers
+	p.pool = newFlowPool(workers, stagedRxSlots*2*workers, p.decryptBusy)
+	for range stagedRxSlots {
+		// Two chunks per worker, so the writer can start committing while the workers are still on the rest
+		b := &rxBatch{items: make([]rxItem, 0, 256), chunks: make([]rxChunk, 2*workers)}
+		for i := range b.chunks {
+			b.chunks[i].done = make(chan struct{}, 1)
+		}
+		p.free = append(p.free, b)
+	}
+	p.cur = p.free[len(p.free)-1]
+	p.free = p.free[:len(p.free)-1]
+	go p.writer()
+	return p
+}
+
+// listen runs the conn's receive loop through the pipeline until it fails, then drains the pipeline. The conn hands
+// its buffers over, so a batch moves between stages without a copy.
+func (p *stagedRx) listen(li handoffConn) error {
+	defer p.close()
+	return li.ListenOutHandoff(p.deliver, p.handoff)
+}
+
+// deliver is the listener's per-packet work: classify the packet and add it to the current batch. Nothing here
+// decrypts or touches the tun.
+func (p *stagedRx) deliver(via netip.AddrPort, packet []byte) {
+	if p.started.IsZero() {
+		p.started = time.Now()
+	}
+	b := p.cur
+	it := rxItem{packet: packet, via: via}
+	p.classify(&it)
+	if it.ci != nil {
+		b.data++
+	}
+	b.items = append(b.items, it)
+}
+
+// classify fills in the tunnel for a plain data packet of an established, non-relayed tunnel, leaving anything else
+// for the writer to handle whole. The replay pre-check waits for the end of the batch, see precheck.
+func (p *stagedRx) classify(it *rxItem) {
+	f := p.f
+	h := p.h
+	packet := it.packet
+	if h.Parse(packet) != nil || h.Version != header.Version || h.Type != header.Message || h.Subtype != header.MessageNone {
+		return
+	}
+	if f.myVpnNetworksTable.Contains(it.via.Addr()) {
+		return
+	}
+	hostinfo := f.hostMap.QueryIndexCached(h.RemoteIndex, p.hostmapCache)
+	if hostinfo == nil || hostinfo.ConnectionState == nil {
+		return
+	}
+	ci := hostinfo.ConnectionState
+	if len(packet) < header.Len+ci.dKey.Overhead() {
+		return
+	}
+	it.hostinfo, it.ci, it.counter = hostinfo, ci, h.MessageCounter
+}
+
+// precheck marks obvious replays before anything decrypts them, like ConnectionState.Decrypt does, taking each
+// tunnel's lock once per run of its packets rather than once per packet: the writer takes the same lock for its
+// window updates, from another core. The writer checks again when it commits. The window only moves forward in
+// between, so a stale check here can only let a replay through to that second check. Marked packets keep their
+// place: an earlier RecvError may delete their tunnel, and inline receive then replies with a RecvError rather than
+// dropping a replay.
+func (p *stagedRx) precheck(b *rxBatch) {
+	f := p.f
+	items := b.items
+	for i := 0; i < len(items); {
+		ci := items[i].ci
+		if ci == nil {
+			i++
+			continue
+		}
+		ci.decryptLock.Lock()
+		for ; i < len(items) && items[i].ci == ci; i++ {
+			if ci.window.Check(f.l, items[i].counter) {
+				continue
+			}
+			b.data--
+			items[i].err = ErrAlreadySeen
+		}
+		ci.decryptLock.Unlock()
+	}
+}
+
+// handoff ends a batch: the batch takes bufs, and the conn reads the next batch into the buffers of the next free
+// batch.
+func (p *stagedRx) handoff(bufs [][]byte) [][]byte {
+	p.cur.bufs = bufs
+	p.endBatch()
+	if !p.started.IsZero() {
+		p.readBusy.Inc(int64(time.Since(p.started)))
+		p.started = time.Time{}
+	}
+	b := p.cur
+	if len(b.bufs) != len(bufs) || len(b.bufs[0]) < len(bufs[0]) {
+		b.bufs = make([][]byte, len(bufs))
+		for i := range b.bufs {
+			b.bufs[i] = make([]byte, len(bufs[0]))
+		}
+	}
+	return b.bufs
+}
+
+// endBatch hands the batch to the workers and the ordered writer, then takes a free batch.
+func (p *stagedRx) endBatch() {
+	clear(p.hostmapCache)
+	p.reap()
+	b := p.cur
+	p.precheck(b)
+	n := len(b.items)
+	if n == 0 {
+		b.reset()
+		return
+	}
+
+	if p.inFlight == 0 && n < p.f.staged.inlineBelow {
+		// The writer is idle with nothing in flight, so the listener may commit with its context and flush the tun
+		p.metricInline.Inc(1)
+		p.process(b, p.wrxc)
+		b.reset()
+		return
+	}
+
+	parts := min(len(b.chunks), b.data/rxChunkMin)
+	p.metricBatches.Inc(1)
+	if parts > 0 {
+		size := (n + parts - 1) / parts
+		for off := 0; off < n; off += size {
+			c := &b.chunks[b.nChunks]
+			c.items = b.items[off:min(off+size, n)]
+			b.nChunks++
+		}
+	}
+	// The jobs go out before the batch does: once the writer has it, it may finish and reset it at any moment
+	for i := range b.nChunks {
+		p.pool.jobs <- &b.chunks[i]
+	}
+	p.inFlight++
+	p.ordered <- b
+
+	if len(p.free) == 0 {
+		// Every batch is in flight, wait for the writer: the socket buffer absorbs the backlog meanwhile. The wait is
+		// downstream backpressure, not listener busy time.
+		t := time.Now()
+		p.release(<-p.returned)
+		p.started = p.started.Add(time.Since(t))
+	}
+	p.cur = p.free[len(p.free)-1]
+	p.free = p.free[:len(p.free)-1]
+}
+
+// reap takes back every batch the writer has finished with, without waiting.
+func (p *stagedRx) reap() {
+	for {
+		select {
+		case b := <-p.returned:
+			p.release(b)
+		default:
+			return
+		}
+	}
+}
+
+func (p *stagedRx) release(b *rxBatch) {
+	p.inFlight--
+	p.free = append(p.free, b)
+}
+
+// process is the ordered stage for one batch: decrypt whatever no worker is decrypting, commit every item in arrival
+// order, and flush the tun batcher. It reports how long it waited on workers.
+func (p *stagedRx) process(b *rxBatch, rxc *rxContext) time.Duration {
+	var waited time.Duration
+	if b.nChunks == 0 {
+		decryptRxItems(b.items, rxc.nb)
+		p.commit(b.items, rxc)
+	} else {
+		// Commit chunk by chunk, so the commit overlaps the workers still on later chunks
+		for i := range b.nChunks {
+			c := &b.chunks[i]
+			select {
+			case <-c.done:
+			default:
+				t := time.Now()
+				<-c.done
+				waited += time.Since(t)
+			}
+			p.commit(c.items, rxc)
+		}
+	}
+	p.f.flushOutside(rxc)
+	return waited
+}
+
+// commit runs the ordered stage over items in arrival order. A non-data packet goes through readOutsidePackets in
+// its place, a run of data packets on one tunnel through commitRun.
+func (p *stagedRx) commit(items []rxItem, rxc *rxContext) {
+	for i := 0; i < len(items); {
+		ci := items[i].ci
+		if ci == nil {
+			p.f.readOutsidePackets(ViaSender{UdpAddr: items[i].via}, items[i].packet, rxc)
+			i++
+			continue
+		}
+		j := i + 1
+		for j < len(items) && items[j].ci == ci {
+			j++
+		}
+		p.commitRun(items[i:j], rxc)
+		i = j
+	}
+}
+
+// commitRun commits a run of decrypted data packets on one tunnel. The window updates go first, under one lock, then
+// the rest; handling a data packet never touches the window, so that is the same as one packet at a time.
+func (p *stagedRx) commitRun(run []rxItem, rxc *rxContext) {
+	f := p.f
+	ci := run[0].ci
+	// Classification is speculative: earlier control packets may have removed this tunnel. Resolve it through the
+	// writer's batch-local cache, exactly where inline processing would, before touching replay state.
+	index := binary.BigEndian.Uint32(run[0].packet[4:8])
+	hostinfo := f.hostMap.QueryIndexCached(index, rxc.hostmapCache)
+	if hostinfo != run[0].hostinfo || hostinfo.ConnectionState != ci {
+		if hostinfo == nil {
+			for i := range run {
+				f.maybeSendRecvError(run[i].via, index)
+			}
+		}
+		return
+	}
+
+	ci.decryptLock.Lock()
+	for i := range run {
+		// A packet that failed to decrypt does not touch the window, like ConnectionState.Decrypt
+		if run[i].err == nil && !ci.window.Update(f.l, run[i].counter) {
+			run[i].err = ErrAlreadySeen
+		}
+	}
+	ci.decryptLock.Unlock()
+	for i := range run {
+		it := &run[i]
+		if f.finishOutsideDecrypt(it.hostinfo, ViaSender{UdpAddr: it.via}, it.counter, it.err) {
+			f.handleOutsideMessagePacket(it.hostinfo, it.counter, it.out, rxc)
+		}
+	}
+}
+
+// writer is the ordered receive stage: window, firewall, coalescer and tun write, one batch at a time in arrival
+// order.
+func (p *stagedRx) writer() {
+	defer close(p.writerDone)
+	for b := range p.ordered {
+		if p.f.closed.Load() {
+			// The tun is going away, a write would only log errors. Still wait out the workers on the batch.
+			for i := range b.nChunks {
+				<-b.chunks[i].done
+			}
+		} else {
+			t := time.Now()
+			waited := p.process(b, p.wrxc)
+			p.writeBusy.Inc(int64(time.Since(t) - waited))
+		}
+		b.reset()
+		p.returned <- b
+	}
+}
+
+// close commits anything still in flight, in order, and stops the writer and the workers.
+func (p *stagedRx) close() {
+	if len(p.cur.items) > 0 {
+		p.endBatch()
+	}
+	close(p.ordered)
+	<-p.writerDone
+	p.pool.close()
+}

flow_rx_test.go

@@ -0,0 +1,237 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"fmt"
+	"net/netip"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/flynn/noise"
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/noiseutil"
+	"github.com/slackhq/nebula/udp"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+)
+
+// TestFlowStagedRXOrderAndReplay feeds every packet behind a forgery of it, from another address, and a replay of
+// it, then a test packet that jumps the window and the whole stream again, with shuffled workers and a slow tun,
+// through small and large batches: the tun sees each segment once, in order, the window marks them all, the forgeries
+// neither consume a counter nor roam, and the jump does not get the earlier data packets of its batch rejected.
+func TestFlowStagedRXOrderAndReplay(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	for _, cipher := range []noise.CipherFunc{noiseutil.CipherAESGCM, noise.CipherChaChaPoly} {
+		for _, batchLen := range []int{8, 1000} {
+			t.Run(fmt.Sprintf("%s/batch=%d", cipher.CipherName(), batchLen), func(t *testing.T) {
+				sender, receiver := newFlowTestPair(t, cipher, stagedOn(4))
+				delayCrypto(receiver, func(n uint64) {
+					if n%32 == 0 {
+						time.Sleep(time.Millisecond)
+					}
+				})
+				receiver.tun.delay = 100 * time.Microsecond
+				batches, want := flowTestBatches(30, 4, 1, 10)
+				wire := wireFor(t, sender, batches)
+				var feed [][]byte
+				for i, p := range wire {
+					// Corrupt the tag or the payload
+					bad := append([]byte(nil), p...)
+					bad[[]int{len(p) - 1, header.Len + 3}[i%2]] ^= 1
+					feed = append(feed, bad, p, p)
+				}
+				// The jump lands in a batch with fresh data before it, at batch=1000
+				sender.hostinfo.ConnectionState.messageCounter.Add(2 * ReplayWindow)
+				feed = append(append(feed, sendControl(sender, header.Test, header.TestReply)), wire...)
+				forged := func(i int) netip.AddrPort {
+					if i < 3*len(wire) && i%3 == 0 {
+						return netip.MustParseAddrPort("198.51.100.9:4242")
+					}
+					return flowTestUnder
+				}
+				rx0 := rxBatches.Count()
+				conn := &rxConn{wire: feed, batchLen: batchLen, from: forged}
+				receiver.receive(conn)
+				require.Positive(t, rxBatches.Count()-rx0)
+				require.Equal(t, (len(feed)+batchLen-1)/batchLen, conn.swaps, "every batch is handed off")
+				require.Equal(t, want, receiver.tun.segments(t), "every packet exactly once, in order")
+				require.Equal(t, flowTestUnder, receiver.hostinfo.GetRemote())
+				for _, p := range wire {
+					h := &header.H{}
+					require.NoError(t, h.Parse(p))
+					require.False(t, receiver.hostinfo.ConnectionState.window.Check(receiver.f.l, h.MessageCounter), "counter %d must be marked seen", h.MessageCounter)
+				}
+			})
+		}
+	}
+}
+
+// TestFlowStagedRXClose: a CloseTunnel in the middle of a batch is handled in its place, by the writer. The data after
+// it in the batch is still delivered from the batch's hostmap cache, as inline receive does.
+func TestFlowStagedRXClose(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	batches, want := flowTestBatches(4, 4, 1, 8)
+	wire := wireFor(t, sender, batches)
+	half := len(wire) / 2
+	feed := append(append(append([][]byte{}, wire[:half]...), sendControl(sender, header.CloseTunnel, 0)), wire[half:]...)
+	receiver.receive(&rxConn{wire: feed, batchLen: len(feed)})
+	require.Equal(t, want, receiver.tun.segments(t))
+	require.Nil(t, receiver.f.hostMap.QueryIndex(receiver.hostinfo.localIndexId))
+}
+
+// sendControl returns the wire packet of a control message from sender.
+func sendControl(sender *flowTestPeer, typ header.MessageType, st header.MessageSubType) []byte {
+	sender.f.sendNoMetrics(typ, st, sender.hostinfo.ConnectionState, sender.hostinfo, flowTestUnder, nil, make([]byte, 12), make([]byte, mtu), 0)
+	wire := sender.conn.captured()
+	sender.conn.pkts, sender.conn.addrs = nil, nil
+	return wire[len(wire)-1]
+}
+
+// receivePaths is inline, staged and, with the threshold, adaptive receive.
+func receivePaths() []stagedConfig {
+	return []stagedConfig{{}, stagedOn(2), adaptive(2)}
+}
+
+// TestFlowStagedRXRecvError feeds a RecvError that deletes the tunnel, then a known replay and new data for it, in one
+// batch. Every receive path must drop the data through its own view of the hostmap, without touching the replay window,
+// and answer each packet with a RecvError, as inline receive does, without decrypting the replay again.
+func TestFlowStagedRXRecvError(t *testing.T) {
+	for _, sc := range receivePaths() {
+		t.Run(stagedName(sc), func(t *testing.T) {
+			sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+			receiver.f.messageMetrics = newMessageMetrics()
+			receiver.f.acceptRecvErrorConfig = recvErrorAlways
+			receiver.f.sendRecvErrorConfig = recvErrorAlways
+			batches, want := flowTestBatches(1, 1, 1, 33)
+			wire := wireFor(t, sender, batches)
+			rxc := newRxContext(receiver.f, 0)
+			receiver.f.readOutsidePackets(ViaSender{UdpAddr: flowTestUnder}, append([]byte(nil), wire[0]...), rxc)
+			receiver.f.flushOutside(rxc)
+			h := &header.H{}
+			require.NoError(t, h.Parse(wire[0]))
+			replay := h.MessageCounter
+			var decrypts atomic.Int64
+			delayCrypto(receiver, func(n uint64) {
+				if n == replay {
+					decrypts.Add(1)
+				}
+			})
+			ctl := header.Encode(make([]byte, header.Len), header.Version, header.RecvError, 0, receiver.hostinfo.remoteIndexId, 0)
+			conn := &rxConn{wire: append([][]byte{ctl}, wire...), batchLen: len(wire) + 1}
+			receiver.receive(conn)
+			require.Nil(t, receiver.f.hostMap.QueryIndex(receiver.hostinfo.localIndexId))
+			require.Equal(t, want[:1], receiver.tun.segments(t))
+			require.Zero(t, decrypts.Load(), "known replays must not be decrypted again")
+			replies := conn.captured()
+			require.Len(t, replies, len(wire))
+			for i, reply := range replies {
+				require.NoError(t, h.Parse(reply))
+				require.Equal(t, header.RecvError, h.Type)
+				require.NoError(t, h.Parse(wire[i]))
+				require.True(t, i == 0 || receiver.hostinfo.ConnectionState.window.Check(nil, h.MessageCounter), "dropped data must not update replay state")
+			}
+		})
+	}
+}
+
+// TestFlowStagedRXReplacedTunnel classifies a data packet, then a close and a new handshake replace its tunnel at the
+// same index before the writer commits it, as when the listener runs ahead: it must be dropped untouched.
+func TestFlowStagedRXReplacedTunnel(t *testing.T) {
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	receiver.f.messageMetrics = newMessageMetrics()
+	batches, _ := flowTestBatches(1, 1, 1, 1)
+	wire := wireFor(t, sender, batches)
+	closePacket := sendControl(sender, header.CloseTunnel, 0)
+	p := &stagedRx{f: receiver.f, h: &header.H{}, hostmapCache: map[uint32]*HostInfo{}}
+	item := rxItem{packet: wire[0], via: flowTestUnder}
+	p.classify(&item)
+	require.NotNil(t, item.ci)
+	rxc := newRxContext(receiver.f, 0)
+	p.process(&rxBatch{items: []rxItem{{packet: closePacket, via: flowTestUnder}}}, rxc)
+	require.Nil(t, receiver.f.hostMap.QueryIndex(receiver.hostinfo.localIndexId))
+	_, replacement := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	receiver.f.hostMap.Lock()
+	receiver.f.hostMap.unlockedAddHostInfo(replacement.hostinfo, receiver.f)
+	receiver.f.hostMap.Unlock()
+	p.process(&rxBatch{items: []rxItem{item}}, rxc)
+	require.Empty(t, receiver.tun.writes)
+	require.True(t, item.ci.window.Check(nil, item.counter))
+	require.True(t, replacement.hostinfo.ConnectionState.window.Check(nil, item.counter))
+}
+
+// TestFlowStagedRXRealSocket exercises buffer handoff through the real Linux UDP receive loop.
+func TestFlowStagedRXRealSocket(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(3))
+	batches, want := flowTestBatches(60, 4, 1, 8)
+	wire := wireFor(t, sender, batches)
+
+	l := receiver.f.l
+	conn, err := udp.NewListener(l, udp.Settings{Listen: netip.MustParseAddrPort("127.0.0.1:0"), Batch: 16, Offloads: true})
+	require.NoError(t, err)
+	addr, err := conn.LocalAddr()
+	require.NoError(t, err)
+	out, err := udp.NewListener(l, udp.Settings{Listen: netip.MustParseAddrPort("127.0.0.1:0"), Batch: 16, Offloads: true})
+	require.NoError(t, err)
+	defer out.Close()
+	// The sender's address is the tunnel's remote, so no roaming
+	from, err := out.LocalAddr()
+	require.NoError(t, err)
+	receiver.hostinfo.remote.Store(&from)
+
+	done := make(chan struct{})
+	go func() { receiver.receive(conn); close(done) }()
+	// Send in bursts small enough for the default socket buffer
+	for off := 0; off < len(wire); off += 32 {
+		end := min(off+32, len(wire))
+		addrs := make([]netip.AddrPort, end-off)
+		for i := range addrs {
+			addrs[i] = addr
+		}
+		_, err := out.WriteBatch(wire[off:end], addrs)
+		require.NoError(t, err)
+		require.Eventually(t, func() bool { return len(receiver.tun.segments(t)) >= end }, 5*time.Second, 100*time.Microsecond)
+	}
+	receiver.f.closed.Store(true)
+	require.NoError(t, conn.Close())
+	<-done
+	require.Equal(t, want, receiver.tun.segments(t))
+}
+
+// TestFlowStagedRXCloseFlushesPartial: packets the listener holds in a batch it never got to end when the conn fails
+// are still committed, in order, on a live interface.
+func TestFlowStagedRXCloseFlushesPartial(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	batches, want := flowTestBatches(11, 4, 1, 8)
+	wire := wireFor(t, sender, batches)
+	p := newStagedRx(receiver.f, 0, newRxContext(receiver.f, 0))
+	for i, packet := range wire {
+		p.deliver(flowTestUnder, packet)
+		if (i+1)%64 == 0 {
+			p.endBatch()
+		}
+	}
+	p.close()
+	require.Equal(t, want, receiver.tun.segments(t))
+}
+
+// TestFlowStagedRXInlineBelow: small idle batches stay on the listener, larger batches still use the ordered writer.
+func TestFlowStagedRXInlineBelow(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sc := adaptive(3)
+	sc.inlineBelow = 16
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+	batches, want := flowTestBatches(40, 4, 1, 16)
+	wire := wireFor(t, sender, batches)
+	inline := metrics.GetOrRegisterCounter("flow.rx.inline", nil)
+	inline0, rx0 := inline.Count(), rxBatches.Count()
+	receiver.receive(&rxConn{wire: wire, sizes: []int{8, 8, 8, 128, 128, 8, 8, 128, 4, 128}})
+	require.Equal(t, want, receiver.tun.segments(t))
+	require.Positive(t, inline.Count()-inline0)
+	require.Positive(t, rxBatches.Count()-rx0)
+}

flow_test.go

@@ -0,0 +1,491 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"encoding/binary"
+	"io"
+	"math"
+	"net/netip"
+	"os"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/flynn/noise"
+	"github.com/gaissmai/bart"
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/cert"
+	"github.com/slackhq/nebula/config"
+	"github.com/slackhq/nebula/firewall"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/iputil"
+	"github.com/slackhq/nebula/noiseutil"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/overlay/tio"
+	"github.com/slackhq/nebula/test"
+	"github.com/slackhq/nebula/udp"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+)
+
+// flowTestMSS is the segment payload size of the synthetic TSO superpackets.
+const flowTestMSS = 1200
+
+var (
+	flowTestLocal  = netip.MustParseAddr("10.0.0.1")
+	flowTestRemote = netip.MustParseAddr("10.0.0.2")
+	flowTestUnder  = netip.MustParseAddrPort("192.0.2.2:4242")
+)
+
+// flowTCPSuperpacket builds an IPv4/TCP TSO superpacket from local to remote carrying segs segments. Segment i's
+// payload starts with the flow id, the superpacket seq and i, so a receiver can check what it got and in which order.
+func flowTCPSuperpacket(flow uint16, seq uint32, segs int) tio.Packet {
+	const hdr = 40
+	pay := segs * flowTestMSS
+	b := make([]byte, hdr+pay)
+	b[0] = 0x45
+	binary.BigEndian.PutUint16(b[2:4], uint16(min(hdr+pay, 65535)))
+	b[8] = 64
+	b[9] = iputil.IPProtocolTCP
+	copy(b[12:16], flowTestLocal.AsSlice())
+	copy(b[16:20], flowTestRemote.AsSlice())
+	binary.BigEndian.PutUint16(b[20:22], 10000+flow)
+	binary.BigEndian.PutUint16(b[22:24], 80)
+	binary.BigEndian.PutUint32(b[24:28], seq*uint32(pay))
+	b[32] = 0x50
+	b[33] = 0x10 // ACK
+	binary.BigEndian.PutUint16(b[34:36], 65535)
+	for i := range segs {
+		p := b[hdr+i*flowTestMSS:]
+		binary.BigEndian.PutUint16(p[0:2], flow)
+		binary.BigEndian.PutUint32(p[2:6], seq)
+		binary.BigEndian.PutUint16(p[6:8], uint16(i))
+	}
+	if segs == 1 {
+		return tio.Packet{Bytes: b}
+	}
+	return tio.Packet{Bytes: b, GSO: tio.GSOInfo{Size: flowTestMSS, HdrLen: hdr, CsumStart: 20, Proto: tio.GSOProtoTCP}}
+}
+
+// decodedSeg identifies one segment by the marker flowTCPSuperpacket put in its payload.
+type decodedSeg struct {
+	flow uint16
+	seq  uint32
+	idx  uint16
+}
+
+func decodeSeg(p []byte) decodedSeg {
+	return decodedSeg{flow: binary.BigEndian.Uint16(p[0:2]), seq: binary.BigEndian.Uint32(p[2:6]), idx: binary.BigEndian.Uint16(p[6:8])}
+}
+
+// flowTestBatches builds n reads of per superpackets of segs segments, across flows interleaved flows, and the
+// segments they carry in order.
+func flowTestBatches(n, per, flows, segs int) ([][]tio.Packet, []decodedSeg) {
+	var batches [][]tio.Packet
+	var want []decodedSeg
+	seq := make([]uint32, flows)
+	for range n {
+		var b []tio.Packet
+		for j := range per {
+			flow := uint16(j % flows)
+			b = append(b, flowTCPSuperpacket(flow, seq[flow], segs))
+			for s := range segs {
+				want = append(want, decodedSeg{flow: flow, seq: seq[flow], idx: uint16(s)})
+			}
+			seq[flow]++
+		}
+		batches = append(batches, b)
+	}
+	return batches, want
+}
+
+// flowTestPeer is one end of a tunnel: the interface, the hostinfo for the other end, and its capture points.
+type flowTestPeer struct {
+	f        *Interface
+	hostinfo *HostInfo
+	conn     *flowCaptureConn
+	tun      *flowTunCapture
+}
+
+// newFlowTestPair builds two interfaces joined by one real tunnel (keys from an actual handshake). sender sends from
+// flowTestLocal to flowTestRemote, receiver is the far side.
+func newFlowTestPair(t testing.TB, cipher noise.CipherFunc, sc stagedConfig) (sender, receiver *flowTestPeer) {
+	t.Helper()
+	initR, respR := runTestHandshakeWith(t, cipher)
+	sci, err := newConnectionStateFromResult(initR)
+	require.NoError(t, err)
+	rci, err := newConnectionStateFromResult(respR)
+	require.NoError(t, err)
+	sender = newFlowTestPeer(t, flowTestLocal, flowTestRemote, sci, 100, 200)
+	receiver = newFlowTestPeer(t, flowTestRemote, flowTestLocal, rci, 200, 100)
+	sender.f.staged, receiver.f.staged = sc, sc
+	return sender, receiver
+}
+
+func newFlowTestPeer(t testing.TB, me, them netip.Addr, ci *ConnectionState, localIdx, remoteIdx uint32) *flowTestPeer {
+	l := test.NewLogger()
+	hm := newHostMap(l)
+	hm.preferredRanges.Store(&[]netip.Prefix{})
+
+	myNets, myAddrs := &bart.Lite{}, &bart.Lite{}
+	myNets.Insert(netip.PrefixFrom(me, 24))
+	myAddrs.Insert(netip.PrefixFrom(me, 32))
+	myCert := &dummyCert{version: cert.Version1, networks: []netip.Prefix{netip.PrefixFrom(me, 24)}}
+	theirCert := &dummyCert{version: cert.Version1, networks: []netip.Prefix{netip.PrefixFrom(them, 24)}, name: "peer"}
+	ci.peerCert = &cert.CachedCertificate{Certificate: theirCert, InvertedGroups: map[string]struct{}{}}
+
+	fw := NewFirewall(l, time.Minute, time.Minute, time.Minute, myCert)
+	require.NoError(t, fw.AddRule(false, firewall.ProtoAny, 0, 0, []string{"any"}, "", "", "", "", ""))
+	require.NoError(t, fw.AddRule(true, firewall.ProtoAny, 0, 0, []string{"any"}, "", "", "", "", ""))
+
+	conn, tun := &flowCaptureConn{}, &flowTunCapture{}
+	lh := newTestLighthouse()
+	lh.remoteAllowList.Store(&RemoteAllowList{})
+
+	f := &Interface{
+		hostMap:            hm,
+		outside:            conn,
+		firewall:           fw,
+		lightHouse:         lh,
+		pki:                &PKI{},
+		myVpnNetworksTable: myNets,
+		myVpnAddrsTable:    myAddrs,
+		routines:           1,
+		writers:            []udp.Conn{conn},
+		messageMetrics:     &MessageMetrics{txExhausted: metrics.NewCounter()},
+		metricTxDropped:    metrics.NilCounter{},
+		l:                  l,
+	}
+	f.ctx = t.Context()
+	f.pki.caPool.Store(cert.NewCAPool())
+	// Never promote or requery, those would put extra test packets on the wire
+	f.tryPromoteEvery.Store(math.MaxUint32)
+	f.reQueryEvery.Store(math.MaxUint32)
+	f.batchers = []*batch.MultiCoalescer{batch.NewMultiCoalescer(tun, l)}
+	f.connectionManager = &connectionManager{
+		intf:          f,
+		relayUsed:     map[uint32]struct{}{},
+		relayUsedLock: &sync.RWMutex{},
+		trafficTimer:  NewLockingTimerWheel[uint32](time.Second, time.Minute),
+	}
+	f.handshakeManager = NewHandshakeManager(l, hm, lh, conn, defaultHandshakeConfig)
+	f.handshakeManager.f = f
+
+	hostinfo := &HostInfo{
+		ConnectionState: ci,
+		vpnAddrs:        []netip.Addr{them},
+		localIndexId:    localIdx,
+		remoteIndexId:   remoteIdx,
+		relayState:      RelayState{relayForByAddr: map[netip.Addr]*Relay{}, relayForByIdx: map[uint32]*Relay{}},
+	}
+	hostinfo.remote.Store(&flowTestUnder)
+	hostinfo.buildNetworks(myNets, theirCert)
+	hm.unlockedAddHostInfo(hostinfo, f)
+
+	return &flowTestPeer{f: f, hostinfo: hostinfo, conn: conn, tun: tun}
+}
+
+// flowCaptureConn records every packet sent, in order, as copies.
+type flowCaptureConn struct {
+	udp.NoopConn
+	mu    sync.Mutex
+	pkts  [][]byte
+	addrs []netip.AddrPort
+	// delay, when set, slows every WriteBatch down so batches pile up in flight
+	delay time.Duration
+}
+
+func (c *flowCaptureConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	time.Sleep(c.delay)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for i, b := range bufs {
+		c.pkts = append(c.pkts, append([]byte(nil), b...))
+		c.addrs = append(c.addrs, addrs[i])
+	}
+	return len(bufs), nil
+}
+
+func (c *flowCaptureConn) WriteTo(b []byte, addr netip.AddrPort) error {
+	_, err := c.WriteBatch([][]byte{b}, []netip.AddrPort{addr})
+	return err
+}
+
+func (c *flowCaptureConn) captured() [][]byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([][]byte(nil), c.pkts...)
+}
+
+// flowTunCapture is a tun queue that records every plaintext write.
+type flowTunCapture struct {
+	mu     sync.Mutex
+	writes [][]byte
+	// delay, when set, slows every write down
+	delay time.Duration
+}
+
+func (c *flowTunCapture) Read() ([]tio.Packet, error) { return nil, io.EOF }
+func (c *flowTunCapture) Close() error                { return nil }
+func (c *flowTunCapture) Write(b []byte) (int, error) {
+	time.Sleep(c.delay)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.writes = append(c.writes, append([]byte(nil), b...))
+	return len(b), nil
+}
+
+// segments decodes what reached the tun. The coalescer may have merged segments back into superpackets, so this
+// walks each payload in MSS steps.
+func (c *flowTunCapture) segments(t testing.TB) []decodedSeg {
+	t.Helper()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	var out []decodedSeg
+	for _, w := range c.writes {
+		pay := w[40:]
+		require.Zero(t, len(pay)%flowTestMSS, "tun write payload %d bytes", len(pay))
+		for off := 0; off < len(pay); off += flowTestMSS {
+			out = append(out, decodeSeg(pay[off:]))
+		}
+	}
+	return out
+}
+
+// decryptWire checks every wire packet decrypts on the receiver, in order, with strictly increasing counters, and
+// returns what each one carried. It uses its own replay window and copies, so the wire can still be fed to the
+// receiver afterwards.
+func decryptWire(t testing.TB, receiver *flowTestPeer, wire [][]byte) []decodedSeg {
+	t.Helper()
+	peer := receiver.hostinfo.ConnectionState
+	ci := &ConnectionState{dKey: peer.dKey, window: NewBits(peer.window.length)}
+	nb := make([]byte, 12)
+	var last uint64
+	out := make([]decodedSeg, 0, len(wire))
+	for i, pkt := range wire {
+		h := &header.H{}
+		require.NoError(t, h.Parse(pkt))
+		require.Equal(t, header.Message, h.Type)
+		require.Equal(t, receiver.hostinfo.localIndexId, h.RemoteIndex)
+		require.Greater(t, h.MessageCounter, last, "wire packet %d: counter %d after %d", i, h.MessageCounter, last)
+		last = h.MessageCounter
+		plain, err := ci.Decrypt(receiver.f.l, h.MessageCounter, append([]byte(nil), pkt...), nb)
+		require.NoError(t, err, "wire packet %d", i)
+		out = append(out, decodeSeg(plain[40:]))
+	}
+	return out
+}
+
+// flowBatchQueue is a tun queue that hands out scripted Read batches, then fails with io.EOF. Every batch is copied
+// into the current receive buffer, so a test sees the buffer ownership the Offload queue has.
+type flowBatchQueue struct {
+	batches [][]tio.Packet
+	rxBuf   []byte
+	out     []tio.Packet
+	closed  chan struct{}
+	// reads counts the batches handed out
+	reads atomic.Int64
+}
+
+func newFlowBatchQueue(batches [][]tio.Packet) *flowBatchQueue {
+	return &flowBatchQueue{batches: batches, rxBuf: make([]byte, 8*65536), closed: make(chan struct{})}
+}
+
+func (q *flowBatchQueue) Read() ([]tio.Packet, error) {
+	select {
+	case <-q.closed:
+		return nil, os.ErrClosed
+	default:
+	}
+	n := int(q.reads.Load())
+	if n == len(q.batches) {
+		return nil, io.EOF
+	}
+	q.out = q.out[:0]
+	off := 0
+	for _, p := range q.batches[n] {
+		l := copy(q.rxBuf[off:], p.Bytes)
+		q.out = append(q.out, tio.Packet{Bytes: q.rxBuf[off : off+l : off+l], GSO: p.GSO})
+		off += l
+	}
+	q.reads.Add(1)
+	return q.out, nil
+}
+
+func (q *flowBatchQueue) Write(b []byte) (int, error) { return len(b), nil }
+
+func (q *flowBatchQueue) Close() error { close(q.closed); return nil }
+
+// delayedCipher changes scheduling through the real cipher interface, only in tests.
+type delayedCipher struct {
+	noiseutil.CipherState
+	before func(uint64)
+}
+
+func (c delayedCipher) EncryptDanger(out, ad, plain []byte, n uint64, nb []byte) ([]byte, error) {
+	c.before(n)
+	return c.CipherState.EncryptDanger(out, ad, plain, n, nb)
+}
+
+func (c delayedCipher) DecryptDanger(out, ad, ciphertext []byte, n uint64, nb []byte) ([]byte, error) {
+	c.before(n)
+	return c.CipherState.DecryptDanger(out, ad, ciphertext, n, nb)
+}
+
+func delayCrypto(peer *flowTestPeer, before func(uint64)) {
+	ci := peer.hostinfo.ConnectionState
+	ci.eKey = delayedCipher{ci.eKey, before}
+	ci.dKey = delayedCipher{ci.dKey, before}
+}
+
+// stagedOn enables staging with n workers.
+func stagedOn(n int) stagedConfig {
+	return stagedConfig{workers: n}
+}
+
+// adaptive is staging with the inline threshold parseStagedConfig sets.
+func adaptive(n int) stagedConfig {
+	return stagedConfig{workers: n, inlineBelow: 256}
+}
+
+func stagedName(sc stagedConfig) string {
+	switch {
+	case sc.workers == 0:
+		return "inline"
+	case sc.inlineBelow > 0:
+		return "adaptive"
+	default:
+		return "staged"
+	}
+}
+
+// TestFlowStagedFallback: a conn that cannot hand off its batches keeps receive on the inline path.
+func TestFlowStagedFallback(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(3))
+	batches, want := flowTestBatches(10, 2, 1, 4)
+	wire := wireFor(t, sender, batches)
+	require.Equal(t, want, decryptWire(t, receiver, wire))
+	rx0 := rxBatches.Count()
+	receiver.receive(plainConn{&rxConn{wire: wire, batchLen: 16}})
+	require.Equal(t, want, receiver.tun.segments(t))
+	require.Equal(t, rx0, rxBatches.Count())
+}
+
+// wireFor produces the wire packets the sender emits for batches.
+func wireFor(t testing.TB, sender *flowTestPeer, batches [][]tio.Packet) [][]byte {
+	t.Helper()
+	sender.f.listenIn(newFlowBatchQueue(batches), 0)
+	wire := sender.conn.captured()
+	sender.conn.pkts, sender.conn.addrs = nil, nil
+	return wire
+}
+
+// rxConn plays wire back to listenOut through the buffer handoff, in batches of batchLen or of the scripted sizes
+// (the last one repeating), and captures what is sent back. Every buffer handed back is scribbled over before it is
+// read into, so a stage still using one shows corruption. from, when set, picks each packet's source address by its
+// index in wire.
+type rxConn struct {
+	flowCaptureConn
+	wire     [][]byte
+	batchLen int
+	sizes    []int
+	from     func(i int) netip.AddrPort
+	swaps    int
+}
+
+func (c *rxConn) ListenOut(r udp.EncReader, flush func()) error {
+	return c.ListenOutHandoff(r, func(bufs [][]byte) [][]byte {
+		flush()
+		return bufs
+	})
+}
+
+func (c *rxConn) ListenOutHandoff(r udp.EncReader, handoff func(bufs [][]byte) [][]byte) error {
+	n := c.batchLen
+	for _, size := range c.sizes {
+		n = max(n, size)
+	}
+	bufs := make([][]byte, n)
+	for i := range bufs {
+		bufs[i] = make([]byte, udp.MTU)
+	}
+	for off, k := 0, 0; off < len(c.wire); k++ {
+		size := c.batchLen
+		if len(c.sizes) > 0 {
+			size = c.sizes[min(k, len(c.sizes)-1)]
+		}
+		end := min(off+size, len(c.wire))
+		for i, p := range c.wire[off:end] {
+			n := copy(bufs[i], p)
+			from := flowTestUnder
+			if c.from != nil {
+				from = c.from(off + i)
+			}
+			r(from, bufs[i][:n:n])
+		}
+		next := handoff(bufs)
+		if len(next) != len(bufs) {
+			panic("handoff returned the wrong number of buffers")
+		}
+		bufs = next
+		for _, b := range bufs {
+			for j := range b[:udp.MTU] {
+				b[j] = 0x5a
+			}
+		}
+		c.swaps++
+		off = end
+	}
+	return os.ErrClosed
+}
+
+// plainConn hides ListenOutHandoff, so it cannot hand off its batches.
+type plainConn struct{ udp.Conn }
+
+var rxBatches = metrics.GetOrRegisterCounter("flow.rx.batches", nil)
+
+// receive runs the peer's listenOut over conn until conn runs out.
+func (p *flowTestPeer) receive(conn udp.Conn) {
+	p.f.writers = []udp.Conn{conn}
+	p.f.outside = conn
+	p.f.listenOut(0)
+}
+
+// TestFlowStagedFirewallParity: both receive paths apply the same inbound policy. An authenticated inbound denial
+// still consumes its replay counter, exactly as inline Decrypt does.
+func TestFlowStagedFirewallParity(t *testing.T) {
+	for _, sc := range []stagedConfig{{}, stagedOn(2)} {
+		t.Run(stagedName(sc), func(t *testing.T) {
+			batches, _ := flowTestBatches(4, 3, 1, 8)
+			sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+			wire := wireFor(t, sender, batches)
+			receiver.f.firewall.InRules = newFirewallTable()
+			receiver.receive(&rxConn{wire: wire, batchLen: 64})
+			require.Empty(t, receiver.tun.writes)
+			for _, pkt := range wire {
+				h := &header.H{}
+				require.NoError(t, h.Parse(pkt))
+				require.False(t, receiver.hostinfo.ConnectionState.window.Check(receiver.f.l, h.MessageCounter))
+			}
+		})
+	}
+}
+
+// TestFlowStagedAESGCM runs the parsed configuration with the real AES cipher. Receive decryption has no nonce
+// ordering requirement, so it stays staged in BoringCrypto and FIPS builds too.
+func TestFlowStagedAESGCM(t *testing.T) {
+	l := test.NewLogger()
+	c := config.NewC(l)
+	c.Settings = map[string]any{"tun": map[string]any{"staged": map[string]any{"workers": 2}}}
+	sender, receiver := newFlowTestPair(t, noiseutil.CipherAESGCM, parseStagedConfig(c, l))
+	require.Equal(t, 2, receiver.f.staged.workers)
+	batches, want := flowTestBatches(4, 8, 1, 45)
+	wire := wireFor(t, sender, batches)
+	rx0 := rxBatches.Count()
+	receiver.receive(&rxConn{wire: wire, batchLen: 512})
+	require.Greater(t, rxBatches.Count(), rx0)
+	require.Equal(t, want, receiver.tun.segments(t))
+}

interface.go

@@ -65,6 +65,7 @@ type InterfaceConfig struct {
 	// goroutine's sendmmsg on one XPS-selected NIC TX ring so per-flow
 	// packets stay ordered on the wire.
 	PinThreads bool
+	staged     stagedConfig
 
 	l *slog.Logger
 }
@@ -99,6 +100,7 @@ type Interface struct {
 	// a CPU at all (tun.pin_threads, default true). When false, threads are
 	// left free to migrate as on stock nebula.
 	pinThreads   bool
+	staged       stagedConfig
 	relayManager *relayManager
 
 	tryPromoteEvery atomic.Uint32
@@ -248,6 +250,7 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		conntrackCacheTimeout: c.ConntrackCacheTimeout,
 		cpuAffinity:           c.CpuAffinity,
 		pinThreads:            c.PinThreads,
+		staged:                c.staged,
 		udpTxErrLimiter:       udpTxErrLimiter,
 		udpRelayTxErrLimiter:  udpRelayTxErrLimiter,
 		txOversizeLimiter:     txOversizeLimiter,
@@ -413,15 +416,19 @@ func (f *Interface) listenOut(i int) {
 		f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, payload, rxc)
 	}
 
-	flusher := func() {
-		if err := f.batchers[i].Flush(); err != nil {
-			f.l.Error("Failed to flush tun coalescer", "error", err)
+	flusher := func() { f.flushOutside(rxc) }
+
+	var err error
+	if hc, ok := li.(handoffConn); ok && f.staged.workers > 0 {
+		// Every stage on its own goroutine, see flow_rx.go. rxc becomes the ordered writer's context.
+		err = newStagedRx(f, i, rxc).listen(hc)
+	} else {
+		if f.staged.workers > 0 {
+			f.l.Info("udp conn cannot hand off receive batches, receive stays inline", "reader", i)
 		}
-		clear(rxc.hostmapCache)
+		err = li.ListenOut(listener, flusher)
 	}
 
-	err := li.ListenOut(listener, flusher)
-
 	// An error after teardown began is shutdown noise, the closed flag covers resources
 	// Close releases itself and the cancelled ctx covers ones torn down by their owners
 	// reacting to it, like the user device pipes

main.go

@@ -244,6 +244,7 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 
 	pinThreads := c.GetBool("tun.pin_threads", true)
 	cpuAffinity := parseCpuAffinity(c, l, routines)
+	staged := parseStagedConfig(c, l)
 	if pinThreads && routines > 1 && len(cpuAffinity) == 0 && !configTest {
 		// The operator didn't choose pin CPUs, so pick a default set that
 		// prefers performance cores and doesn't stack co-located instances
@@ -296,6 +297,7 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 		ConntrackCacheTimeout: conntrackCacheTimeout,
 		CpuAffinity:           cpuAffinity,
 		PinThreads:            pinThreads,
+		staged:                staged,
 		l:                     l,
 	}
 

main_test.go

@@ -5,6 +5,7 @@ import (
 	"net/netip"
 	"os"
 	"path/filepath"
+	"runtime"
 	"testing"
 	"time"
 
@@ -123,3 +124,23 @@ firewall:
 	require.NoError(t, err)
 	require.Nil(t, ctrl, "a config test hands back nothing to stop, so it must stop itself")
 }
+
+func TestParseStagedConfig(t *testing.T) {
+	l := test.NewLogger()
+	staged := func(workers any) stagedConfig {
+		c := config.NewC(l)
+		if workers != nil {
+			c.Settings["tun"] = map[string]any{"staged": map[string]any{"workers": workers}}
+		}
+		return parseStagedConfig(c, l)
+	}
+	for _, workers := range []any{nil, -1, 0} {
+		require.Equal(t, stagedConfig{}, staged(workers))
+	}
+	if runtime.GOOS != "linux" {
+		require.Equal(t, stagedConfig{}, staged(2))
+		return
+	}
+	want := stagedConfig{workers: 2, inlineBelow: 256}
+	require.Equal(t, want, staged(2))
+}

outside.go

@@ -130,17 +130,10 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 	}
 
 	out, err := hostinfo.ConnectionState.Decrypt(f.l, h.MessageCounter, packet, rxc.nb)
-	if err != nil {
-		if f.l.Enabled(context.Background(), slog.LevelDebug) {
-			hostinfo.logger(f.l).Debug("Failed to decrypt packet", "error", err, "from", via, "header", h)
-		}
+	if !f.finishOutsideDecrypt(hostinfo, via, h.MessageCounter, err) {
 		return
 	}
 
-	// Roam before we respond
-	f.handleHostRoaming(hostinfo, via)
-	f.connectionManager.In(hostinfo)
-
 	switch h.Type {
 	case header.Message:
 		switch h.Subtype {
@@ -473,6 +466,29 @@ func parseV4(data []byte, incoming bool, fp *firewall.ParsedPacket) error {
 	return nil
 }
 
+// finishOutsideDecrypt logs a failed decrypt, or roams and records the traffic for an authenticated packet before it
+// is handled. It reports whether the packet authenticated.
+func (f *Interface) finishOutsideDecrypt(hostinfo *HostInfo, via ViaSender, counter uint64, err error) bool {
+	if err != nil {
+		if f.l.Enabled(context.Background(), slog.LevelDebug) {
+			hostinfo.logger(f.l).Debug("Failed to decrypt packet", "error", err, "from", via, "counter", counter)
+		}
+		return false
+	}
+	// Roam before we respond
+	f.handleHostRoaming(hostinfo, via)
+	f.connectionManager.In(hostinfo)
+	return true
+}
+
+// flushOutside ends a receive batch: it writes out the tun coalescer and drops the batch's hostmap cache.
+func (f *Interface) flushOutside(rxc *rxContext) {
+	if err := f.batchers[rxc.q].Flush(); err != nil {
+		f.l.Error("Failed to flush tun coalescer", "error", err)
+	}
+	clear(rxc.hostmapCache)
+}
+
 func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, messageCounter uint64, out []byte, rxc *rxContext) {
 	err := newPacket(out, true, rxc.fwPacket)
 	if err != nil {

udp/udp_linux.go

@@ -246,6 +246,18 @@ func getFrom(names [][]byte, i int, isV4 bool) netip.AddrPort {
 }
 
 func (u *StdConn) ListenOut(r EncReader, flush func()) error {
+	return u.listenOut(r, flush, nil)
+}
+
+// ListenOutHandoff is ListenOut for a reader that keeps the packets: after each batch, handoff takes the batch's
+// buffers and returns replacements, of the same count and at least the same size, to receive the next one into.
+func (u *StdConn) ListenOutHandoff(r EncReader, handoff func(bufs [][]byte) [][]byte) error {
+	return u.listenOut(r, nil, handoff)
+}
+
+// listenOut is the receive loop. After each batch it calls flush, or with handoff set, hands the batch's buffers
+// over and points the recvmmsg iovecs at the ones it gets back.
+func (u *StdConn) listenOut(r EncReader, flush func(), handoff func(bufs [][]byte) [][]byte) error {
 	bufSize := MTU
 	cmsgSpace := 0
 	if u.groSupported {
@@ -283,7 +295,21 @@ func (u *StdConn) ListenOut(r EncReader, flush func()) error {
 			deliverSegments(r, from, payload, segSize)
 		}
 
-		flush()
+		if handoff == nil {
+			flush()
+			continue
+		}
+		next := handoff(buffers)
+		if len(next) != len(buffers) {
+			panic(fmt.Sprintf("udp: handoff returned %d buffers, need %d", len(next), len(buffers)))
+		}
+		for i := range next {
+			if len(next[i]) < bufSize {
+				panic(fmt.Sprintf("udp: handoff buffer is %d bytes, need %d", len(next[i]), bufSize))
+			}
+			msgs[i].Hdr.Iov.Base = &next[i][0]
+		}
+		buffers = next
 	}
 }
 

udp/udp_tester.go

@@ -183,15 +183,24 @@ func (u *TesterConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, err
 }
 
 func (u *TesterConn) ListenOut(r EncReader, flush func()) error {
+	// The batcher borrows plaintext decrypted in place inside p.Data
+	// until Flush, so the packet must stay alive across flush()
+	return u.listenOut(r, func(p *Packet) { flush() })
+}
+
+// ListenOutHandoff hands each packet's buffer over, the replacement goes back to the pool.
+func (u *TesterConn) ListenOutHandoff(r EncReader, handoff func(bufs [][]byte) [][]byte) error {
+	return u.listenOut(r, func(p *Packet) { p.Data = handoff([][]byte{p.Data})[0] })
+}
+
+func (u *TesterConn) listenOut(r EncReader, finish func(*Packet)) error {
 	for {
 		select {
 		case <-u.done:
 			return os.ErrClosed
 		case p := <-u.RxPackets:
 			r(p.From, p.Data[:len(p.Data):len(p.Data)])
-			// The batcher borrows plaintext decrypted in place inside p.Data
-			// until Flush, so the packet must stay alive across flush()
-			flush()
+			finish(p)
 			p.Release()
 		}
 	}