1use std::collections::HashMap;
10use std::time::{Duration, Instant};
11
12pub const DEFAULT_TTL: Duration = Duration::from_secs(120);
15
16pub struct PaginationStore {
17 session: String,
19 counter: u64,
20 ttl: Duration,
21 entries: HashMap<u64, (Vec<u8>, Instant)>,
22}
23
24impl PaginationStore {
25 pub fn new(session: u64, ttl: Duration) -> Self {
26 Self {
27 session: format!("{:08x}", session as u32),
30 counter: 0,
31 ttl,
32 entries: HashMap::new(),
33 }
34 }
35
36 pub fn store(&mut self, blob: Vec<u8>, now: Instant) -> String {
38 self.evict(now);
39 self.counter += 1;
40 let id = self.counter;
41 self.entries.insert(id, (blob, now));
42 format!("{}{id:x}", self.session)
43 }
44
45 pub fn take(&mut self, token: &str, now: Instant) -> Option<Vec<u8>> {
48 self.evict(now);
49 let rest = token.strip_prefix(&self.session)?;
50 let id: u64 = u64::from_str_radix(rest, 16).ok()?;
51 let (blob, minted) = self.entries.remove(&id)?;
52 (now.duration_since(minted) < self.ttl).then_some(blob)
53 }
54
55 fn evict(&mut self, now: Instant) {
56 let ttl = self.ttl;
57 self.entries
58 .retain(|_, (_, minted)| now.duration_since(*minted) < ttl);
59 }
60
61 #[cfg(test)]
62 fn len(&self) -> usize {
63 self.entries.len()
64 }
65}
66
67#[cfg(test)]
68mod tests {
69 use super::*;
70
71 #[test]
72 fn store_then_take_roundtrips_and_is_single_use() {
73 let mut s = PaginationStore::new(0xABCD, DEFAULT_TTL);
74 let t0 = Instant::now();
75 let tok = s.store(b"hello".to_vec(), t0);
76 assert!(tok.starts_with("0000abcd"));
77 assert_eq!(s.take(&tok, t0).as_deref(), Some(&b"hello"[..]));
78 assert_eq!(s.take(&tok, t0), None);
80 assert_eq!(s.len(), 0);
81 }
82
83 #[test]
84 fn expired_token_and_foreign_session_miss() {
85 let mut s = PaginationStore::new(1, Duration::from_secs(120));
86 let t0 = Instant::now();
87 let tok = s.store(b"x".to_vec(), t0);
88 assert_eq!(s.take(&tok, t0 + Duration::from_secs(121)), None);
89
90 let other = PaginationStore::new(2, DEFAULT_TTL).store(b"y".to_vec(), t0);
92 let mut s = PaginationStore::new(1, DEFAULT_TTL);
93 assert_eq!(s.take(&other, t0), None);
94 assert_eq!(s.take("not-a-token", t0), None);
95 }
96}