use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;
use std::time::{Duration, SystemTime};
use crate::identity::ServiceIdentity;
use crate::proto::Did;
use crate::wire::handshake::SessionNonce;
use crate::wire::nonce::{NonceFreshness, NonceTrackerError};
pub const MAX_HANDSHAKE_NONCE_REPLAY_WINDOW: Duration =
Duration::from_secs(24 * 3600);
pub const MAX_HANDSHAKE_NONCE_TRACKER_ENTRIES: usize = 1_000_000;
pub trait HandshakeNonceTracker: Send + Sync {
fn check_and_record(
&self,
initiator: &ServiceIdentity,
nonce: &SessionNonce,
observed_at: SystemTime,
) -> Result<NonceFreshness, NonceTrackerError>;
fn replay_window(&self) -> Duration;
}
type FifoEntry = (Did, [u8; 32]);
struct State {
seen: HashMap<FifoEntry, SystemTime>,
order: VecDeque<FifoEntry>,
}
pub struct DefaultHandshakeNonceTracker {
inner: Mutex<State>,
replay_window: Duration,
cap: usize,
}
impl DefaultHandshakeNonceTracker {
#[must_use]
pub fn new() -> Self {
Self::with_config(
MAX_HANDSHAKE_NONCE_REPLAY_WINDOW,
MAX_HANDSHAKE_NONCE_TRACKER_ENTRIES,
)
}
#[must_use]
pub fn with_config(replay_window: Duration, cap: usize) -> Self {
DefaultHandshakeNonceTracker {
inner: Mutex::new(State {
seen: HashMap::new(),
order: VecDeque::new(),
}),
replay_window,
cap,
}
}
}
impl Default for DefaultHandshakeNonceTracker {
fn default() -> Self {
Self::new()
}
}
impl HandshakeNonceTracker for DefaultHandshakeNonceTracker {
fn check_and_record(
&self,
initiator: &ServiceIdentity,
nonce: &SessionNonce,
observed_at: SystemTime,
) -> Result<NonceFreshness, NonceTrackerError> {
let mut guard = self
.inner
.lock()
.map_err(|_| NonceTrackerError::BackendUnavailable)?;
let cutoff = observed_at
.checked_sub(self.replay_window)
.unwrap_or(SystemTime::UNIX_EPOCH);
while let Some(front_key) = guard.order.front() {
let stale = matches!(
guard.seen.get(front_key),
Some(first_seen) if *first_seen < cutoff
);
if !stale {
break;
}
let evicted = guard.order.pop_front().expect("front existed");
guard.seen.remove(&evicted);
}
let key: FifoEntry = (initiator.service_did().clone(), *nonce.as_bytes());
if let Some(first_seen) = guard.seen.get(&key) {
return Ok(NonceFreshness::Replay {
first_seen_at: *first_seen,
});
}
if guard.seen.len() >= self.cap {
if let Some(evicted) = guard.order.pop_front() {
guard.seen.remove(&evicted);
}
}
guard.seen.insert(key.clone(), observed_at);
guard.order.push_back(key);
Ok(NonceFreshness::Fresh)
}
fn replay_window(&self) -> Duration {
self.replay_window
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::identity::{KeyId, PublicKey, SignatureAlgorithm};
fn identity(seed: u8) -> ServiceIdentity {
let did_str = format!("did:plc:{seed:02x}sample0000000000");
ServiceIdentity::new_internal(
Did::new(&did_str).unwrap(),
KeyId::from_bytes([seed; 32]),
PublicKey {
algorithm: SignatureAlgorithm::Ed25519,
bytes: [seed.wrapping_add(1); 32],
},
None,
)
}
#[test]
fn defaults_pinned_per_7_5() {
assert_eq!(
MAX_HANDSHAKE_NONCE_REPLAY_WINDOW,
Duration::from_secs(24 * 3600)
);
assert_eq!(MAX_HANDSHAKE_NONCE_TRACKER_ENTRIES, 1_000_000);
}
#[test]
fn fresh_nonce_returns_fresh() {
let t = DefaultHandshakeNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
let r = t
.check_and_record(&identity(1), &SessionNonce::from_bytes([0xAA; 32]), now)
.unwrap();
assert_eq!(r, NonceFreshness::Fresh);
}
#[test]
fn replay_within_window_returns_replay() {
let t = DefaultHandshakeNonceTracker::new();
let id = identity(1);
let nonce = SessionNonce::from_bytes([0xBB; 32]);
let first = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
t.check_and_record(&id, &nonce, first).unwrap();
let later = first + Duration::from_secs(60);
let r = t.check_and_record(&id, &nonce, later).unwrap();
assert_eq!(
r,
NonceFreshness::Replay {
first_seen_at: first
}
);
}
#[test]
fn replay_past_window_returns_fresh() {
let t = DefaultHandshakeNonceTracker::with_config(Duration::from_secs(60), 100);
let id = identity(1);
let nonce = SessionNonce::from_bytes([0xCC; 32]);
let first = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
t.check_and_record(&id, &nonce, first).unwrap();
let later = first + Duration::from_secs(120);
let r = t.check_and_record(&id, &nonce, later).unwrap();
assert_eq!(r, NonceFreshness::Fresh);
}
#[test]
fn same_nonce_under_different_initiators_is_independent() {
let t = DefaultHandshakeNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
let nonce = SessionNonce::from_bytes([0xDD; 32]);
let r1 = t.check_and_record(&identity(1), &nonce, now).unwrap();
let r2 = t.check_and_record(&identity(2), &nonce, now).unwrap();
assert_eq!(r1, NonceFreshness::Fresh);
assert_eq!(r2, NonceFreshness::Fresh);
}
#[test]
fn same_did_with_different_key_id_shares_partition() {
let t = DefaultHandshakeNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
let nonce = SessionNonce::from_bytes([0xEE; 32]);
let did = Did::new("did:plc:samedidsamedidsamedid").unwrap();
let id_k1 = ServiceIdentity::new_internal(
did.clone(),
KeyId::from_bytes([0x01; 32]),
PublicKey {
algorithm: SignatureAlgorithm::Ed25519,
bytes: [0x02; 32],
},
None,
);
let id_k2 = ServiceIdentity::new_internal(
did,
KeyId::from_bytes([0x99; 32]),
PublicKey {
algorithm: SignatureAlgorithm::Ed25519,
bytes: [0x9A; 32],
},
None,
);
let r1 = t.check_and_record(&id_k1, &nonce, now).unwrap();
assert_eq!(r1, NonceFreshness::Fresh);
let r2 = t.check_and_record(&id_k2, &nonce, now).unwrap();
assert!(
matches!(r2, NonceFreshness::Replay { .. }),
"same DID + same nonce must be Replay regardless of key rotation"
);
}
#[test]
fn fifo_eviction_drops_oldest_at_cap() {
let t = DefaultHandshakeNonceTracker::with_config(Duration::from_secs(3600), 4);
let id = identity(1);
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
let nonces: [[u8; 32]; 4] = [
[0x01; 32],
[0x02; 32],
[0x03; 32],
[0x04; 32],
];
for n in &nonces {
assert_eq!(
t.check_and_record(&id, &SessionNonce::from_bytes(*n), now).unwrap(),
NonceFreshness::Fresh
);
}
let overflow = SessionNonce::from_bytes([0x05; 32]);
assert_eq!(
t.check_and_record(&id, &overflow, now).unwrap(),
NonceFreshness::Fresh
);
for live in &[nonces[1], nonces[2], nonces[3], [0x05; 32]] {
let r = t
.check_and_record(&id, &SessionNonce::from_bytes(*live), now)
.unwrap();
assert!(
matches!(r, NonceFreshness::Replay { .. }),
"non-evicted nonce must still be Replay"
);
}
let evicted_replay = t
.check_and_record(&id, &SessionNonce::from_bytes(nonces[0]), now)
.unwrap();
assert_eq!(
evicted_replay,
NonceFreshness::Fresh,
"evicted nonce must be insertable again (LRU degradation)"
);
}
#[test]
fn replay_window_round_trips() {
let t = DefaultHandshakeNonceTracker::with_config(Duration::from_secs(123), 100);
assert_eq!(t.replay_window(), Duration::from_secs(123));
}
}