agentguard_auth/
jti.rs

1//! jti tracker for replay protection.
2//!
3//! Seen JTIs are stored with monotonic timestamps. Expired entries are
4//! periodically reaped, while duplicate lookups and insertions stay atomic
5//! under one mutex. A hard capacity bounds memory; when full, new proofs are
6//! rejected instead of discarding live replay history.
7
8use crate::error::Result;
9use parking_lot::Mutex;
10use std::collections::HashMap;
11use std::time::{Duration, Instant};
12
13/// Maximum number of live proof identifiers retained by one tracker.
14const DEFAULT_MAX_ENTRIES: usize = 262_144;
15/// Opportunistic expiry sweep interval as a fraction of the replay TTL.
16const SWEEP_DIVISOR: u32 = 4;
17
18struct JtiState {
19    seen: HashMap<[u8; 16], Instant>,
20    last_sweep: Instant,
21}
22
23/// In-memory tracker of seen `jti` values.
24pub struct JtiTracker {
25    state: Mutex<JtiState>,
26    ttl: Duration,
27    max_entries: usize,
28    sweep_interval: Duration,
29}
30
31impl JtiTracker {
32    pub fn new(ttl: Duration) -> Self {
33        Self::build(ttl, DEFAULT_MAX_ENTRIES)
34    }
35
36    /// Construct a tracker with an explicit live-entry limit.
37    ///
38    /// The default [`Self::new`] limit is 262,144 identifiers. Capacity
39    /// exhaustion rejects new proofs; it never evicts unexpired replay state.
40    pub fn with_capacity(ttl: Duration, max_entries: usize) -> Result<Self> {
41        if max_entries == 0 {
42            return Err(crate::error::AuthError::DpopInvalid(
43                "JTI tracker capacity must be positive".into(),
44            ));
45        }
46        Ok(Self::build(ttl, max_entries))
47    }
48
49    fn build(ttl: Duration, max_entries: usize) -> Self {
50        let now = Instant::now();
51        Self {
52            state: Mutex::new(JtiState {
53                seen: HashMap::new(),
54                last_sweep: now,
55            }),
56            ttl,
57            max_entries,
58            sweep_interval: ttl / SWEEP_DIVISOR,
59        }
60    }
61
62    /// Record `jti`. Returns an error for replays or when capacity is exhausted.
63    pub fn check_and_record(&self, jti: &[u8; 16]) -> Result<()> {
64        let now = Instant::now();
65        let mut state = self.state.lock();
66
67        if state
68            .seen
69            .get(jti)
70            .is_some_and(|seen_at| now.duration_since(*seen_at) < self.ttl)
71        {
72            return Err(crate::error::AuthError::DpopReplay(hex::encode(jti)));
73        }
74        state.seen.remove(jti);
75
76        if now.duration_since(state.last_sweep) >= self.sweep_interval {
77            state
78                .seen
79                .retain(|_, seen_at| now.duration_since(*seen_at) < self.ttl);
80            state.last_sweep = now;
81        }
82
83        if state.seen.len() >= self.max_entries {
84            return Err(crate::error::AuthError::DpopCapacityExceeded);
85        }
86
87        state.seen.insert(*jti, now);
88        Ok(())
89    }
90
91    pub fn len(&self) -> usize {
92        self.state.lock().seen.len()
93    }
94
95    pub fn is_empty(&self) -> bool {
96        self.len() == 0
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::error::AuthError;
104
105    #[test]
106    fn first_occurrence_ok_replay_blocked() {
107        let t = JtiTracker::new(Duration::from_secs(60));
108        let jti = [1u8; 16];
109        t.check_and_record(&jti).unwrap();
110        assert!(t.check_and_record(&jti).is_err());
111    }
112
113    #[test]
114    fn distinct_jtis_accepted() {
115        let t = JtiTracker::new(Duration::from_secs(60));
116        t.check_and_record(&[1u8; 16]).unwrap();
117        t.check_and_record(&[2u8; 16]).unwrap();
118        assert_eq!(t.len(), 2);
119    }
120
121    #[test]
122    fn expired_entries_are_reaped() {
123        // Use a 0-second TTL so any entry expires immediately.
124        let t = JtiTracker::new(Duration::from_secs(0));
125        t.check_and_record(&[1u8; 16]).unwrap();
126        // Sleep briefly so the entry is "old enough" to expire.
127        std::thread::sleep(Duration::from_millis(10));
128        // The reaping is opportunistic — a new identifier triggers a sweep.
129        t.check_and_record(&[2u8; 16]).unwrap();
130        // Old entry should have been reaped.
131        assert_eq!(t.len(), 1);
132    }
133
134    #[test]
135    fn new_is_empty() {
136        let t = JtiTracker::new(Duration::from_secs(60));
137        assert!(t.is_empty());
138        assert_eq!(t.len(), 0);
139    }
140
141    #[test]
142    fn replay_blocked_returns_descriptive_error() {
143        // The error must include the offending jti (hex-encoded) so the
144        // operator can identify the duplicate in their logs.
145        let t = JtiTracker::new(Duration::from_secs(60));
146        let jti = [0xab; 16];
147        t.check_and_record(&jti).unwrap();
148        let err = t.check_and_record(&jti).unwrap_err();
149        match err {
150            AuthError::DpopReplay(s) => {
151                assert!(s.contains("abababababababab"), "got: {}", s);
152            }
153            other => panic!("expected DpopReplay, got: {:?}", other),
154        }
155    }
156
157    #[test]
158    fn capacity_exhaustion_does_not_discard_live_replay_history() {
159        let tracker = JtiTracker::with_capacity(Duration::from_secs(60), 2).unwrap();
160        let first = [1u8; 16];
161        tracker.check_and_record(&first).unwrap();
162        tracker.check_and_record(&[2u8; 16]).unwrap();
163
164        assert!(matches!(
165            tracker.check_and_record(&[3u8; 16]),
166            Err(AuthError::DpopCapacityExceeded)
167        ));
168        assert!(matches!(
169            tracker.check_and_record(&first),
170            Err(AuthError::DpopReplay(_))
171        ));
172        assert_eq!(tracker.len(), 2);
173    }
174
175    #[test]
176    fn expired_entries_are_reclaimed_before_capacity_rejection() {
177        let tracker = JtiTracker::with_capacity(Duration::from_millis(20), 1).unwrap();
178        tracker.check_and_record(&[1u8; 16]).unwrap();
179        std::thread::sleep(Duration::from_millis(25));
180
181        tracker.check_and_record(&[2u8; 16]).unwrap();
182        assert_eq!(tracker.len(), 1);
183    }
184
185    #[test]
186    fn explicit_capacity_must_be_positive() {
187        assert!(matches!(
188            JtiTracker::with_capacity(Duration::from_secs(60), 0),
189            Err(AuthError::DpopInvalid(_))
190        ));
191    }
192
193    #[test]
194    fn concurrent_replays_are_recorded_atomically() {
195        let tracker = std::sync::Arc::new(JtiTracker::new(Duration::from_secs(60)));
196        let barrier = std::sync::Arc::new(std::sync::Barrier::new(16));
197        let threads = (0..16)
198            .map(|_| {
199                let tracker = std::sync::Arc::clone(&tracker);
200                let barrier = std::sync::Arc::clone(&barrier);
201                std::thread::spawn(move || {
202                    barrier.wait();
203                    tracker.check_and_record(&[42u8; 16]).is_ok()
204                })
205            })
206            .collect::<Vec<_>>();
207        let accepted = threads
208            .into_iter()
209            .map(|thread| thread.join().unwrap())
210            .filter(|accepted| *accepted)
211            .count();
212
213        assert_eq!(accepted, 1);
214        assert_eq!(tracker.len(), 1);
215    }
216}
agentguard_auth/jti.rs

agentguard_auth/
jti.rs
