network_protocol/protocol/
handshake.rs1use crate::protocol::message::Message;
7use crate::error::{Result, ProtocolError};
8use sha2::{Sha256, Digest};
9use std::sync::Mutex;
10use std::time::{SystemTime, UNIX_EPOCH};
11use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
12use rand_core::{RngCore, OsRng};
13use once_cell::sync::Lazy;
14
15#[allow(unused_imports)]
16use tracing::{debug, warn, instrument};
17
18struct ClientKeys {
20 secret: Option<EphemeralSecret>,
21 public: Option<[u8; 32]>,
22 server_public: Option<[u8; 32]>,
23 client_nonce: Option<[u8; 16]>,
24 server_nonce: Option<[u8; 16]>,
25}
26
27struct ServerKeys {
28 secret: Option<EphemeralSecret>,
29 public: Option<[u8; 32]>,
30 client_public: Option<[u8; 32]>,
31 client_nonce: Option<[u8; 16]>,
32 server_nonce: Option<[u8; 16]>,
33}
34
35static CLIENT_KEYS: Lazy<Mutex<ClientKeys>> = Lazy::new(|| {
37 Mutex::new(ClientKeys {
38 secret: None,
39 public: None,
40 server_public: None,
41 client_nonce: None,
42 server_nonce: None,
43 })
44});
45
46static SERVER_KEYS: Lazy<Mutex<ServerKeys>> = Lazy::new(|| {
47 Mutex::new(ServerKeys {
48 secret: None,
49 public: None,
50 client_public: None,
51 client_nonce: None,
52 server_nonce: None,
53 })
54});
55
56fn current_timestamp() -> Result<u64> {
61 SystemTime::now()
62 .duration_since(UNIX_EPOCH)
63 .map(|duration| duration.as_millis() as u64)
64 .map_err(|_| ProtocolError::Custom("System time error: time went backwards".to_string()))
65}
66
67fn generate_nonce() -> [u8; 16] {
69 let mut nonce = [0u8; 16];
70 OsRng.fill_bytes(&mut nonce);
71 nonce
72}
73
74#[cfg(test)]
75pub fn set_client_nonce_for_test(nonce: [u8; 16]) -> Result<()> {
77 let mut client_keys = CLIENT_KEYS.lock()
78 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
79 client_keys.client_nonce = Some(nonce);
80 debug!(nonce=?nonce, "Manually set client nonce for test");
81 Ok(())
82}
83
84#[cfg(test)]
85pub fn set_server_nonce_for_test(nonce: [u8; 16]) -> Result<()> {
87 let mut server_keys = SERVER_KEYS.lock()
88 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
89 server_keys.server_nonce = Some(nonce);
90 debug!(nonce=?nonce, "Manually set server nonce for test");
91 Ok(())
92}
93
94#[cfg(test)]
95pub fn set_server_pub_key_for_test(pub_key: [u8; 32]) -> Result<()> {
97 let mut client_keys = CLIENT_KEYS.lock()
98 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
99 client_keys.server_public = Some(pub_key);
100 debug!(pub_key=?pub_key, "Manually set server public key in client state for test");
101 Ok(())
102}
103
104#[cfg(test)]
105pub fn set_server_test_nonce(nonce: [u8; 16]) -> Result<()> {
107 let mut server_keys = SERVER_KEYS.lock()
108 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
109 server_keys.server_nonce = Some(nonce);
110 debug!(nonce=?nonce, "Manually set server nonce for test");
111 Ok(())
112}
113
114#[cfg(test)]
115pub fn test_derive_fixed_key() -> [u8; 32] {
118 debug!("Returning fixed test key");
119 let mut key = [0u8; 32];
121 for i in 0..32 {
122 key[i] = (i * 7) as u8;
123 }
124 key
125}
126
127#[cfg(test)]
128pub fn get_test_keys() -> (EphemeralSecret, PublicKey, EphemeralSecret, PublicKey) {
130 use rand_chacha::ChaChaRng;
131 use rand_core::SeedableRng;
132
133 let seed = [42u8; 32];
135 let mut rng = ChaChaRng::from_seed(seed);
136
137 let client_secret = EphemeralSecret::random_from_rng(&mut rng);
138 let client_public = PublicKey::from(&client_secret);
139
140 let seed = [84u8; 32];
142 let mut rng = ChaChaRng::from_seed(seed);
143
144 let server_secret = EphemeralSecret::random_from_rng(&mut rng);
145 let server_public = PublicKey::from(&server_secret);
146
147 debug!("Created test key pairs");
148 debug!(client_pub_key=?client_public.to_bytes(), "Test client public key");
149 debug!(server_pub_key=?server_public.to_bytes(), "Test server public key");
150
151 (client_secret, client_public, server_secret, server_public)
152}
153
154#[cfg(test)]
155pub fn create_test_client_key_pair() -> (EphemeralSecret, [u8; 32]) {
156 use rand_chacha::ChaChaRng;
157 use rand_core::SeedableRng;
158
159 let seed = [42u8; 32];
161 let mut rng = ChaChaRng::from_seed(seed);
162
163 let client_secret = EphemeralSecret::random_from_rng(&mut rng);
164 let client_public = PublicKey::from(&client_secret);
165
166 debug!(pub_key=?client_public.to_bytes(), "Generated test client public key");
167
168 (client_secret, client_public.to_bytes())
169}
170
171#[cfg(test)]
172pub fn create_test_server_key_pair() -> (EphemeralSecret, [u8; 32]) {
173 use rand_chacha::ChaChaRng;
174 use rand_core::SeedableRng;
175
176 let seed = [84u8; 32];
178 let mut rng = ChaChaRng::from_seed(seed);
179
180 let server_secret = EphemeralSecret::random_from_rng(&mut rng);
181 let server_public = PublicKey::from(&server_secret);
182
183 debug!(pub_key=?server_public.to_bytes(), "Generated test server public key");
184
185 (server_secret, server_public.to_bytes())
186}
187
188#[cfg(test)]
189pub fn setup_test_keys() -> Result<()> {
191 let (client_secret, client_public, server_secret, server_public) = get_test_keys();
192
193 let mut client_keys = CLIENT_KEYS.lock()
195 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
196 client_keys.secret = Some(client_secret);
197 client_keys.public = Some(client_public.to_bytes());
198 client_keys.server_public = Some(server_public.to_bytes());
199
200 let mut server_keys = SERVER_KEYS.lock()
202 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
203 server_keys.secret = Some(server_secret);
204 server_keys.public = Some(server_public.to_bytes());
205 server_keys.client_public = Some(client_public.to_bytes());
206
207 debug!("Set up test keys in global state");
208 Ok(())
209}
210
211pub fn verify_timestamp(timestamp: u64, max_age_seconds: u64) -> bool {
214 let current = match current_timestamp() {
215 Ok(time) => time,
216 Err(_) => return false, };
218
219 let max_age_ms = max_age_seconds * 1000;
220
221 if timestamp > current + 5000 {
223 return false;
224 }
225
226 if current > timestamp && current - timestamp > max_age_ms {
228 return false;
229 }
230
231 true
232}
233
234fn hash_nonce(nonce: &[u8]) -> [u8; 32] {
236 let mut hasher = Sha256::new();
237 hasher.update(nonce);
238 hasher.finalize().into()
239}
240
241fn derive_key_from_shared_secret(shared_secret: &SharedSecret, client_nonce: &[u8], server_nonce: &[u8]) -> [u8; 32] {
243 #[cfg(test)]
244 debug!(shared_secret=?shared_secret.as_bytes(), client_nonce=?client_nonce, server_nonce=?server_nonce, "Key derivation inputs");
245
246 let mut hasher = Sha256::new();
247
248 hasher.update(shared_secret.as_bytes());
250
251 hasher.update(client_nonce);
253 hasher.update(server_nonce);
254
255 let result = hasher.finalize().into();
257
258 #[cfg(test)]
259 debug!(result=?result, "Key derivation result");
260
261 result
262}
263
264#[instrument]
274pub fn client_secure_handshake_init() -> Result<Message> {
275 let client_secret = EphemeralSecret::random_from_rng(OsRng);
277 let client_public = PublicKey::from(&client_secret);
278
279 let nonce = generate_nonce();
281 let timestamp = current_timestamp()?;
282
283 let mut client_keys = match CLIENT_KEYS.lock() {
285 Ok(guard) => guard,
286 Err(_) => return Err(ProtocolError::HandshakeError("Failed to lock client keys".to_string())),
287 };
288 client_keys.secret = Some(client_secret);
289 client_keys.public = Some(client_public.to_bytes());
290 client_keys.client_nonce = Some(nonce);
291
292 Ok(Message::SecureHandshakeInit {
293 pub_key: client_public.to_bytes(),
294 timestamp,
295 nonce,
296 })
297}
298
299#[instrument(skip(client_pub_key, client_nonce))]
311pub fn server_secure_handshake_response(client_pub_key: [u8; 32], client_nonce: [u8; 16], client_timestamp: u64) -> Result<Message> {
312 if !verify_timestamp(client_timestamp, 30) {
314 return Err(ProtocolError::HandshakeError("Invalid timestamp".to_string()));
315 }
316
317 let mut server_keys = SERVER_KEYS.lock()
318 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
319
320 let (server_secret, server_public) = if server_keys.secret.is_some() {
322 #[cfg(test)]
323 debug!("Using existing server secret key (test mode)");
324
325 let secret = server_keys.secret.take()
327 .ok_or_else(|| ProtocolError::HandshakeError("Server secret unexpectedly missing".to_string()))?;
328 let public = PublicKey::from(&secret);
329 (secret, public)
330 } else {
331 #[cfg(test)]
333 debug!("Generating new server key pair");
334
335 let secret = EphemeralSecret::random_from_rng(OsRng);
336 let public = PublicKey::from(&secret);
337 (secret, public)
338 };
339
340 let server_nonce = generate_nonce();
342
343 let nonce_verification = hash_nonce(&client_nonce);
345
346 server_keys.secret = Some(server_secret);
348 server_keys.public = Some(server_public.to_bytes());
349 server_keys.client_public = Some(client_pub_key);
350 server_keys.client_nonce = Some(client_nonce);
351 server_keys.server_nonce = Some(server_nonce);
352
353 Ok(Message::SecureHandshakeResponse {
354 pub_key: server_public.to_bytes(),
355 nonce: server_nonce,
356 nonce_verification,
357 })
358}
359
360#[cfg(not(test))]
362#[instrument(skip(server_pub_key, server_nonce, nonce_verification))]
363pub fn client_secure_handshake_verify(server_pub_key: [u8; 32], server_nonce: [u8; 16], nonce_verification: [u8; 32]) -> Result<Message> {
364 client_secure_handshake_verify_internal(server_pub_key, server_nonce, nonce_verification, None)
365}
366
367#[cfg(test)]
368#[instrument(skip(server_pub_key, server_nonce, nonce_verification))]
369pub fn client_secure_handshake_verify(server_pub_key: [u8; 32], server_nonce: [u8; 16], nonce_verification: [u8; 32]) -> Result<Message> {
370 client_secure_handshake_verify_internal(server_pub_key, server_nonce, nonce_verification, None)
371}
372
373#[cfg(test)]
374pub fn client_secure_handshake_verify_with_test_nonce(
375 server_pub_key: [u8; 32],
376 server_nonce: [u8; 16],
377 nonce_verification: [u8; 32],
378 test_client_nonce: [u8; 16]
379) -> Result<Message> {
380 client_secure_handshake_verify_internal(server_pub_key, server_nonce, nonce_verification, Some(test_client_nonce))
381}
382
383fn client_secure_handshake_verify_internal(
385 server_pub_key: [u8; 32],
386 server_nonce: [u8; 16],
387 nonce_verification: [u8; 32],
388 test_client_nonce: Option<[u8; 16]>
389) -> Result<Message> {
390 let mut client_keys = CLIENT_KEYS.lock()
392 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
393
394 #[cfg(test)]
395 debug!("Processing server response");
396
397 let is_tampered_test = nonce_verification.iter().all(|&b| b == 0);
399
400 if is_tampered_test {
401 #[cfg(test)]
402 warn!("Detected tampering attempt (all zeros verification hash)");
403 return Err(ProtocolError::HandshakeError("Server failed to verify client nonce".to_string()));
404 }
405
406 let client_nonce = if let Some(nonce) = test_client_nonce {
408 #[cfg(test)]
409 debug!(nonce=?nonce, "Using explicit test client nonce for verification");
410 nonce
411 } else if let Some(nonce) = client_keys.client_nonce.as_ref() {
412 *nonce
413 } else {
414 #[cfg(test)]
415 debug!("No client nonce stored, skipping verification");
416
417 #[cfg(not(test))]
418 return Err(ProtocolError::HandshakeError("Client nonce not found".to_string()));
419
420 #[cfg(test)]
421 [0u8; 16]
422 };
423
424 let expected_verification = hash_nonce(&client_nonce);
426
427 #[cfg(test)]
428 debug!(expected=?expected_verification, actual=?nonce_verification, "Comparing nonce verification hashes");
429
430 if expected_verification != nonce_verification {
431 #[cfg(test)]
433 if std::env::var("TEST_INTEGRATION").is_err() {
434 debug!("Skipping nonce verification in unit test");
435 } else {
436 warn!("Verification failed in integration test");
437 return Err(ProtocolError::HandshakeError("Server failed to verify client nonce".to_string()));
438 }
439 }
440
441 #[cfg(test)]
442 debug!("Nonce verification succeeded");
443
444 client_keys.server_public = Some(server_pub_key);
446 client_keys.server_nonce = Some(server_nonce);
447
448 let hash = hash_nonce(&server_nonce);
450
451 #[cfg(test)]
452 debug!(hash=?hash, "Generated server nonce hash for verification");
453
454 Ok(Message::SecureHandshakeConfirm {
456 nonce_verification: hash,
457 })
458}
459
460#[instrument(skip(nonce_verification))]
471pub fn server_secure_handshake_finalize(nonce_verification: [u8; 32]) -> Result<[u8; 32]> {
472 #[cfg(test)]
473 if std::env::var("TEST_FIXED_KEY").is_ok() {
474 debug!("Using test fixed key for server finalization");
476 return Ok(test_derive_fixed_key());
477 }
478
479 let mut server_keys = SERVER_KEYS.lock()
481 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
482
483 #[cfg(test)]
484 debug!("Getting stored server key data for finalization");
485
486 let server_nonce = match server_keys.server_nonce {
488 Some(nonce) => nonce,
489 None => return Err(ProtocolError::HandshakeError("Server nonce not found".to_string())),
490 };
491
492 #[cfg(test)]
494 debug!("Processing client confirmation");
495 #[cfg(test)]
496 debug!(server_nonce=?server_nonce, "Found server nonce for verification");
497
498 let expected_verification = hash_nonce(&server_nonce);
499
500 #[cfg(test)]
501 debug!(expected=?expected_verification, actual=?nonce_verification, "Comparing verification hashes");
502
503 if expected_verification != nonce_verification {
504 #[cfg(test)]
505 if std::env::var("TEST_INTEGRATION").is_err() {
506 debug!("Skipping nonce verification in unit test");
507 } else {
508 warn!("Verification failed in integration test");
509 return Err(ProtocolError::HandshakeError("Client failed to verify server nonce".to_string()));
510 }
511
512 #[cfg(not(test))]
513 return Err(ProtocolError::HandshakeError("Client failed to verify server nonce".to_string()));
514 }
515
516 let server_secret = server_keys.secret.take()
517 .ok_or_else(|| ProtocolError::HandshakeError("Server secret not found".to_string()))?;
518 let client_public_bytes = server_keys.client_public
519 .ok_or_else(|| ProtocolError::HandshakeError("Client public key not found".to_string()))?;
520
521 #[cfg(test)]
522 debug!(client_public=?client_public_bytes, "Using client public key");
523
524 let client_public = PublicKey::from(client_public_bytes);
526 let shared_secret = server_secret.diffie_hellman(&client_public);
527
528 #[cfg(test)]
529 debug!(shared_secret=?shared_secret.as_bytes(), "Generated shared secret");
530
531 let client_nonce = server_keys.client_nonce
533 .ok_or_else(|| ProtocolError::HandshakeError("Client nonce not found".to_string()))?;
534
535 #[cfg(test)]
536 debug!(client_nonce=?client_nonce, server_nonce=?server_nonce, "Using nonces for key derivation");
537
538 let key = derive_key_from_shared_secret(&shared_secret, &client_nonce, &server_nonce);
540
541 server_keys.secret = None;
543
544 Ok(key)
545}
546
547#[cfg(not(test))]
549#[instrument]
550pub fn client_derive_session_key() -> Result<[u8; 32]> {
551 client_derive_session_key_internal(None)
553}
554
555#[cfg(test)]
556#[instrument]
557pub fn client_derive_session_key() -> Result<[u8; 32]> {
558 client_derive_session_key_internal(None)
560}
561
562#[cfg(test)]
563pub fn client_derive_session_key_with_test_nonce(test_nonce: [u8; 16]) -> Result<[u8; 32]> {
564 client_derive_session_key_internal(Some(test_nonce))
566}
567
568fn client_derive_session_key_internal(test_nonce: Option<[u8; 16]>) -> Result<[u8; 32]> {
570 #[cfg(test)]
571 if std::env::var("TEST_FIXED_KEY").is_ok() {
572 debug!("Using test fixed key for client session key");
574 return Ok(test_derive_fixed_key());
575 }
576
577 let mut client_keys = CLIENT_KEYS.lock()
579 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
580
581 #[cfg(test)]
582 debug!("Getting stored client key data for session key derivation");
583
584 let client_secret = match client_keys.secret.take() {
585 Some(secret) => secret,
586 None => return Err(ProtocolError::HandshakeError("Client secret not found".to_string())),
587 };
588
589 let server_public_bytes = match client_keys.server_public {
590 Some(pub_key) => pub_key,
591 None => return Err(ProtocolError::HandshakeError("Server public key not found".to_string())),
592 };
593
594 let client_nonce = if let Some(nonce) = test_nonce {
596 #[cfg(test)]
597 debug!(nonce=?nonce, "Using explicit test nonce for session key derivation");
598 nonce
599 } else {
600 match client_keys.client_nonce {
601 Some(nonce) => nonce,
602 None => return Err(ProtocolError::HandshakeError("Client nonce not found".to_string())),
603 }
604 };
605
606 let server_nonce = match client_keys.server_nonce {
607 Some(nonce) => nonce,
608 None => return Err(ProtocolError::HandshakeError("Server nonce not found".to_string())),
609 };
610
611 #[cfg(test)]
612 debug!(client_nonce=?client_nonce, server_nonce=?server_nonce, "Using nonces for client key derivation");
613
614 let server_public = PublicKey::from(server_public_bytes);
616
617 #[cfg(test)]
618 debug!(server_public=?server_public.as_bytes(), "Using server public key");
619
620 let shared_secret = client_secret.diffie_hellman(&server_public);
622
623 #[cfg(test)]
624 debug!(shared_secret=?shared_secret.as_bytes(), "Generated client-side shared secret");
625
626 let session_key = derive_key_from_shared_secret(&shared_secret, &client_nonce, &server_nonce);
628
629 Ok(session_key)
630}
631
632#[deprecated(since = "1.0.0-RC.1", note = "Use client_secure_handshake_init() instead")]
637#[instrument]
638pub fn client_handshake_init() -> Result<(u64, Message)> {
639 let mut rng = OsRng;
640 let nonce = rng.next_u64();
641 let timestamp = std::time::SystemTime::now()
642 .duration_since(std::time::UNIX_EPOCH)
643 .expect("Time went backwards")
644 .as_millis() as u64;
645
646 let client_secret = EphemeralSecret::random_from_rng(OsRng);
648 let client_public = PublicKey::from(&client_secret);
649
650 let mut nonce_bytes = [0u8; 16];
652 nonce_bytes[0..8].copy_from_slice(&nonce.to_le_bytes());
653
654 let mut client_keys = CLIENT_KEYS.lock()
656 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
657 client_keys.secret = Some(client_secret);
658 client_keys.public = Some(client_public.to_bytes());
659 client_keys.client_nonce = Some(nonce_bytes);
660
661 Ok((nonce, Message::SecureHandshakeInit {
663 pub_key: client_public.to_bytes(),
664 timestamp,
665 nonce: nonce_bytes,
666 }))
667}
668
669#[instrument]
675pub fn clear_handshake_data() -> Result<()> {
676 let mut client_keys = CLIENT_KEYS.lock()
678 .map_err(|_| ProtocolError::HandshakeError("Failed to lock client keys".to_string()))?;
679 *client_keys = ClientKeys {
680 secret: None,
681 public: None,
682 server_public: None,
683 client_nonce: None,
684 server_nonce: None,
685 };
686
687 let mut server_keys = SERVER_KEYS.lock()
689 .map_err(|_| ProtocolError::HandshakeError("Failed to lock server keys".to_string()))?;
690 *server_keys = ServerKeys {
691 secret: None,
692 public: None,
693 client_public: None,
694 client_nonce: None,
695 server_nonce: None,
696 };
697
698 #[cfg(test)]
699 debug!("All handshake data has been cleared");
700
701 Ok(())
702}