udp: Add ordered staged receive processing

With tun.staged.workers set, each Linux UDP listener classifies a recvmmsg
batch, hands its data packets to that many decrypt workers, and an ordered
writer commits the batch in arrival order: replay window, roaming, firewall,
control packets and tun writes go through the same helpers inline receive
uses. Six batches per routine bound memory and backpressure. Conns that cannot
hand off their buffers stay inline.

The listener's classification is speculative. Before touching replay state the
writer resolves each data packet through its own batch-local hostmap cache,
exactly where inline receive would, so data after a close in the same batch is
still delivered while data after a removal or replacement between batches is
dropped, with the same RecvError reply inline receive sends. Known replays are
marked by a precheck and skip decryption, but keep their place so that reply
stays identical when an earlier RecvError deleted their tunnel.

New API: udp.StdConn.ListenOutHandoff, and the same on TesterConn for e2e,
runs the receive loop while handing each batch's buffers to the caller in
exchange for replacements. It is found through an unexported interface in
nebula, so udp.Conn and its other implementations are unchanged.

Metrics: flow.rx.batches, and flow.stage.{rx_read,rx_decrypt,rx_write}.q<N>.busy_ns.

Tests cover both ciphers, ordering, duplicate and corrupted packets from other
addresses, window jumps, control packets mid-batch, RecvError then replay,
tunnel replacement, firewall parity, fallback, real UDP sockets, partial
batches at close, and a handshake and traffic through e2e that checks pipeline
use on Linux and the inline fallback elsewhere. No isolated receive throughput
claim is made.

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)

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,30 @@ func TestCloseTunnelAuthenticated(t *testing.T) {
 	myControl.Stop()
 	theirControl.Stop()
 }
+
+// TestStagedTunnel completes a handshake and traffic both ways with staged receive. The test tun cannot hand off
+// batches, so transmit stays inline; 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.batches", 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

@@ -278,6 +278,7 @@ tun:
   # Linux only. staged.workers enables staged packet processing: each routine hands its batches to this many crypto
   # workers per direction, then sends or writes them in order. 0, the default, keeps processing inline. Not reloadable.
   # Transmit needs use_offloads; BoringCrypto and FIPS builds keep it inline because AES-GCM there needs ordered nonces.
+  # Receive needs a listener that can hand off its batches, which Linux has.
   # IMPORTANT: staged senders require peers with at least the 8192-entry replay window of this release. Released
   # v1.11.x peers use 1024 and are incompatible. The window is not negotiated.
   #staged:

flow.go

@@ -12,21 +12,24 @@ import (
 	"github.com/slackhq/nebula/noiseutil"
 )
 
-// Staged transmit (tun.staged.workers) moves encryption off each tun reader onto worker goroutines, behind which an
-// ordered sender flushes every batch in read order. It reuses the inline filter and encryption: staging changes
-// ownership and scheduling, not policy. Each routine has its own workers and sender, so order is per routine. Full
-// queues block the reader. A read error drains the batches in flight, while shutdown waits for the crypto and then
-// drops packets whose sockets have closed. Eight slots take about 9 MiB per routine; small-MSS arenas can grow.
+// Staged processing (tun.staged.workers) moves crypto off each tun reader and udp listener onto worker goroutines,
+// behind which an ordered stage sends or writes every batch in read order. It reuses the inline packet handlers:
+// staging changes ownership and scheduling, not policy. Each routine has its own workers and ordered stages, so order
+// is per routine. Full queues block the reader. A read error drains the batches in flight, while shutdown waits for
+// the crypto and then drops packets whose destination has closed. TX uses eight slots and RX six; at listen.batch=64
+// with GRO their buffers take about 9 MiB and 24 MiB per routine. Small-MSS TX arenas can grow.
 //
-// Metrics are registered only for active pipelines. flow.tx.batches counts reads handed to the sender.
-// flow.stage.<stage>.q<queue>.busy_ns locates the limiting stage: tx_read, tx_encrypt or tx_send. 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.
+// Metrics are registered only for active pipelines. flow.{tx,rx}.batches counts reads handed to the ordered stage.
+// flow.stage.<stage>.q<queue>.busy_ns locates the limiting stage: tx_read, tx_encrypt, tx_send, 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 RX batches sent to the
+// writer include decryption in writer time.
 
 // stagedConfig is tun.staged. workers is per routine and direction; zero keeps processing inline.
 type stagedConfig struct {
 	workers int
 	tx      bool
+	rx      bool
 }
 
 // parseStagedConfig enables staging on Linux only; queues and conns that cannot hand off batches also stay inline.
@@ -39,12 +42,12 @@ func parseStagedConfig(c *config.C, l *slog.Logger) stagedConfig {
 		l.Warn("tun.staged is only supported on Linux; packet processing stays inline")
 		return stagedConfig{}
 	}
-	sc.tx = true
+	sc.tx, sc.rx = true, true
 	if noiseutil.EncryptLockNeeded {
 		l.Warn("tun.staged cannot encrypt in parallel with boringcrypto or fips140; transmit stays inline")
 		sc.tx = false
 	}
-	l.Info("Staged packet processing enabled", "workers", sc.workers, "tx", sc.tx)
+	l.Info("Staged packet processing enabled", "workers", sc.workers, "tx", sc.tx, "rx", sc.rx)
 	return sc
 }
 

flow_rx.go

@@ -0,0 +1,388 @@
+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
+}
+
+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),
+	}
+	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
+	}
+
+	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,220 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"fmt"
+	"net/netip"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/flynn/noise"
+	"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)}
+}
+
+// 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))
+}

flow_test.go

@@ -17,6 +17,7 @@ import (
 	"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"
@@ -236,6 +237,23 @@ func (c *flowTunCapture) Write(b []byte) (int, error) {
 	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.
@@ -337,7 +355,7 @@ func delayCrypto(peer *flowTestPeer, before func(uint64)) {
 
 // stagedOn enables staging with n workers.
 func stagedOn(n int) stagedConfig {
-	return stagedConfig{workers: n, tx: true}
+	return stagedConfig{workers: n, tx: true, rx: true}
 }
 
 func stagedName(sc stagedConfig) string {
@@ -356,12 +374,143 @@ func TestFlowStagedFallback(t *testing.T) {
 	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(3))
 	batches, want := flowTestBatches(10, 2, 1, 4)
 	tx := metrics.GetOrRegisterCounter("flow.tx.batches", nil)
-	tx0 := tx.Count()
+	tx0, rx0 := tx.Count(), rxBatches.Count()
 	sender.f.listenIn(plainQueue{newFlowBatchQueue(batches)}, 0)
 	wire := sender.conn.captured()
 	require.Equal(t, want, decryptWire(t, receiver, wire))
 	require.Equal(t, tx0, tx.Count())
+	receiver.receive(plainConn{&rxConn{wire: wire, batchLen: 16}})
+	require.Equal(t, want, receiver.tun.segments(t))
+	require.Equal(t, rx0, rxBatches.Count())
 }
 
 // plainQueue and plainConn hide SwapRxBuf and ListenOutHandoff, so they cannot hand off their batches.
 type plainQueue struct{ tio.Queue }
+
+// wireFor produces the wire packets the sender emits for batches, inline.
+func wireFor(t testing.TB, sender *flowTestPeer, batches [][]tio.Packet) [][]byte {
+	t.Helper()
+	sc := sender.f.staged
+	sender.f.staged = stagedConfig{}
+	sender.f.listenIn(newFlowBatchQueue(batches), 0)
+	sender.f.staged = sc
+	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
+}
+
+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 paths apply the same outbound and 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, _ := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+			sender.f.firewall.OutRules = newFirewallTable()
+			sender.f.listenIn(newFlowBatchQueue(batches), 0)
+			require.Empty(t, sender.conn.captured())
+
+			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))
+			}
+		})
+	}
+}
+
+// TestFlowStagedCipherFallback runs the parsed configuration with the real AES cipher: transmit falls back to inline
+// when the build needs nonce-ordered encryption (BoringCrypto, FIPS), receive decryption stays staged.
+func TestFlowStagedCipherFallback(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, !noiseutil.EncryptLockNeeded, sender.f.staged.tx)
+	require.True(t, sender.f.staged.rx)
+	batches, want := flowTestBatches(4, 8, 1, 45)
+	tx := metrics.GetOrRegisterCounter("flow.tx.batches", nil)
+	tx0, rx0 := tx.Count(), rxBatches.Count()
+	sender.f.listenIn(newFlowBatchQueue(batches), 0)
+	if noiseutil.EncryptLockNeeded {
+		require.Equal(t, tx0, tx.Count())
+	}
+	receiver.receive(&rxConn{wire: sender.conn.captured(), batchLen: 512})
+	require.Greater(t, rxBatches.Count(), rx0)
+	require.Equal(t, want, receiver.tun.segments(t))
+}

interface.go

@@ -422,15 +422,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.rx {
+		// 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.rx {
+			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_test.go

@@ -142,6 +142,6 @@ func TestParseStagedConfig(t *testing.T) {
 		require.Equal(t, stagedConfig{}, staged(2))
 		return
 	}
-	want := stagedConfig{workers: 2, tx: !noiseutil.EncryptLockNeeded}
+	want := stagedConfig{workers: 2, tx: !noiseutil.EncryptLockNeeded, rx: true}
 	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()
 		}
 	}