use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, SystemTime};
use super::nonce::{
NonceFreshness, NonceIssuerKey, NonceKind, NonceTracker, NonceTrackerError,
};
pub const DEFAULT_NONCE_RETENTION: Duration = Duration::from_secs(3700);
pub const DEFAULT_PER_PARTITION_CAP: usize = 16_384;
type Partition = HashMap<[u8; 16], SystemTime>;
type PartitionMap = HashMap<(NonceKind, NonceIssuerKey), Partition>;
pub struct DefaultNonceTracker {
inner: Mutex<PartitionMap>,
retention: Duration,
per_partition_cap: usize,
}
impl DefaultNonceTracker {
#[must_use]
pub fn new() -> Self {
Self::with_config(DEFAULT_NONCE_RETENTION, DEFAULT_PER_PARTITION_CAP)
}
#[must_use]
pub fn with_config(retention: Duration, per_partition_cap: usize) -> Self {
DefaultNonceTracker {
inner: Mutex::new(HashMap::new()),
retention,
per_partition_cap,
}
}
}
impl Default for DefaultNonceTracker {
fn default() -> Self {
Self::new()
}
}
impl NonceTracker for DefaultNonceTracker {
fn record(
&self,
kind: NonceKind,
issuer: &NonceIssuerKey,
nonce_bytes: &[u8; 16],
observed_at: SystemTime,
) -> Result<NonceFreshness, NonceTrackerError> {
let mut guard = self
.inner
.lock()
.map_err(|_| NonceTrackerError::BackendUnavailable)?;
let key = (kind, issuer.clone());
let partition = guard.entry(key).or_default();
let cutoff = observed_at
.checked_sub(self.retention)
.unwrap_or(SystemTime::UNIX_EPOCH);
partition.retain(|_, first_seen| *first_seen >= cutoff);
if let Some(first_seen) = partition.get(nonce_bytes) {
return Ok(NonceFreshness::Replay {
first_seen_at: *first_seen,
});
}
if partition.len() >= self.per_partition_cap {
return Err(NonceTrackerError::OverCapacity);
}
partition.insert(*nonce_bytes, observed_at);
Ok(NonceFreshness::Fresh)
}
fn retention_window(&self) -> Duration {
self.retention
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::identity::KeyId;
use crate::proto::Did;
use crate::wire::nonce::NoncePrincipal;
fn issuer(byte: u8) -> NonceIssuerKey {
NonceIssuerKey {
principal: NoncePrincipal::Service(Did::new("did:plc:tracker").unwrap()),
key_id: KeyId::from_bytes([byte; 32]),
}
}
#[test]
fn defaults_meet_4_8_retention_lower_bound() {
assert!(DEFAULT_NONCE_RETENTION >= Duration::from_secs(3630));
}
#[test]
fn fresh_nonce_returns_fresh() {
let t = DefaultNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
let r = t
.record(NonceKind::CapabilityClaim, &issuer(1), &[0xAA; 16], now)
.unwrap();
assert_eq!(r, NonceFreshness::Fresh);
}
#[test]
fn replayed_nonce_within_retention_returns_replay_with_first_seen() {
let t = DefaultNonceTracker::new();
let first = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
t.record(NonceKind::CapabilityClaim, &issuer(1), &[0xBB; 16], first)
.unwrap();
let later = first + Duration::from_secs(60);
let r = t
.record(NonceKind::CapabilityClaim, &issuer(1), &[0xBB; 16], later)
.unwrap();
assert_eq!(
r,
NonceFreshness::Replay {
first_seen_at: first
}
);
}
#[test]
fn replayed_nonce_after_retention_returns_fresh() {
let t = DefaultNonceTracker::with_config(Duration::from_secs(60), 100);
let first = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
t.record(NonceKind::CapabilityClaim, &issuer(1), &[0xCC; 16], first)
.unwrap();
let later = first + Duration::from_secs(120);
let r = t
.record(NonceKind::CapabilityClaim, &issuer(1), &[0xCC; 16], later)
.unwrap();
assert_eq!(r, NonceFreshness::Fresh);
}
#[test]
fn same_nonce_under_different_partitions_is_independent() {
let t = DefaultNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
let nonce = [0xDD; 16];
let r1 = t.record(NonceKind::CapabilityClaim, &issuer(1), &nonce, now).unwrap();
let r2 = t.record(NonceKind::CapabilityClaim, &issuer(2), &nonce, now).unwrap();
assert_eq!(r1, NonceFreshness::Fresh);
assert_eq!(r2, NonceFreshness::Fresh);
}
#[test]
fn same_nonce_across_kinds_is_independent() {
let t = DefaultNonceTracker::new();
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
let nonce = [0xEE; 16];
let r1 = t.record(NonceKind::CapabilityClaim, &issuer(1), &nonce, now).unwrap();
let r2 = t.record(NonceKind::Jwt, &issuer(1), &nonce, now).unwrap();
assert_eq!(r1, NonceFreshness::Fresh);
assert_eq!(r2, NonceFreshness::Fresh);
}
#[test]
fn capacity_exhaustion_returns_overcapacity() {
let t = DefaultNonceTracker::with_config(Duration::from_secs(3600), 4);
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
for i in 0..4u8 {
let mut nonce = [0u8; 16];
nonce[0] = i;
let r = t
.record(NonceKind::CapabilityClaim, &issuer(1), &nonce, now)
.unwrap();
assert_eq!(r, NonceFreshness::Fresh);
}
let mut overflow = [0u8; 16];
overflow[0] = 99;
let err = t
.record(NonceKind::CapabilityClaim, &issuer(1), &overflow, now)
.unwrap_err();
assert_eq!(err, NonceTrackerError::OverCapacity);
}
#[test]
fn replay_check_still_works_at_capacity() {
let t = DefaultNonceTracker::with_config(Duration::from_secs(3600), 1);
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000);
let nonce = [0xFF; 16];
let r1 = t.record(NonceKind::CapabilityClaim, &issuer(1), &nonce, now).unwrap();
assert_eq!(r1, NonceFreshness::Fresh);
let r2 = t.record(NonceKind::CapabilityClaim, &issuer(1), &nonce, now).unwrap();
assert!(matches!(r2, NonceFreshness::Replay { .. }));
}
#[test]
fn retention_window_round_trips() {
let t = DefaultNonceTracker::with_config(Duration::from_secs(123), 100);
assert_eq!(t.retention_window(), Duration::from_secs(123));
}
}