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