JackDoan 3 месяцев назад
Родитель
Сommit
f2b04d2e9d

+ 10 - 7
inside.go

@@ -279,27 +279,30 @@ func (f *Interface) rejectInside(packet []byte, out []byte, q int) {
 	}
 }
 
-func (f *Interface) rejectOutside(packet []byte, ci *ConnectionState, hostinfo *HostInfo, nb, out []byte, q int) {
+func (f *Interface) rejectOutside(packet []byte, ci *ConnectionState, hostinfo *HostInfo, nb, rejectBuf []byte, q int) {
 	if !f.firewall.InboundSendReject {
 		return
 	}
 
-	out = iputil.CreateRejectPacket(packet, out)
+	// split rejectBuf to make sure we have room to write the plaintext rejection, then encrypt it, without trampling anything
+	// we can't re-use packet, if we need to send an icmp reject, it won't be long enough.
+	half := len(rejectBuf) / 2
+	encryptBuf := rejectBuf[0:0:half] //the first half of rejectBuf's capacity, len set to 0
+	buildBuf := rejectBuf[half:]
+
+	out := iputil.CreateRejectPacket(packet, buildBuf)
 	if len(out) == 0 {
 		return
 	}
 
 	if len(out) > iputil.MaxRejectPacketSize {
 		if f.l.Enabled(context.Background(), slog.LevelInfo) {
-			f.l.Info("rejectOutside: packet too big, not sending",
-				"packet", packet,
-				"outPacket", out,
-			)
+			f.l.Info("rejectOutside: packet too big, not sending", "packet", packet, "outPacket", out)
 		}
 		return
 	}
 
-	f.sendNoMetrics(header.Message, 0, ci, hostinfo, netip.AddrPort{}, out, nb, packet, q)
+	f.sendNoMetrics(header.Message, 0, ci, hostinfo, netip.AddrPort{}, out, nb, encryptBuf, q)
 }
 
 // Handshake will attempt to initiate a tunnel with the provided vpn address. This is a no-op if the tunnel is already established or being established

+ 9 - 9
interface.go

@@ -120,9 +120,11 @@ type Interface struct {
 	ctx     context.Context
 	writers []udp.Conn
 	queues  []tio.Queue
-	// batchers is one per tun queue, wrapping queues[i].
-	// decryptToTun sends plaintext into the batch.RxBatcher;
-	// listenOut calls its Flush at the end of each UDP recvmmsg batch.
+	// batchers is one per tun queue, wrapping queues[i]. readOutsidePackets
+	// commits plaintext into the batch.RxBatcher; the plaintext is decrypted
+	// in place inside the UDP receive buffers, so listenOut must call Flush
+	// at the end of each UDP recvmmsg batch, before those buffers are
+	// reused (every udp.Conn ListenOut guarantees that ordering).
 	batchers []batch.RxBatcher
 	wg       sync.WaitGroup
 
@@ -303,11 +305,9 @@ func (f *Interface) activate() error {
 			// is on, everything else (and either lane disabled) falls
 			// through to passthrough so non-IP / non-TCP-UDP traffic still
 			// reaches the TUN.
-			arena := batch.NewArena(batch.DefaultMultiArenaCap)
-			f.batchers[i] = batch.NewMultiCoalescer(f.queues[i], f.l, arena, caps.TSO, caps.USO)
+			f.batchers[i] = batch.NewMultiCoalescer(f.queues[i], f.l, caps.TSO, caps.USO)
 		} else {
-			arena := batch.NewArena(batch.DefaultPassthroughArenaCap)
-			f.batchers[i] = batch.NewPassthrough(f.queues[i], arena.Reserve, arena.Reset)
+			f.batchers[i] = batch.NewPassthrough(f.queues[i])
 		}
 	}
 
@@ -369,10 +369,10 @@ func (f *Interface) listenOut(i int) {
 	h := &header.H{}
 	fwPacket := &firewall.Packet{}
 	nb := make([]byte, 12, 12)
+	scratch := make([]byte, mtu)
 
 	listener := func(fromUdpAddr netip.AddrPort, payload []byte, meta udp.RxMeta) {
-		plaintext := f.batchers[i].Reserve(len(payload))
-		f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, plaintext[:0], payload, h, fwPacket, lhh, nb, i, ctCache.Get(), meta)
+		f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, scratch, payload, h, fwPacket, lhh, nb, i, ctCache.Get(), meta)
 	}
 
 	flusher := func() {

+ 45 - 0
noiseutil/cipher_state_test.go

@@ -164,3 +164,48 @@ func TestCipherStateNilSafety(t *testing.T) {
 	assert.Empty(t, out)
 	assert.Equal(t, 0, cc.Overhead())
 }
+
+func TestCipherStateAESGCMInPlaceDecrypt(t *testing.T) {
+	enc, dec := buildCipherStates(t, CipherAESGCM)
+	inPlaceDecrypt(t, NewCipherStateAESGCM(enc), NewCipherStateAESGCM(dec))
+}
+
+func TestCipherStateChaChaPolyInPlaceDecrypt(t *testing.T) {
+	enc, dec := buildCipherStates(t, noise.CipherChaChaPoly)
+	inPlaceDecrypt(t, NewCipherStateChaChaPoly(enc), NewCipherStateChaChaPoly(dec))
+}
+
+func inPlaceDecrypt(t *testing.T, enc, dec CipherState) {
+	t.Helper()
+	const hdrLen = 16
+	plaintext := []byte("in-place decrypt should replace the ciphertext bytes")
+	nb := make([]byte, 12)
+
+	// packet = [16-byte header | ciphertext+tag], like a nebula Message.
+	packet := make([]byte, hdrLen, hdrLen+len(plaintext)+enc.Overhead())
+	for i := range packet {
+		packet[i] = byte(i)
+	}
+	packet, err := enc.EncryptDanger(packet, packet[:hdrLen], plaintext, 1, nb)
+	require.NoError(t, err)
+
+	// Simulate a GRO row: [packet | next segment]. A failed auth on packet
+	// may zero packet's plaintext region but must not touch the header, the
+	// tag, or the neighboring segment.
+	neighbor := []byte("next coalesced segment, must stay intact")
+	row := append(append([]byte(nil), packet...), neighbor...)
+	tampered := row[:len(packet)]
+	tampered[hdrLen] ^= 0x01
+	_, err = dec.DecryptDanger(tampered[hdrLen:hdrLen], tampered[:hdrLen], tampered[hdrLen:], 1, nb)
+	require.Error(t, err)
+	assert.Equal(t, packet[:hdrLen], tampered[:hdrLen], "failed auth must not touch the header")
+	assert.Equal(t, packet[len(packet)-dec.Overhead():], tampered[len(tampered)-dec.Overhead():],
+		"failed auth must not touch the tag")
+	assert.Equal(t, neighbor, row[len(packet):], "failed auth must not touch the next segment")
+
+	out, err := dec.DecryptDanger(packet[hdrLen:hdrLen], packet[:hdrLen], packet[hdrLen:], 1, nb)
+	require.NoError(t, err)
+	assert.Equal(t, plaintext, out)
+	// The plaintext must be IN the packet buffer, not a fresh allocation.
+	assert.Equal(t, &packet[hdrLen], &out[0], "plaintext must alias the packet buffer")
+}

+ 36 - 23
outside.go

@@ -23,7 +23,10 @@ const (
 
 var ErrOutOfWindow = errors.New("out of window packet")
 
-func (f *Interface) readOutsidePackets(via ViaSender, out []byte, packet []byte, h *header.H, fwPacket *firewall.Packet, lhf *LightHouseHandler, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
+// readOutsidePackets processes one received underlay packet.
+// Message payloads are decrypted IN PLACE, so packet must stay untouched
+// by the caller until the batcher for queue q has been flushed
+func (f *Interface) readOutsidePackets(via ViaSender, scratch []byte, packet []byte, h *header.H, fwPacket *firewall.Packet, lhf *LightHouseHandler, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
 	err := h.Parse(packet)
 	if err != nil {
 		// Hole punch packets are 0 or 1 byte big, so lets ignore printing those errors
@@ -111,11 +114,11 @@ func (f *Interface) readOutsidePackets(via ViaSender, out []byte, packet []byte,
 
 	// Relay packets are special
 	if isMessageRelay {
-		f.handleOutsideRelayPacket(hostinfo, via, out, packet, h, fwPacket, lhf, nb, q, localCache, meta)
+		f.handleOutsideRelayPacket(hostinfo, via, scratch, packet, h, fwPacket, lhf, nb, q, localCache, meta)
 		return
 	}
 
-	out, err = f.decrypt(hostinfo, h.MessageCounter, out, packet, h, nb)
+	out, err := f.decrypt(hostinfo, h.MessageCounter, packet, nb)
 	if err != nil {
 		if f.l.Enabled(context.Background(), slog.LevelDebug) {
 			hostinfo.logger(f.l).Debug("Failed to decrypt packet",
@@ -135,7 +138,7 @@ func (f *Interface) readOutsidePackets(via ViaSender, out []byte, packet []byte,
 	case header.Message:
 		switch h.Subtype {
 		case header.MessageNone:
-			f.handleOutsideMessagePacket(hostinfo, out, packet, fwPacket, nb, q, localCache, meta)
+			f.handleOutsideMessagePacket(hostinfo, out, scratch, fwPacket, nb, q, localCache, meta)
 		default:
 			hostinfo.logger(f.l).Error("IsValidSubType was true, but unexpected message subtype seen", "from", via, "header", h)
 			return
@@ -150,8 +153,15 @@ func (f *Interface) readOutsidePackets(via ViaSender, out []byte, packet []byte,
 		case header.TestReply:
 			// No-op, useful for the Roaming and connectionManager side-effects above
 		case header.TestRequest:
-			//recycle the input packet ciphertext as our output buffer
-			f.send(header.Test, header.TestReply, ci, hostinfo, out, nb, packet)
+			const maxCipherOverhead = 16 //todo we use this too often, needs a real importable const
+			const maxOverhead = header.Len + header.Len + maxCipherOverhead + maxCipherOverhead
+			if maxOverhead+len(out) <= len(scratch) {
+				f.send(header.Test, header.TestReply, ci, hostinfo, out, nb, scratch[:0])
+				return
+			} else if f.l.Enabled(context.Background(), slog.LevelDebug) {
+				hostinfo.logger(f.l).Debug("dropping oversized test request", "payloadLen", len(out), "from", via)
+				return
+			}
 		default:
 			hostinfo.logger(f.l).Error("IsValidSubType was true, but unexpected test subtype seen", "from", via, "header", h)
 			return
@@ -169,7 +179,7 @@ func (f *Interface) readOutsidePackets(via ViaSender, out []byte, packet []byte,
 	}
 }
 
-func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender, out []byte, packet []byte, h *header.H, fwPacket *firewall.Packet, lhf *LightHouseHandler, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
+func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender, scratch []byte, packet []byte, h *header.H, fwPacket *firewall.Packet, lhf *LightHouseHandler, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
 	// The entire body is sent as AD, not encrypted.
 	// The packet consists of a 16-byte parsed Nebula header, Associated Data-protected payload, and a trailing 16-byte AEAD signature value.
 	// The packet is guaranteed to be at least 16 bytes at this point, b/c it got past the h.Parse() call above. If it's
@@ -177,9 +187,7 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 	// which will gracefully fail in the DecryptDanger call.
 	signedPayload := packet[:len(packet)-hostinfo.ConnectionState.dKey.Overhead()]
 	signatureValue := packet[len(packet)-hostinfo.ConnectionState.dKey.Overhead():]
-	var err error
-	out, err = hostinfo.ConnectionState.dKey.DecryptDanger(out, signedPayload, signatureValue, h.MessageCounter, nb)
-	if err != nil {
+	if _, err := hostinfo.ConnectionState.dKey.DecryptDanger(nil, signedPayload, signatureValue, h.MessageCounter, nb); err != nil {
 		return
 	}
 	// Advance the replay window now that the frame is authenticated
@@ -217,7 +225,7 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 			relay:     relay,
 			IsRelayed: true,
 		}
-		f.readOutsidePackets(via, out[:0], signedPayload, h, fwPacket, lhf, nb, q, localCache, meta)
+		f.readOutsidePackets(via, scratch, signedPayload, h, fwPacket, lhf, nb, q, localCache, meta)
 	case ForwardingType:
 		// Find the target HostInfo relay object
 		targetHI, targetRelay, err := f.hostMap.QueryVpnAddrsRelayFor(hostinfo.vpnAddrs, relay.PeerAddr)
@@ -234,9 +242,11 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 		if targetRelay.State == Established {
 			switch targetRelay.Type {
 			case ForwardingType:
-				// Forward this packet through the relay tunnel
-				// Find the target HostInfo //todo it would potentially be nice to batch these
-				f.SendVia(targetHI, targetRelay, signedPayload, nb, out, false)
+				// Forward this packet through the relay tunnel, rebuilding it in place.
+				// Encode overwrites the old outer header, and the new AEAD tag lands where the old one was
+				fwdBuf := packet[:0:len(packet)] // Cap to len(packet) to protect memory from a larger parent buffer
+				//todo it would potentially be nice to batch these
+				f.SendVia(targetHI, targetRelay, signedPayload, nb, fwdBuf, true)
 			case TerminalType:
 				hostinfo.logger(f.l).Error("Unexpected Relay Type of Terminal")
 				return
@@ -503,9 +513,16 @@ func parseV4(data []byte, incoming bool, fp *firewall.Packet) error {
 	return nil
 }
 
-func (f *Interface) decrypt(hostinfo *HostInfo, mc uint64, out []byte, packet []byte, h *header.H, nb []byte) ([]byte, error) {
-	var err error
-	out, err = hostinfo.ConnectionState.dKey.DecryptDanger(out, packet[:header.Len], packet[header.Len:], mc, nb)
+// decrypt authenticates and decrypts a Message packet IN PLACE: the dst is
+// packet[header.Len:header.Len], the exact-alias append pattern
+// (ciphertext[:0]) that crypto/cipher.AEAD.Open documents, so the plaintext
+// lands where the ciphertext sat and no separate plaintext buffer exists.
+// On a failed auth the AEAD zeroes the would-be plaintext region (both
+// AES-GCM and ChaCha20-Poly1305 do) but never writes outside it, so the
+// other segments of a shared GRO receive row stay intact; nothing reads a
+// packet after its decrypt fails. The returned slice aliases packet.
+func (f *Interface) decrypt(hostinfo *HostInfo, mc uint64, packet []byte, nb []byte) ([]byte, error) {
+	out, err := hostinfo.ConnectionState.dKey.DecryptDanger(packet[header.Len:header.Len], packet[:header.Len], packet[header.Len:], mc, nb)
 	if err != nil {
 		return nil, err
 	}
@@ -580,7 +597,7 @@ func applyOuterECN(pkt []byte, outerECN byte, hostinfo *HostInfo, l *slog.Logger
 	}
 }
 
-func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, out []byte, packet []byte, fwPacket *firewall.Packet, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
+func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, out []byte, scratch []byte, fwPacket *firewall.Packet, nb []byte, q int, localCache firewall.ConntrackCache, meta udp.RxMeta) {
 	// RFC 6040 normal-mode combine: fold any outer CE mark stamped by the
 	// underlay into the inner header before firewall + TUN write. Other
 	// outer codepoints are advisory only — we keep the inner unchanged.
@@ -599,11 +616,7 @@ func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, out []byte, p
 
 	dropReason := f.firewall.Drop(*fwPacket, true, hostinfo, f.pki.GetCAPool(), localCache)
 	if dropReason != nil {
-		// NOTE: We give `packet` as the `out` here since we already decrypted from it and we don't need it anymore
-		// This gives us a buffer to build the reject packet in. With UDP GRO this is a single segment of a shared
-		// recvmmsg row whose capacity runs to the end of the whole row, so cap it to its own length (cap==len) to
-		// keep the reject builder from writing past this segment into the next, not-yet-processed coalesced segment.
-		f.rejectOutside(out, hostinfo.ConnectionState, hostinfo, nb, packet[:len(packet):len(packet)], q)
+		f.rejectOutside(out, hostinfo.ConnectionState, hostinfo, nb, scratch, q)
 		if f.l.Enabled(context.Background(), slog.LevelDebug) {
 			hostinfo.logger(f.l).Debug("dropping inbound packet",
 				"fwPacket", fwPacket,

+ 2 - 4
overlay/batch/batch.go

@@ -3,13 +3,11 @@ package batch
 import "net/netip"
 
 type RxBatcher interface {
-	// Reserve creates a pkt to borrow
-	Reserve(sz int) []byte
-	// Commit borrows pkt. The caller must keep pkt valid until the next Flush
+	// Commit commits pkt to be flushed by the batch. The caller must keep pkt valid until the next Flush, and not re-use it.
 	Commit(pkt []byte) error
 	// Flush emits every queued packet in arrival order.
 	// Returns the first error observed; keeps draining so one bad packet doesn't hold up the rest.
-	// After Flush returns, borrowed payload slices may be recycled.
+	// After Flush returns, committed payload slices may be recycled.
 	Flush() error
 }
 

+ 0 - 7
overlay/batch/coalesce_core.go

@@ -172,10 +172,3 @@ func (a *Arena) Reserve(sz int) []byte {
 func (a *Arena) Reset() {
 	a.buf = a.buf[:0]
 }
-
-// Reserver hands out an sz-byte slice valid until its Resetter runs.
-type Reserver func(sz int) []byte
-
-// Resetter clears all reservations held by a Reserver. Only the arena's
-// owner holds one; lanes inside a MultiCoalescer get nil.
-type Resetter func()

+ 7 - 28
overlay/batch/multi_coalesce.go

@@ -7,8 +7,7 @@ import (
 )
 
 // MultiCoalescer fans plaintext packets out to lane-specific batchers based
-// on the IP/L4 protocol of the packet, sharing a single Reserve arena
-// across lanes so the caller's allocation pattern is unchanged.
+// on the IP/L4 protocol of the packet.
 //
 // Lanes are processed independently: the TCP coalescer only sees TCP, the
 // UDP coalescer only sees UDP, and the passthrough lane handles everything
@@ -19,7 +18,7 @@ import (
 // This is acceptable because the carrier-side recvmmsg path already
 // stable-sorts by (peer, message counter) before delivering plaintext
 // here, so replay-window invariants are unaffected, and apps observe
-// correct per-flow ordering — which is all the IP layer guarantees anyway.
+// correct per-flow ordering; which is all the IP layer guarantees anyway.
 // Do not "fix" this by interleaving lane outputs at flush time; that
 // negates the entire point of coalescing (each lane needs to see runs of
 // adjacent same-flow packets to coalesce them).
@@ -27,43 +26,26 @@ type MultiCoalescer struct {
 	tcp *TCPCoalescer
 	udp *UDPCoalescer
 	pt  *Passthrough
-	// arena is owned by the Multi: lanes get only its Reserve (nil Resetter)
-	// and Flush resets it exactly once after every lane has drained.
-	arena *Arena
 }
 
-// DefaultMultiArenaCap is the recommended arena capacity for a Multi-lane
-// batcher: 64 slots × 65535 bytes ≈ 4 MiB, enough to hold one recvmmsg
-// burst worth of MTU-sized packets without the arena growing.
-const DefaultMultiArenaCap = initialSlots * 65535
-
 // NewMultiCoalescer builds a multi-lane batcher. tcpEnabled lets the caller
 // opt out of TCP coalescing (e.g. when the queue can't do TSO); udpEnabled
 // likewise gates UDP coalescing (only enable when USO was negotiated).
 // Either lane disabled redirects its traffic into the passthrough lane.
-// arena is the single backing slab shared across every lane; the caller
-// pre-sizes it via NewArena so the hot path never allocates.
-func NewMultiCoalescer(w io.Writer, l *slog.Logger, arena *Arena, tcpEnabled, udpEnabled bool) *MultiCoalescer {
+func NewMultiCoalescer(w io.Writer, l *slog.Logger, tcpEnabled, udpEnabled bool) *MultiCoalescer {
 	m := &MultiCoalescer{
-		pt:    NewPassthrough(w, arena.Reserve, nil),
-		arena: arena,
+		pt: NewPassthrough(w),
 	}
 	if tcpEnabled {
-		m.tcp = NewTCPCoalescer(w, l, arena.Reserve, nil)
+		m.tcp = NewTCPCoalescer(w, l)
 	}
 	if udpEnabled {
-		m.udp = NewUDPCoalescer(w, arena.Reserve, nil)
+		m.udp = NewUDPCoalescer(w)
 	}
 	return m
 }
 
-func (m *MultiCoalescer) Reserve(sz int) []byte {
-	return m.arena.Reserve(sz)
-}
-
-// Commit dispatches pkt to the appropriate lane based on IP version + L4
-// proto. Borrowed slice contract is identical to the single-lane batchers,
-// pkt must remain valid until the next Flush.
+// 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
@@ -110,8 +92,6 @@ func (m *MultiCoalescer) Commit(pkt []byte) error {
 	return m.pt.Commit(pkt)
 }
 
-// Flush drains every lane in a fixed order, then resets the shared arena once.
-// A lane error doesn't stop the remaining lanes; the joined errors are returned.
 func (m *MultiCoalescer) Flush() error {
 	var errs []error
 	if m.tcp != nil {
@@ -127,6 +107,5 @@ func (m *MultiCoalescer) Flush() error {
 	if err := m.pt.Flush(); err != nil {
 		errs = append(errs, err)
 	}
-	m.arena.Reset()
 	return errors.Join(errs...)
 }

+ 3 - 3
overlay/batch/multi_coalesce_test.go

@@ -11,7 +11,7 @@ import (
 // else (ICMP here) falls through to plain Write.
 func TestMultiCoalescerRoutesByProto(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	m := NewMultiCoalescer(w, test.NewLogger(), NewArena(0), true, true)
+	m := NewMultiCoalescer(w, test.NewLogger(), true, true)
 
 	tcpPay := make([]byte, 1200)
 	udpPay := make([]byte, 1200)
@@ -53,7 +53,7 @@ func TestMultiCoalescerRoutesByProto(t *testing.T) {
 // the kernel via the passthrough lane rather than being lost.
 func TestMultiCoalescerDisabledUDPFallsThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	m := NewMultiCoalescer(w, test.NewLogger(), NewArena(0), true, false) // TSO on, USO off
+	m := NewMultiCoalescer(w, test.NewLogger(), true, false) // TSO on, USO off
 
 	if err := m.Commit(buildUDPv4(1000, 53, make([]byte, 800))); err != nil {
 		t.Fatal(err)
@@ -75,7 +75,7 @@ func TestMultiCoalescerDisabledUDPFallsThrough(t *testing.T) {
 // TestMultiCoalescerDisabledTCPFallsThrough mirrors the TSO=off case.
 func TestMultiCoalescerDisabledTCPFallsThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	m := NewMultiCoalescer(w, test.NewLogger(), NewArena(0), false, true) // TSO off, USO on
+	m := NewMultiCoalescer(w, test.NewLogger(), false, true) // TSO off, USO on
 
 	pay := make([]byte, 1200)
 	if err := m.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {

+ 5 - 32
overlay/batch/passthrough.go

@@ -2,54 +2,27 @@ package batch
 
 import (
 	"io"
-
-	"github.com/slackhq/nebula/udp"
 )
 
 // Passthrough is a RxBatcher that doesn't batch anything, it just accumulates and then sends packets.
 type Passthrough struct {
-	out      io.Writer
-	slots    [][]byte
-	reserver Reserver
-	resetter Resetter
-	cursor   int
+	out   io.Writer
+	slots [][]byte
 }
 
-const passthroughBaseNumSlots = 128
-
-// DefaultPassthroughArenaCap is the recommended arena capacity for a
-// standalone Passthrough batcher: 128 slots × udp.MTU ≈ 1.1 MiB.
-const DefaultPassthroughArenaCap = passthroughBaseNumSlots * udp.MTU
-
-func NewPassthrough(w io.Writer, reserver Reserver, resetter Resetter) *Passthrough {
+func NewPassthrough(w io.Writer) *Passthrough {
 	return &Passthrough{
-		out:      w,
-		slots:    make([][]byte, 0, passthroughBaseNumSlots),
-		reserver: reserver,
-		resetter: resetter,
+		out:   w,
+		slots: make([][]byte, 0, 128),
 	}
 }
 
-func (p *Passthrough) Reserve(sz int) []byte {
-	return p.reserver(sz)
-}
-
 func (p *Passthrough) Commit(pkt []byte) error {
 	p.slots = append(p.slots, pkt)
 	return nil
 }
 
-// Flush drains every queued packet and calls the configured Resetter
 func (p *Passthrough) Flush() error {
-	firstErr := p.drain()
-	if p.resetter != nil {
-		p.resetter()
-	}
-	return firstErr
-}
-
-// drain writes out every queued packet and clears the slot list.
-func (p *Passthrough) drain() error {
 	var firstErr error
 	for _, s := range p.slots {
 		_, err := p.out.Write(s)

+ 1 - 21
overlay/batch/tcp_coalesce.go

@@ -79,19 +79,15 @@ type TCPCoalescer struct {
 	// at is removed/sealed.
 	lastSlot *coalesceSlot
 	pool     []*coalesceSlot // free list for reuse
-	reserver Reserver
-	resetter Resetter
 	l        *slog.Logger
 }
 
-func NewTCPCoalescer(w io.Writer, l *slog.Logger, reserver Reserver, resetter Resetter) *TCPCoalescer {
+func NewTCPCoalescer(w io.Writer, l *slog.Logger) *TCPCoalescer {
 	c := &TCPCoalescer{
 		plainW:    w,
 		slots:     make([]*coalesceSlot, 0, initialSlots),
 		openSlots: make(map[flowKey]*coalesceSlot, initialSlots),
 		pool:      make([]*coalesceSlot, 0, initialSlots),
-		reserver:  reserver,
-		resetter:  resetter,
 		l:         l,
 	}
 	if gw, ok := tio.SupportsGSO(w, tio.GSOProtoTCP); ok {
@@ -170,11 +166,6 @@ func (p parsedTCP) coalesceable() bool {
 	return p.payLen > 0
 }
 
-func (c *TCPCoalescer) Reserve(sz int) []byte {
-	return c.reserver(sz)
-}
-
-// Commit borrows pkt. The caller must keep pkt valid until the next Flush.
 func (c *TCPCoalescer) Commit(pkt []byte) error {
 	if c.gsoW == nil {
 		c.addPassthrough(pkt)
@@ -240,18 +231,7 @@ func (c *TCPCoalescer) commitParsed(pkt []byte, info parsedTCP) error {
 	return nil
 }
 
-// Flush emits every queued event in (per-flow) seq order.
 func (c *TCPCoalescer) Flush() error {
-	first := c.drain()
-	if c.resetter != nil {
-		c.resetter()
-	}
-	return first
-}
-
-// drain emits every queued slot (reordering/merging coalesced runs first)
-// and clears the slot state.
-func (c *TCPCoalescer) drain() error {
 	c.reorderForFlush()
 	var first error
 	for _, s := range c.slots {

+ 2 - 3
overlay/batch/tcp_coalesce_bench_test.go

@@ -71,8 +71,7 @@ func buildICMPv4() []byte {
 // between batches, and reports per-packet cost.
 func runCommitBench(b *testing.B, pkts [][]byte, batchSize int) {
 	b.Helper()
-	arena := NewArena(0)
-	c := NewTCPCoalescer(nopTunWriter{}, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(nopTunWriter{}, test.NewLogger())
 	b.ReportAllocs()
 	b.SetBytes(int64(len(pkts[0])))
 	b.ResetTimer()
@@ -141,7 +140,7 @@ func BenchmarkCommitNonCoalesceableTCP(b *testing.B) {
 // is the bench that shows the savings of skipping the lane's re-parse.
 func runMultiCommitBench(b *testing.B, pkts [][]byte, batchSize int) {
 	b.Helper()
-	m := NewMultiCoalescer(nopTunWriter{}, test.NewLogger(), NewArena(0), true, true)
+	m := NewMultiCoalescer(nopTunWriter{}, test.NewLogger(), true, true)
 	b.ReportAllocs()
 	b.SetBytes(int64(len(pkts[0])))
 	b.ResetTimer()

+ 28 - 56
overlay/batch/tcp_coalesce_test.go

@@ -128,8 +128,7 @@ const (
 
 func TestCoalescerPassthroughWhenGSOUnavailable(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: false}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pkt := buildTCPv4(1000, tcpAck, []byte("hello"))
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -148,8 +147,7 @@ func TestCoalescerPassthroughWhenGSOUnavailable(t *testing.T) {
 
 func TestCoalescerNonTCPPassthrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pkt := make([]byte, 28)
 	pkt[0] = 0x45
 	binary.BigEndian.PutUint16(pkt[2:4], 28)
@@ -169,8 +167,7 @@ func TestCoalescerNonTCPPassthrough(t *testing.T) {
 
 func TestCoalescerSeedThenFlushAlone(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pkt := buildTCPv4(1000, tcpAck, make([]byte, 1000))
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -197,8 +194,7 @@ func TestCoalescerSeedThenFlushAlone(t *testing.T) {
 
 func TestCoalescerCoalescesAdjacentACKs(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -238,8 +234,7 @@ func TestCoalescerCoalescesAdjacentACKs(t *testing.T) {
 
 func TestCoalescerRejectsSeqGap(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -258,8 +253,7 @@ func TestCoalescerRejectsSeqGap(t *testing.T) {
 
 func TestCoalescerRejectsFlagMismatch(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -280,8 +274,7 @@ func TestCoalescerRejectsFlagMismatch(t *testing.T) {
 
 func TestCoalescerRejectsFIN(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	fin := buildTCPv4(1000, tcpAck|tcpFin, []byte("x"))
 	if err := c.Commit(fin); err != nil {
 		t.Fatal(err)
@@ -297,8 +290,7 @@ func TestCoalescerRejectsFIN(t *testing.T) {
 
 func TestCoalescerShortLastSegmentClosesChain(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	full := make([]byte, 1200)
 	half := make([]byte, 500)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, full)); err != nil {
@@ -333,8 +325,7 @@ func TestCoalescerShortLastSegmentClosesChain(t *testing.T) {
 
 func TestCoalescerPSHFinalizesChain(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -364,8 +355,7 @@ func TestCoalescerPSHFinalizesChain(t *testing.T) {
 // coalescer drops it the sender's push signal never reaches the receiver.
 func TestCoalescerPropagatesPSHFromAppended(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// Seed has no PSH; second segment carries PSH and seals the chain.
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
@@ -393,8 +383,7 @@ func TestCoalescerPropagatesPSHFromAppended(t *testing.T) {
 
 func TestCoalescerRejectsDifferentFlow(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	p1 := buildTCPv4(1000, tcpAck, pay)
 	p2 := buildTCPv4(2200, tcpAck, pay)
@@ -416,8 +405,7 @@ func TestCoalescerRejectsDifferentFlow(t *testing.T) {
 
 func TestCoalescerRejectsIPOptions(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 500)
 	pkt := buildTCPv4(1000, tcpAck, pay)
 	// Bump IHL to 6 to simulate 4 bytes of IP options. Don't actually add
@@ -437,8 +425,7 @@ func TestCoalescerRejectsIPOptions(t *testing.T) {
 
 func TestCoalescerCapBySegments(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 512)
 	seq := uint32(1000)
 	for i := 0; i < tcpCoalesceMaxSegs+5; i++ {
@@ -462,8 +449,7 @@ func TestCoalescerCapBySegments(t *testing.T) {
 // flows coalesce independently in a single Flush.
 func TestCoalescerMultipleFlowsInSameBatch(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 
 	// Flow A: sport 1000. Flow B: sport 3000.
@@ -520,8 +506,7 @@ func TestCoalescerMultipleFlowsInSameBatch(t *testing.T) {
 // writing passthrough packets synchronously.
 func TestCoalescerPreservesArrivalOrder(t *testing.T) {
 	w := &orderedFakeWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	// Sequence: coalesceable TCP, ICMP (passthrough), coalesceable TCP on
 	// a different flow. Expected emit order: gso(X), plain(ICMP), gso(Y).
 	pay := make([]byte, 1200)
@@ -589,8 +574,7 @@ func stringSliceEq(a, b []string) bool {
 // packet (SYN) mid-flow only flushes its own flow, not others.
 func TestCoalescerInterleavedFlowsPreserveOrdering(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 
 	// Flow A two segments.
@@ -695,8 +679,7 @@ func buildTCPv6(tcLow byte, seq uint32, flags byte, payload []byte) []byte {
 // retains ECE on the wire.
 func TestCoalescerCoalescesEceFlow(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	flags := byte(tcpAck | tcpEce)
 	if err := c.Commit(buildTCPv4(1000, flags, pay)); err != nil {
@@ -725,8 +708,7 @@ func TestCoalescerCoalescesEceFlow(t *testing.T) {
 // in-flow segment seeds a new slot rather than extending the prior burst.
 func TestCoalescerCwrSealsFlow(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -759,8 +741,7 @@ func TestCoalescerCwrSealsFlow(t *testing.T) {
 // a CE-echoing window or none.
 func TestCoalescerEceMismatchReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4(1000, tcpAck|tcpEce, pay)); err != nil {
 		t.Fatal(err)
@@ -790,8 +771,7 @@ func TestCoalescerEceMismatchReseeds(t *testing.T) {
 // across the whole burst.
 func TestCoalescerDifferingECNReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4WithToS(ecnECT0, 1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -836,8 +816,7 @@ func TestCoalescerDifferingECNReseeds(t *testing.T) {
 // codepoint, and neither may end up CE-marked.
 func TestCoalescerECT0ThenECT1NoCE(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	if err := c.Commit(buildTCPv4WithToS(ecnECT0, 1000, tcpAck, pay)); err != nil {
 		t.Fatal(err)
@@ -867,8 +846,7 @@ func TestCoalescerECT0ThenECT1NoCE(t *testing.T) {
 // six DSCP bits must match too.
 func TestCoalescerDscpMismatchReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// Same ECN (Not-ECT), different DSCP (0x10 vs 0x20 in upper 6 bits).
 	tosA := byte(0x10<<2) | ecnNotECT
@@ -891,8 +869,7 @@ func TestCoalescerDscpMismatchReseeds(t *testing.T) {
 // TestCoalescerCoalescesEceFlow.
 func TestCoalescerIPv6CoalescesEceFlow(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	flags := byte(tcpAck | tcpEce)
 	if err := c.Commit(buildTCPv6(0, 1000, flags, pay)); err != nil {
@@ -923,8 +900,7 @@ func TestCoalescerIPv6CoalescesEceFlow(t *testing.T) {
 // seen had the wire never reordered.
 func TestCoalescerSortsReorderedSeedsAndMerges(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// Arrival order: seq 1000, 3400, 2200. The 3400 seeds a separate slot
 	// because 3400 != nextSeq=2200, then 2200 fails to extend the 3400 slot
@@ -960,8 +936,7 @@ func TestCoalescerSortsReorderedSeedsAndMerges(t *testing.T) {
 // without any cross-flow contamination.
 func TestCoalescerSortAcrossFlowsMergesEachIndependently(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// Flow A (sport 1000) seq 100, 1300; flow B (sport 3000) seq 500, 1700.
 	// Arrival: A.1300, B.1700, A.100, B.500 — every flow reordered.
@@ -1012,8 +987,7 @@ func TestCoalescerSortAcrossFlowsMergesEachIndependently(t *testing.T) {
 // boundary by an arbitrary number of segments.
 func TestCoalescerSortKeepsPSHBoundary(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// Seq 1000 (no PSH) + 2200 (PSH) → seal one slot with PSH set.
 	// Seq 3400 (no PSH) is contiguous to 3400 from seq 2200+1200; without
@@ -1041,8 +1015,7 @@ func TestCoalescerSortKeepsPSHBoundary(t *testing.T) {
 // is sorted/merged independently.
 func TestCoalescerSortKeepsPassthroughBarrier(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// First two segments seed S1 (then a 3400 reorder seeds S2).
 	if err := c.Commit(buildTCPv4(1000, tcpAck, pay)); err != nil {
@@ -1076,8 +1049,7 @@ func TestCoalescerSortKeepsPassthroughBarrier(t *testing.T) {
 // 0x30, so ipHeadersMatch (comparing byte 1 fully) still splits them.
 func TestCoalescerIPv6DifferingECNReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewTCPCoalescer(w, test.NewLogger(), arena.Reserve, arena.Reset)
+	c := NewTCPCoalescer(w, test.NewLogger())
 	pay := make([]byte, 1200)
 	// tcLow is the low 4 bits of TC; ECN occupies the bottom 2 of those.
 	if err := c.Commit(buildTCPv6(ecnECT0, 1000, tcpAck, pay)); err != nil {

+ 4 - 41
overlay/batch/udp_coalesce.go

@@ -46,13 +46,7 @@ type udpSlot struct {
 // concurrent flows and emits each flow's run as a single GSO_UDP_L4
 // superpacket via tio.GSOWriter. Falls back to per-packet writes when the
 // underlying writer doesn't support USO.
-//
-// All output — coalesced or not — is deferred until Flush so per-flow
-// arrival order is preserved on the wire. Cross-flow order is NOT preserved
-// across the TCP/UDP/passthrough split when this coalescer runs alongside
-// others — see multi_coalesce.go. Per-flow order is preserved because a
-// single 5-tuple only ever lands in one lane and each lane preserves its
-// own slot order.
+// Preserves the in-flow order of packets as they are Commit-ed
 //
 // Owns no locks; one coalescer per TUN write queue.
 type UDPCoalescer struct {
@@ -62,8 +56,6 @@ type UDPCoalescer struct {
 	slots     []*udpSlot
 	openSlots map[flowKey]*udpSlot
 	pool      []*udpSlot
-	reserver  Reserver
-	resetter  Resetter
 }
 
 // NewUDPCoalescer wraps w. The caller is responsible for only constructing
@@ -71,14 +63,12 @@ type UDPCoalescer struct {
 // the kernel may reject GSO_UDP_L4 writes. If w does not implement
 // tio.GSOWriter at all (single-packet Queue), the coalescer degrades to
 // plain Writes — same defensive shape as the TCP coalescer.
-func NewUDPCoalescer(w io.Writer, reserver Reserver, resetter Resetter) *UDPCoalescer {
+func NewUDPCoalescer(w io.Writer) *UDPCoalescer {
 	c := &UDPCoalescer{
 		plainW:    w,
 		slots:     make([]*udpSlot, 0, initialSlots),
 		openSlots: make(map[flowKey]*udpSlot, initialSlots),
 		pool:      make([]*udpSlot, 0, initialSlots),
-		reserver:  reserver,
-		resetter:  resetter,
 	}
 	if gw, ok := tio.SupportsGSO(w, tio.GSOProtoUDP); ok {
 		c.gsoW = gw
@@ -123,10 +113,6 @@ func parseUDP(pkt []byte) (parsedUDP, bool) {
 	return p, true
 }
 
-func (c *UDPCoalescer) Reserve(sz int) []byte {
-	return c.reserver(sz)
-}
-
 // Commit borrows pkt. The caller must keep pkt valid until the next Flush.
 func (c *UDPCoalescer) Commit(pkt []byte) error {
 	if c.gsoW == nil {
@@ -175,19 +161,7 @@ func (c *UDPCoalescer) commitParsed(pkt []byte, info parsedUDP) error {
 	return nil
 }
 
-// Flush drains every queued slot and calls the configured Resetter.
 func (c *UDPCoalescer) Flush() error {
-	first := c.drain()
-	if c.resetter != nil {
-		c.resetter()
-	}
-	return first
-}
-
-// drain emits every queued slot in arrival order and clears the slot state.
-// It does NOT reset the arena: borrowed payload slices stay valid until the
-// arena's owner recycles it.
-func (c *UDPCoalescer) drain() error {
 	var first error
 	for _, s := range c.slots {
 		var err error
@@ -295,15 +269,7 @@ func (c *UDPCoalescer) release(s *udpSlot) {
 // flushSlot patches the IP header total length / IPv6 payload length and
 // the UDP length to the *total* across all coalesced segments, then seeds
 // the UDP checksum field with the pseudo-header partial (single-fold, not
-// inverted) per virtio NEEDS_CSUM. The kernel's ip_rcv_core (v4) and
-// ip6_rcv_core (v6) trim the skb to those length fields, so per-segment
-// values would silently drop everything but the first segment. The kernel
-// then walks each segment in __udp_gso_segment, recomputing per-segment
-// uh->len / iph->tot_len / IPv6 plen and adjusting the checksum via
-// `check = csum16_add(csum16_sub(uh->check, uh->len), newlen)` — meaning
-// our seed's uh->check must be consistent with the seed's uh->len, which
-// is what passing the total to both pseudoSum and the UDP length field
-// guarantees.
+// inverted) per virtio NEEDS_CSUM.
 func (c *UDPCoalescer) flushSlot(s *udpSlot) error {
 	hdr := s.hdrBuf[:s.hdrLen]
 	total := s.hdrLen + s.totalPay // full IP+UDP+all_payloads bytes
@@ -334,10 +300,7 @@ func (c *UDPCoalescer) flushSlot(s *udpSlot) error {
 }
 
 // udpHeadersMatch compares two IP+UDP header prefixes for byte-equality on
-// every field that must be identical across coalesced segments. Length
-// fields are masked out (flushSlot rewrites them), but the IP-level ECN
-// codepoint is compared (via ipHeadersMatch) so segments with differing ECN
-// don't coalesce, matching kernel GRO.
+// every field that must be identical across coalesced segments
 func udpHeadersMatch(a, b []byte, isV6 bool, ipHdrLen int) bool {
 	if len(a) != len(b) {
 		return false

+ 16 - 32
overlay/batch/udp_coalesce_test.go

@@ -60,8 +60,7 @@ func buildUDPv6(sport, dport uint16, payload []byte) []byte {
 
 func TestUDPCoalescerPassthroughWhenGSOUnavailable(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: false}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv4(1000, 53, make([]byte, 100))
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -79,8 +78,7 @@ func TestUDPCoalescerPassthroughWhenGSOUnavailable(t *testing.T) {
 
 func TestUDPCoalescerNonUDPPassthrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	// ICMP packet
 	pkt := make([]byte, 28)
 	pkt[0] = 0x45
@@ -101,8 +99,7 @@ func TestUDPCoalescerNonUDPPassthrough(t *testing.T) {
 
 func TestUDPCoalescerSeedThenFlushAlone(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv4(1000, 53, make([]byte, 800))
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -119,8 +116,7 @@ func TestUDPCoalescerSeedThenFlushAlone(t *testing.T) {
 
 func TestUDPCoalescerCoalescesEqualSized(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 1200)
 	for i := 0; i < 3; i++ {
 		if err := c.Commit(buildUDPv4(1000, 53, pay)); err != nil {
@@ -160,8 +156,7 @@ func TestUDPCoalescerCoalescesEqualSized(t *testing.T) {
 // Last segment may be shorter, sealing the chain.
 func TestUDPCoalescerShortLastSegmentSeals(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	full := make([]byte, 1200)
 	tail := make([]byte, 600)
 	if err := c.Commit(buildUDPv4(1000, 53, full)); err != nil {
@@ -194,8 +189,7 @@ func TestUDPCoalescerShortLastSegmentSeals(t *testing.T) {
 // A larger-than-gsoSize packet cannot extend the slot — it reseeds.
 func TestUDPCoalescerLargerThanSeedReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	if err := c.Commit(buildUDPv4(1000, 53, make([]byte, 800))); err != nil {
 		t.Fatal(err)
 	}
@@ -213,8 +207,7 @@ func TestUDPCoalescerLargerThanSeedReseeds(t *testing.T) {
 // Different 5-tuples must not coalesce.
 func TestUDPCoalescerDifferentFlowsKeepSeparate(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 800)
 	if err := c.Commit(buildUDPv4(1000, 53, pay)); err != nil {
 		t.Fatal(err)
@@ -245,8 +238,7 @@ func TestUDPCoalescerDifferentFlowsKeepSeparate(t *testing.T) {
 // Caps at udpCoalesceMaxSegs.
 func TestUDPCoalescerCapsAtMaxSegs(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 100)
 	for i := 0; i < udpCoalesceMaxSegs+5; i++ {
 		if err := c.Commit(buildUDPv4(1000, 53, pay)); err != nil {
@@ -275,8 +267,7 @@ func TestUDPCoalescerCapsAtMaxSegs(t *testing.T) {
 // trailing Not-ECT datagram seeds another.
 func TestUDPCoalescerDifferingECNReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 800)
 	pkt0 := buildUDPv4(1000, 53, pay) // ECN=00 (Not-ECT)
 	pkt1 := buildUDPv4(1000, 53, pay)
@@ -307,8 +298,7 @@ func TestUDPCoalescerDifferingECNReseeds(t *testing.T) {
 // IPv6 path: same flow, equal-sized → coalesced.
 func TestUDPCoalescerIPv6Coalesces(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 1200)
 	for i := 0; i < 3; i++ {
 		if err := c.Commit(buildUDPv6(1000, 53, pay)); err != nil {
@@ -344,8 +334,7 @@ func TestUDPCoalescerIPv6Coalesces(t *testing.T) {
 // DSCP differences must reseed: udpHeadersMatch compares the full ToS byte.
 func TestUDPCoalescerDSCPMismatchReseeds(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pay := make([]byte, 800)
 	pkt0 := buildUDPv4(1000, 53, pay)
 	pkt1 := buildUDPv4(1000, 53, pay)
@@ -367,8 +356,7 @@ func TestUDPCoalescerDSCPMismatchReseeds(t *testing.T) {
 // Fragmented IPv4 must not be coalesced.
 func TestUDPCoalescerFragmentedIPv4PassesThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv4(1000, 53, make([]byte, 200))
 	binary.BigEndian.PutUint16(pkt[6:8], 0x2000) // MF=1
 	if err := c.Commit(pkt); err != nil {
@@ -389,8 +377,7 @@ func TestUDPCoalescerFragmentedIPv4PassesThrough(t *testing.T) {
 // reach the GSO path. Regression: must not panic and must be written.
 func TestUDPCoalescerZeroLengthPayloadPassesThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv4(1000, 53, nil) // UDP length 8, zero payload
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -409,8 +396,7 @@ func TestUDPCoalescerZeroLengthPayloadPassesThrough(t *testing.T) {
 // IPv6 zero-length UDP datagram: same passthrough contract as v4.
 func TestUDPCoalescerZeroLengthPayloadIPv6PassesThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv6(1000, 53, nil) // UDP length 8, zero payload
 	if err := c.Commit(pkt); err != nil {
 		t.Fatal(err)
@@ -431,8 +417,7 @@ func TestUDPCoalescerZeroLengthPayloadIPv6PassesThrough(t *testing.T) {
 // wire — per-flow arrival order (full, empty, full) must be preserved.
 func TestUDPCoalescerZeroLengthMidFlowSealsAndPreservesOrder(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	full := make([]byte, 800)
 	if err := c.Commit(buildUDPv4(1000, 53, full)); err != nil {
 		t.Fatal(err)
@@ -456,8 +441,7 @@ func TestUDPCoalescerZeroLengthMidFlowSealsAndPreservesOrder(t *testing.T) {
 // IPv4 with options is not admissible (we require IHL=5).
 func TestUDPCoalescerIPv4WithOptionsPassesThrough(t *testing.T) {
 	w := &fakeTunWriter{gsoEnabled: true}
-	arena := NewArena(0)
-	c := NewUDPCoalescer(w, arena.Reserve, arena.Reset)
+	c := NewUDPCoalescer(w)
 	pkt := buildUDPv4(1000, 53, make([]byte, 200))
 	pkt[0] = 0x46 // IHL = 6 (24-byte IPv4 header — has options)
 	if err := c.Commit(pkt); err != nil {

+ 3 - 1
udp/udp_tester.go

@@ -173,8 +173,10 @@ func (u *TesterConn) ListenOut(r EncReader, flush func()) error {
 			return os.ErrClosed
 		case p := <-u.RxPackets:
 			r(p.From, p.Data, RxMeta{})
-			p.Release()
+			// The batcher borrows plaintext decrypted in place inside p.Data
+			// until Flush, so the packet must stay alive across flush()
 			flush()
+			p.Release()
 		}
 	}
 }