From 5bba4de3c64bc626eadd08bed7dfcd1b46c5cdb9 Mon Sep 17 00:00:00 2001
From: JackDoan <me@jackdoan.com>
Date: Thu, 15 Oct 2026 17:13:30 -0500
Subject: [PATCH] Check staged receive replays on the listener

---
 connection_state.go | 13 ++++---------
 flow_rx.go          | 40 ++++++++++++++++++++++++++++------------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/connection_state.go b/connection_state.go
index e7d03f57..b9108567 100644
--- a/connection_state.go
+++ b/connection_state.go
@@ -97,21 +97,16 @@ func (cs *ConnectionState) Curve() cert.Curve {
 }
 
 func (cs *ConnectionState) Decrypt(l *slog.Logger, messageCounter uint64, packet []byte, nb []byte) ([]byte, error) {
-	if !cs.checkReplay(l, messageCounter) {
+	cs.decryptLock.Lock()
+	result := cs.window.Check(l, messageCounter)
+	cs.decryptLock.Unlock()
+	if !result {
 		return nil, ErrAlreadySeen
 	}
 	out, err := cs.decryptPacket(messageCounter, packet, nb)
 	return cs.acceptDecrypted(l, messageCounter, out, err)
 }
 
-// checkReplay reports whether the replay window would still accept messageCounter, without recording it.
-func (cs *ConnectionState) checkReplay(l *slog.Logger, messageCounter uint64) bool {
-	cs.decryptLock.Lock()
-	result := cs.window.Check(l, messageCounter)
-	cs.decryptLock.Unlock()
-	return result
-}
-
 // 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)
diff --git a/flow_rx.go b/flow_rx.go
index 8ac16a0b..305c4470 100644
--- a/flow_rx.go
+++ b/flow_rx.go
@@ -1,7 +1,6 @@
 package nebula
 
 import (
-	"log/slog"
 	"net/netip"
 	"sync"
 
@@ -37,17 +36,12 @@ type rxChunk struct {
 	done  chan struct{}
 }
 
-// decryptRxItems checks the data packets among items against the replay window and decrypts the ones it accepts in
-// place, as Decrypt does. Recording the counters is left to readOutsidePackets on the writer, which also handles
-// the packets without a tunnel.
-func decryptRxItems(l *slog.Logger, items []rxItem, nb []byte) {
+// decryptRxItems decrypts the data packets among items in place, skipping the replays checkReplays marked. Recording
+// the counters is left to readOutsidePackets on the writer, which also handles the packets without a tunnel.
+func decryptRxItems(items []rxItem, nb []byte) {
 	for i := range items {
 		it := &items[i]
-		if it.ci == nil {
-			continue
-		}
-		if !it.ci.checkReplay(l, it.counter) {
-			it.err = ErrAlreadySeen
+		if it.ci == nil || it.err != nil {
 			continue
 		}
 		it.out, it.err = it.ci.decryptPacket(it.counter, it.packet, nb)
@@ -135,7 +129,7 @@ func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
 func (p *stagedRx) decrypt() {
 	nb := make([]byte, 12)
 	for c := range p.jobs {
-		decryptRxItems(p.f.l, c.items, nb)
+		decryptRxItems(c.items, nb)
 		c.done <- struct{}{}
 	}
 }
@@ -205,6 +199,7 @@ func (p *stagedRx) endBatch() {
 		b.reset()
 		return
 	}
+	p.checkReplays(b.items)
 
 	if p.inFlight == 0 && n < p.f.staged.inlineBelow {
 		// The writer is idle with nothing in flight, so the listener may commit with its context and flush the tun
@@ -239,6 +234,27 @@ func (p *stagedRx) endBatch() {
 	p.free = p.free[:len(p.free)-1]
 }
 
+// checkReplays marks the replays among items before anything decrypts them, as Decrypt checks before it decrypts,
+// taking each tunnel's lock once per run of its packets rather than once per packet. The window only moves forward
+// until the writer records these counters, so a stale check can only let a replay through to acceptDecrypted, which
+// rejects it.
+func (p *stagedRx) checkReplays(items []rxItem) {
+	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(p.f.l, items[i].counter) {
+				items[i].err = ErrAlreadySeen
+			}
+		}
+		ci.decryptLock.Unlock()
+	}
+}
+
 // reap takes back every batch the writer has finished with, without waiting.
 func (p *stagedRx) reap() {
 	for {
@@ -260,7 +276,7 @@ func (p *stagedRx) release(b *rxBatch) {
 // order, and flush the tun batcher.
 func (p *stagedRx) process(b *rxBatch, rxc *rxContext) {
 	if b.nChunks == 0 {
-		decryptRxItems(p.f.l, b.items, rxc.nb)
+		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