use std::collections::HashSet;
use std::sync::Mutex;
use base64::Engine;
use rand::RngCore;
const NONCE_BYTES: usize = 24;
const MAX_OUTSTANDING: usize = 4096;
pub struct NonceStore {
issued: Mutex<HashSet<String>>,
}
impl NonceStore {
pub fn new() -> Self {
Self {
issued: Mutex::new(HashSet::new()),
}
}
pub fn issue(&self) -> String {
let mut bytes = [0u8; NONCE_BYTES];
rand::rng().fill_bytes(&mut bytes);
let nonce = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
let mut set = self.issued.lock().unwrap_or_else(|e| e.into_inner());
if set.len() >= MAX_OUTSTANDING {
set.clear();
}
set.insert(nonce.clone());
nonce
}
pub fn redeem(&self, nonce: &str) -> bool {
let mut set = self.issued.lock().unwrap_or_else(|e| e.into_inner());
set.remove(nonce)
}
}
impl Default for NonceStore {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn issued_nonce_redeems_once() {
let store = NonceStore::new();
let n = store.issue();
assert!(store.redeem(&n), "freshly issued nonce must redeem");
assert!(
!store.redeem(&n),
"a redeemed nonce must NOT redeem again (replay)"
);
}
#[test]
fn unknown_nonce_does_not_redeem() {
let store = NonceStore::new();
assert!(!store.redeem("never-issued"));
}
#[test]
fn nonces_are_unique() {
let store = NonceStore::new();
let a = store.issue();
let b = store.issue();
assert_ne!(a, b, "each issued nonce must be distinct");
}
#[test]
fn concurrent_issue_redeem_is_safe() {
use std::sync::Arc;
use std::thread;
let store = Arc::new(NonceStore::new());
let mut handles = Vec::new();
for _ in 0..8 {
let s = Arc::clone(&store);
handles.push(thread::spawn(move || {
for _ in 0..100 {
let n = s.issue();
assert!(s.redeem(&n));
}
}));
}
for h in handles {
h.join().unwrap();
}
}
}