Skip to main content

reifydb_auth/
challenge.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::{collections::HashMap, sync::RwLock, time::Duration};
5
6use reifydb_runtime::context::{
7	clock::{Clock, Instant},
8	rng::Rng,
9};
10use uuid::Builder;
11
12/// A pending authentication challenge.
13struct ChallengeEntry {
14	pub identifier: String,
15	pub method: String,
16	pub payload: HashMap<String, String>,
17	pub created_at: Instant,
18}
19
20/// Stored challenge info returned when consuming a challenge.
21pub struct ChallengeInfo {
22	pub identifier: String,
23	pub method: String,
24	pub payload: HashMap<String, String>,
25}
26
27/// In-memory store for pending authentication challenges.
28///
29/// Challenges are created during multi-step authentication (e.g., wallet signing)
30/// and consumed on the client's response. Each challenge is one-time-use and
31/// expires after the configured TTL.
32pub 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	/// Create a new challenge and return its ID.
46	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	/// Consume a challenge by ID. Returns the challenge data if valid and not expired.
69	/// The challenge is removed after consumption (one-time use).
70	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	/// Remove all expired entries.
86	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()); // second attempt fails
126	}
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}