Skip to main content

reifydb_auth/
challenge.rs

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