ant_quic/relay/
authenticator.rs1use crate::relay::{RelayError, RelayResult};
11use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
12use rand::rngs::OsRng;
13use std::collections::HashSet;
14use std::sync::{Arc, Mutex};
15use std::time::{SystemTime, UNIX_EPOCH};
16
17#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct AuthToken {
20 pub nonce: u64,
22 pub timestamp: u64,
24 pub bandwidth_limit: u32,
26 pub timeout_seconds: u32,
28 pub signature: [u8; 64],
30}
31
32#[derive(Debug)]
34pub struct RelayAuthenticator {
35 signing_key: SigningKey,
37 verifying_key: VerifyingKey,
39 used_nonces: Arc<Mutex<HashSet<u64>>>,
41 max_token_age: u64,
43 replay_window_size: u64,
45}
46
47impl AuthToken {
48 pub fn new(
50 bandwidth_limit: u32,
51 timeout_seconds: u32,
52 signing_key: &SigningKey,
53 ) -> RelayResult<Self> {
54 let nonce = Self::generate_nonce();
55 let timestamp = Self::current_timestamp()?;
56
57 let mut token = Self {
58 nonce,
59 timestamp,
60 bandwidth_limit,
61 timeout_seconds,
62 signature: [0; 64],
63 };
64
65 let signature_bytes = signing_key.sign(&token.signable_data()).to_bytes();
67 token.signature = signature_bytes;
68
69 Ok(token)
70 }
71
72 fn generate_nonce() -> u64 {
74 use rand::Rng;
75 OsRng.r#gen()
76 }
77
78 fn current_timestamp() -> RelayResult<u64> {
80 SystemTime::now()
81 .duration_since(UNIX_EPOCH)
82 .map(|d| d.as_secs())
83 .map_err(|_| RelayError::AuthenticationFailed {
84 reason: "System time before Unix epoch".to_string(),
85 })
86 }
87
88 fn signable_data(&self) -> Vec<u8> {
90 let mut data = Vec::new();
91 data.extend_from_slice(&self.nonce.to_le_bytes());
92 data.extend_from_slice(&self.timestamp.to_le_bytes());
93 data.extend_from_slice(&self.bandwidth_limit.to_le_bytes());
94 data.extend_from_slice(&self.timeout_seconds.to_le_bytes());
95 data
96 }
97
98 pub fn verify(&self, verifying_key: &VerifyingKey) -> RelayResult<()> {
100 let signature = Signature::from_bytes(&self.signature);
101
102 verifying_key
103 .verify(&self.signable_data(), &signature)
104 .map_err(|_| RelayError::AuthenticationFailed {
105 reason: "Signature verification failed".to_string(),
106 })
107 }
108
109 pub fn is_expired(&self, max_age_seconds: u64) -> RelayResult<bool> {
111 let current_time = Self::current_timestamp()?;
112 Ok(current_time > self.timestamp + max_age_seconds)
113 }
114}
115
116impl RelayAuthenticator {
117 pub fn new() -> Self {
119 let signing_key = SigningKey::generate(&mut OsRng);
120 let verifying_key = signing_key.verifying_key();
121
122 Self {
123 signing_key,
124 verifying_key,
125 used_nonces: Arc::new(Mutex::new(HashSet::new())),
126 max_token_age: 300, replay_window_size: 1000,
128 }
129 }
130
131 pub fn with_key(signing_key: SigningKey) -> Self {
133 let verifying_key = signing_key.verifying_key();
134
135 Self {
136 signing_key,
137 verifying_key,
138 used_nonces: Arc::new(Mutex::new(HashSet::new())),
139 max_token_age: 300,
140 replay_window_size: 1000,
141 }
142 }
143
144 pub fn verifying_key(&self) -> &VerifyingKey {
146 &self.verifying_key
147 }
148
149 pub fn create_token(
151 &self,
152 bandwidth_limit: u32,
153 timeout_seconds: u32,
154 ) -> RelayResult<AuthToken> {
155 AuthToken::new(bandwidth_limit, timeout_seconds, &self.signing_key)
156 }
157
158 #[allow(clippy::expect_used)]
160 pub fn verify_token(
161 &self,
162 token: &AuthToken,
163 peer_verifying_key: &VerifyingKey,
164 ) -> RelayResult<()> {
165 token.verify(peer_verifying_key)?;
167
168 if token.is_expired(self.max_token_age)? {
170 return Err(RelayError::AuthenticationFailed {
171 reason: "Token expired".to_string(),
172 });
173 }
174
175 let mut used_nonces = self
177 .used_nonces
178 .lock()
179 .expect("Mutex poisoning is unexpected in normal operation");
180
181 if used_nonces.contains(&token.nonce) {
182 return Err(RelayError::AuthenticationFailed {
183 reason: "Token replay detected".to_string(),
184 });
185 }
186
187 if used_nonces.len() >= self.replay_window_size as usize {
189 let to_remove: Vec<_> = used_nonces.iter().take(100).cloned().collect();
191 for nonce in to_remove {
192 used_nonces.remove(&nonce);
193 }
194 }
195
196 used_nonces.insert(token.nonce);
197
198 Ok(())
199 }
200
201 pub fn set_max_token_age(&mut self, max_age_seconds: u64) {
203 self.max_token_age = max_age_seconds;
204 }
205
206 pub fn max_token_age(&self) -> u64 {
208 self.max_token_age
209 }
210
211 #[allow(clippy::unwrap_used, clippy::expect_used)]
213 pub fn clear_nonces(&self) {
214 let mut used_nonces = self
215 .used_nonces
216 .lock()
217 .expect("Mutex poisoning is unexpected in normal operation");
218 used_nonces.clear();
219 }
220
221 #[allow(clippy::unwrap_used, clippy::expect_used)]
223 pub fn nonce_count(&self) -> usize {
224 let used_nonces = self
225 .used_nonces
226 .lock()
227 .expect("Mutex poisoning is unexpected in normal operation");
228 used_nonces.len()
229 }
230}
231
232impl Default for RelayAuthenticator {
233 fn default() -> Self {
234 Self::new()
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241 use std::thread;
242 use std::time::Duration;
243
244 #[test]
245 fn test_auth_token_creation_and_verification() {
246 let authenticator = RelayAuthenticator::new();
247 let token = authenticator.create_token(1024, 300).unwrap();
248
249 assert!(token.bandwidth_limit == 1024);
250 assert!(token.timeout_seconds == 300);
251 assert!(token.nonce != 0);
252 assert!(token.timestamp > 0);
253
254 assert!(token.verify(authenticator.verifying_key()).is_ok());
256 }
257
258 #[test]
259 fn test_token_verification_with_wrong_key() {
260 let authenticator1 = RelayAuthenticator::new();
261 let authenticator2 = RelayAuthenticator::new();
262
263 let token = authenticator1.create_token(1024, 300).unwrap();
264
265 assert!(token.verify(authenticator2.verifying_key()).is_err());
267 }
268
269 #[test]
270 fn test_token_expiration() {
271 let mut authenticator = RelayAuthenticator::new();
272 authenticator.set_max_token_age(1); let token = authenticator.create_token(1024, 300).unwrap();
275
276 let max_age = authenticator.max_token_age();
278 assert!(!token.is_expired(max_age).unwrap());
279
280 thread::sleep(Duration::from_secs(2)); assert!(token.is_expired(max_age).unwrap());
285 }
286
287 #[test]
288 fn test_anti_replay_protection() {
289 let authenticator = RelayAuthenticator::new();
290 let token = authenticator.create_token(1024, 300).unwrap();
291
292 assert!(
294 authenticator
295 .verify_token(&token, authenticator.verifying_key())
296 .is_ok()
297 );
298
299 assert!(
301 authenticator
302 .verify_token(&token, authenticator.verifying_key())
303 .is_err()
304 );
305 }
306
307 #[test]
308 fn test_nonce_uniqueness() {
309 let authenticator = RelayAuthenticator::new();
310 let mut nonces = HashSet::new();
311
312 for _ in 0..1000 {
314 let token = authenticator.create_token(1024, 300).unwrap();
315 assert!(!nonces.contains(&token.nonce), "Duplicate nonce detected");
316 nonces.insert(token.nonce);
317 }
318 }
319
320 #[test]
321 fn test_token_signable_data() {
322 let authenticator = RelayAuthenticator::new();
323 let token1 = authenticator.create_token(1024, 300).unwrap();
324 let token2 = authenticator.create_token(1024, 300).unwrap();
325
326 assert_ne!(token1.signable_data(), token2.signable_data());
328 }
329
330 #[test]
331 fn test_nonce_window_management() {
332 let authenticator = RelayAuthenticator::new();
333
334 for _ in 0..1000 {
336 let token = authenticator.create_token(1024, 300).unwrap();
337 let _ = authenticator.verify_token(&token, authenticator.verifying_key());
338 }
339
340 assert_eq!(authenticator.nonce_count(), 1000);
341
342 let token = authenticator.create_token(1024, 300).unwrap();
344 let _ = authenticator.verify_token(&token, authenticator.verifying_key());
345
346 assert!(authenticator.nonce_count() <= 1000);
348 }
349
350 #[test]
351 fn test_clear_nonces() {
352 let authenticator = RelayAuthenticator::new();
353 let token = authenticator.create_token(1024, 300).unwrap();
354
355 let _ = authenticator.verify_token(&token, authenticator.verifying_key());
357 assert!(authenticator.nonce_count() > 0);
358
359 authenticator.clear_nonces();
361 assert_eq!(authenticator.nonce_count(), 0);
362
363 assert!(
365 authenticator
366 .verify_token(&token, authenticator.verifying_key())
367 .is_ok()
368 );
369 }
370
371 #[test]
372 fn test_with_specific_key() {
373 let signing_key = SigningKey::generate(&mut OsRng);
374 let authenticator = RelayAuthenticator::with_key(signing_key);
375
376 let token = authenticator.create_token(1024, 300).unwrap();
377 assert!(token.verify(authenticator.verifying_key()).is_ok());
378 }
379}