Skip to main content

phantom_protocol/security/
replay_protection.rs

1use anyhow::{bail, Result};
2use parking_lot::RwLock;
3use std::collections::HashMap;
4use std::time::{Duration, Instant};
5
6/// Replay protection using nonce tracking
7///
8/// Tracks seen nonces within a time window to prevent replay attacks.
9/// Old nonces are automatically cleaned up to prevent memory exhaustion.
10pub struct ReplayProtection {
11    seen_nonces: RwLock<HashMap<u32, Instant>>,
12    window: Duration,
13    max_entries: usize,
14}
15
16impl ReplayProtection {
17    /// Create new replay protection with default settings
18    ///
19    /// Default window: 120 seconds (2x the timestamp validation window)
20    /// Max entries: 100,000 nonces
21    pub fn new() -> Self {
22        Self {
23            seen_nonces: RwLock::new(HashMap::new()),
24            window: Duration::from_secs(120),
25            max_entries: 100_000,
26        }
27    }
28
29    /// Create with custom window and max entries
30    pub fn with_config(window: Duration, max_entries: usize) -> Self {
31        Self {
32            seen_nonces: RwLock::new(HashMap::new()),
33            window,
34            max_entries,
35        }
36    }
37
38    /// Check if nonce has been seen before and record it
39    ///
40    /// Returns:
41    /// - Ok(()) if nonce is new
42    /// - Err if nonce is duplicate (replay attack detected)
43    pub fn check_and_record(&self, nonce: u32) -> Result<()> {
44        let mut seen = self.seen_nonces.write();
45
46        // Check if nonce exists
47        if let Some(first_seen) = seen.get(&nonce) {
48            // Check if still within window
49            if first_seen.elapsed() < self.window {
50                bail!(
51                    "Replay attack detected: duplicate nonce {} (first seen {:?} ago)",
52                    nonce,
53                    first_seen.elapsed()
54                );
55            }
56            // Nonce is old, can be reused (remove old entry)
57            seen.remove(&nonce);
58        }
59
60        // Check if we need to cleanup
61        if seen.len() >= self.max_entries {
62            self.cleanup_old_nonces(&mut seen);
63
64            // If still full after cleanup, reject
65            if seen.len() >= self.max_entries {
66                bail!("Replay protection cache full (possible DoS attack)");
67            }
68        }
69
70        // Record nonce
71        seen.insert(nonce, Instant::now());
72
73        Ok(())
74    }
75
76    /// Cleanup nonces older than the window
77    fn cleanup_old_nonces(&self, seen: &mut HashMap<u32, Instant>) {
78        let now = Instant::now();
79        seen.retain(|_, first_seen| now.duration_since(*first_seen) < self.window);
80    }
81
82    /// Get current cache statistics
83    pub fn stats(&self) -> ReplayProtectionStats {
84        let seen = self.seen_nonces.read();
85        ReplayProtectionStats {
86            total_nonces: seen.len(),
87            window_seconds: self.window.as_secs(),
88            max_entries: self.max_entries,
89        }
90    }
91
92    /// Clear all tracked nonces (for testing)
93    #[cfg(test)]
94    pub fn clear(&self) {
95        self.seen_nonces.write().clear();
96    }
97}
98
99impl Default for ReplayProtection {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105#[derive(Debug, Clone)]
106pub struct ReplayProtectionStats {
107    pub total_nonces: usize,
108    pub window_seconds: u64,
109    pub max_entries: usize,
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn test_duplicate_nonce_rejected() {
118        let rp = ReplayProtection::new();
119
120        // First nonce should succeed
121        assert!(rp.check_and_record(12345).is_ok());
122
123        // Duplicate nonce should fail
124        assert!(rp.check_and_record(12345).is_err());
125    }
126
127    #[test]
128    fn test_different_nonces_accepted() {
129        let rp = ReplayProtection::new();
130
131        assert!(rp.check_and_record(1).is_ok());
132        assert!(rp.check_and_record(2).is_ok());
133        assert!(rp.check_and_record(3).is_ok());
134    }
135
136    #[test]
137    fn test_nonce_expires() {
138        let rp = ReplayProtection::with_config(Duration::from_millis(100), 1000);
139
140        assert!(rp.check_and_record(999).is_ok());
141
142        // Wait for expiration
143        std::thread::sleep(Duration::from_millis(150));
144
145        // Same nonce should be accepted after expiration
146        assert!(rp.check_and_record(999).is_ok());
147    }
148
149    #[test]
150    fn test_cleanup() {
151        let rp = ReplayProtection::with_config(Duration::from_millis(50), 10);
152
153        // Fill cache
154        for i in 0..10 {
155            assert!(rp.check_and_record(i).is_ok());
156        }
157
158        // Wait for expiration
159        std::thread::sleep(Duration::from_millis(100));
160
161        // Adding new nonce should trigger cleanup
162        assert!(rp.check_and_record(999).is_ok());
163
164        // Old nonces should be cleared
165        let stats = rp.stats();
166        assert_eq!(stats.total_nonces, 1);
167    }
168}