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 pending_public_key: Option<String>,
21	pub created_at: Instant,
22}
23
24pub struct ChallengeInfo {
25	pub identifier: String,
26	pub method: String,
27	pub payload: HashMap<String, String>,
28	pub pending_public_key: Option<String>,
29}
30
31pub struct ChallengeStore {
32	entries: RwLock<HashMap<String, ChallengeEntry>>,
33	ttl: Duration,
34}
35
36impl ChallengeStore {
37	pub fn new(ttl: Duration) -> Self {
38		Self {
39			entries: RwLock::new(HashMap::new()),
40			ttl,
41		}
42	}
43
44	pub fn create(
45		&self,
46		identifier: String,
47		method: String,
48		payload: HashMap<String, String>,
49		pending_public_key: Option<String>,
50		clock: &Clock,
51		rng: &Rng,
52	) -> String {
53		let millis = clock.now().to_millis();
54		let random_bytes = rng.infra_bytes_10();
55		let challenge_id = Builder::from_unix_timestamp_millis(millis, &random_bytes).into_uuid().to_string();
56		let entry = ChallengeEntry {
57			identifier,
58			method,
59			payload,
60			pending_public_key,
61			created_at: clock.instant(),
62		};
63		let mut entries = self.entries.write();
64		entries.insert(challenge_id.clone(), entry);
65		challenge_id
66	}
67
68	pub fn consume(&self, challenge_id: &str) -> Option<ChallengeInfo> {
69		let mut entries = self.entries.write();
70		let entry = entries.remove(challenge_id)?;
71
72		if entry.created_at.elapsed() > self.ttl.to_std() {
73			return None;
74		}
75
76		Some(ChallengeInfo {
77			identifier: entry.identifier,
78			method: entry.method,
79			payload: entry.payload,
80			pending_public_key: entry.pending_public_key,
81		})
82	}
83
84	pub fn cleanup_expired(&self) {
85		let ttl = self.ttl.to_std();
86		let mut entries = self.entries.write();
87		entries.retain(|_, e| e.created_at.elapsed() <= ttl);
88	}
89}
90
91#[cfg(test)]
92mod tests {
93	use reifydb_runtime::context::clock::MockClock;
94
95	use super::*;
96
97	fn test_clock_and_rng() -> (Clock, MockClock, Rng) {
98		let mock = MockClock::from_millis(1000);
99		(Clock::Mock(mock.clone()), mock, Rng::seeded(42))
100	}
101
102	#[test]
103	fn test_create_and_consume() {
104		let (clock, _, rng) = test_clock_and_rng();
105		let store = ChallengeStore::new(Duration::from_seconds(60).unwrap());
106		let data = HashMap::from([("nonce".to_string(), "abc123".to_string())]);
107
108		let id = store.create("alice".to_string(), "solana".to_string(), data, None, &clock, &rng);
109		let info = store.consume(&id).unwrap();
110
111		assert_eq!(info.identifier, "alice");
112		assert_eq!(info.method, "solana");
113		assert_eq!(info.payload.get("nonce").unwrap(), "abc123");
114	}
115
116	#[test]
117	fn test_one_time_use() {
118		let (clock, _, rng) = test_clock_and_rng();
119		let store = ChallengeStore::new(Duration::from_seconds(60).unwrap());
120		let id = store.create("alice".to_string(), "solana".to_string(), HashMap::new(), None, &clock, &rng);
121
122		assert!(store.consume(&id).is_some());
123		assert!(store.consume(&id).is_none()); // second attempt fails
124	}
125
126	#[test]
127	fn test_unknown_challenge() {
128		let store = ChallengeStore::new(Duration::from_seconds(60).unwrap());
129		assert!(store.consume("nonexistent").is_none());
130	}
131
132	#[test]
133	fn test_pending_public_key_round_trips_outside_the_payload() {
134		// The pending key must stay server-side; in the payload a client could swap it.
135		let (clock, _, rng) = test_clock_and_rng();
136		let store = ChallengeStore::new(Duration::from_seconds(60).unwrap());
137
138		let id = store.create(
139			"wallet".to_string(),
140			"solana".to_string(),
141			HashMap::from([("message".to_string(), "sign me".to_string())]),
142			Some("PubKey111".to_string()),
143			&clock,
144			&rng,
145		);
146		let info = store.consume(&id).unwrap();
147
148		assert_eq!(info.pending_public_key.as_deref(), Some("PubKey111"));
149		assert!(
150			!info.payload.contains_key("public_key"),
151			"the pending key must never reach the client payload"
152		);
153	}
154
155	#[test]
156	fn test_expired_challenge() {
157		let (clock, mock, rng) = test_clock_and_rng();
158		let store = ChallengeStore::new(Duration::from_milliseconds(1).unwrap());
159		let id = store.create("alice".to_string(), "solana".to_string(), HashMap::new(), None, &clock, &rng);
160
161		mock.advance_millis(10);
162		assert!(store.consume(&id).is_none());
163	}
164}
reifydb_auth/challenge.rs

reifydb_auth/
challenge.rs
