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