reifydb_auth/
challenge.rs1use std::{collections::HashMap, sync::RwLock, time::Duration};
9
10use reifydb_runtime::context::{
11 clock::{Clock, Instant},
12 rng::Rng,
13};
14use uuid::Builder;
15
16struct ChallengeEntry {
18 pub identifier: String,
19 pub method: String,
20 pub payload: HashMap<String, String>,
21 pub created_at: Instant,
22}
23
24pub struct ChallengeInfo {
26 pub identifier: String,
27 pub method: String,
28 pub payload: HashMap<String, String>,
29}
30
31pub struct ChallengeStore {
37 entries: RwLock<HashMap<String, ChallengeEntry>>,
38 ttl: Duration,
39}
40
41impl ChallengeStore {
42 pub fn new(ttl: Duration) -> Self {
43 Self {
44 entries: RwLock::new(HashMap::new()),
45 ttl,
46 }
47 }
48
49 pub fn create(
51 &self,
52 identifier: String,
53 method: String,
54 payload: HashMap<String, String>,
55 clock: &Clock,
56 rng: &Rng,
57 ) -> String {
58 let millis = clock.now_millis();
59 let random_bytes = rng.infra_bytes_10();
60 let challenge_id = Builder::from_unix_timestamp_millis(millis, &random_bytes).into_uuid().to_string();
61 let entry = ChallengeEntry {
62 identifier,
63 method,
64 payload,
65 created_at: clock.instant(),
66 };
67 let mut entries = self.entries.write().unwrap();
68 entries.insert(challenge_id.clone(), entry);
69 challenge_id
70 }
71
72 pub fn consume(&self, challenge_id: &str) -> Option<ChallengeInfo> {
75 let mut entries = self.entries.write().unwrap();
76 let entry = entries.remove(challenge_id)?;
77
78 if entry.created_at.elapsed() > self.ttl {
79 return None;
80 }
81
82 Some(ChallengeInfo {
83 identifier: entry.identifier,
84 method: entry.method,
85 payload: entry.payload,
86 })
87 }
88
89 pub fn cleanup_expired(&self) {
91 let ttl = self.ttl;
92 let mut entries = self.entries.write().unwrap();
93 entries.retain(|_, e| e.created_at.elapsed() <= ttl);
94 }
95}
96
97#[cfg(test)]
98mod tests {
99 use reifydb_runtime::context::clock::MockClock;
100
101 use super::*;
102
103 fn test_clock_and_rng() -> (Clock, MockClock, Rng) {
104 let mock = MockClock::from_millis(1000);
105 (Clock::Mock(mock.clone()), mock, Rng::seeded(42))
106 }
107
108 #[test]
109 fn test_create_and_consume() {
110 let (clock, _, rng) = test_clock_and_rng();
111 let store = ChallengeStore::new(Duration::from_secs(60));
112 let data = HashMap::from([("nonce".to_string(), "abc123".to_string())]);
113
114 let id = store.create("alice".to_string(), "solana".to_string(), data, &clock, &rng);
115 let info = store.consume(&id).unwrap();
116
117 assert_eq!(info.identifier, "alice");
118 assert_eq!(info.method, "solana");
119 assert_eq!(info.payload.get("nonce").unwrap(), "abc123");
120 }
121
122 #[test]
123 fn test_one_time_use() {
124 let (clock, _, rng) = test_clock_and_rng();
125 let store = ChallengeStore::new(Duration::from_secs(60));
126 let id = store.create("alice".to_string(), "solana".to_string(), HashMap::new(), &clock, &rng);
127
128 assert!(store.consume(&id).is_some());
129 assert!(store.consume(&id).is_none()); }
131
132 #[test]
133 fn test_unknown_challenge() {
134 let store = ChallengeStore::new(Duration::from_secs(60));
135 assert!(store.consume("nonexistent").is_none());
136 }
137
138 #[test]
139 fn test_expired_challenge() {
140 let (clock, mock, rng) = test_clock_and_rng();
141 let store = ChallengeStore::new(Duration::from_millis(1));
142 let id = store.create("alice".to_string(), "solana".to_string(), HashMap::new(), &clock, &rng);
143
144 mock.advance_millis(10);
145 assert!(store.consume(&id).is_none());
146 }
147}