batch stuff

iputil/packet.go

@@ -199,7 +199,7 @@ func ipv4CreateRejectTCPPacket(packet []byte, out []byte) []byte {
 }
 
 func ipv6CreateRejectPacket(packet []byte, out []byte) []byte {
-	proto, offset, isFragment := ipv6FindUpperProtocol(packet)
+	proto, offset, isFragment := IPv6FindUpperProtocol(packet)
 	if isFragment {
 		return nil
 	}
@@ -333,11 +333,34 @@ func ipv6CreateRejectTCPPacket(packet []byte, out []byte, offset int) []byte {
 	return out
 }
 
-func ipv6FindUpperProtocol(packet []byte) (nextHeader uint8, offset int, isFragment bool) {
+// maxIPv6ExtHeaders caps the extension-header walk in IPv6FindUpperProtocol.
+// RFC 8200 legal chains are shorter (each header at most once, Destination
+// Options at most twice), so the cap only bites crafted packets, which would
+// otherwise make us walk their whole payload 8 bytes at a time.
+const maxIPv6ExtHeaders = 8
+
+// IPv6FindUpperProtocol walks packet's IPv6 extension-header chain and
+// returns the terminal (upper-layer) protocol number, the byte offset where
+// that protocol's header begins, and whether the packet is a non-first
+// fragment. It steps over Hop-by-Hop (0), Routing (43), Fragment (44),
+// AH (51), and Destination Options (60); anything else — including ESP,
+// whose payload is encrypted — terminates the walk.
+//
+// For a non-first fragment, nextHeader still names the flow's upper
+// protocol (copied from the fragment header) but offset points at fragment
+// payload, not a real transport header: consult isFragment before
+// dereferencing offset. If the chain is truncated, over-long, or the packet
+// is shorter than an IPv6 header, the walk stops early and nextHeader is
+// whatever it stopped on (59, IPPROTO_NONE, for the too-short case) —
+// callers treat any non-transport result as unclassifiable.
+func IPv6FindUpperProtocol(packet []byte) (nextHeader uint8, offset int, isFragment bool) {
+	if len(packet) < ipv6.HeaderLen {
+		return 59, 0, false // IPPROTO_NONE: nothing to classify
+	}
 	nextHeader = packet[6]
 	offset = ipv6.HeaderLen
 
-	for {
+	for range maxIPv6ExtHeaders {
 		switch nextHeader {
 		case 0, 43, 60: // Hop-by-Hop, Routing, Destination
 			if len(packet) < offset+2 {
@@ -367,6 +390,7 @@ func ipv6FindUpperProtocol(packet []byte) (nextHeader uint8, offset int, isFragm
 			return nextHeader, offset, isFragment
 		}
 	}
+	return nextHeader, offset, isFragment
 }
 
 func CreateICMPEchoResponse(packet, out []byte) []byte {

iputil/packet_test.go

@@ -515,3 +515,121 @@ func TestCreateICMPEchoResponse_IPv6_NotICMPv6(t *testing.T) {
 	result := CreateICMPEchoResponse(packet, out)
 	assert.Nil(t, result)
 }
+
+func TestIPv6FindUpperProtocol(t *testing.T) {
+	src := net.ParseIP("fd00::1")
+	dst := net.ParseIP("fd00::2")
+
+	// extHdr builds one 8-byte-unit extension header: next, hdrExtLen
+	// ((extra+1)*8 bytes total), padded to size.
+	extHdr := func(next uint8, extra int) []byte {
+		b := make([]byte, (extra+1)*8)
+		b[0] = next
+		b[1] = uint8(extra)
+		return b
+	}
+
+	t.Run("no extension headers", func(t *testing.T) {
+		for _, proto := range []uint8{6, 17, 58} {
+			nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, proto, make([]byte, 20)))
+			assert.Equal(t, proto, nh)
+			assert.Equal(t, ipv6.HeaderLen, offset)
+			assert.False(t, frag)
+		}
+	})
+
+	t.Run("hop-by-hop then TCP", func(t *testing.T) {
+		payload := append(extHdr(6, 0), make([]byte, 20)...)
+		nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 0, payload))
+		assert.Equal(t, uint8(6), nh)
+		assert.Equal(t, ipv6.HeaderLen+8, offset)
+		assert.False(t, frag)
+	})
+
+	t.Run("chained headers honor length units", func(t *testing.T) {
+		// Hop-by-Hop (8B) -> Dest Options (16B) -> Routing (8B) -> UDP.
+		payload := extHdr(60, 0)
+		payload = append(payload, extHdr(43, 1)...)
+		payload = append(payload, extHdr(17, 0)...)
+		payload = append(payload, make([]byte, 8)...)
+		nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 0, payload))
+		assert.Equal(t, uint8(17), nh)
+		assert.Equal(t, ipv6.HeaderLen+8+16+8, offset)
+		assert.False(t, frag)
+	})
+
+	t.Run("AH length is in 4-byte units plus 2", func(t *testing.T) {
+		// AH payload-len byte 4 -> (4+2)*4 = 24 bytes on the wire.
+		ah := make([]byte, 24)
+		ah[0] = 6
+		ah[1] = 4
+		payload := append(ah, make([]byte, 20)...)
+		nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 51, payload))
+		assert.Equal(t, uint8(6), nh)
+		assert.Equal(t, ipv6.HeaderLen+24, offset)
+		assert.False(t, frag)
+	})
+
+	t.Run("first fragment walks to the transport header", func(t *testing.T) {
+		frag := make([]byte, 8)
+		frag[0] = 17
+		binary.BigEndian.PutUint16(frag[2:4], 0x0001) // offset 0, M=1
+		payload := append(frag, make([]byte, 8)...)
+		nh, offset, isFrag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 44, payload))
+		assert.Equal(t, uint8(17), nh)
+		assert.Equal(t, ipv6.HeaderLen+8, offset)
+		assert.False(t, isFrag, "first fragment carries the real transport header")
+	})
+
+	t.Run("non-first fragment is flagged", func(t *testing.T) {
+		frag := make([]byte, 8)
+		frag[0] = 17
+		binary.BigEndian.PutUint16(frag[2:4], 1<<3) // offset 1, M=0
+		payload := append(frag, make([]byte, 8)...)
+		nh, _, isFrag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 44, payload))
+		assert.Equal(t, uint8(17), nh, "fragment header still names the flow's L4")
+		assert.True(t, isFrag, "offset points at fragment payload, not a header")
+	})
+
+	t.Run("ESP terminates the walk", func(t *testing.T) {
+		nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 50, make([]byte, 16)))
+		assert.Equal(t, uint8(50), nh)
+		assert.Equal(t, ipv6.HeaderLen, offset)
+		assert.False(t, frag)
+	})
+
+	t.Run("unknown protocol terminates the walk", func(t *testing.T) {
+		nh, offset, _ := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 132, make([]byte, 16))) // SCTP
+		assert.Equal(t, uint8(132), nh)
+		assert.Equal(t, ipv6.HeaderLen, offset)
+	})
+
+	t.Run("truncated extension header stops the walk", func(t *testing.T) {
+		// Next header says Hop-by-Hop but the packet ends at the IPv6 header.
+		nh, offset, frag := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 0, nil))
+		assert.Equal(t, uint8(0), nh, "unresolvable chain returns the extension header it stopped on")
+		assert.Equal(t, ipv6.HeaderLen, offset)
+		assert.False(t, frag)
+	})
+
+	t.Run("crafted over-long chain hits the cap", func(t *testing.T) {
+		// Ten chained Hop-by-Hop headers, then TCP. Illegal per RFC 8200
+		// (Hop-by-Hop may only appear first); the cap must stop the walk
+		// before it resolves rather than crawling arbitrary crafted chains.
+		var payload []byte
+		for i := 0; i < 9; i++ {
+			payload = append(payload, extHdr(0, 0)...)
+		}
+		payload = append(payload, extHdr(6, 0)...)
+		payload = append(payload, make([]byte, 20)...)
+		nh, _, _ := IPv6FindUpperProtocol(makeIPv6Packet(src, dst, 0, payload))
+		assert.Equal(t, uint8(0), nh, "walk must stop at the cap, not resolve to TCP")
+	})
+
+	t.Run("packet shorter than an IPv6 header", func(t *testing.T) {
+		nh, offset, frag := IPv6FindUpperProtocol(make([]byte, 39))
+		assert.Equal(t, uint8(59), nh) // IPPROTO_NONE
+		assert.Equal(t, 0, offset)
+		assert.False(t, frag)
+	})
+}

overlay/batch/coalesce_core.go

@@ -136,6 +136,30 @@ func ipHeadersMatch(a, b []byte, isV6 bool) bool {
 	return true
 }
 
+// ipv4FlagDF is the Don't Fragment bit in the IPv4 flags byte (header byte 6).
+const ipv4FlagDF = 0x40
+
+// ipv4CanCoalesceID reports whether an IPv4 packet whose header starts at
+// nextHdr may join a chain whose seed header is seedHdr as segment index seg
+// (the seed is segment 0). Kernel GSO re-stamps outgoing segment IDs as
+// seed_id+n, so coalescing is only transparent when that re-stamp is either
+// harmless (DF set: RFC 6864 atomic datagrams, the ID carries no meaning) or
+// reproduces the original IDs exactly (DF clear + IDs already sequential —
+// the same admission rule kernel GRO applies). Without this, a DF=0 sender
+// with non-sequential IDs (e.g. OpenBSD's randomized IDs) could have IDs
+// rewritten into ranges that collide across superpackets, corrupting
+// reassembly if the packets are fragmented after the TUN write.
+//
+// DF itself is guaranteed uniform across a chain by ipHeadersMatch (byte 6
+// is inside its compared range), so checking the seed's copy suffices.
+func ipv4CanCoalesceID(seedHdr, nextHdr []byte, seg int) bool {
+	if seedHdr[6]&ipv4FlagDF != 0 {
+		return true
+	}
+	expect := binary.BigEndian.Uint16(seedHdr[4:6]) + uint16(seg)
+	return binary.BigEndian.Uint16(nextHdr[4:6]) == expect
+}
+
 // Arena is an injectable byte-slab that hands out non-overlapping borrowed
 // slices via Reserve and releases them in bulk via Reset.
 type Arena struct {

overlay/batch/multi_coalesce.go

@@ -4,6 +4,8 @@ import (
 	"errors"
 	"io"
 	"log/slog"
+
+	"github.com/slackhq/nebula/iputil"
 )
 
 // MultiCoalescer fans plaintext packets out to lane-specific batchers based
@@ -13,6 +15,10 @@ import (
 // UDP coalescer only sees UDP, and the passthrough lane handles everything else.
 // Per-flow delivery order is preserved because a single 5-tuple only
 // ever lands in one lane and each lane preserves its own slot order.
+// Routing follows the flow, not the coalesceability: IPv4 fragments keep
+// their L4 proto visible and IPv6 extension chains are walked to the
+// terminal proto, so a flow's non-coalesceable shapes ride its lane as
+// in-lane passthroughs rather than falling to the later-flushed pt lane.
 //
 // Cross-lane order is intentionally NOT preserved across the TCP/UDP/passthrough split.
 type MultiCoalescer struct {
@@ -34,6 +40,28 @@ func NewMultiCoalescer(w io.Writer, l *slog.Logger) RxBatcher {
 	return m
 }
 
+// IANA protocol numbers for the IPv6 extension headers
+// iputil.IPv6FindUpperProtocol can step over. The set here must match what
+// that walker walks: it is the hot path's cheap pre-guard, so the walk is
+// only paid when it can actually make progress.
+const (
+	ipProtoHopByHop = 0
+	ipProtoRouting  = 43
+	ipProtoFragment = 44
+	ipProtoAH       = 51
+	ipProtoDestOpts = 60
+)
+
+// isIPv6ExtHeader reports whether nh is an extension header the terminal-
+// protocol walk knows how to step over.
+func isIPv6ExtHeader(nh byte) bool {
+	switch nh {
+	case ipProtoHopByHop, ipProtoRouting, ipProtoFragment, ipProtoAH, ipProtoDestOpts:
+		return true
+	}
+	return false
+}
+
 // Commit dispatches pkt to the appropriate lane based on IP version + L4 proto.
 // On the success path the IP/TCP-or-UDP parse happens here once and the
 // parsed struct is handed to the lane via commitParsed so the lane doesn't re-walk the header.
@@ -51,6 +79,22 @@ func (m *MultiCoalescer) Commit(pkt []byte) error {
 			return m.pt.Commit(pkt)
 		}
 		proto = pkt[6]
+		if isIPv6ExtHeader(proto) {
+			// Walk to the terminal protocol so the packet routes to its
+			// flow's lane. It stays non-coalesceable — the lane's parser
+			// rejects the ext-header shape and emits it as an in-lane
+			// passthrough — but landing in the right lane preserves
+			// per-flow order, exactly like IPv4 fragments (whose header
+			// keeps the L4 proto visible) already do. Fragments are the
+			// case that matters: every fragment names the flow's L4, so a
+			// fragmented datagram travels with its flow's unfragmented
+			// siblings instead of the passthrough lane, which flushes
+			// after every coalescer lane and would emit it behind data
+			// that arrived later. An unresolved walk (truncated or crafted
+			// over-long chain) yields a non-transport number and falls to
+			// the pt lane below.
+			proto, _, _ = iputil.IPv6FindUpperProtocol(pkt)
+		}
 	default:
 		return m.pt.Commit(pkt)
 	}

overlay/batch/multi_coalesce_test.go

@@ -2,6 +2,7 @@ package batch
 
 import (
 	"bytes"
+	"encoding/binary"
 	"io"
 	"testing"
 
@@ -131,6 +132,72 @@ func TestMultiCoalescerNoOffloadsIsPassthrough(t *testing.T) {
 	}
 }
 
+// buildUDPv6Fragment builds an IPv6 packet whose extension chain is a
+// single fragment header (NH=44) naming UDP as the terminal protocol —
+// a first fragment (offset 0, MF set) carrying the UDP header and a
+// partial payload.
+func buildUDPv6Fragment(sport, dport uint16, payload []byte) []byte {
+	const ipHdrLen = 40
+	const fragHdrLen = 8
+	const udpHdrLen = 8
+	total := ipHdrLen + fragHdrLen + udpHdrLen + len(payload)
+	pkt := make([]byte, total)
+
+	pkt[0] = 0x60
+	binary.BigEndian.PutUint16(pkt[4:6], uint16(total-ipHdrLen))
+	pkt[6] = 44 // fragment extension header
+	pkt[7] = 64
+	pkt[8] = 0xfe
+	pkt[9] = 0x80
+	pkt[23] = 1
+	pkt[24] = 0xfe
+	pkt[25] = 0x80
+	pkt[39] = 2
+
+	pkt[40] = ipProtoUDP                              // fragment's next header
+	binary.BigEndian.PutUint16(pkt[42:44], 0x0001)    // offset 0, MF set
+	binary.BigEndian.PutUint32(pkt[44:48], 0x1badf00) // identification
+
+	binary.BigEndian.PutUint16(pkt[48:50], sport)
+	binary.BigEndian.PutUint16(pkt[50:52], dport)
+	binary.BigEndian.PutUint16(pkt[52:54], uint16(udpHdrLen+len(payload)))
+	copy(pkt[56:], payload)
+	return pkt
+}
+
+// TestMultiCoalescerIPv6FragmentStaysInLane locks in extension-header
+// routing: a fragment whose chain terminates in UDP must ride the UDP lane
+// as an in-lane passthrough — emitted ahead of later same-flow datagrams —
+// not the passthrough lane, which flushes after every coalescer lane and
+// would reorder it behind data that arrived after it.
+func TestMultiCoalescerIPv6FragmentStaysInLane(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	m := newTestMultiCoalescer(t, w)
+
+	if err := m.Commit(buildUDPv6Fragment(2000, 53, make([]byte, 512))); err != nil {
+		t.Fatal(err)
+	}
+	if err := m.Commit(buildUDPv6(2000, 53, make([]byte, 800))); err != nil {
+		t.Fatal(err)
+	}
+	if err := m.Commit(buildUDPv6(2000, 53, make([]byte, 800))); err != nil {
+		t.Fatal(err)
+	}
+	if err := m.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.writes) != 1 {
+		t.Fatalf("want the fragment as 1 plain write, got %d", len(w.writes))
+	}
+	if len(w.gsoWrites) != 1 {
+		t.Fatalf("want the two whole datagrams coalesced into 1 gso write, got %d", len(w.gsoWrites))
+	}
+	// Arrival order was fragment-then-data; same-lane routing must keep it.
+	if w.order[0] != "write" {
+		t.Fatalf("fragment must be emitted before later data (in-lane passthrough), order=%v", w.order)
+	}
+}
+
 // TestMultiCoalescerNoTSOFallsThrough mirrors the no-TSO case.
 func TestMultiCoalescerNoTSOFallsThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true, noTSO: true}

overlay/batch/tcp_coalesce.go

@@ -71,10 +71,11 @@ type TCPCoalescer struct {
 	// removed from this map when they close (PSH or short-last-segment),
 	// when a non-admissible packet for that flow arrives, or in Flush.
 	openSlots map[flowKey]*coalesceSlot
-	// lastSlot caches the most recently touched open slot. Steady-state
-	// bulk traffic is dominated by a single flow, so comparing the
-	// incoming key against the cached slot's own fk lets the hot path
-	// skip the map lookup (and the aeshash of a 38-byte key) entirely.
+	// lastSlot caches the most recently touched open slot. Bulk traffic
+	// arrives in same-flow runs (single-flow steady state, or GRO bursts
+	// under multi-flow), so comparing the incoming key against the cached
+	// slot's own fk lets the hot path skip the map lookup (and the aeshash
+	// of a 38-byte key) for the length of each run.
 	// Kept in lockstep with openSlots: nil whenever the slot it pointed
 	// at is removed/sealed.
 	lastSlot *coalesceSlot
@@ -183,21 +184,26 @@ func (c *TCPCoalescer) commitParsed(pkt []byte, info parsedTCP) error {
 	if !info.coalesceable() {
 		// TCP but not admissible (SYN/FIN/RST/URG/CWR or zero-payload).
 		// Seal this flow's open slot so later in-flow packets don't extend
-		// it and accidentally reorder past this passthrough.
-		if last := c.lastSlot; last != nil && last.fk == info.fk {
-			c.lastSlot = nil
+		// it and accidentally reorder past this passthrough. The len guard
+		// skips hashing the 38-byte key on ack-dominant queues, where the
+		// map is almost always empty.
+		if len(c.openSlots) != 0 {
+			if last := c.lastSlot; last != nil && last.fk == info.fk {
+				c.lastSlot = nil
+			}
+			delete(c.openSlots, info.fk)
 		}
-		delete(c.openSlots, info.fk)
 		c.addPassthrough(pkt)
 		return nil
 	}
 
-	// Single-flow fast path: with only one open flow the cache hits every
-	// packet, and len(openSlots)==1 lets us skip the 38-byte fk compare
-	// when there are multiple flows in flight (where the hit rate would
-	// be ~0 and the compare is pure overhead).
+	// Cached-slot fast path. Arrival isn't per-packet interleaved even with
+	// many flows: wire-side GRO delivers runs of same-flow packets
+	// (deliverSegments splits a superdatagram into up to 64), so the cache
+	// hits for the length of each run and a miss costs one fk compare
+	// before the map lookup carries the weight.
 	var open *coalesceSlot
-	if last := c.lastSlot; last != nil && len(c.openSlots) == 1 && last.fk == info.fk {
+	if last := c.lastSlot; last != nil && last.fk == info.fk {
 		open = last
 	} else {
 		open = c.openSlots[info.fk]
@@ -313,6 +319,9 @@ func (c *TCPCoalescer) canAppend(s *coalesceSlot, pkt []byte, info parsedTCP) bo
 	if (seedFlags^info.flags)&tcpFlagEce != 0 {
 		return false
 	}
+	if !s.isV6 && !ipv4CanCoalesceID(s.hdrBuf[:], pkt, s.numSeg) {
+		return false
+	}
 	if !headersMatch(s.hdrBuf[:s.hdrLen], pkt[:info.hdrLen], s.isV6, s.ipHdrLen) {
 		return false
 	}
@@ -352,6 +361,15 @@ func (c *TCPCoalescer) release(s *coalesceSlot) {
 	s.numSeg = 0
 	s.totalPay = 0
 	s.sealed = false
+	// Zero the identity fields too: addPassthrough doesn't set them, so a
+	// pooled slot reused as a passthrough must not carry a stale flow key
+	// that a future refactor could mistake for real.
+	s.fk = flowKey{}
+	s.hdrLen = 0
+	s.ipHdrLen = 0
+	s.isV6 = false
+	s.gsoSize = 0
+	s.nextSeq = 0
 	c.pool = append(c.pool, s)
 }
 
@@ -620,6 +638,13 @@ func canMergeSlots(prev, s *coalesceSlot) bool {
 	if (prevFlags^sFlags)&tcpFlagEce != 0 {
 		return false
 	}
+	// Same IPv4 ID rule as canAppend: s becomes segment prev.numSeg of the
+	// merged chain, so its seed ID must continue prev's sequence (or DF must
+	// make the IDs meaningless). s's own interior segments already passed
+	// this check against s's seed when they were appended.
+	if !prev.isV6 && !ipv4CanCoalesceID(prev.hdrBuf[:], s.hdrBuf[:], prev.numSeg) {
+		return false
+	}
 	if !headersMatch(prev.hdrBuf[:prev.hdrLen], s.hdrBuf[:s.hdrLen], prev.isV6, prev.ipHdrLen) {
 		return false
 	}

overlay/batch/tcp_coalesce_bench_test.go

@@ -55,6 +55,30 @@ func buildTCPv4Interleaved(nFlows, perFlow, payloadLen int) [][]byte {
 	return pkts
 }
 
+// buildTCPv4RunInterleaved returns nFlows*perFlow packets delivered in
+// runs of runLen per flow — the arrival pattern wire-side GRO actually
+// produces (deliverSegments splits each superdatagram into up to 64
+// same-flow packets back to back). Contrast with buildTCPv4Interleaved's
+// per-packet round-robin, the adversarial worst case for a last-slot cache.
+func buildTCPv4RunInterleaved(nFlows, perFlow, runLen, payloadLen int) [][]byte {
+	pay := make([]byte, payloadLen)
+	seqs := make([]uint32, nFlows)
+	for i := range seqs {
+		seqs[i] = uint32(1000 + i*1000000)
+	}
+	pkts := make([][]byte, 0, nFlows*perFlow)
+	for done := 0; done < perFlow; done += runLen {
+		for f := range nFlows {
+			sport := uint16(10000 + f)
+			for range runLen {
+				pkts = append(pkts, buildTCPv4Ports(sport, 2000, seqs[f], tcpAck, pay))
+				seqs[f] += uint32(payloadLen)
+			}
+		}
+	}
+	return pkts
+}
+
 // buildICMPv4 returns a minimal non-TCP packet that takes the passthrough
 // branch in Commit.
 func buildICMPv4() []byte {
@@ -112,6 +136,15 @@ func BenchmarkCommitInterleaved16(b *testing.B) {
 	runCommitBench(b, pkts, len(pkts))
 }
 
+// BenchmarkCommitRunInterleaved4 is 4 concurrent flows arriving in
+// GRO-burst runs of 16 — the realistic multi-flow pattern. A last-slot
+// cache hits for the length of each run; the per-packet round-robin
+// benches above are its worst case.
+func BenchmarkCommitRunInterleaved4(b *testing.B) {
+	pkts := buildTCPv4RunInterleaved(4, tcpCoalesceMaxSegs, 16, 1200)
+	runCommitBench(b, pkts, len(pkts))
+}
+
 // BenchmarkCommitPassthrough exercises the non-TCP branch: parseTCPBase
 // bails early and addPassthrough is the only work.
 func BenchmarkCommitPassthrough(b *testing.B) {

overlay/batch/tcp_coalesce_test.go

@@ -1,6 +1,7 @@
 package batch
 
 import (
+	"bytes"
 	"encoding/binary"
 	"io"
 	"testing"
@@ -21,6 +22,9 @@ type fakeTunWriter struct {
 	noUSO      bool
 	writes     [][]byte
 	gsoWrites  []fakeGSOWrite
+	// order records the interleaving of Write ("write") and WriteGSO ("gso")
+	// calls for tests that assert cross-call emission order.
+	order []string
 }
 
 // fakeGSOWrite captures one WriteGSO call. hdr is the concatenation of the
@@ -56,6 +60,7 @@ func (w *fakeTunWriter) Write(p []byte) (int, error) {
 	buf := make([]byte, len(p))
 	copy(buf, p)
 	w.writes = append(w.writes, buf)
+	w.order = append(w.order, "write")
 	return len(p), nil
 }
 
@@ -81,6 +86,7 @@ func (w *fakeTunWriter) WriteGSO(hdr []byte, transportHdr []byte, pays [][]byte,
 		isV6:      isV6,
 		csumStart: uint16(len(hdr)),
 	})
+	w.order = append(w.order, "gso")
 	return nil
 }
 
@@ -132,6 +138,18 @@ const (
 	tcpAckPsh = tcpAck | tcpPsh
 )
 
+// setIPv4ID stamps an IPv4 ID and DF state onto a builder packet. The
+// builders default to DF=1/ID=0 (an atomic datagram); the ID-admission
+// tests use this to fabricate non-atomic (DF=0) senders.
+func setIPv4ID(pkt []byte, id uint16, df bool) {
+	binary.BigEndian.PutUint16(pkt[4:6], id)
+	var flags uint16
+	if df {
+		flags = 0x4000
+	}
+	binary.BigEndian.PutUint16(pkt[6:8], flags)
+}
+
 // newTestTCPCoalescer builds a coalescer over w and fails the test if w can't
 // do TSO. Every test but TestNewTCPCoalescerRefusesWhenGSOUnavailable wants the
 // GSO path, and the constructor now hands back a nil coalescer otherwise.
@@ -1213,3 +1231,131 @@ func TestCoalescerMergePreservesRealPSH(t *testing.T) {
 		t.Errorf("merged header flags=%#x: real PSH lost in merge", flags)
 	}
 }
+
+// TestCoalescerSeqWrapAroundSortsAndMerges pins the serial-number
+// arithmetic through the sort-and-merge path: a chain that crosses the
+// 2^32 seq wrap must still sort pre-wrap before post-wrap and merge into
+// one superpacket when contiguous.
+func TestCoalescerSeqWrapAroundSortsAndMerges(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestTCPCoalescer(t, w)
+
+	payA := bytes.Repeat([]byte{'A'}, 32)
+	payB := bytes.Repeat([]byte{'B'}, 32)
+	seqA := uint32(0xffffffe0) // 32 before the wrap: nextSeq lands exactly on 0
+
+	// The post-wrap segment arrives first — wire reorder across a batch
+	// boundary, the case reorderForFlush exists for.
+	if err := c.Commit(buildTCPv4(0, tcpAck, payB)); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Commit(buildTCPv4(seqA, tcpAck, payA)); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 1 {
+		t.Fatalf("want 1 merged gso write across the wrap, got %d (plain=%d)", len(w.gsoWrites), len(w.writes))
+	}
+	g := w.gsoWrites[0]
+	const ipHdrLen = 20
+	if seedSeq := binary.BigEndian.Uint32(g.hdr[ipHdrLen+4 : ipHdrLen+8]); seedSeq != seqA {
+		t.Errorf("merged seed seq=%#x want %#x (pre-wrap segment first)", seedSeq, seqA)
+	}
+	if len(g.pays) != 2 {
+		t.Fatalf("merged segs=%d want 2", len(g.pays))
+	}
+	if !bytes.Equal(g.pays[0], payA) || !bytes.Equal(g.pays[1], payB) {
+		t.Errorf("payload order wrong across the wrap: got %q then %q", g.pays[0][:1], g.pays[1][:1])
+	}
+}
+
+// TestCoalescerNonAtomicSequentialIDsCoalesce: with DF clear, coalescing
+// is allowed when the IPv4 IDs already run seed+1 per segment — kernel
+// TSO's re-stamp then reproduces the originals exactly (the kernel GRO
+// admission rule).
+func TestCoalescerNonAtomicSequentialIDsCoalesce(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestTCPCoalescer(t, w)
+	pay := make([]byte, 1200)
+
+	seq := uint32(1000)
+	for i := range 3 {
+		pkt := buildTCPv4(seq, tcpAck, pay)
+		setIPv4ID(pkt, uint16(700+i), false)
+		if err := c.Commit(pkt); err != nil {
+			t.Fatal(err)
+		}
+		seq += uint32(len(pay))
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 1 || len(w.gsoWrites[0].pays) != 3 {
+		t.Fatalf("sequential-ID DF=0 chain must coalesce: gso=%d", len(w.gsoWrites))
+	}
+	if id := binary.BigEndian.Uint16(w.gsoWrites[0].hdr[4:6]); id != 700 {
+		t.Errorf("superpacket seed ID=%d want 700", id)
+	}
+}
+
+// TestCoalescerNonAtomicIDGapDoesNotCoalesce: with DF clear and an ID jump
+// mid-flow, neither the append path nor the flush-time merge may combine
+// the segments — TSO would re-stamp seed+n and rewrite the second
+// packet's ID, which is meaningful on non-atomic datagrams.
+func TestCoalescerNonAtomicIDGapDoesNotCoalesce(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestTCPCoalescer(t, w)
+	pay := make([]byte, 1200)
+
+	p1 := buildTCPv4(1000, tcpAck, pay)
+	setIPv4ID(p1, 700, false)
+	p2 := buildTCPv4(1000+uint32(len(pay)), tcpAck, pay)
+	setIPv4ID(p2, 900, false)
+
+	if err := c.Commit(p1); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Commit(p2); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 2 {
+		t.Fatalf("ID gap on DF=0 must not coalesce (append or merge): gso=%d", len(w.gsoWrites))
+	}
+	for i, want := range []uint16{700, 900} {
+		if id := binary.BigEndian.Uint16(w.gsoWrites[i].hdr[4:6]); id != want {
+			t.Errorf("write %d: ID=%d want %d (must be preserved)", i, id, want)
+		}
+	}
+}
+
+// TestCoalescerAtomicRandomIDsCoalesce guards the other direction: DF set
+// makes the datagram atomic (RFC 6864), so arbitrary IDs must not block
+// coalescing.
+func TestCoalescerAtomicRandomIDsCoalesce(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestTCPCoalescer(t, w)
+	pay := make([]byte, 1200)
+
+	p1 := buildTCPv4(1000, tcpAck, pay)
+	setIPv4ID(p1, 0x1234, true)
+	p2 := buildTCPv4(1000+uint32(len(pay)), tcpAck, pay)
+	setIPv4ID(p2, 0x0007, true)
+
+	if err := c.Commit(p1); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Commit(p2); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 1 || len(w.gsoWrites[0].pays) != 2 {
+		t.Fatalf("DF=1 chain with arbitrary IDs must coalesce: gso=%d", len(w.gsoWrites))
+	}
+}

overlay/batch/udp_coalesce.go

@@ -51,7 +51,14 @@ type UDPCoalescer struct {
 	w         tio.GSOWriter
 	slots     []*udpSlot
 	openSlots map[flowKey]*udpSlot
-	pool      []*udpSlot
+	// lastSlot caches the most recently touched open slot; see the
+	// TCPCoalescer field of the same name. Single-flow QUIC bulk is the
+	// dominant USO workload, and multi-flow arrival comes in GRO runs, so
+	// the fk compare beats the map's 38-byte key hash on most packets.
+	// Kept in lockstep with openSlots: nil whenever the slot it pointed at
+	// is removed/sealed.
+	lastSlot *udpSlot
+	pool     []*udpSlot
 }
 
 func NewUDPCoalescer(w io.Writer) *UDPCoalescer {
@@ -119,22 +126,41 @@ func (c *UDPCoalescer) Commit(pkt []byte) error {
 // avoid re-walking the IP/UDP header.
 func (c *UDPCoalescer) commitParsed(pkt []byte, info parsedUDP) error {
 	// A zero-length UDP datagram (UDP `length` == 8) is legal and must still
-	// reach the TUN, but it can't be coalesced.
+	// reach the TUN, but it can't be coalesced. The len guard skips hashing
+	// the key when no flow is open.
 	if info.payLen == 0 {
-		delete(c.openSlots, info.fk)
+		if len(c.openSlots) != 0 {
+			if last := c.lastSlot; last != nil && last.fk == info.fk {
+				c.lastSlot = nil
+			}
+			delete(c.openSlots, info.fk)
+		}
 		c.addPassthrough(pkt)
 		return nil
 	}
-	if open := c.openSlots[info.fk]; open != nil {
+	// Cached-slot fast path; see the TCPCoalescer equivalent.
+	var open *udpSlot
+	if last := c.lastSlot; last != nil && last.fk == info.fk {
+		open = last
+	} else {
+		open = c.openSlots[info.fk]
+	}
+	if open != nil {
 		if c.canAppend(open, pkt, info) {
 			c.appendPayload(open, pkt, info)
 			if open.sealed {
 				delete(c.openSlots, info.fk)
+				c.lastSlot = nil
+			} else {
+				c.lastSlot = open
 			}
 			return nil
 		}
 		// Can't extend. Seal it and fall through to seed a fresh slot.
 		delete(c.openSlots, info.fk)
+		if c.lastSlot == open {
+			c.lastSlot = nil
+		}
 	}
 	c.seed(pkt, info)
 	return nil
@@ -157,6 +183,7 @@ func (c *UDPCoalescer) Flush() error {
 	clear(c.slots)
 	c.slots = c.slots[:0]
 	clear(c.openSlots)
+	c.lastSlot = nil
 	return first
 }
 
@@ -187,6 +214,7 @@ func (c *UDPCoalescer) seed(pkt []byte, info parsedUDP) {
 	s.payIovs = append(s.payIovs[:0], pkt[info.hdrLen:info.hdrLen+info.payLen])
 	c.slots = append(c.slots, s)
 	c.openSlots[info.fk] = s
+	c.lastSlot = s
 }
 
 // canAppend reports whether info's packet extends the slot's seed.
@@ -208,6 +236,9 @@ func (c *UDPCoalescer) canAppend(s *udpSlot, pkt []byte, info parsedUDP) bool {
 	if s.hdrLen+s.totalPay+info.payLen > udpCoalesceBufSize {
 		return false
 	}
+	if !s.isV6 && !ipv4CanCoalesceID(s.hdrBuf[:], pkt, s.numSeg) {
+		return false
+	}
 	if !udpHeadersMatch(s.hdrBuf[:s.hdrLen], pkt[:info.hdrLen], s.isV6, s.ipHdrLen) {
 		return false
 	}
@@ -242,6 +273,12 @@ func (c *UDPCoalescer) release(s *udpSlot) {
 	s.numSeg = 0
 	s.totalPay = 0
 	s.sealed = false
+	// Zero the identity fields too; see TCPCoalescer.release.
+	s.fk = flowKey{}
+	s.hdrLen = 0
+	s.ipHdrLen = 0
+	s.isV6 = false
+	s.gsoSize = 0
 	c.pool = append(c.pool, s)
 }
 

overlay/batch/udp_coalesce_bench_test.go

@@ -0,0 +1,72 @@
+package batch
+
+import (
+	"testing"
+)
+
+// buildUDPv4BulkFlow returns n equal-size datagrams on one flow — the
+// steady state for single-flow QUIC bulk, the workload USO exists for.
+func buildUDPv4BulkFlow(n, payloadLen int) [][]byte {
+	pay := make([]byte, payloadLen)
+	pkts := make([][]byte, n)
+	for i := range pkts {
+		pkts[i] = buildUDPv4(40000, 443, pay)
+	}
+	return pkts
+}
+
+// buildUDPv4RunInterleaved mirrors buildTCPv4RunInterleaved: nFlows*perFlow
+// datagrams arriving in GRO-burst runs of runLen per flow.
+func buildUDPv4RunInterleaved(nFlows, perFlow, runLen, payloadLen int) [][]byte {
+	pay := make([]byte, payloadLen)
+	pkts := make([][]byte, 0, nFlows*perFlow)
+	for done := 0; done < perFlow; done += runLen {
+		for f := range nFlows {
+			sport := uint16(40000 + f)
+			for range runLen {
+				pkts = append(pkts, buildUDPv4(sport, 443, pay))
+			}
+		}
+	}
+	return pkts
+}
+
+// runUDPCommitBench drives UDPCoalescer.Commit over pkts batchSize at a
+// time, flushing between batches, and reports per-packet cost.
+func runUDPCommitBench(b *testing.B, pkts [][]byte, batchSize int) {
+	b.Helper()
+	c := newTestUDPCoalescer(b, nopTunWriter{})
+	b.ReportAllocs()
+	b.SetBytes(int64(len(pkts[0])))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		pkt := pkts[i%len(pkts)]
+		if err := c.Commit(pkt); err != nil {
+			b.Fatal(err)
+		}
+		if (i+1)%batchSize == 0 {
+			if err := c.Flush(); err != nil {
+				b.Fatal(err)
+			}
+		}
+	}
+	_ = c.Flush()
+}
+
+// BenchmarkUDPCommitSingleFlow is the single-flow bulk steady state.
+func BenchmarkUDPCommitSingleFlow(b *testing.B) {
+	pkts := buildUDPv4BulkFlow(udpCoalesceMaxSegs, 1200)
+	runUDPCommitBench(b, pkts, udpCoalesceMaxSegs)
+}
+
+// BenchmarkUDPCommitInterleaved4 is the adversarial per-packet round-robin.
+func BenchmarkUDPCommitInterleaved4(b *testing.B) {
+	pkts := buildUDPv4RunInterleaved(4, udpCoalesceMaxSegs, 1, 1200)
+	runUDPCommitBench(b, pkts, len(pkts))
+}
+
+// BenchmarkUDPCommitRunInterleaved4 is 4 flows in GRO-burst runs of 16.
+func BenchmarkUDPCommitRunInterleaved4(b *testing.B) {
+	pkts := buildUDPv4RunInterleaved(4, udpCoalesceMaxSegs, 16, 1200)
+	runUDPCommitBench(b, pkts, len(pkts))
+}

overlay/batch/udp_coalesce_test.go

@@ -459,3 +459,57 @@ func TestUDPCoalescerIPv4WithOptionsPassesThrough(t *testing.T) {
 		t.Fatalf("ipv4-with-options must pass through plain, got writes=%d gso=%d", len(w.writes), len(w.gsoWrites))
 	}
 }
+
+// TestUDPCoalescerNonAtomicSequentialIDsCoalesce mirrors the TCP rule: DF
+// clear is fine as long as the IDs already run seed+1 per datagram, so
+// kernel USO's re-stamp reproduces them.
+func TestUDPCoalescerNonAtomicSequentialIDsCoalesce(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestUDPCoalescer(t, w)
+	pay := make([]byte, 1200)
+
+	for i := range 2 {
+		pkt := buildUDPv4(40000, 443, pay)
+		setIPv4ID(pkt, uint16(40+i), false)
+		if err := c.Commit(pkt); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 1 || len(w.gsoWrites[0].pays) != 2 {
+		t.Fatalf("sequential-ID DF=0 datagrams must coalesce: gso=%d", len(w.gsoWrites))
+	}
+}
+
+// TestUDPCoalescerNonAtomicIDGapReseeds: an ID jump on a DF=0 flow breaks
+// the chain; each datagram must keep its own (meaningful) ID.
+func TestUDPCoalescerNonAtomicIDGapReseeds(t *testing.T) {
+	w := &fakeTunWriter{gsoEnabled: true}
+	c := newTestUDPCoalescer(t, w)
+	pay := make([]byte, 1200)
+
+	p1 := buildUDPv4(40000, 443, pay)
+	setIPv4ID(p1, 40, false)
+	p2 := buildUDPv4(40000, 443, pay)
+	setIPv4ID(p2, 50, false)
+
+	if err := c.Commit(p1); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Commit(p2); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(w.gsoWrites) != 2 {
+		t.Fatalf("ID gap on DF=0 must reseed: gso=%d", len(w.gsoWrites))
+	}
+	for i, want := range []uint16{40, 50} {
+		if id := binary.BigEndian.Uint16(w.gsoWrites[i].hdr[4:6]); id != want {
+			t.Errorf("write %d: ID=%d want %d (must be preserved)", i, id, want)
+		}
+	}
+}