1use std::net::SocketAddr;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11use crate::bytes::{Reader, Writer};
12use crate::crypto::{self, XNONCE_BYTES};
13use crate::{
14 CONNECT_TOKEN_BYTES, Error, KEY_BYTES, Key, MAX_SERVERS_PER_CONNECT, USER_DATA_BYTES, UserData,
15 VERSION_INFO,
16};
17
18pub(crate) const CONNECT_TOKEN_PRIVATE_BYTES: usize = 1024;
19pub(crate) const CONNECT_TOKEN_NONCE_BYTES: usize = XNONCE_BYTES;
20pub(crate) const CHALLENGE_TOKEN_BYTES: usize = 300;
21
22fn connect_token_additional_data(protocol_id: u64, expire_timestamp: u64) -> [u8; 13 + 8 + 8] {
26 let mut additional_data = [0u8; 13 + 8 + 8];
27 let mut writer = Writer::new(&mut additional_data);
28 writer.write_bytes(&VERSION_INFO);
29 writer.write_u64(protocol_id);
30 writer.write_u64(expire_timestamp);
31 additional_data
32}
33
34#[cfg_attr(fuzzing, derive(Debug, PartialEq))]
39pub(crate) struct PrivateConnectToken {
40 pub client_id: u64,
41 pub timeout_seconds: i32,
42 pub server_addresses: Vec<SocketAddr>,
43 pub client_to_server_key: Key,
44 pub server_to_client_key: Key,
45 pub user_data: UserData,
46}
47
48impl PrivateConnectToken {
49 pub fn generate(
50 client_id: u64,
51 timeout_seconds: i32,
52 server_addresses: &[SocketAddr],
53 user_data: &UserData,
54 ) -> Self {
55 debug_assert!(!server_addresses.is_empty());
56 debug_assert!(server_addresses.len() <= MAX_SERVERS_PER_CONNECT);
57 Self {
58 client_id,
59 timeout_seconds,
60 server_addresses: server_addresses.to_vec(),
61 client_to_server_key: crypto::generate_key(),
62 server_to_client_key: crypto::generate_key(),
63 user_data: *user_data,
64 }
65 }
66
67 pub fn write(&self, buffer: &mut [u8; CONNECT_TOKEN_PRIVATE_BYTES]) {
68 let mut writer = Writer::new(buffer);
69 writer.write_u64(self.client_id);
70 writer.write_u32(self.timeout_seconds as u32);
71 writer.write_u32(self.server_addresses.len() as u32);
72 for &address in &self.server_addresses {
73 writer.write_address(address);
74 }
75 writer.write_bytes(&self.client_to_server_key);
76 writer.write_bytes(&self.server_to_client_key);
77 writer.write_bytes(&self.user_data);
78 writer.zero_pad_to_end();
79 }
80
81 pub fn read(buffer: &[u8]) -> Result<Self, Error> {
82 if buffer.len() < CONNECT_TOKEN_PRIVATE_BYTES {
83 return Err(Error::InvalidConnectToken);
84 }
85
86 let mut reader = Reader::new(buffer);
87 (|| {
88 let client_id = reader.read_u64()?;
89 let timeout_seconds = reader.read_u32()? as i32;
90
91 let num_server_addresses = reader.read_u32()? as usize;
92 if !(1..=MAX_SERVERS_PER_CONNECT).contains(&num_server_addresses) {
93 return None;
94 }
95
96 let mut server_addresses = Vec::with_capacity(num_server_addresses);
97 for _ in 0..num_server_addresses {
98 server_addresses.push(reader.read_address()?);
99 }
100
101 Some(Self {
102 client_id,
103 timeout_seconds,
104 server_addresses,
105 client_to_server_key: reader.read_bytes::<KEY_BYTES>()?,
106 server_to_client_key: reader.read_bytes::<KEY_BYTES>()?,
107 user_data: reader.read_bytes::<USER_DATA_BYTES>()?,
108 })
109 })()
110 .ok_or(Error::InvalidConnectToken)
111 }
112}
113
114pub(crate) fn encrypt_connect_token_private(
117 buffer: &mut [u8; CONNECT_TOKEN_PRIVATE_BYTES],
118 protocol_id: u64,
119 expire_timestamp: u64,
120 nonce: &[u8; CONNECT_TOKEN_NONCE_BYTES],
121 key: &Key,
122) -> Result<(), Error> {
123 let additional_data = connect_token_additional_data(protocol_id, expire_timestamp);
124 crypto::encrypt_aead_big_nonce(buffer, &additional_data, nonce, key)
125}
126
127pub(crate) fn decrypt_connect_token_private(
130 buffer: &mut [u8; CONNECT_TOKEN_PRIVATE_BYTES],
131 protocol_id: u64,
132 expire_timestamp: u64,
133 nonce: &[u8; CONNECT_TOKEN_NONCE_BYTES],
134 key: &Key,
135) -> Result<(), Error> {
136 let additional_data = connect_token_additional_data(protocol_id, expire_timestamp);
137 crypto::decrypt_aead_big_nonce(buffer, &additional_data, nonce, key)
138}
139
140pub(crate) struct ChallengeToken {
143 pub client_id: u64,
144 pub user_data: UserData,
145}
146
147impl ChallengeToken {
148 pub fn write(&self, buffer: &mut [u8; CHALLENGE_TOKEN_BYTES]) {
149 let mut writer = Writer::new(buffer);
150 writer.write_u64(self.client_id);
151 writer.write_bytes(&self.user_data);
152 writer.zero_pad_to_end();
153 }
154
155 pub fn read(buffer: &[u8; CHALLENGE_TOKEN_BYTES]) -> Self {
156 let mut reader = Reader::new(buffer);
157 Self {
158 client_id: reader.read_u64().unwrap(),
159 user_data: reader.read_bytes::<USER_DATA_BYTES>().unwrap(),
160 }
161 }
162}
163
164pub(crate) fn encrypt_challenge_token(
167 buffer: &mut [u8; CHALLENGE_TOKEN_BYTES],
168 sequence: u64,
169 key: &Key,
170) -> Result<(), Error> {
171 crypto::encrypt_aead(buffer, &[], &crypto::sequence_nonce(sequence), key)
172}
173
174pub(crate) fn decrypt_challenge_token(
175 buffer: &mut [u8; CHALLENGE_TOKEN_BYTES],
176 sequence: u64,
177 key: &Key,
178) -> Result<(), Error> {
179 crypto::decrypt_aead(buffer, &[], &crypto::sequence_nonce(sequence), key)
180}
181
182#[cfg_attr(fuzzing, derive(Debug, PartialEq))]
187pub(crate) struct ConnectToken {
188 pub protocol_id: u64,
189 pub create_timestamp: u64,
190 pub expire_timestamp: u64,
191 pub nonce: [u8; CONNECT_TOKEN_NONCE_BYTES],
192 pub private_data: Box<[u8; CONNECT_TOKEN_PRIVATE_BYTES]>,
193 pub timeout_seconds: i32,
194 pub server_addresses: Vec<SocketAddr>,
195 pub client_to_server_key: Key,
196 pub server_to_client_key: Key,
197}
198
199impl ConnectToken {
200 pub fn write(&self, buffer: &mut [u8; CONNECT_TOKEN_BYTES]) {
201 let mut writer = Writer::new(buffer);
202 writer.write_bytes(&VERSION_INFO);
203 writer.write_u64(self.protocol_id);
204 writer.write_u64(self.create_timestamp);
205 writer.write_u64(self.expire_timestamp);
206 writer.write_bytes(&self.nonce);
207 writer.write_bytes(&self.private_data[..]);
208 writer.write_u32(self.timeout_seconds as u32);
209 writer.write_u32(self.server_addresses.len() as u32);
210 for &address in &self.server_addresses {
211 writer.write_address(address);
212 }
213 writer.write_bytes(&self.client_to_server_key);
214 writer.write_bytes(&self.server_to_client_key);
215 writer.zero_pad_to_end();
216 }
217
218 pub fn read(buffer: &[u8; CONNECT_TOKEN_BYTES]) -> Result<Self, Error> {
219 let mut reader = Reader::new(buffer);
220 (|| {
221 if reader.read_bytes::<13>()? != VERSION_INFO {
222 return None;
223 }
224
225 let protocol_id = reader.read_u64()?;
226 let create_timestamp = reader.read_u64()?;
227 let expire_timestamp = reader.read_u64()?;
228 if create_timestamp > expire_timestamp {
229 return None;
230 }
231
232 let nonce = reader.read_bytes::<CONNECT_TOKEN_NONCE_BYTES>()?;
233 let private_data = Box::new(reader.read_bytes::<CONNECT_TOKEN_PRIVATE_BYTES>()?);
234 let timeout_seconds = reader.read_u32()? as i32;
235
236 let num_server_addresses = reader.read_u32()? as usize;
237 if !(1..=MAX_SERVERS_PER_CONNECT).contains(&num_server_addresses) {
238 return None;
239 }
240
241 let mut server_addresses = Vec::with_capacity(num_server_addresses);
242 for _ in 0..num_server_addresses {
243 server_addresses.push(reader.read_address()?);
244 }
245
246 Some(Self {
247 protocol_id,
248 create_timestamp,
249 expire_timestamp,
250 nonce,
251 private_data,
252 timeout_seconds,
253 server_addresses,
254 client_to_server_key: reader.read_bytes::<KEY_BYTES>()?,
255 server_to_client_key: reader.read_bytes::<KEY_BYTES>()?,
256 })
257 })()
258 .ok_or(Error::InvalidConnectToken)
259 }
260}
261
262pub(crate) fn unix_timestamp() -> u64 {
263 SystemTime::now().duration_since(UNIX_EPOCH).expect("system clock is before 1970").as_secs()
264}
265
266#[allow(clippy::too_many_arguments)]
282pub fn generate_connect_token(
283 public_server_addresses: &[SocketAddr],
284 internal_server_addresses: &[SocketAddr],
285 expire_seconds: i32,
286 timeout_seconds: i32,
287 client_id: u64,
288 protocol_id: u64,
289 private_key: &Key,
290 user_data: &UserData,
291) -> Result<[u8; CONNECT_TOKEN_BYTES], Error> {
292 if public_server_addresses.is_empty()
293 || public_server_addresses.len() > MAX_SERVERS_PER_CONNECT
294 || public_server_addresses.len() != internal_server_addresses.len()
295 {
296 return Err(Error::InvalidServerAddresses);
297 }
298
299 let create_timestamp = unix_timestamp();
300 let expire_timestamp =
301 if expire_seconds >= 0 { create_timestamp + expire_seconds as u64 } else { u64::MAX };
302
303 let mut nonce = [0u8; CONNECT_TOKEN_NONCE_BYTES];
304 crypto::random_bytes(&mut nonce);
305
306 let private_token = PrivateConnectToken::generate(
307 client_id,
308 timeout_seconds,
309 internal_server_addresses,
310 user_data,
311 );
312
313 let mut private_data = Box::new([0u8; CONNECT_TOKEN_PRIVATE_BYTES]);
314 private_token.write(&mut private_data);
315 encrypt_connect_token_private(
316 &mut private_data,
317 protocol_id,
318 expire_timestamp,
319 &nonce,
320 private_key,
321 )?;
322
323 let connect_token = ConnectToken {
324 protocol_id,
325 create_timestamp,
326 expire_timestamp,
327 nonce,
328 private_data,
329 timeout_seconds,
330 server_addresses: public_server_addresses.to_vec(),
331 client_to_server_key: private_token.client_to_server_key,
332 server_to_client_key: private_token.server_to_client_key,
333 };
334
335 let mut buffer = [0u8; CONNECT_TOKEN_BYTES];
336 connect_token.write(&mut buffer);
337 Ok(buffer)
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343 use crate::generate_key;
344
345 const TEST_PROTOCOL_ID: u64 = 0x1122334455667788;
346
347 fn test_user_data() -> UserData {
348 let mut user_data = [0u8; USER_DATA_BYTES];
349 crypto::random_bytes(&mut user_data);
350 user_data
351 }
352
353 #[test]
354 fn private_connect_token_round_trip() {
355 let server_addresses: Vec<SocketAddr> =
356 vec!["127.0.0.1:40000".parse().unwrap(), "[::1]:50000".parse().unwrap()];
357 let user_data = test_user_data();
358 let token = PrivateConnectToken::generate(0x1234, 10, &server_addresses, &user_data);
359
360 let key = generate_key();
361 let mut nonce = [0u8; CONNECT_TOKEN_NONCE_BYTES];
362 crypto::random_bytes(&mut nonce);
363 let expire_timestamp = unix_timestamp() + 30;
364
365 let mut buffer = Box::new([0u8; CONNECT_TOKEN_PRIVATE_BYTES]);
366 token.write(&mut buffer);
367 encrypt_connect_token_private(
368 &mut buffer,
369 TEST_PROTOCOL_ID,
370 expire_timestamp,
371 &nonce,
372 &key,
373 )
374 .unwrap();
375 decrypt_connect_token_private(
376 &mut buffer,
377 TEST_PROTOCOL_ID,
378 expire_timestamp,
379 &nonce,
380 &key,
381 )
382 .unwrap();
383
384 let output = PrivateConnectToken::read(&buffer[..]).unwrap();
385 assert_eq!(output.client_id, token.client_id);
386 assert_eq!(output.timeout_seconds, token.timeout_seconds);
387 assert_eq!(output.server_addresses, token.server_addresses);
388 assert_eq!(output.client_to_server_key, token.client_to_server_key);
389 assert_eq!(output.server_to_client_key, token.server_to_client_key);
390 assert_eq!(output.user_data, token.user_data);
391 }
392
393 #[test]
394 fn private_connect_token_rejects_modified_associated_data() {
395 let server_addresses: Vec<SocketAddr> = vec!["127.0.0.1:40000".parse().unwrap()];
396 let user_data = test_user_data();
397 let token = PrivateConnectToken::generate(0x1234, 10, &server_addresses, &user_data);
398
399 let key = generate_key();
400 let mut nonce = [0u8; CONNECT_TOKEN_NONCE_BYTES];
401 crypto::random_bytes(&mut nonce);
402 let expire_timestamp = unix_timestamp() + 30;
403
404 let mut buffer = Box::new([0u8; CONNECT_TOKEN_PRIVATE_BYTES]);
405 token.write(&mut buffer);
406 encrypt_connect_token_private(
407 &mut buffer,
408 TEST_PROTOCOL_ID,
409 expire_timestamp,
410 &nonce,
411 &key,
412 )
413 .unwrap();
414
415 assert!(
417 decrypt_connect_token_private(
418 &mut buffer.clone(),
419 TEST_PROTOCOL_ID + 1,
420 expire_timestamp,
421 &nonce,
422 &key
423 )
424 .is_err()
425 );
426 assert!(
427 decrypt_connect_token_private(
428 &mut buffer.clone(),
429 TEST_PROTOCOL_ID,
430 expire_timestamp + 1,
431 &nonce,
432 &key
433 )
434 .is_err()
435 );
436 }
437
438 #[test]
439 fn challenge_token_round_trip() {
440 let token = ChallengeToken { client_id: 0xDEADBEEF, user_data: test_user_data() };
441
442 let key = generate_key();
443 let sequence = 42;
444
445 let mut buffer = [0u8; CHALLENGE_TOKEN_BYTES];
446 token.write(&mut buffer);
447 encrypt_challenge_token(&mut buffer, sequence, &key).unwrap();
448 decrypt_challenge_token(&mut buffer, sequence, &key).unwrap();
449
450 let output = ChallengeToken::read(&buffer);
451 assert_eq!(output.client_id, token.client_id);
452 assert_eq!(output.user_data, token.user_data);
453 }
454
455 #[test]
456 fn connect_token_round_trip() {
457 let server_address: SocketAddr = "127.0.0.1:40000".parse().unwrap();
458 let private_key = generate_key();
459 let user_data = test_user_data();
460
461 let buffer = generate_connect_token(
462 &[server_address],
463 &[server_address],
464 30,
465 5,
466 0x1234,
467 TEST_PROTOCOL_ID,
468 &private_key,
469 &user_data,
470 )
471 .unwrap();
472
473 let token = ConnectToken::read(&buffer).unwrap();
474 assert_eq!(token.protocol_id, TEST_PROTOCOL_ID);
475 assert_eq!(token.expire_timestamp, token.create_timestamp + 30);
476 assert_eq!(token.timeout_seconds, 5);
477 assert_eq!(token.server_addresses, vec![server_address]);
478
479 let mut private_data = token.private_data.clone();
481 decrypt_connect_token_private(
482 &mut private_data,
483 TEST_PROTOCOL_ID,
484 token.expire_timestamp,
485 &token.nonce,
486 &private_key,
487 )
488 .unwrap();
489 let private_token = PrivateConnectToken::read(&private_data[..]).unwrap();
490 assert_eq!(private_token.client_id, 0x1234);
491 assert_eq!(private_token.timeout_seconds, 5);
492 assert_eq!(private_token.server_addresses, vec![server_address]);
493 assert_eq!(private_token.client_to_server_key, token.client_to_server_key);
494 assert_eq!(private_token.server_to_client_key, token.server_to_client_key);
495 assert_eq!(private_token.user_data, user_data);
496 }
497
498 #[test]
499 fn connect_token_rejects_bad_version_info() {
500 let server_address: SocketAddr = "127.0.0.1:40000".parse().unwrap();
501 let mut buffer = generate_connect_token(
502 &[server_address],
503 &[server_address],
504 30,
505 5,
506 1,
507 TEST_PROTOCOL_ID,
508 &generate_key(),
509 &[0u8; USER_DATA_BYTES],
510 )
511 .unwrap();
512
513 buffer[0] = b'X';
514 assert!(ConnectToken::read(&buffer).is_err());
515 }
516
517 #[test]
518 fn connect_token_rejects_create_after_expire() {
519 let server_address: SocketAddr = "127.0.0.1:40000".parse().unwrap();
520 let mut buffer = generate_connect_token(
521 &[server_address],
522 &[server_address],
523 30,
524 5,
525 1,
526 TEST_PROTOCOL_ID,
527 &generate_key(),
528 &[0u8; USER_DATA_BYTES],
529 )
530 .unwrap();
531
532 let create: [u8; 8] = buffer[21..29].try_into().unwrap();
534 let expire: [u8; 8] = buffer[29..37].try_into().unwrap();
535 buffer[21..29].copy_from_slice(&expire);
536 buffer[29..37].copy_from_slice(&create);
537 assert!(ConnectToken::read(&buffer).is_err());
538 }
539}