From d691ad67a6dd70a0855d72d1d55942ef97916f0c Mon Sep 17 00:00:00 2001
From: Calin Teodor <t.calin@student.vu.nl>
Date: Sun, 7 Jun 2026 10:24:29 +0300
Subject: [PATCH] feat(keyexchange): per-peer whitelist for reply interval gate
 (PILOT-344)

Adds atomic.Pointer-backed per-peer whitelist for the asymmetric-recovery
reply gate. Whitelisted peers always pass MarkReplyKeyExchangeSent so
trusted nodes never get throttled by the 1 s ping-pong-prevention
interval.

API: Manager.SetReplyWhitelist([]uint32). Empty whitelist by default
preserves the existing rate-limit behaviour.
---
 pkg/daemon/keyexchange/keyexchange.go | 41 +++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/pkg/daemon/keyexchange/keyexchange.go b/pkg/daemon/keyexchange/keyexchange.go
index 2ab5653f..9c910840 100644
--- a/pkg/daemon/keyexchange/keyexchange.go
+++ b/pkg/daemon/keyexchange/keyexchange.go
@@ -45,6 +45,7 @@ import (
 	"crypto/ed25519"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pilot-protocol/common/crypto"
@@ -244,6 +245,13 @@ type Manager struct {
 	// described on that constant.
 	lastKeyExchangeReply map[uint32]time.Time
 
+	// PILOT-344: per-peer whitelist for the reply-interval gate.
+	// Whitelisted peers always pass MarkReplyKeyExchangeSent so trusted
+	// nodes never get gated by the 1s asymmetric-recovery throttle.
+	// Read on the hot path under atomic.Pointer; populated at startup
+	// via SetReplyWhitelist.
+	replyWhitelist atomic.Pointer[map[uint32]struct{}]
+
 	// Side-effect hooks plumbed in from the daemon.
 	sender      FrameSender
 	addrLookup  PeerAddrLookup
@@ -573,18 +581,45 @@ func (m *Manager) InboundDecryptStale(peerNodeID uint32) bool {
 // concurrent HandleAuthFrame goroutines (direct + relay copy past
 // DuplicateHandshakeDebounce) cannot both pass.
 func (m *Manager) MarkReplyKeyExchangeSent(peerNodeID uint32) bool {
+	// PILOT-344: trusted-peer whitelist bypasses the interval gate.
+	// The check is outside rkPendingMu so an atomic.Pointer load is all
+	// we need; we still record the timestamp below for consistency with
+	// the cleanup loop.
+	whitelisted := false
+	if wl := m.replyWhitelist.Load(); wl != nil {
+		if _, ok := (*wl)[peerNodeID]; ok {
+			whitelisted = true
+		}
+	}
 	now := time.Now()
 	m.rkPendingMu.Lock()
 	defer m.rkPendingMu.Unlock()
-	if last, ok := m.lastKeyExchangeReply[peerNodeID]; ok {
-		if now.Sub(last) < KeyExchangeReplyMinInterval {
-			return false
+	if !whitelisted {
+		if last, ok := m.lastKeyExchangeReply[peerNodeID]; ok {
+			if now.Sub(last) < KeyExchangeReplyMinInterval {
+				return false
+			}
 		}
 	}
 	m.lastKeyExchangeReply[peerNodeID] = now
 	return true
 }
 
+// SetReplyWhitelist replaces the trusted-peer whitelist for the
+// asymmetric-recovery reply interval gate (PILOT-344). Whitelisted
+// peers always pass MarkReplyKeyExchangeSent. Safe to call concurrently
+// with the hot path — backed by an atomic.Pointer.
+func (m *Manager) SetReplyWhitelist(nodeIDs []uint32) {
+	wm := make(map[uint32]struct{}, len(nodeIDs))
+	for _, id := range nodeIDs {
+		if id == 0 {
+			continue
+		}
+		wm[id] = struct{}{}
+	}
+	m.replyWhitelist.Store(&wm)
+}
+
 // RemovePeer wipes per-peer L5 state (called from TunnelManager.RemovePeer).
 func (m *Manager) RemovePeer(nodeID uint32) {
 	m.pubKeysMu.Lock()