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

reifydb_auth/
challenge.rs
