Skip to main content

reifydb_auth/
challenge.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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()); // second attempt fails
119	}
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}