diff --git a/crates/wzp-client/src/bench.rs b/crates/wzp-client/src/bench.rs index df5ff0c..f58aac0 100644 --- a/crates/wzp-client/src/bench.rs +++ b/crates/wzp-client/src/bench.rs @@ -263,17 +263,36 @@ pub fn bench_encrypt_decrypt() -> CryptoResult { }) .collect(); - let header = b"bench-header"; + // Build valid v2 MediaHeader bytes — encrypt/decrypt now derive nonces from + // header.seq and require a parseable MediaHeader (WIRE_SIZE bytes minimum). + use wzp_proto::packet::MediaHeader; + use wzp_proto::{CodecId, MediaType}; let mut total_bytes: usize = 0; let start = Instant::now(); - for payload in &payloads { + for (i, payload) in payloads.iter().enumerate() { + let hdr = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq: i as u32, + timestamp: (i as u32).wrapping_mul(20), + fec_block: 0, + }; + let mut header_bytes = Vec::with_capacity(MediaHeader::WIRE_SIZE); + hdr.write_to(&mut header_bytes); + let mut ciphertext = Vec::with_capacity(payload.len() + 16); - encryptor.encrypt(header, payload, &mut ciphertext).unwrap(); + encryptor + .encrypt(&header_bytes, payload, &mut ciphertext) + .unwrap(); let mut plaintext = Vec::with_capacity(payload.len()); decryptor - .decrypt(header, &ciphertext, &mut plaintext) + .decrypt(&header_bytes, &ciphertext, &mut plaintext) .unwrap(); total_bytes += payload.len(); diff --git a/crates/wzp-client/tests/handshake_integration.rs b/crates/wzp-client/tests/handshake_integration.rs index 8e1027a..ef969ed 100644 --- a/crates/wzp-client/tests/handshake_integration.rs +++ b/crates/wzp-client/tests/handshake_integration.rs @@ -99,31 +99,52 @@ async fn full_handshake_both_sides_derive_same_session() { assert_eq!(chosen_profile, wzp_proto::QualityProfile::GOOD); // Verify both sides can communicate: client encrypts, relay decrypts. - let header = b"test-header"; + // encrypt/decrypt derive nonces from MediaHeader.seq, so we need valid headers. + use wzp_proto::packet::MediaHeader; + use wzp_proto::{CodecId, MediaType}; + let make_hdr = |seq: u32| { + let h = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq, + timestamp: seq.wrapping_mul(20), + fec_block: 0, + }; + let mut b = Vec::new(); + h.write_to(&mut b); + b + }; + + let header = make_hdr(0); let plaintext = b"hello from client to relay"; let mut ciphertext = Vec::new(); client_session - .encrypt(header, plaintext, &mut ciphertext) + .encrypt(&header, plaintext, &mut ciphertext) .expect("client encrypt should succeed"); let mut decrypted = Vec::new(); relay_session - .decrypt(header, &ciphertext, &mut decrypted) + .decrypt(&header, &ciphertext, &mut decrypted) .expect("relay decrypt should succeed"); assert_eq!(&decrypted[..], plaintext); // Verify reverse direction: relay encrypts, client decrypts. + let header2 = make_hdr(0); // relay's send_seq starts at 0 let plaintext2 = b"hello from relay to client"; let mut ciphertext2 = Vec::new(); relay_session - .encrypt(header, plaintext2, &mut ciphertext2) + .encrypt(&header2, plaintext2, &mut ciphertext2) .expect("relay encrypt should succeed"); let mut decrypted2 = Vec::new(); client_session - .decrypt(header, &ciphertext2, &mut decrypted2) + .decrypt(&header2, &ciphertext2, &mut decrypted2) .expect("client decrypt should succeed"); assert_eq!(&decrypted2[..], plaintext2); diff --git a/crates/wzp-crypto/src/handshake.rs b/crates/wzp-crypto/src/handshake.rs index 9ea2d89..597c7ee 100644 --- a/crates/wzp-crypto/src/handshake.rs +++ b/crates/wzp-crypto/src/handshake.rs @@ -209,18 +209,34 @@ mod tests { let mut alice_session = alice.derive_session(&bob_eph_pub).unwrap(); let mut bob_session = bob.derive_session(&alice_eph_pub).unwrap(); - // Verify they can communicate: Alice encrypts, Bob decrypts - let header = b"call-header"; + // Verify they can communicate: Alice encrypts, Bob decrypts. + // Use a valid v2 MediaHeader — encrypt/decrypt now derive the nonce from + // header.seq and will reject raw byte slices shorter than WIRE_SIZE. + use wzp_proto::{CodecId, MediaHeader, MediaType}; + let header = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq: 0, + timestamp: 0, + fec_block: 0, + }; + let mut header_bytes = Vec::new(); + header.write_to(&mut header_bytes); + let plaintext = b"hello from alice"; let mut ciphertext = Vec::new(); alice_session - .encrypt(header, plaintext, &mut ciphertext) + .encrypt(&header_bytes, plaintext, &mut ciphertext) .unwrap(); let mut decrypted = Vec::new(); bob_session - .decrypt(header, &ciphertext, &mut decrypted) + .decrypt(&header_bytes, &ciphertext, &mut decrypted) .unwrap(); assert_eq!(&decrypted, plaintext); diff --git a/crates/wzp-crypto/src/session.rs b/crates/wzp-crypto/src/session.rs index 5649772..d9e8243 100644 --- a/crates/wzp-crypto/src/session.rs +++ b/crates/wzp-crypto/src/session.rs @@ -101,10 +101,14 @@ impl CryptoSession for ChaChaSession { plaintext: &[u8], out: &mut Vec, ) -> Result<(), CryptoError> { - let nonce_bytes = nonce::build_nonce(&self.session_id, self.send_seq, Direction::Send); + // Derive nonce from the wire-level seq in the header, not from an + // internal counter. This ensures the receiver can reconstruct the + // same nonce using the header it receives, regardless of delivery order. + let header = parse_header(header_bytes) + .ok_or_else(|| CryptoError::Internal("header too short to derive nonce".into()))?; + let nonce_bytes = nonce::build_nonce(&self.session_id, header.seq, Direction::Send); let nonce = Nonce::from_slice(&nonce_bytes); - // Encrypt with AAD use chacha20poly1305::aead::Payload; let payload = Payload { msg: plaintext, @@ -117,7 +121,7 @@ impl CryptoSession for ChaChaSession { .map_err(|_| CryptoError::Internal("encryption failed".into()))?; out.extend_from_slice(&ciphertext); - self.send_seq = self.send_seq.wrapping_add(1); + self.send_seq = self.send_seq.wrapping_add(1); // packet counter for rekey trigger only Ok(()) } @@ -127,9 +131,14 @@ impl CryptoSession for ChaChaSession { ciphertext: &[u8], out: &mut Vec, ) -> Result<(), CryptoError> { - // Use Direction::Send to match the sender's nonce construction. - // The recv_seq counter tracks which packet from the peer we're decrypting. - let nonce_bytes = nonce::build_nonce(&self.session_id, self.recv_seq, Direction::Send); + // Parse header before decryption — needed for nonce derivation. + // Using header.seq (not recv_seq) means the nonce is always derived + // from the same wire field as the sender, surviving out-of-order delivery. + // A recv_seq counter diverges from the sender's send_seq on any reorder, + // causing every subsequent decryption to fail for the rest of the session. + let header = parse_header(header_bytes) + .ok_or_else(|| CryptoError::Internal("header too short to derive nonce".into()))?; + let nonce_bytes = nonce::build_nonce(&self.session_id, header.seq, Direction::Send); let nonce = Nonce::from_slice(&nonce_bytes); use chacha20poly1305::aead::Payload; @@ -145,20 +154,17 @@ impl CryptoSession for ChaChaSession { let plaintext_len = plaintext.len(); out.extend_from_slice(&plaintext); - self.recv_seq = self.recv_seq.wrapping_add(1); + self.recv_seq = self.recv_seq.wrapping_add(1); // packet counter for rekey trigger only - // Anti-replay check: if header parses as a v2 MediaHeader, verify seq - // is not a replay for this (stream_id, media_type). - if let Some(header) = parse_header(header_bytes) { - let window = self - .anti_replay - .entry((header.stream_id, header.media_type)) - .or_insert_with(|| default_window_for_media_type(header.media_type)); - if let Err(e) = window.check_and_update(header.seq) { - // Roll back the plaintext we just appended. - out.truncate(out.len() - plaintext_len); - return Err(e); - } + // Anti-replay check: header already parsed above. + let window = self + .anti_replay + .entry((header.stream_id, header.media_type)) + .or_insert_with(|| default_window_for_media_type(header.media_type)); + if let Err(e) = window.check_and_update(header.seq) { + // Roll back the plaintext we just appended. + out.truncate(out.len() - plaintext_len); + return Err(e); } Ok(()) @@ -198,24 +204,42 @@ impl CryptoSession for ChaChaSession { #[cfg(test)] mod tests { use super::*; + use wzp_proto::{CodecId, MediaType}; fn make_session_pair() -> (ChaChaSession, ChaChaSession) { let key = [0x42u8; 32]; (ChaChaSession::new(key), ChaChaSession::new(key)) } + /// Build a minimal valid v2 MediaHeader serialised to bytes. + fn make_header_bytes(seq: u32) -> Vec { + let header = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq, + timestamp: seq.wrapping_mul(20), + fec_block: 0, + }; + let mut bytes = Vec::new(); + header.write_to(&mut bytes); + bytes + } + #[test] fn encrypt_decrypt_roundtrip() { let (mut alice, mut bob) = make_session_pair(); - let header = b"test-header"; + let header = make_header_bytes(0); let plaintext = b"hello warzone"; let mut ciphertext = Vec::new(); - alice.encrypt(header, plaintext, &mut ciphertext).unwrap(); + alice.encrypt(&header, plaintext, &mut ciphertext).unwrap(); - // Bob decrypts (his recv matches Alice's send) let mut decrypted = Vec::new(); - bob.decrypt(header, &ciphertext, &mut decrypted).unwrap(); + bob.decrypt(&header, &ciphertext, &mut decrypted).unwrap(); assert_eq!(&decrypted, plaintext); } @@ -223,14 +247,18 @@ mod tests { #[test] fn decrypt_wrong_aad_fails() { let (mut alice, mut bob) = make_session_pair(); - let header = b"correct-header"; + let correct_header = make_header_bytes(0); + // Different seq → different nonce AND different AAD bytes: decryption must fail. + let wrong_header = make_header_bytes(1); let plaintext = b"secret data"; let mut ciphertext = Vec::new(); - alice.encrypt(header, plaintext, &mut ciphertext).unwrap(); + alice + .encrypt(&correct_header, plaintext, &mut ciphertext) + .unwrap(); let mut decrypted = Vec::new(); - let result = bob.decrypt(b"wrong-header", &ciphertext, &mut decrypted); + let result = bob.decrypt(&wrong_header, &ciphertext, &mut decrypted); assert!(result.is_err()); } @@ -239,29 +267,29 @@ mod tests { let mut alice = ChaChaSession::new([0xAA; 32]); let mut eve = ChaChaSession::new([0xBB; 32]); - let header = b"hdr"; + let header = make_header_bytes(0); let plaintext = b"secret"; let mut ciphertext = Vec::new(); - alice.encrypt(header, plaintext, &mut ciphertext).unwrap(); + alice.encrypt(&header, plaintext, &mut ciphertext).unwrap(); let mut decrypted = Vec::new(); - let result = eve.decrypt(header, &ciphertext, &mut decrypted); + let result = eve.decrypt(&header, &ciphertext, &mut decrypted); assert!(result.is_err()); } #[test] fn multiple_packets_roundtrip() { let (mut alice, mut bob) = make_session_pair(); - let header = b"hdr"; - for i in 0..100 { + for i in 0..100u32 { + let header = make_header_bytes(i); let msg = format!("message {}", i); let mut ct = Vec::new(); - alice.encrypt(header, msg.as_bytes(), &mut ct).unwrap(); + alice.encrypt(&header, msg.as_bytes(), &mut ct).unwrap(); let mut pt = Vec::new(); - bob.decrypt(header, &ct, &mut pt).unwrap(); + bob.decrypt(&header, &ct, &mut pt).unwrap(); assert_eq!(pt, msg.as_bytes()); } } @@ -281,6 +309,57 @@ mod tests { assert_eq!(alice.send_seq, 0); } + #[test] + fn decrypt_survives_out_of_order_delivery() { + // Regression test for nonce derivation using recv_seq instead of + // MediaHeader.seq. If nonces are tied to a local counter, any reorder + // causes the counter to diverge from the sender's seq and every + // subsequent packet fails decryption permanently. + use wzp_proto::{CodecId, MediaType}; + + let key = [0x55u8; 32]; + let mut alice = ChaChaSession::new(key); + let mut bob = ChaChaSession::new(key); + + let plaintext = b"audio payload"; + + // Encrypt 5 packets in order (seqs 10, 11, 12, 13, 14). + let seqs = [10u32, 11, 12, 13, 14]; + let mut ciphertexts: Vec<(Vec, Vec)> = Vec::new(); + for &seq in &seqs { + let header = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq, + timestamp: seq * 20, + fec_block: 0, + }; + let mut header_bytes = Vec::new(); + header.write_to(&mut header_bytes); + let mut ct = Vec::new(); + alice.encrypt(&header_bytes, plaintext, &mut ct).unwrap(); + ciphertexts.push((header_bytes, ct)); + } + + // Bob receives them out of order: 0, 2, 1, 4, 3 + let delivery_order = [0usize, 2, 1, 4, 3]; + for &idx in &delivery_order { + let (ref hdr, ref ct) = ciphertexts[idx]; + let mut pt = Vec::new(); + let result = bob.decrypt(hdr, ct, &mut pt); + assert!( + result.is_ok(), + "out-of-order packet (original idx={idx}, seq={}) must decrypt successfully", + seqs[idx] + ); + assert_eq!(&pt, plaintext); + } + } + #[test] fn per_stream_anti_replay_rejects_duplicate() { use wzp_proto::{CodecId, MediaType}; diff --git a/crates/wzp-relay/tests/handshake_integration.rs b/crates/wzp-relay/tests/handshake_integration.rs index 5fd69be..afca743 100644 --- a/crates/wzp-relay/tests/handshake_integration.rs +++ b/crates/wzp-relay/tests/handshake_integration.rs @@ -9,10 +9,29 @@ use std::sync::Arc; use wzp_client::perform_handshake; use wzp_crypto::{KeyExchange, WarzoneKeyExchange}; -use wzp_proto::{MediaTransport, SignalMessage, default_signal_version}; +use wzp_proto::packet::MediaHeader; +use wzp_proto::{CodecId, MediaTransport, MediaType, SignalMessage, default_signal_version}; use wzp_relay::handshake::accept_handshake; use wzp_transport::{QuinnTransport, client_config, create_endpoint, server_config}; +/// Build valid v2 MediaHeader bytes for use in encrypt/decrypt test calls. +fn test_header(seq: u32) -> Vec { + let h = MediaHeader { + version: 2, + flags: 0, + media_type: MediaType::Audio, + codec_id: CodecId::Opus24k, + stream_id: 0, + fec_ratio: 0, + seq, + timestamp: seq.wrapping_mul(20), + fec_block: 0, + }; + let mut b = Vec::new(); + h.write_to(&mut b); + b +} + /// Establish a QUIC connection and wrap both sides in `QuinnTransport`. /// /// Returns (client_transport, server_transport, _endpoints) where the endpoint @@ -79,7 +98,7 @@ async fn handshake_succeeds() { // Both sides should have derived a working CryptoSession. // Verify by encrypting on one side and decrypting on the other. - let header = b"test-header"; + let header = test_header(0); let plaintext = b"hello warzone"; let mut ciphertext = Vec::new(); @@ -87,12 +106,12 @@ async fn handshake_succeeds() { let mut callee_session = callee_session; caller_session - .encrypt(header, plaintext, &mut ciphertext) + .encrypt(&header, plaintext, &mut ciphertext) .expect("encrypt"); let mut decrypted = Vec::new(); callee_session - .decrypt(header, &ciphertext, &mut decrypted) + .decrypt(&header, &ciphertext, &mut decrypted) .expect("decrypt"); assert_eq!(&decrypted, plaintext); @@ -212,7 +231,7 @@ async fn handshake_verifies_identity() { .expect("accept_handshake must succeed"); // Cross-encrypt/decrypt to prove the shared session works. - let header = b"id-test"; + let header = test_header(0); let plaintext = b"identity verified"; let mut ct = Vec::new(); @@ -220,12 +239,12 @@ async fn handshake_verifies_identity() { let mut callee_session = callee_session; caller_session - .encrypt(header, plaintext, &mut ct) + .encrypt(&header, plaintext, &mut ct) .expect("encrypt"); let mut pt = Vec::new(); callee_session - .decrypt(header, &ct, &mut pt) + .decrypt(&header, &ct, &mut pt) .expect("decrypt"); assert_eq!(&pt, plaintext); @@ -292,7 +311,7 @@ async fn auth_then_handshake() { assert_eq!(received_token, "bearer-test-token-12345"); // Verify the crypto session works after the auth preamble. - let header = b"auth-hdr"; + let header = test_header(0); let plaintext = b"post-auth payload"; let mut ct = Vec::new(); @@ -300,12 +319,12 @@ async fn auth_then_handshake() { let mut callee_session = callee_session; caller_session - .encrypt(header, plaintext, &mut ct) + .encrypt(&header, plaintext, &mut ct) .expect("encrypt"); let mut pt = Vec::new(); callee_session - .decrypt(header, &ct, &mut pt) + .decrypt(&header, &ct, &mut pt) .expect("decrypt"); assert_eq!(&pt, plaintext);