From 63582dd4e1f97b5304f2cb287bdb816ca9c3a46c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Wed, 18 Feb 2026 16:37:37 +0000 Subject: [PATCH] LpSession cleanup --- Cargo.toml | 14 +- common/nym-kkt/src/keys.rs | 76 +- common/nym-lp-transport/src/traits.rs | 19 + common/nym-lp/Cargo.toml | 11 - common/nym-lp/src/codec.rs | 2204 ++++++++---------- common/nym-lp/src/error.rs | 36 +- common/nym-lp/src/kkt_orchestrator.rs | 492 ---- common/nym-lp/src/lib.rs | 211 +- common/nym-lp/src/message.rs | 668 +----- common/nym-lp/src/noise_protocol.rs | 337 --- common/nym-lp/src/packet.rs | 199 +- common/nym-lp/src/peer.rs | 28 +- common/nym-lp/src/psk.rs | 788 ------- common/nym-lp/src/psq/helpers.rs | 49 - common/nym-lp/src/psq/initiator.rs | 69 +- common/nym-lp/src/psq/mod.rs | 40 +- common/nym-lp/src/psq/responder.rs | 60 +- common/nym-lp/src/session.rs | 708 +----- common/nym-lp/src/session_integration/mod.rs | 1398 ++++++----- common/nym-lp/src/session_manager.rs | 69 +- common/nym-lp/src/state_machine.rs | 1079 ++------- gateway/src/node/lp_listener/handler.rs | 22 +- 22 files changed, 2367 insertions(+), 6210 deletions(-) delete mode 100644 common/nym-lp/src/kkt_orchestrator.rs delete mode 100644 common/nym-lp/src/noise_protocol.rs delete mode 100644 common/nym-lp/src/psk.rs diff --git a/Cargo.toml b/Cargo.toml index 4e28dda315..e58bf37a29 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -397,13 +397,13 @@ prometheus = { version = "0.14.0" } # libcrux -libcrux-kem = { git = "https://github.com/cryspen/libcrux" } -libcrux-ecdh = { git = "https://github.com/cryspen/libcrux" } -libcrux-chacha20poly1305 = { git = "https://github.com/cryspen/libcrux" } -libcrux-psq = { git = "https://github.com/cryspen/libcrux" } -libcrux-ml-kem = { git = "https://github.com/cryspen/libcrux" } -libcrux-sha3 = { git = "https://github.com/cryspen/libcrux" } -libcrux-traits = { git = "https://github.com/cryspen/libcrux" } +libcrux-kem = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-ecdh = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-chacha20poly1305 = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-psq = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-ml-kem = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-sha3 = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } +libcrux-traits = { git = "https://github.com/jstuczyn/libcrux", branch = "jstuczyn/exposed-transport" } # Workspace dep definitions required by crates.io publication - we need a workspace version since `cargo workspaces` doesn't work with path imports from crate manifests nym-api-requests = { version = "1.20.4", path = "nym-api/nym-api-requests" } diff --git a/common/nym-kkt/src/keys.rs b/common/nym-kkt/src/keys.rs index ebcb0d9f2c..c5b70c8dcb 100644 --- a/common/nym-kkt/src/keys.rs +++ b/common/nym-kkt/src/keys.rs @@ -7,11 +7,15 @@ use libcrux_psq::classic_mceliece; use libcrux_psq::handshake::types::PQEncapsulationKey; use nym_kkt_ciphersuite::{KEM, mceliece}; use std::fmt::{Debug, Formatter}; +use std::sync::Arc; /// Wrapper around keys used for the KEM exchange +/// with cheap clones thanks to Arc wrappers pub struct KEMKeys { - mc_eliece: classic_mceliece::KeyPair, - ml_kem768: MlKem768KeyPair, + mc_eliece_pk: Arc, + mc_eliece_sk: Arc, + ml_kem768_pk: Arc, + ml_kem768_sk: Arc, } impl Debug for KEMKeys { @@ -25,40 +29,61 @@ impl Debug for KEMKeys { impl KEMKeys { pub fn new(mc_eliece: classic_mceliece::KeyPair, ml_kem768: MlKem768KeyPair) -> Self { + let (ml_kem768_sk, ml_kem768_pk) = ml_kem768.into_parts(); Self { - mc_eliece, - ml_kem768, + mc_eliece_pk: Arc::new(mc_eliece.pk), + mc_eliece_sk: Arc::new(mc_eliece.sk), + ml_kem768_pk: Arc::new(ml_kem768_pk), + ml_kem768_sk: Arc::new(ml_kem768_sk), } } pub fn encoded_encapsulation_key(&self, kem: KEM) -> Option<&[u8]> { match kem { - KEM::McEliece => Some(self.mc_eliece.pk.as_ref()), - KEM::MlKem768 => Some(self.ml_kem768.pk()), + KEM::McEliece => Some(self.mc_eliece_pk.as_ref().as_ref()), + KEM::MlKem768 => Some(self.ml_kem768_pk.as_slice()), + // _ => None, + } + } + + pub fn encapsulation_key(&self, kem: KEM) -> Option { + match kem { + KEM::McEliece => Some(EncapsulationKey::McEliece(self.mc_eliece_pk.clone())), + KEM::MlKem768 => Some(EncapsulationKey::MlKem768(self.ml_kem768_pk.clone())), // _ => None, } } pub fn mc_eliece_encapsulation_key(&self) -> &classic_mceliece::PublicKey { - &self.mc_eliece.pk + &self.mc_eliece_pk } pub fn ml_kem768_encapsulation_key(&self) -> &MlKem768PublicKey { - self.ml_kem768.public_key() + self.ml_kem768_pk.as_ref() } pub fn mc_eliece_decapsulation_key(&self) -> &classic_mceliece::SecretKey { - &self.mc_eliece.sk + &self.mc_eliece_sk } pub fn ml_kem768_decapsulation_key(&self) -> &MlKem768PrivateKey { - self.ml_kem768.private_key() + &self.ml_kem768_sk } } +#[derive(Clone)] pub enum EncapsulationKey { - McEliece(classic_mceliece::PublicKey), - MlKem768(MlKem768PublicKey), + McEliece(Arc), + MlKem768(Arc), +} + +impl Debug for EncapsulationKey { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + EncapsulationKey::McEliece(_) => write!(f, "EncapsulationKey::McEliece"), + EncapsulationKey::MlKem768(_) => write!(f, "EncapsulationKey::MlKem768"), + } + } } impl EncapsulationKey { @@ -78,9 +103,11 @@ impl EncapsulationKey { pub fn try_from_bytes(bytes: Vec, kem: KEM) -> Result { match kem { - KEM::MlKem768 => Ok(EncapsulationKey::MlKem768( - MlKem768PublicKey::try_from(bytes.as_slice()).unwrap(), - )), + KEM::MlKem768 => Ok(EncapsulationKey::MlKem768(Arc::new( + MlKem768PublicKey::try_from(bytes.as_slice()).map_err(|_| KKTError::KEMError { + info: "mlkem768 key of invalid length", + })?, + ))), KEM::McEliece => { let boxed_array: Box<[u8; mceliece::PUBLIC_KEY_LENGTH]> = bytes .into_boxed_slice() @@ -89,26 +116,17 @@ impl EncapsulationKey { info: "mceliece key of invalid length", })?; - Ok(EncapsulationKey::McEliece( - classic_mceliece::PublicKey::try_from(boxed_array).unwrap(), - )) + Ok(EncapsulationKey::McEliece(Arc::new( + classic_mceliece::PublicKey::from(boxed_array), + ))) } } } pub fn as_bytes(&self) -> &[u8] { match self { - EncapsulationKey::McEliece(k) => k.as_ref(), - EncapsulationKey::MlKem768(k) => k.as_ref(), - } - } -} - -impl Debug for EncapsulationKey { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self { - EncapsulationKey::McEliece(_) => write!(f, "McEliece"), - EncapsulationKey::MlKem768(_) => write!(f, "MlKem768"), + EncapsulationKey::McEliece(k) => k.as_ref().as_ref(), + EncapsulationKey::MlKem768(k) => k.as_ref().as_ref(), } } } diff --git a/common/nym-lp-transport/src/traits.rs b/common/nym-lp-transport/src/traits.rs index 045bf29498..069adc308c 100644 --- a/common/nym-lp-transport/src/traits.rs +++ b/common/nym-lp-transport/src/traits.rs @@ -8,6 +8,25 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; use tracing::debug; +// only used in internal code (and tests) +#[allow(async_fn_in_trait)] +pub trait LpChannel: Sized { + /// Sends a serialised acket over the data stream. + /// + /// # Arguments + /// * `packet_data` - The serialised packet to send + /// + /// # Errors + /// Returns an error on network transmission fails. + async fn send_serialised_packet(&mut self, packet_data: &[u8]) -> std::io::Result<()>; + + /// Receives a data chunk of the set length from the data stream. + /// + /// # Errors + /// Returns an error on network transmission fails. + async fn receive_raw_packet(&mut self, len: usize) -> std::io::Result>; +} + // only used in internal code (and tests) #[allow(async_fn_in_trait)] pub trait LpTransport: Sized { diff --git a/common/nym-lp/Cargo.toml b/common/nym-lp/Cargo.toml index e53cdf8190..ff0e7cbb2d 100644 --- a/common/nym-lp/Cargo.toml +++ b/common/nym-lp/Cargo.toml @@ -7,17 +7,11 @@ publish = false [dependencies] thiserror = { workspace = true } -parking_lot = { workspace = true } -snow = { workspace = true } bs58 = { workspace = true } -serde = { workspace = true } bytes = { workspace = true } -dashmap = { workspace = true } -sha2 = { workspace = true } tracing = { workspace = true } rand = { workspace = true } rand09 = { workspace = true } -rand_core09 = { package = "rand_core", version = "=0.9.2" } nym-crypto = { path = "../crypto", features = ["hashing", "asymmetric"] } nym-kkt = { path = "../nym-kkt" } @@ -27,10 +21,7 @@ nym-lp-transport = { path = "../nym-lp-transport" } # libcrux dependencies for PSQ (Post-Quantum PSK derivation) libcrux-psq = { workspace = true, features = ["test-utils"] } -libcrux-kem = { workspace = true } -tls_codec = { workspace = true } num_enum = { workspace = true } -chacha20poly1305 = { workspace = true } zeroize = { workspace = true, features = ["zeroize_derive"] } # needed for the 'mock 'feature @@ -38,8 +29,6 @@ nym-test-utils = { workspace = true, optional = true } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } -#rand_chacha = "0.3" -mock_instant = { workspace = true } nym-crypto = { path = "../crypto", features = ["rand"] } nym-test-utils = { workspace = true } anyhow = { workspace = true } diff --git a/common/nym-lp/src/codec.rs b/common/nym-lp/src/codec.rs index 95a6e2da3a..5be334f1e0 100644 --- a/common/nym-lp/src/codec.rs +++ b/common/nym-lp/src/codec.rs @@ -1,101 +1,13 @@ // Copyright 2025 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::LpError; -use crate::message::{LpMessage, MessageType}; -use crate::packet::{LpHeader, LpPacket, OuterHeader, TRAILER_LEN}; -use bytes::{BufMut, BytesMut}; -use chacha20poly1305::{ - ChaCha20Poly1305, Key, Nonce, Tag, - aead::{AeadInPlace, KeyInit}, -}; -use zeroize::{Zeroize, ZeroizeOnDrop}; +use crate::packet::{EncryptedLpPacket, InnerHeader, LpHeader, LpPacket, OuterHeader}; +use crate::{LpError, LpMessage}; +use bytes::BytesMut; +use libcrux_psq::Channel; +use tracing::error; -/// Size of outer header (receiver_idx + counter) - always cleartext -pub const OUTER_HEADER_SIZE: usize = OuterHeader::SIZE; // 12 bytes - -// georgio: maybe remove this? -/// Size of inner prefix (proto + reserved) - cleartext or encrypted depending on mode -const INNER_PREFIX_SIZE: usize = 4; // proto(1) + reserved(3) - -/// Outer AEAD key for LP packet encryption. -/// -/// Derived from PSK using Blake3 KDF with domain separation. -/// Used for opportunistic encryption: before PSK packets are cleartext, -/// after PSK packets have encrypted payload and authenticated header. -/// -/// # Security: Nonce Reuse Prevention -/// -/// ChaCha20-Poly1305 requires unique nonces per key. The counter starts at 0 -/// for each session, which is safe because: -/// -/// 1. **PSK is always fresh**: Each handshake uses PSQ -/// with a client-generated random salt. This ensures a unique -/// PSK for every session, even between the same client-gateway pair. -/// -/// 2. **Key derivation**: `outer_key = Blake3_KDF("lp-outer-aead", PSK)`. -/// Different PSK → different outer_key → nonce reuse impossible. -/// -/// 3. **No PSK persistence**: PSK handles are not stored/reused across sessions. -/// Each connection performs fresh KKT+PSQ handshake. -/// -// noiserm -#[derive(Clone, Zeroize, ZeroizeOnDrop)] -pub struct OuterAeadKey { - key: [u8; 32], -} - -impl OuterAeadKey { - /// KDF context for outer AEAD key derivation (domain separation) - const KDF_CONTEXT: &'static str = "lp-outer-aead"; - - /// Derive outer AEAD key from PSK. - /// - /// Uses Blake3 KDF with domain separation to avoid key reuse - /// between the outer AEAD layer and the inner Noise layer. - pub fn from_psk(psk: &[u8; 32]) -> Self { - let key = nym_crypto::hkdf::blake3::derive_key_blake3(Self::KDF_CONTEXT, psk, &[]); - Self { key } - } - - /// Get reference to the raw key bytes. - pub fn as_bytes(&self) -> &[u8; 32] { - &self.key - } -} - -impl std::fmt::Debug for OuterAeadKey { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OuterAeadKey") - .field("key", &"[REDACTED]") - .finish() - } -} - -// noiserm -/// Build 12-byte nonce from 8-byte counter (zero-padded). -/// -/// Format: counter (8 bytes LE) || 0x00000000 (4 bytes) -fn build_nonce(counter: u64) -> [u8; 12] { - let mut nonce = [0u8; 12]; - nonce[..8].copy_from_slice(&counter.to_le_bytes()); - // bytes 8..12 remain zero (zero-padding) - nonce -} - -/// Parse message from raw type and content bytes. -/// -/// Used when decrypting outer-encrypted packets where the message type -/// was encrypted along with the content. -fn parse_message_from_type_and_content( - msg_type_raw: u32, - content: &[u8], -) -> Result { - let message_type = MessageType::from_u32(msg_type_raw) - .ok_or_else(|| LpError::invalid_message_type(msg_type_raw))?; - - LpMessage::decode_content(content, message_type) -} +pub(crate) const CIPHERTEXT_OVERHEAD: usize = 25; /// Parse only the outer header from raw packet bytes. /// @@ -111,1220 +23,954 @@ pub fn parse_lp_header_only(src: &[u8]) -> Result { OuterHeader::parse(src) } -/// Parses a complete Lewes Protocol packet from a byte slice (e.g., a UDP datagram payload). -/// -/// ## Unified Packet Format -/// -/// Both cleartext and encrypted packets have the same structure: -/// - Outer header (12B): receiver_idx(4) + counter(8) - always cleartext -/// - Inner payload: proto(1) + reserved(3) + msg_type(4) + content - cleartext or encrypted -/// - Trailer (16B): zeros (cleartext) or AEAD tag (encrypted) -/// -/// # Arguments -/// * `src` - Raw packet bytes -/// * `outer_key` - None for cleartext parsing, Some for AEAD decryption -/// -/// # Errors -/// * `LpError::AeadTagMismatch` - Tag verification failed (when outer_key provided) -/// * `LpError::InsufficientBufferSize` - Packet too small -pub fn parse_lp_packet(src: &[u8], outer_key: Option<&OuterAeadKey>) -> Result { - // Minimum size check: OuterHeader + InnerPrefix + MsgType + Trailer (for 0-payload message) - // 12 + 4 + 2 + 16 = 34 bytes - let min_size = OUTER_HEADER_SIZE + INNER_PREFIX_SIZE + 2 + TRAILER_LEN; - if src.len() < min_size { +// /// Parses a complete Lewes Protocol packet from a byte slice (e.g., a UDP datagram payload). +// /// +// /// # Arguments +// /// * `outer_header` - The parsed OuterHeader from the underlying stream +// /// * `plaintext` - the decrypted plaintext of the remainer of the packet +// /// +// /// # Errors +// /// * `LpError::InsufficientBufferSize` - Packet too small +// pub fn parse_lp_packet(outer_header: OuterHeader, plaintext: &[u8]) -> Result { +// todo!() +// // if plaintext.len() < InnerHeader::SIZE { +// // return Err(LpError::InsufficientBufferSize); +// // } +// // +// // let inner_header = InnerHeader::parse(plaintext)?; +// // let payload = &plaintext[InnerHeader::SIZE..]; +// // let message = LpMessage::decode_content(payload, inner_header.message_type)?; +// // +// // Ok(LpPacket { +// // header: LpHeader { +// // outer: outer_header, +// // inner: inner_header, +// // }, +// // message, +// // }) +// } + +pub fn encrypt_lp_packet( + packet: LpPacket, + transport: &mut libcrux_psq::Transport, +) -> Result { + let mut plaintext = BytesMut::with_capacity(InnerHeader::SIZE + packet.message.len()); + packet.header.inner.encode(&mut plaintext); + packet.message.encode_content(&mut plaintext); + + let mut ciphertext = vec![0u8; plaintext.len() + CIPHERTEXT_OVERHEAD]; + let n = transport.write_message(&*plaintext, &mut ciphertext)?; + + if ciphertext.len() != n { + error!("inconsistent ciphertext overhead") + } + + ciphertext.truncate(n); + + Ok(EncryptedLpPacket { + outer_header: packet.header.outer, + ciphertext, + }) +} + +pub fn decrypt_lp_packet( + packet: EncryptedLpPacket, + transport: &mut libcrux_psq::Transport, +) -> Result { + if packet.ciphertext.len() < InnerHeader::SIZE + CIPHERTEXT_OVERHEAD { return Err(LpError::InsufficientBufferSize); } - // Parse outer header (always cleartext at bytes 0-12) - let outer_header = OuterHeader::parse(src)?; - - // Extract trailer (potential AEAD tag) - let trailer_start = src.len() - TRAILER_LEN; - let mut trailer = [0u8; TRAILER_LEN]; - trailer.copy_from_slice(&src[trailer_start..]); - - // Inner payload is everything between outer header and trailer - let inner_bytes = &src[OUTER_HEADER_SIZE..trailer_start]; - - // Handle decryption if outer key provided - match outer_key { - None => { - // Cleartext mode - parse inner directly - // Inner format: proto(1) + reserved(3) + msg_type(4) + content - if inner_bytes.len() < INNER_PREFIX_SIZE + 4 { - return Err(LpError::InsufficientBufferSize); - } - - let protocol_version = inner_bytes[0]; - // reserved bytes [1..4] are ignored - let msg_type = u32::from_le_bytes([ - inner_bytes[4], - inner_bytes[5], - inner_bytes[6], - inner_bytes[7], - ]); - let message_content = &inner_bytes[8..]; - - let header = LpHeader { - protocol_version, - reserved: [0u8; 3], - receiver_idx: outer_header.receiver_idx, - counter: outer_header.counter, - }; - - let message = parse_message_from_type_and_content(msg_type, message_content)?; - - Ok(LpPacket { - header, - message, - trailer, - }) - } - // noiserm (potentially) - Some(key) => { - // AEAD decryption mode - // AAD is the outer header (12 bytes) - let nonce = build_nonce(outer_header.counter); - let aad = &src[..OUTER_HEADER_SIZE]; - - // Copy inner payload for in-place decryption - let mut decrypted = inner_bytes.to_vec(); - - // Convert trailer to Tag - let tag = Tag::from_slice(&trailer); - - // Decrypt and verify - let cipher = ChaCha20Poly1305::new(Key::from_slice(key.as_bytes())); - cipher - .decrypt_in_place_detached(Nonce::from_slice(&nonce), aad, &mut decrypted, tag) - .map_err(|_| LpError::AeadTagMismatch)?; - - // Decrypted format: proto(1) + reserved(3) + msg_type(4) + content - if decrypted.len() < INNER_PREFIX_SIZE + 4 { - return Err(LpError::InsufficientBufferSize); - } - - let protocol_version = decrypted[0]; - // reserved bytes [1..4] are ignored - let msg_type = - u32::from_le_bytes([decrypted[4], decrypted[5], decrypted[6], decrypted[7]]); - let message_content = &decrypted[8..]; - - let header = LpHeader { - protocol_version, - reserved: [0u8; 3], - receiver_idx: outer_header.receiver_idx, - counter: outer_header.counter, - }; - - let message = parse_message_from_type_and_content(msg_type, message_content)?; - - Ok(LpPacket { - header, - message, - trailer, - }) - } + let mut plaintext = Vec::with_capacity(packet.ciphertext.len() - CIPHERTEXT_OVERHEAD); + let (_, n) = transport.read_message(&packet.ciphertext, &mut plaintext)?; + if n != packet.ciphertext.len() - CIPHERTEXT_OVERHEAD { + error!("inconsistent ciphertext overhead") } + + let inner_header = InnerHeader::parse(&plaintext)?; + let payload = &plaintext[InnerHeader::SIZE..]; + let message = LpMessage::decode_content(payload, inner_header.message_type)?; + + Ok(LpPacket { + header: LpHeader { + outer: packet.outer_header, + inner: inner_header, + }, + message, + }) } -// georgio: start with no outer_key /// Serializes an LpPacket into the provided BytesMut buffer. /// /// ## Unified Packet Format /// /// Both cleartext and encrypted packets have the same structure: /// - Outer header (12B): receiver_idx(4) + counter(8) - always cleartext -/// - Inner payload: proto(1) + reserved(3) + msg_type(4) + content - cleartext or encrypted -/// - Trailer (16B): zeros (cleartext) or AEAD tag (encrypted) +/// - Inner payload: proto(1) + reserved(3) + msg_type(4) + content - encrypted /// /// # Arguments /// * `item` - Packet to serialize /// * `dst` - Output buffer -/// * `outer_key` - None for cleartext, Some for AEAD encryption +/// * `transport` - AEAD encryption channel pub fn serialize_lp_packet( - item: &LpPacket, + item: LpPacket, dst: &mut BytesMut, - outer_key: Option<&OuterAeadKey>, + transport: &mut libcrux_psq::Transport, ) -> Result<(), LpError> { - // Total size: outer_header(12) + inner_prefix(4) + msg_type(4) + content + trailer(16) - let total_size = OUTER_HEADER_SIZE + INNER_PREFIX_SIZE + 4 + item.message.len() + TRAILER_LEN; - dst.reserve(total_size); + // 1. encrypt the inner header and payload + let encrypted_packet = encrypt_lp_packet(item, transport)?; - // 1. Write outer header (always cleartext) - 12 bytes - let outer_header = OuterHeader::new(item.header.receiver_idx, item.header.counter); - outer_header.encode_into(dst); + // 2. Write outer header (always cleartext) followed by the ciphertext + encrypted_packet.encode(dst); - match outer_key { - None => { - // Cleartext mode - // 2. Write inner prefix: proto(1) + reserved(3) - dst.put_u8(item.header.protocol_version); - dst.put_slice(&item.header.reserved); // reserved - - // 3. Write message type (4B) + content - dst.put_slice(&(item.message.typ() as u32).to_le_bytes()); - item.message.encode_content(dst); - - // 4. Write zeros trailer - dst.put_slice(&[0u8; TRAILER_LEN]); - - Ok(()) - } - Some(key) => { - // AEAD encryption mode - // AAD is the outer header (first 12 bytes) - let aad = outer_header.encode(); - - // 2. Build plaintext: proto(1) + reserved(3) + msg_type(4) + content - let mut plaintext = BytesMut::new(); - plaintext.put_u8(item.header.protocol_version); - plaintext.put_slice(&item.header.reserved); // reserved - plaintext.put_slice(&(item.message.typ() as u32).to_le_bytes()); - item.message.encode_content(&mut plaintext); - - // 3. Copy plaintext to dst for in-place encryption - let payload_start = dst.len(); - dst.put_slice(&plaintext); - - // 4. Build nonce from counter - let nonce = build_nonce(item.header.counter); - - // 5. Encrypt payload in-place - let cipher = ChaCha20Poly1305::new(Key::from_slice(key.as_bytes())); - let tag = cipher - .encrypt_in_place_detached( - Nonce::from_slice(&nonce), - &aad, - &mut dst[payload_start..], - ) - .map_err(|_| LpError::Internal("AEAD encryption failed".to_string()))?; - - // 6. Append tag as trailer - dst.put_slice(&tag); - - Ok(()) - } - } -} - -// Add a new error variant for invalid message types (Moved from previous impl LpError block) -impl LpError { - pub fn invalid_message_type(message_type: u32) -> Self { - LpError::InvalidMessageType(message_type) - } + Ok(()) } #[cfg(test)] mod tests { - use std::time::{SystemTime, UNIX_EPOCH}; - - // Import standalone functions - use super::{OuterAeadKey, parse_lp_packet, serialize_lp_packet}; - // Keep necessary imports - use crate::LpError; - use crate::message::{EncryptedDataPayload, LpMessage, MessageType, PSQRequestData}; - use crate::packet::{LpHeader, LpPacket, TRAILER_LEN}; - use bytes::BytesMut; - use nym_crypto::asymmetric::{ed25519, x25519}; - use rand::thread_rng; - - // With unified format, outer header (receiver_idx + counter) is always first - // and is the only cleartext portion for encrypted packets - const OUTER_HDR: usize = super::OUTER_HEADER_SIZE; // 12 bytes // === Cleartext Encode/Decode Tests === #[test] - fn test_serialize_parse_busy() { - let mut dst = BytesMut::new(); - - // Create a Busy packet - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 123, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - // Serialize the packet (cleartext) - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse the packet (cleartext) - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify the packet fields - assert_eq!(decoded.header.protocol_version, 1); - assert_eq!(decoded.header.receiver_idx, 42); - assert_eq!(decoded.header.counter, 123); - assert!(matches!(decoded.message, LpMessage::Busy)); - assert_eq!(decoded.trailer, [0; TRAILER_LEN]); + fn restore_below_tests() { + todo!() } - - #[test] - fn test_serialize_parse_handshake() { - let mut dst = BytesMut::new(); - - // Create a Handshake message packet - let payload = vec![42u8; 80]; // Example payload size - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 123, - }, - message: LpMessage::PSQRequest(PSQRequestData(payload.clone())), - trailer: [0; TRAILER_LEN], - }; - - // Serialize the packet (cleartext) - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse the packet (cleartext) - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify the packet fields - assert_eq!(decoded.header.protocol_version, 1); - assert_eq!(decoded.header.receiver_idx, 42); - assert_eq!(decoded.header.counter, 123); - - // Verify message type and data - match decoded.message { - LpMessage::PSQRequest(decoded_payload) => { - assert_eq!(decoded_payload, PSQRequestData(payload)); - } - _ => panic!("Expected Handshake message"), - } - assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - } - - #[test] - fn test_serialize_parse_encrypted_data() { - let mut dst = BytesMut::new(); - - // Create an EncryptedData message packet - let payload = vec![43u8; 124]; // Example payload size - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 123, - }, - message: LpMessage::EncryptedData(EncryptedDataPayload(payload.clone())), - trailer: [0; TRAILER_LEN], - }; - - // Serialize the packet (cleartext) - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse the packet (cleartext) - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify the packet fields - assert_eq!(decoded.header.protocol_version, 1); - assert_eq!(decoded.header.receiver_idx, 42); - assert_eq!(decoded.header.counter, 123); - - // Verify message type and data - match decoded.message { - LpMessage::EncryptedData(decoded_payload) => { - assert_eq!(decoded_payload, EncryptedDataPayload(payload)); - } - _ => panic!("Expected EncryptedData message"), - } - assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - } - - // === Incomplete Data Tests === - - #[test] - fn test_parse_incomplete_header() { - // Create a buffer with incomplete header - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Only 4 bytes, not enough for LpHeader::SIZE - - // Attempt to parse - expect error - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - assert!(matches!( - result.unwrap_err(), - LpError::InsufficientBufferSize - )); - } - - #[test] - fn test_parse_incomplete_message_type() { - // Create a buffer with complete header but incomplete message type - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&[0]); // Only 1 byte of message type (need 2) - - // Buffer length = 16 + 1 = 17. Min size = 16 + 2 + 16 = 34. - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - assert!(matches!( - result.unwrap_err(), - LpError::InsufficientBufferSize - )); - } - - #[test] - fn test_parse_incomplete_message_data() { - // Create a buffer simulating Handshake but missing trailer and maybe partial payload - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type - buf.extend_from_slice(&[42; 40]); // 40 bytes of payload data - - // Buffer length = 16 + 2 + 40 = 58. Min size = 16 + 2 + 16 = 34. - // Payload size calculated as 58 - 34 = 24. - // Trailer expected at index 16 + 2 + 24 = 42. - // Trailer read attempts src[42..58]. - // This *should* parse successfully based on the logic, but the trailer is garbage. - // Let's rethink: parse_lp_packet assumes the *entire slice* is the packet. - // If the slice doesn't end exactly where the trailer should, it's an error. - // In this case, total length is 58. OuterHdr(12) + InnerPrefix(4) + Type(2) + Trailer(16) = 34. Payload = 58-34=24. - // Trailer starts at 16+2+24 = 42. Ends at 42+16=58. It fits exactly. - // This test *still* doesn't test incompleteness correctly for the datagram parser. - - // Let's test a buffer that's *too short* even for header+type+trailer+min_payload - // Note: Buffer order doesn't matter for this test since we fail on minimum size check - let mut buf_too_short = BytesMut::new(); - buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx (outer header) - buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // counter (outer header) - buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // version + reserved (inner prefix) - buf_too_short.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // msg type - // No payload, no trailer. Length = 12+4+2=18. Min size = 34. - let result_too_short = parse_lp_packet(&buf_too_short, None); - assert!(result_too_short.is_err()); - assert!(matches!( - result_too_short.unwrap_err(), - LpError::InsufficientBufferSize - )); - - // Test a buffer missing PART of the trailer - let mut buf_partial_trailer = BytesMut::new(); - buf_partial_trailer.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf_partial_trailer.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf_partial_trailer.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf_partial_trailer.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type - let payload = vec![42u8; 20]; // Assume 20 byte payload - buf_partial_trailer.extend_from_slice(&payload); - buf_partial_trailer.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer - - // Total length = 16 + 2 + 20 + 15 = 53. Min size = 34. This passes. - // Payload size = 53 - 34 = 19. <--- THIS IS WRONG. The parser assumes the length dictates payload. - // Let's fix the parser logic slightly. - - // The point is, parse_lp_packet expects a COMPLETE datagram. Providing less bytes - // than LpHeader + Type + Trailer should fail. Providing *more* is also an issue unless - // the length calculation works out perfectly. The most direct test is just < min_size. - // Renaming test to reflect this. - } - - #[test] - fn test_parse_buffer_smaller_than_minimum() { - // Test a buffer that's smaller than the smallest possible packet (LpHeader+Type+Trailer) - let mut buf_too_short = BytesMut::new(); - buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf_too_short.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Type - buf_too_short.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer - // Length = 16 + 2 + 15 = 33. Min Size = 34. - let result_too_short = parse_lp_packet(&buf_too_short, None); - assert!( - result_too_short.is_err(), - "Expected error for buffer size 33, min 34" - ); - assert!(matches!( - result_too_short.unwrap_err(), - LpError::InsufficientBufferSize - )); - } - - #[test] - fn test_parse_invalid_message_type() { - // Create a buffer with invalid message type - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&231u16.to_le_bytes()); // Invalid message type - // Need payload and trailer to meet min_size requirement - let payload_size = 10; // Arbitrary - buf.extend_from_slice(&vec![0u8; payload_size]); // Some data - buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - - // Attempt to parse - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - match result { - Err(LpError::InvalidMessageType(231)) => {} // Expected error - Err(e) => panic!("Expected InvalidMessageType error, got {:?}", e), - Ok(_) => panic!("Expected error, but got Ok"), - } - } - - #[test] - fn test_parse_incorrect_payload_size_for_busy() { - // Create a Busy packet but *with* a payload (which is invalid) - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Busy type - buf.extend_from_slice(&[42; 1]); // <<< Invalid 1-byte payload for Busy - buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - - // Total size = 16 + 2 + 1 + 16 = 35. Min size = 34. - // Calculated payload size = 35 - 34 = 1. - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - assert!(matches!( - result.unwrap_err(), - LpError::InvalidPayloadSize { - expected: 0, - actual: 1 - } - )); - } - - // Test multiple packets simulation isn't relevant for datagram parsing + // // #[test] - // fn test_multiple_packets_in_buffer() { ... } - - // === ClientHello Serialization Tests === - - #[test] - fn test_serialize_parse_client_hello() { - use crate::message::ClientHelloData; - - let mut rng = thread_rng(); - let valid_ed25519 = ed25519::KeyPair::new(&mut rng); - let mut dst = BytesMut::new(); - - // Create ClientHelloData - let client_key = x25519::PublicKey::from_byte_array(&[42u8; 32]); - let client_ed25519_key = *valid_ed25519.public_key(); - let salt = [99u8; 32]; - let hello_data = ClientHelloData { - receiver_index: 12345, - client_lp_public_key: client_key, - client_ed25519_public_key: client_ed25519_key, - salt, - }; - - // Create a ClientHello message packet - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 123, - }, - message: LpMessage::ClientHello(hello_data), - trailer: [0; TRAILER_LEN], - }; - - // Serialize the packet - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse the packet - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify the packet fields - assert_eq!(decoded.header.protocol_version, 1); - assert_eq!(decoded.header.receiver_idx, 42); - assert_eq!(decoded.header.counter, 123); - - // Verify message type and data - match decoded.message { - LpMessage::ClientHello(decoded_data) => { - assert_eq!(decoded_data.client_lp_public_key, client_key); - assert_eq!(decoded_data.salt, salt); - } - _ => panic!("Expected ClientHello message"), - } - assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - } - - #[test] - fn test_serialize_parse_client_hello_with_fresh_salt() { - use crate::message::ClientHelloData; - - let mut dst = BytesMut::new(); - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("System time before UNIX epoch") - .as_secs(); - - // Create ClientHelloData with fresh salt - let mut rng = thread_rng(); - let valid_ed25519 = ed25519::KeyPair::new(&mut rng); - - let client_key = x25519::PublicKey::from_byte_array(&[7u8; 32]); - let client_ed25519_key = *valid_ed25519.public_key(); - let hello_data = - ClientHelloData::new_with_fresh_salt(client_key, client_ed25519_key, timestamp); - - // Create a ClientHello message packet - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 100, - counter: 200, - }, - message: LpMessage::ClientHello(hello_data), - trailer: [55; TRAILER_LEN], - }; - - // Serialize the packet - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse the packet - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify message type and data - match decoded.message { - LpMessage::ClientHello(decoded_data) => { - assert_eq!(decoded_data.client_lp_public_key, client_key); - assert_eq!(decoded_data.salt, hello_data.salt); - - // Verify timestamp can be extracted - let timestamp = decoded_data.extract_timestamp(); - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - // Timestamp should be within 2 seconds of now - assert!((timestamp as i64 - now as i64).abs() <= 2); - } - _ => panic!("Expected ClientHello message"), - } - } - - #[test] - fn test_parse_client_hello_malformed_data() { - // Create a buffer with ClientHello message type but invalid bincode data - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&MessageType::ClientHello.to_u32().to_le_bytes()); // ClientHello type - - // Add data that does not equal to ClientHelloData::LEN - buf.extend_from_slice(&[0xFF; 50]); // invalid data - buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - - // Attempt to parse - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - match result { - Err(LpError::DeserializationError(_)) => {} // Expected error - Err(e) => panic!("Expected DeserializationError, got {:?}", e), - Ok(_) => panic!("Expected error, but got Ok"), - } - } - - #[test] - fn test_parse_client_hello_incomplete_bincode() { - // Create a buffer with ClientHello but truncated bincode data - let mut buf = BytesMut::new(); - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - buf.extend_from_slice(&MessageType::ClientHello.to_u32().to_le_bytes()); // ClientHello type - - // Add incomplete bincode data (only partial ClientHelloData) - buf.extend_from_slice(&[0; 20]); // Too few bytes for full ClientHelloData - buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - - // Attempt to parse - let result = parse_lp_packet(&buf, None); - assert!(result.is_err()); - match result { - Err(LpError::DeserializationError(_)) => {} // Expected error - Err(e) => panic!("Expected DeserializationError, got {:?}", e), - Ok(_) => panic!("Expected error, but got Ok"), - } - } - - #[test] - fn test_forward_packet_encode_decode_roundtrip_v4() { - let mut dst = BytesMut::new(); - - let forward_data = crate::message::ForwardPacketData { - target_gateway_identity: [77u8; 32], - target_lp_address: "1.2.3.4:41264".parse().unwrap(), - inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 999, - counter: 555, - }, - message: LpMessage::ForwardPacket(forward_data), - trailer: [0xff; TRAILER_LEN], - }; - - // Serialize - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse back - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify LP protocol handling works correctly - assert_eq!(decoded.header.receiver_idx, 999); - assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); - - if let LpMessage::ForwardPacket(data) = decoded.message { - assert_eq!(data.target_gateway_identity, [77u8; 32]); - assert_eq!(data.target_lp_address, "1.2.3.4:41264".parse().unwrap()); - assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); - } else { - panic!("Expected ForwardPacket message"); - } - } - - #[test] - fn test_forward_packet_encode_decode_roundtrip_v6() { - let mut dst = BytesMut::new(); - - let forward_data = crate::message::ForwardPacketData { - target_gateway_identity: [77u8; 32], - target_lp_address: "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap(), - inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 999, - counter: 555, - }, - message: LpMessage::ForwardPacket(forward_data), - trailer: [0xff; TRAILER_LEN], - }; - - // Serialize - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - - // Parse back - let decoded = parse_lp_packet(&dst, None).unwrap(); - - // Verify LP protocol handling works correctly - assert_eq!(decoded.header.receiver_idx, 999); - assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); - - if let LpMessage::ForwardPacket(data) = decoded.message { - assert_eq!(data.target_gateway_identity, [77u8; 32]); - assert_eq!( - data.target_lp_address, - "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap() - ); - assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); - } else { - panic!("Expected ForwardPacket message"); - } - } - - // === Outer AEAD Tests === - - #[test] - fn test_aead_roundtrip_with_key() { - // Test that encrypt/decrypt roundtrip works with an AEAD key - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 999, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - // Parse back with the same key - let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - - assert_eq!(decoded.header.protocol_version, 1); - assert_eq!(decoded.header.receiver_idx, 12345); - assert_eq!(decoded.header.counter, 999); - assert!(matches!(decoded.message, LpMessage::Busy)); - } - - #[test] - fn test_aead_ciphertext_differs_from_plaintext() { - // Verify that encrypted payload differs from plaintext - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 999, - }, - message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload(vec![ - 0xAA, 0xBB, 0xCC, 0xDD, - ])), - trailer: [0; TRAILER_LEN], - }; - - let mut cleartext = BytesMut::new(); - serialize_lp_packet(&packet, &mut cleartext, None).unwrap(); - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - // Outer header (receiver_idx + counter) should be the same - always cleartext - assert_eq!(&cleartext[..OUTER_HDR], &encrypted[..OUTER_HDR]); - - // Inner payload (proto + reserved + msg_type + content) should differ (encrypted) - let payload_start = OUTER_HDR; - let payload_end_cleartext = cleartext.len() - TRAILER_LEN; - let payload_end_encrypted = encrypted.len() - TRAILER_LEN; - - assert_ne!( - &cleartext[payload_start..payload_end_cleartext], - &encrypted[payload_start..payload_end_encrypted], - "Encrypted payload should differ from plaintext" - ); - - // Trailer should differ (zeros vs AEAD tag) - assert_ne!( - &cleartext[payload_end_cleartext..], - &encrypted[payload_end_encrypted..], - "Encrypted trailer should be a tag, not zeros" - ); - } - - #[test] - fn test_aead_tampered_tag_fails() { - // Verify that tampering with the tag causes decryption failure - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 999, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - // Tamper with the tag (last byte) - let last_idx = encrypted.len() - 1; - encrypted[last_idx] ^= 0xFF; - - // Parsing should fail with AeadTagMismatch - let result = parse_lp_packet(&encrypted, Some(&outer_key)); - assert!(matches!(result, Err(LpError::AeadTagMismatch))); - } - - #[test] - fn test_aead_tampered_header_fails() { - // Verify that tampering with the header (AAD) causes decryption failure - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 999, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - // Tamper with the outer header AAD (flip a bit in counter at byte 4) - // New format: [receiver_idx(0-3), counter(4-11)], so byte 4 is counter's LSB - encrypted[4] ^= 0x01; - - // Parsing should fail with AeadTagMismatch - let result = parse_lp_packet(&encrypted, Some(&outer_key)); - assert!(matches!(result, Err(LpError::AeadTagMismatch))); - } - - #[test] - fn test_aead_different_counters_produce_different_ciphertext() { - // Verify that different counters (nonces) produce different ciphertexts - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let packet1 = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 1, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let packet2 = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 2, // Different counter - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted1 = BytesMut::new(); - serialize_lp_packet(&packet1, &mut encrypted1, Some(&outer_key)).unwrap(); - - let mut encrypted2 = BytesMut::new(); - serialize_lp_packet(&packet2, &mut encrypted2, Some(&outer_key)).unwrap(); - - // The encrypted inner payloads should differ even though the message is the same - // (because nonce is different). Inner payload starts after outer header. - let payload_start = OUTER_HDR; - assert_ne!( - &encrypted1[payload_start..], - &encrypted2[payload_start..], - "Different counters should produce different ciphertexts" - ); - } - - #[test] - fn test_aead_wrong_key_fails() { - // Verify that decryption with wrong key fails - let psk1 = [42u8; 32]; - let psk2 = [43u8; 32]; // Different PSK - let outer_key1 = OuterAeadKey::from_psk(&psk1); - let outer_key2 = OuterAeadKey::from_psk(&psk2); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 12345, - counter: 999, - }, - message: LpMessage::Busy, - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key1)).unwrap(); - - // Parsing with wrong key should fail - let result = parse_lp_packet(&encrypted, Some(&outer_key2)); - assert!(matches!(result, Err(LpError::AeadTagMismatch))); - } - - #[test] - fn test_aead_encrypted_data_message_roundtrip() { - // Test AEAD with EncryptedData message type (larger payload) - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let payload_data = vec![0xDE; 100]; - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 54321, - counter: 12345678, - }, - message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload( - payload_data.clone(), - )), - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - - match decoded.message { - LpMessage::EncryptedData(data) => { - assert_eq!(data.0, payload_data); - } - _ => panic!("Expected EncryptedData message"), - } - } - - #[test] - fn test_aead_handshake_message_roundtrip() { - // Test AEAD with Handshake message type - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let handshake_data = vec![0x01, 0x02, 0x03, 0x04, 0x05]; - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 99999, - counter: 2, - }, - message: LpMessage::PSQRequest(PSQRequestData(handshake_data.clone())), - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - - match decoded.message { - LpMessage::PSQResponse(data) => { - assert_eq!(data.0, handshake_data); - } - _ => panic!("Expected Handshake message"), - } - } - - // === Subsession Message Tests === - - #[test] - fn test_serialize_parse_subsession_request() { - let mut dst = BytesMut::new(); - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 100, - }, - message: LpMessage::SubsessionRequest, - trailer: [0; TRAILER_LEN], - }; - - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - let decoded = parse_lp_packet(&dst, None).unwrap(); - - assert_eq!(decoded.header.receiver_idx, 42); - assert_eq!(decoded.header.counter, 100); - assert!(matches!(decoded.message, LpMessage::SubsessionRequest)); - } - - #[test] - fn test_serialize_parse_subsession_kk1() { - use crate::message::SubsessionKK1Data; - - let mut dst = BytesMut::new(); - - let kk1_data = SubsessionKK1Data { - payload: vec![0xAA; 50], // 50 bytes KK payload - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 123, - counter: 456, - }, - message: LpMessage::SubsessionKK1(kk1_data.clone()), - trailer: [0; TRAILER_LEN], - }; - - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - let decoded = parse_lp_packet(&dst, None).unwrap(); - - assert_eq!(decoded.header.receiver_idx, 123); - match decoded.message { - LpMessage::SubsessionKK1(data) => { - assert_eq!(data.payload, kk1_data.payload); - } - _ => panic!("Expected SubsessionKK1 message"), - } - } - - #[test] - fn test_serialize_parse_subsession_kk2() { - use crate::message::SubsessionKK2Data; - - let mut dst = BytesMut::new(); - - let kk2_data = SubsessionKK2Data { - payload: vec![0x11; 60], // 60 bytes KK response payload - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 789, - counter: 1000, - }, - message: LpMessage::SubsessionKK2(kk2_data.clone()), - trailer: [0; TRAILER_LEN], - }; - - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - let decoded = parse_lp_packet(&dst, None).unwrap(); - - assert_eq!(decoded.header.receiver_idx, 789); - match decoded.message { - LpMessage::SubsessionKK2(data) => { - assert_eq!(data.payload, kk2_data.payload); - } - _ => panic!("Expected SubsessionKK2 message"), - } - } - - #[test] - fn test_serialize_parse_subsession_ready() { - use crate::message::SubsessionReadyData; - - let mut dst = BytesMut::new(); - - let ready_data = SubsessionReadyData { - receiver_index: 99999, - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 200, - }, - message: LpMessage::SubsessionReady(ready_data.clone()), - trailer: [0; TRAILER_LEN], - }; - - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - let decoded = parse_lp_packet(&dst, None).unwrap(); - - assert_eq!(decoded.header.receiver_idx, 42); - match decoded.message { - LpMessage::SubsessionReady(data) => { - assert_eq!(data.receiver_index, 99999); - } - _ => panic!("Expected SubsessionReady message"), - } - } - - #[test] - fn test_subsession_request_with_payload_fails() { - // SubsessionRequest should have no payload - let mut buf = BytesMut::new(); - buf.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx - buf.extend_from_slice(&123u64.to_le_bytes()); // counter - buf.extend_from_slice(&[1, 0, 0, 0]); // version + reserved - buf.extend_from_slice(&MessageType::SubsessionRequest.to_u32().to_le_bytes()); - buf.extend_from_slice(&[0xFF]); // Invalid payload for SubsessionRequest - buf.extend_from_slice(&[0; TRAILER_LEN]); - - let result = parse_lp_packet(&buf, None); - assert!(matches!( - result, - Err(LpError::InvalidPayloadSize { - expected: 0, - actual: 1 - }) - )); - } - - #[test] - fn test_aead_subsession_roundtrip() { - use crate::message::SubsessionKK1Data; - - let psk = [42u8; 32]; - let outer_key = OuterAeadKey::from_psk(&psk); - - let kk1_data = SubsessionKK1Data { - payload: vec![0xDE; 48], // 48 bytes KK payload - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 54321, - counter: 999, - }, - message: LpMessage::SubsessionKK1(kk1_data.clone()), - trailer: [0; TRAILER_LEN], - }; - - let mut encrypted = BytesMut::new(); - serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - - let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - - match decoded.message { - LpMessage::SubsessionKK1(data) => { - assert_eq!(data.payload, kk1_data.payload); - } - _ => panic!("Expected SubsessionKK1 message"), - } - } - - #[test] - fn test_serialize_parse_error() { - use crate::message::ErrorPacketData; - - let mut dst = BytesMut::new(); - - let error_data = ErrorPacketData { - message: "this is an error".to_string(), - }; - - let packet = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 200, - }, - message: LpMessage::Error(error_data.clone()), - trailer: [0; TRAILER_LEN], - }; - - serialize_lp_packet(&packet, &mut dst, None).unwrap(); - let decoded = parse_lp_packet(&dst, None).unwrap(); - - assert_eq!(decoded.header.receiver_idx, 42); - match decoded.message { - LpMessage::Error(data) => { - assert_eq!(data.message, "this is an error"); - } - _ => panic!("Expected Error message"), - } - } + // fn test_serialize_parse_busy() { + // let mut dst = BytesMut::new(); + // + // // Create a Busy packet + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 123, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // // Serialize the packet (cleartext) + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // + // // Parse the packet (cleartext) + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // // Verify the packet fields + // assert_eq!(decoded.header.protocol_version, 1); + // assert_eq!(decoded.header.receiver_idx, 42); + // assert_eq!(decoded.header.counter, 123); + // assert!(matches!(decoded.message, LpMessage::Busy)); + // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); + // } + // + // #[test] + // fn test_serialize_parse_handshake() { + // let mut dst = BytesMut::new(); + // + // // Create a Handshake message packet + // let payload = vec![42u8; 80]; // Example payload size + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 123, + // }, + // message: LpMessage::PSQRequest(PSQRequestData(payload.clone())), + // trailer: [0; TRAILER_LEN], + // }; + // + // // Serialize the packet (cleartext) + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // + // // Parse the packet (cleartext) + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // // Verify the packet fields + // assert_eq!(decoded.header.protocol_version, 1); + // assert_eq!(decoded.header.receiver_idx, 42); + // assert_eq!(decoded.header.counter, 123); + // + // // Verify message type and data + // match decoded.message { + // LpMessage::PSQRequest(decoded_payload) => { + // assert_eq!(decoded_payload, PSQRequestData(payload)); + // } + // _ => panic!("Expected Handshake message"), + // } + // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); + // } + // + // #[test] + // fn test_serialize_parse_encrypted_data() { + // let mut dst = BytesMut::new(); + // + // // Create an EncryptedData message packet + // let payload = vec![43u8; 124]; // Example payload size + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 123, + // }, + // message: LpMessage::EncryptedData(EncryptedDataPayload(payload.clone())), + // trailer: [0; TRAILER_LEN], + // }; + // + // // Serialize the packet (cleartext) + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // + // // Parse the packet (cleartext) + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // // Verify the packet fields + // assert_eq!(decoded.header.protocol_version, 1); + // assert_eq!(decoded.header.receiver_idx, 42); + // assert_eq!(decoded.header.counter, 123); + // + // // Verify message type and data + // match decoded.message { + // LpMessage::EncryptedData(decoded_payload) => { + // assert_eq!(decoded_payload, EncryptedDataPayload(payload)); + // } + // _ => panic!("Expected EncryptedData message"), + // } + // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); + // } + // + // // === Incomplete Data Tests === + // + // #[test] + // fn test_parse_incomplete_header() { + // // Create a buffer with incomplete header + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&[1, 0, 0, 0]); // Only 4 bytes, not enough for LpHeader::SIZE + // + // // Attempt to parse - expect error + // let result = parse_lp_packet(&buf, None); + // assert!(result.is_err()); + // assert!(matches!( + // result.unwrap_err(), + // LpError::InsufficientBufferSize + // )); + // } + // + // #[test] + // fn test_parse_incomplete_message_type() { + // // Create a buffer with complete header but incomplete message type + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf.extend_from_slice(&[0]); // Only 1 byte of message type (need 2) + // + // // Buffer length = 16 + 1 = 17. Min size = 16 + 2 + 16 = 34. + // let result = parse_lp_packet(&buf, None); + // assert!(result.is_err()); + // assert!(matches!( + // result.unwrap_err(), + // LpError::InsufficientBufferSize + // )); + // } + // + // #[test] + // fn test_parse_incomplete_message_data() { + // // Create a buffer simulating Handshake but missing trailer and maybe partial payload + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type + // buf.extend_from_slice(&[42; 40]); // 40 bytes of payload data + // + // // Buffer length = 16 + 2 + 40 = 58. Min size = 16 + 2 + 16 = 34. + // // Payload size calculated as 58 - 34 = 24. + // // Trailer expected at index 16 + 2 + 24 = 42. + // // Trailer read attempts src[42..58]. + // // This *should* parse successfully based on the logic, but the trailer is garbage. + // // Let's rethink: parse_lp_packet assumes the *entire slice* is the packet. + // // If the slice doesn't end exactly where the trailer should, it's an error. + // // In this case, total length is 58. OuterHdr(12) + InnerPrefix(4) + Type(2) + Trailer(16) = 34. Payload = 58-34=24. + // // Trailer starts at 16+2+24 = 42. Ends at 42+16=58. It fits exactly. + // // This test *still* doesn't test incompleteness correctly for the datagram parser. + // + // // Let's test a buffer that's *too short* even for header+type+trailer+min_payload + // // Note: Buffer order doesn't matter for this test since we fail on minimum size check + // let mut buf_too_short = BytesMut::new(); + // buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx (outer header) + // buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // counter (outer header) + // buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // version + reserved (inner prefix) + // buf_too_short.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // msg type + // // No payload, no trailer. Length = 12+4+2=18. Min size = 34. + // let result_too_short = parse_lp_packet(&buf_too_short, None); + // assert!(result_too_short.is_err()); + // assert!(matches!( + // result_too_short.unwrap_err(), + // LpError::InsufficientBufferSize + // )); + // + // // Test a buffer missing PART of the trailer + // let mut buf_partial_trailer = BytesMut::new(); + // buf_partial_trailer.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf_partial_trailer.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf_partial_trailer.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf_partial_trailer.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type + // let payload = vec![42u8; 20]; // Assume 20 byte payload + // buf_partial_trailer.extend_from_slice(&payload); + // buf_partial_trailer.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer + // + // // Total length = 16 + 2 + 20 + 15 = 53. Min size = 34. This passes. + // // Payload size = 53 - 34 = 19. <--- THIS IS WRONG. The parser assumes the length dictates payload. + // // Let's fix the parser logic slightly. + // + // // The point is, parse_lp_packet expects a COMPLETE datagram. Providing less bytes + // // than LpHeader + Type + Trailer should fail. Providing *more* is also an issue unless + // // the length calculation works out perfectly. The most direct test is just < min_size. + // // Renaming test to reflect this. + // } + // + // #[test] + // fn test_parse_buffer_smaller_than_minimum() { + // // Test a buffer that's smaller than the smallest possible packet (LpHeader+Type+Trailer) + // let mut buf_too_short = BytesMut::new(); + // buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf_too_short.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Type + // buf_too_short.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer + // // Length = 16 + 2 + 15 = 33. Min Size = 34. + // let result_too_short = parse_lp_packet(&buf_too_short, None); + // assert!( + // result_too_short.is_err(), + // "Expected error for buffer size 33, min 34" + // ); + // assert!(matches!( + // result_too_short.unwrap_err(), + // LpError::InsufficientBufferSize + // )); + // } + // + // #[test] + // fn test_parse_invalid_message_type() { + // // Create a buffer with invalid message type + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf.extend_from_slice(&231u16.to_le_bytes()); // Invalid message type + // // Need payload and trailer to meet min_size requirement + // let payload_size = 10; // Arbitrary + // buf.extend_from_slice(&vec![0u8; payload_size]); // Some data + // buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer + // + // // Attempt to parse + // let result = parse_lp_packet(&buf, None); + // assert!(result.is_err()); + // match result { + // Err(LpError::InvalidMessageType(231)) => {} // Expected error + // Err(e) => panic!("Expected InvalidMessageType error, got {:?}", e), + // Ok(_) => panic!("Expected error, but got Ok"), + // } + // } + // + // #[test] + // fn test_parse_incorrect_payload_size_for_busy() { + // // Create a Busy packet but *with* a payload (which is invalid) + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index + // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter + // buf.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Busy type + // buf.extend_from_slice(&[42; 1]); // <<< Invalid 1-byte payload for Busy + // buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer + // + // // Total size = 16 + 2 + 1 + 16 = 35. Min size = 34. + // // Calculated payload size = 35 - 34 = 1. + // let result = parse_lp_packet(&buf, None); + // assert!(result.is_err()); + // assert!(matches!( + // result.unwrap_err(), + // LpError::InvalidPayloadSize { + // expected: 0, + // actual: 1 + // } + // )); + // } + // + // // Test multiple packets simulation isn't relevant for datagram parsing + // // #[test] + // // fn test_multiple_packets_in_buffer() { ... } + // + // + // + // #[test] + // fn test_forward_packet_encode_decode_roundtrip_v4() { + // let mut dst = BytesMut::new(); + // + // let forward_data = crate::message::ForwardPacketData { + // target_gateway_identity: [77u8; 32], + // target_lp_address: "1.2.3.4:41264".parse().unwrap(), + // inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 999, + // counter: 555, + // }, + // message: LpMessage::ForwardPacket(forward_data), + // trailer: [0xff; TRAILER_LEN], + // }; + // + // // Serialize + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // + // // Parse back + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // // Verify LP protocol handling works correctly + // assert_eq!(decoded.header.receiver_idx, 999); + // assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); + // + // if let LpMessage::ForwardPacket(data) = decoded.message { + // assert_eq!(data.target_gateway_identity, [77u8; 32]); + // assert_eq!(data.target_lp_address, "1.2.3.4:41264".parse().unwrap()); + // assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); + // } else { + // panic!("Expected ForwardPacket message"); + // } + // } + // + // #[test] + // fn test_forward_packet_encode_decode_roundtrip_v6() { + // let mut dst = BytesMut::new(); + // + // let forward_data = crate::message::ForwardPacketData { + // target_gateway_identity: [77u8; 32], + // target_lp_address: "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap(), + // inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 999, + // counter: 555, + // }, + // message: LpMessage::ForwardPacket(forward_data), + // trailer: [0xff; TRAILER_LEN], + // }; + // + // // Serialize + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // + // // Parse back + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // // Verify LP protocol handling works correctly + // assert_eq!(decoded.header.receiver_idx, 999); + // assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); + // + // if let LpMessage::ForwardPacket(data) = decoded.message { + // assert_eq!(data.target_gateway_identity, [77u8; 32]); + // assert_eq!( + // data.target_lp_address, + // "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap() + // ); + // assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); + // } else { + // panic!("Expected ForwardPacket message"); + // } + // } + // + // // === Outer AEAD Tests === + // + // #[test] + // fn test_aead_roundtrip_with_key() { + // // Test that encrypt/decrypt roundtrip works with an AEAD key + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 999, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // // Parse back with the same key + // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); + // + // assert_eq!(decoded.header.protocol_version, 1); + // assert_eq!(decoded.header.receiver_idx, 12345); + // assert_eq!(decoded.header.counter, 999); + // assert!(matches!(decoded.message, LpMessage::Busy)); + // } + // + // #[test] + // fn test_aead_ciphertext_differs_from_plaintext() { + // // Verify that encrypted payload differs from plaintext + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 999, + // }, + // message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload(vec![ + // 0xAA, 0xBB, 0xCC, 0xDD, + // ])), + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut cleartext = BytesMut::new(); + // serialize_lp_packet(&packet, &mut cleartext, None).unwrap(); + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // // Outer header (receiver_idx + counter) should be the same - always cleartext + // assert_eq!(&cleartext[..OUTER_HDR], &encrypted[..OUTER_HDR]); + // + // // Inner payload (proto + reserved + msg_type + content) should differ (encrypted) + // let payload_start = OUTER_HDR; + // let payload_end_cleartext = cleartext.len() - TRAILER_LEN; + // let payload_end_encrypted = encrypted.len() - TRAILER_LEN; + // + // assert_ne!( + // &cleartext[payload_start..payload_end_cleartext], + // &encrypted[payload_start..payload_end_encrypted], + // "Encrypted payload should differ from plaintext" + // ); + // + // // Trailer should differ (zeros vs AEAD tag) + // assert_ne!( + // &cleartext[payload_end_cleartext..], + // &encrypted[payload_end_encrypted..], + // "Encrypted trailer should be a tag, not zeros" + // ); + // } + // + // #[test] + // fn test_aead_tampered_tag_fails() { + // // Verify that tampering with the tag causes decryption failure + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 999, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // // Tamper with the tag (last byte) + // let last_idx = encrypted.len() - 1; + // encrypted[last_idx] ^= 0xFF; + // + // // Parsing should fail with AeadTagMismatch + // let result = parse_lp_packet(&encrypted, Some(&outer_key)); + // assert!(matches!(result, Err(LpError::AeadTagMismatch))); + // } + // + // #[test] + // fn test_aead_tampered_header_fails() { + // // Verify that tampering with the header (AAD) causes decryption failure + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 999, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // // Tamper with the outer header AAD (flip a bit in counter at byte 4) + // // New format: [receiver_idx(0-3), counter(4-11)], so byte 4 is counter's LSB + // encrypted[4] ^= 0x01; + // + // // Parsing should fail with AeadTagMismatch + // let result = parse_lp_packet(&encrypted, Some(&outer_key)); + // assert!(matches!(result, Err(LpError::AeadTagMismatch))); + // } + // + // #[test] + // fn test_aead_different_counters_produce_different_ciphertext() { + // // Verify that different counters (nonces) produce different ciphertexts + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let packet1 = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 1, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let packet2 = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 2, // Different counter + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted1 = BytesMut::new(); + // serialize_lp_packet(&packet1, &mut encrypted1, Some(&outer_key)).unwrap(); + // + // let mut encrypted2 = BytesMut::new(); + // serialize_lp_packet(&packet2, &mut encrypted2, Some(&outer_key)).unwrap(); + // + // // The encrypted inner payloads should differ even though the message is the same + // // (because nonce is different). Inner payload starts after outer header. + // let payload_start = OUTER_HDR; + // assert_ne!( + // &encrypted1[payload_start..], + // &encrypted2[payload_start..], + // "Different counters should produce different ciphertexts" + // ); + // } + // + // #[test] + // fn test_aead_wrong_key_fails() { + // // Verify that decryption with wrong key fails + // let psk1 = [42u8; 32]; + // let psk2 = [43u8; 32]; // Different PSK + // let outer_key1 = OuterAeadKey::from_psk(&psk1); + // let outer_key2 = OuterAeadKey::from_psk(&psk2); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 12345, + // counter: 999, + // }, + // message: LpMessage::Busy, + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key1)).unwrap(); + // + // // Parsing with wrong key should fail + // let result = parse_lp_packet(&encrypted, Some(&outer_key2)); + // assert!(matches!(result, Err(LpError::AeadTagMismatch))); + // } + // + // #[test] + // fn test_aead_encrypted_data_message_roundtrip() { + // // Test AEAD with EncryptedData message type (larger payload) + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let payload_data = vec![0xDE; 100]; + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 54321, + // counter: 12345678, + // }, + // message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload( + // payload_data.clone(), + // )), + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); + // + // match decoded.message { + // LpMessage::EncryptedData(data) => { + // assert_eq!(data.0, payload_data); + // } + // _ => panic!("Expected EncryptedData message"), + // } + // } + // + // #[test] + // fn test_aead_handshake_message_roundtrip() { + // // Test AEAD with Handshake message type + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let handshake_data = vec![0x01, 0x02, 0x03, 0x04, 0x05]; + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 99999, + // counter: 2, + // }, + // message: LpMessage::PSQRequest(PSQRequestData(handshake_data.clone())), + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); + // + // match decoded.message { + // LpMessage::PSQResponse(data) => { + // assert_eq!(data.0, handshake_data); + // } + // _ => panic!("Expected Handshake message"), + // } + // } + // + // // === Subsession Message Tests === + // + // #[test] + // fn test_serialize_parse_subsession_request() { + // let mut dst = BytesMut::new(); + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 100, + // }, + // message: LpMessage::SubsessionRequest, + // trailer: [0; TRAILER_LEN], + // }; + // + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // assert_eq!(decoded.header.receiver_idx, 42); + // assert_eq!(decoded.header.counter, 100); + // assert!(matches!(decoded.message, LpMessage::SubsessionRequest)); + // } + // + // #[test] + // fn test_serialize_parse_subsession_kk1() { + // use crate::message::SubsessionKK1Data; + // + // let mut dst = BytesMut::new(); + // + // let kk1_data = SubsessionKK1Data { + // payload: vec![0xAA; 50], // 50 bytes KK payload + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 123, + // counter: 456, + // }, + // message: LpMessage::SubsessionKK1(kk1_data.clone()), + // trailer: [0; TRAILER_LEN], + // }; + // + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // assert_eq!(decoded.header.receiver_idx, 123); + // match decoded.message { + // LpMessage::SubsessionKK1(data) => { + // assert_eq!(data.payload, kk1_data.payload); + // } + // _ => panic!("Expected SubsessionKK1 message"), + // } + // } + // + // #[test] + // fn test_serialize_parse_subsession_kk2() { + // use crate::message::SubsessionKK2Data; + // + // let mut dst = BytesMut::new(); + // + // let kk2_data = SubsessionKK2Data { + // payload: vec![0x11; 60], // 60 bytes KK response payload + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 789, + // counter: 1000, + // }, + // message: LpMessage::SubsessionKK2(kk2_data.clone()), + // trailer: [0; TRAILER_LEN], + // }; + // + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // assert_eq!(decoded.header.receiver_idx, 789); + // match decoded.message { + // LpMessage::SubsessionKK2(data) => { + // assert_eq!(data.payload, kk2_data.payload); + // } + // _ => panic!("Expected SubsessionKK2 message"), + // } + // } + // + // #[test] + // fn test_serialize_parse_subsession_ready() { + // use crate::message::SubsessionReadyData; + // + // let mut dst = BytesMut::new(); + // + // let ready_data = SubsessionReadyData { + // receiver_index: 99999, + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 200, + // }, + // message: LpMessage::SubsessionReady(ready_data.clone()), + // trailer: [0; TRAILER_LEN], + // }; + // + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // assert_eq!(decoded.header.receiver_idx, 42); + // match decoded.message { + // LpMessage::SubsessionReady(data) => { + // assert_eq!(data.receiver_index, 99999); + // } + // _ => panic!("Expected SubsessionReady message"), + // } + // } + // + // #[test] + // fn test_subsession_request_with_payload_fails() { + // // SubsessionRequest should have no payload + // let mut buf = BytesMut::new(); + // buf.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx + // buf.extend_from_slice(&123u64.to_le_bytes()); // counter + // buf.extend_from_slice(&[1, 0, 0, 0]); // version + reserved + // buf.extend_from_slice(&MessageType::SubsessionRequest.to_u32().to_le_bytes()); + // buf.extend_from_slice(&[0xFF]); // Invalid payload for SubsessionRequest + // buf.extend_from_slice(&[0; TRAILER_LEN]); + // + // let result = parse_lp_packet(&buf, None); + // assert!(matches!( + // result, + // Err(LpError::InvalidPayloadSize { + // expected: 0, + // actual: 1 + // }) + // )); + // } + // + // #[test] + // fn test_aead_subsession_roundtrip() { + // use crate::message::SubsessionKK1Data; + // + // let psk = [42u8; 32]; + // let outer_key = OuterAeadKey::from_psk(&psk); + // + // let kk1_data = SubsessionKK1Data { + // payload: vec![0xDE; 48], // 48 bytes KK payload + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 54321, + // counter: 999, + // }, + // message: LpMessage::SubsessionKK1(kk1_data.clone()), + // trailer: [0; TRAILER_LEN], + // }; + // + // let mut encrypted = BytesMut::new(); + // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); + // + // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); + // + // match decoded.message { + // LpMessage::SubsessionKK1(data) => { + // assert_eq!(data.payload, kk1_data.payload); + // } + // _ => panic!("Expected SubsessionKK1 message"), + // } + // } + // + // #[test] + // fn test_serialize_parse_error() { + // use crate::message::ErrorPacketData; + // + // let mut dst = BytesMut::new(); + // + // let error_data = ErrorPacketData { + // message: "this is an error".to_string(), + // }; + // + // let packet = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 200, + // }, + // message: LpMessage::Error(error_data.clone()), + // trailer: [0; TRAILER_LEN], + // }; + // + // serialize_lp_packet(&packet, &mut dst, None).unwrap(); + // let decoded = parse_lp_packet(&dst, None).unwrap(); + // + // assert_eq!(decoded.header.receiver_idx, 42); + // match decoded.message { + // LpMessage::Error(data) => { + // assert_eq!(data.message, "this is an error"); + // } + // _ => panic!("Expected Error message"), + // } + // } } diff --git a/common/nym-lp/src/error.rs b/common/nym-lp/src/error.rs index 066b849a9b..d687099d34 100644 --- a/common/nym-lp/src/error.rs +++ b/common/nym-lp/src/error.rs @@ -2,7 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use crate::message::MessageType; -use crate::{noise_protocol::NoiseError, replay::ReplayError}; +use crate::replay::ReplayError; +use crate::session::SessionId; use libcrux_psq::handshake::HandshakeError; use libcrux_psq::handshake::builders::BuilderError; use libcrux_psq::session::SessionError; @@ -16,27 +17,9 @@ pub enum LpError { #[error("IO Error: {0}")] IoError(#[from] std::io::Error), - // noiserm - #[error("Snow Error: {0}")] - SnowKeyError(#[from] snow::Error), - - // noiserm - #[error("Snow Pattern Error: {0}")] - SnowPatternError(String), - - // noiserm - #[error("Noise Protocol Error: {0}")] - NoiseError(#[from] NoiseError), - - #[error("PSQ Error: {0}")] - PSQError(String), - #[error("Replay detected: {0}")] Replay(#[from] ReplayError), - #[error("Invalid packet format: {0}")] - InvalidPacketFormat(String), - #[error("Invalid message type: {0}")] InvalidMessageType(u32), @@ -81,17 +64,13 @@ pub enum LpError { LpSessionProcessing, /// State machine not found. - #[error("State machine not found for lp_id: {lp_id}")] - StateMachineNotFound { lp_id: u32 }, + #[error("State machine not found for lp_id: {lp_id:?}")] + StateMachineNotFound { lp_id: SessionId }, /// Ed25519 to X25519 conversion error. #[error("Ed25519 key conversion error: {0}")] Ed25519RecoveryError(#[from] Ed25519RecoveryError), - /// Outer AEAD authentication tag verification failed. - #[error("AEAD authentication tag verification failed")] - AeadTagMismatch, - /// Received an LP packet with an incompatible, future, version #[error("incompatible LP packet version. got: {got}, highest supported: {highest_supported}")] IncompatibleFuturePacketVersion { got: u8, highest_supported: u8 }, @@ -135,6 +114,9 @@ pub enum LpError { #[error("failed to run the PSQ session: {inner:?}")] PSQSessionFailure { inner: SessionError }, + #[error("failed to derive a transport channel: {inner:?}")] + TransportDerivationFailure { inner: SessionError }, + #[error("the initiator authenticator is not available after ingesting PSQ msg1")] MissingInitiatorAuthenticator, } @@ -149,6 +131,10 @@ impl LpError { "received unexpected response, got: {got:?}, expected: {expected:?}" )) } + + pub fn invalid_message_type(message_type: u32) -> Self { + LpError::InvalidMessageType(message_type) + } } impl From for LpError { diff --git a/common/nym-lp/src/kkt_orchestrator.rs b/common/nym-lp/src/kkt_orchestrator.rs deleted file mode 100644 index 98e3cdff71..0000000000 --- a/common/nym-lp/src/kkt_orchestrator.rs +++ /dev/null @@ -1,492 +0,0 @@ -// Copyright 2025 - Nym Technologies SA -// SPDX-License-Identifier: Apache-2.0 - -//! KKT (Key Encapsulation Transport) orchestration for nym-lp sessions. -//! -//! This module provides functions to perform KKT key exchange before establishing -//! an nym-lp session. The KKT protocol allows secure distribution of post-quantum -//! KEM public keys, which are then used with PSQ to derive a strong pre-shared key -//! for the Noise protocol. -//! -//! # Protocol Flow -//! -//! 1. **Client (Initiator)**: -//! - Calls `create_request()` to generate a KKT request -//! - Sends `LpMessage::KKTRequest` to gateway -//! - Receives `LpMessage::KKTResponse` from gateway -//! - Calls `process_response()` to validate and extract gateway's KEM key -//! -//! 2. **Gateway (Responder)**: -//! - Receives `LpMessage::KKTRequest` from client -//! - Calls `handle_request()` to validate request and generate response -//! - Sends `LpMessage::KKTResponse` to client -//! -//! # Example -//! -//! ```ignore -//! use nym_lp::kkt_orchestrator::{create_request, process_response, handle_request}; -//! use nym_lp::message::{KKTRequestData, KKTResponseData}; -//! use nym_kkt::ciphersuite::{Ciphersuite, KEM, HashFunction, SignatureScheme, EncapsulationKey}; -//! -//! // Setup ciphersuite -//! let ciphersuite = Ciphersuite::resolve_ciphersuite( -//! KEM::X25519, -//! HashFunction::Blake3, -//! SignatureScheme::Ed25519, -//! None, -//! ).unwrap(); -//! -//! // Client: Create request -//! let (session_secret, client_context, request_data) = create_request( -//! ciphersuite, -//! &client_signing_key, -//! &responder_dh_public_key -//! ).unwrap(); -//! -//! // Gateway: Handle request -//! let response_data = handle_request( -//! &request_data, -//! Some(&client_verification_key), -//! &gateway_signing_key, -//! &gateway_dh_private_key, -//! &gateway_kem_public_key, -//! ).unwrap(); -//! -//! // Client: Process response -//! let gateway_kem_key = process_response( -//! client_context, -//! &session_secret, -//! &gateway_verification_key, -//! &expected_key_hash, -//! &response_data, -//! ).unwrap(); -//! ``` - -use crate::LpError; -use crate::message::{KKTRequestData, KKTResponseData}; -use nym_crypto::asymmetric::{ed25519, x25519}; -use nym_kkt::ciphersuite::{Ciphersuite, EncapsulationKey}; -use nym_kkt::context::KKTContext; -use nym_kkt::encryption::KKTSessionSecret; -use nym_kkt::kkt::{handle_kem_request, request_kem_key, validate_kem_response}; - -/// Creates a KKT request to obtain the responder's KEM public key. -/// -/// This is called by the **client (initiator)** to begin the KKT exchange. -/// The returned context must be used when processing the response. -/// -/// # Arguments -/// * `ciphersuite` - Negotiated ciphersuite (KEM, hash, signature algorithms) -/// * `signing_key` - Client's Ed25519 signing key for authentication -/// * `responder_dh_public_key` - Gateway's x25519 public key (from directory) -/// -/// # Returns -/// * `KKTSessionSecret` - Session secret key to encrypt/decrypt KKT messages for this session -/// * `KKTContext` - Context to use when validating the response -/// * `KKTRequestData` - Serialized KKT request frame to send to gateway -/// -/// # Errors -/// Returns `LpError::KKTError` if KKT request generation fails. -pub fn create_request( - ciphersuite: Ciphersuite, - signing_key: &ed25519::PrivateKey, - responder_dh_public_key: &x25519::PublicKey, -) -> Result<(KKTSessionSecret, KKTContext, KKTRequestData), LpError> { - // Note: Uses rand 0.9's thread_rng() to match nym-kkt's rand version - let mut rng = rand09::rng(); - let (session_secret, context, request_bytes) = - request_kem_key(&mut rng, ciphersuite, signing_key, responder_dh_public_key) - .map_err(|e| LpError::KKTError(e.to_string()))?; - - Ok((session_secret, context, KKTRequestData(request_bytes))) -} - -/// Processes a KKT response and extracts the responder's KEM public key. -/// -/// This is called by the **client (initiator)** after receiving a KKT response -/// from the gateway. It verifies the signature and validates the key hash. -/// -/// # Arguments -/// * `context` - Context from the initial `create_request()` call -/// * `session_secret` - The KKT session secret key from the initial `create_request()` call -/// * `responder_vk` - Responder's Ed25519 verification key (from directory) -/// * `expected_key_hash` - Expected hash of responder's KEM key (from directory) -/// * `response_data` - Serialized KKT response frame from responder -/// -/// # Returns -/// * `EncapsulationKey` - Authenticated KEM public key of the responder -/// -/// # Errors -/// Returns `LpError::KKTError` if: -/// - Response deserialization fails -/// - Signature verification fails -/// - Key hash doesn't match expected value -pub fn process_response<'a>( - mut context: KKTContext, - session_secret: &KKTSessionSecret, - responder_vk: &ed25519::PublicKey, - expected_key_hash: &[u8], - response_data: &KKTResponseData, -) -> Result, LpError> { - validate_kem_response( - &mut context, - session_secret, - responder_vk, - expected_key_hash, - &response_data.0, - ) - .map_err(|e| LpError::KKTError(e.to_string())) -} - -/// Handles a KKT request and generates a signed response with the responder's KEM key. -/// -/// This is called by the **gateway (responder)** when receiving a KKT request -/// from a client. It validates the request signature (if authenticated) and -/// responds with the gateway's KEM public key, signed for authenticity. -/// -/// # Arguments -/// * `request_data` - Serialized KKT request frame from initiator -/// * `initiator_vk` - Initiator's Ed25519 verification key (None for anonymous) -/// * `responder_signing_key` - Gateway's Ed25519 signing key -/// * `responder_dh_private_key` - Gateway's x25519 private key -/// * `responder_kem_key` - Gateway's KEM public key to send -/// -/// # Returns -/// * `KKTResponseData` - Signed response frame containing the KEM public key -/// -/// # Errors -/// Returns `LpError::KKTError` if: -/// - Request deserialization fails -/// - Signature verification fails (if authenticated) -/// - Response generation fails -pub fn handle_request<'a>( - request_data: &KKTRequestData, - initiator_vk: Option<&ed25519::PublicKey>, - responder_signing_key: &ed25519::PrivateKey, - responder_dh_private_key: &x25519::PrivateKey, - responder_kem_key: &EncapsulationKey<'a>, -) -> Result { - let mut rng = rand09::rng(); - // Handle the request and generate response - let response_bytes = handle_kem_request( - &mut rng, - &request_data.0, - initiator_vk, - responder_signing_key, - responder_dh_private_key, - responder_kem_key, - ) - .map_err(|e| LpError::KKTError(e.to_string()))?; - - Ok(KKTResponseData(response_bytes)) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::peer::mock_peers; - use nym_kkt::ciphersuite::{HashFunction, KEM, SignatureScheme}; - use nym_kkt::key_utils::{ - generate_keypair_ed25519, generate_keypair_libcrux, generate_keypair_x25519, - hash_encapsulation_key, - }; - use rand09::RngCore; - - #[test] - fn test_kkt_roundtrip_authenticated() { - let mut rng = rand09::rng(); - - // Generate Ed25519 keypairs for both parties - let initiator_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(0)); - let responder_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(1)); - - let responder_x25519 = generate_keypair_x25519(&mut rng); - - // Generate responder's KEM keypair (X25519 for testing) - let (_, responder_kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - let responder_kem_key = EncapsulationKey::X25519(responder_kem_pk); - - // Create ciphersuite - let ciphersuite = Ciphersuite::resolve_ciphersuite( - KEM::X25519, - HashFunction::Blake3, - SignatureScheme::Ed25519, - None, - ) - .unwrap(); - - // Hash the KEM key (simulating directory storage) - let key_hash = hash_encapsulation_key( - &ciphersuite.hash_function(), - ciphersuite.hash_len(), - &responder_kem_key.encode(), - ); - - // Client: Create request - let (session_secret, context, request_data) = create_request( - ciphersuite, - initiator_ed25519_keypair.private_key(), - responder_x25519.public_key(), - ) - .unwrap(); - - // Gateway: Handle request - let response_data = handle_request( - &request_data, - Some(initiator_ed25519_keypair.public_key()), - responder_ed25519_keypair.private_key(), - responder_x25519.private_key(), - &responder_kem_key, - ) - .unwrap(); - - // Client: Process response - let obtained_key = process_response( - context, - &session_secret, - responder_ed25519_keypair.public_key(), - &key_hash, - &response_data, - ) - .unwrap(); - - // Verify we got the correct KEM key - assert_eq!(obtained_key.encode(), responder_kem_key.encode()); - } - - // #[test] - // fn test_kkt_roundtrip_anonymous() { - // let mut rng = rand09::rng(); - - // // Only responder has keys (anonymous initiator) - // // Generate Ed25519 keypairs for both parties - - // let responder_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(1)); - - // let responder_x25519 = generate_keypair_x25519(&mut rng); - - // let (_, responder_kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let responder_kem_key = EncapsulationKey::X25519(responder_kem_pk); - - // let ciphersuite = Ciphersuite::resolve_ciphersuite( - // KEM::X25519, - // HashFunction::Blake3, - // SignatureScheme::Ed25519, - // None, - // ) - // .unwrap(); - - // let key_hash = hash_encapsulation_key( - // &ciphersuite.hash_function(), - // ciphersuite.hash_len(), - // &responder_kem_key.encode(), - // ); - - // // Anonymous initiator - use anonymous_initiator_process directly - // use nym_kkt::kkt::anonymous_initiator_process; - // let (mut context, request_frame) = - // anonymous_initiator_process(&mut rng, ciphersuite).unwrap(); - // let request_data = KKTRequestData(request_frame.to_bytes()); - - // // Gateway: Handle anonymous request - // let response_data = handle_request( - // &request_data, - // None, - // responder_ed25519_keypair.private_key(), - // &responder_x25519_sk, - // &responder_kem_key, - // ) - // .unwrap(); - - // // Initiator: Validate response - // let obtained_key = initiator_ingest_response( - // &mut context, - // responder_ed25519_keypair.public_key(), - // &key_hash, - // &response_data.0, - // ) - // .unwrap(); - - // assert_eq!(obtained_key.encode(), responder_kem_key.encode()); - // } - - #[test] - fn test_invalid_signature_rejected() { - let mut rng = rand09::rng(); - - // Generate Ed25519 keypairs for both parties - let initiator_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(0)); - let responder_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(1)); - - let responder_x25519 = generate_keypair_x25519(&mut rng); - - // Different keypair for wrong signature - let mut wrong_secret = [0u8; 32]; - rng.fill_bytes(&mut wrong_secret); - let wrong_keypair = ed25519::KeyPair::from_secret(wrong_secret, 2); - - let (_, responder_kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - let responder_kem_key = EncapsulationKey::X25519(responder_kem_pk); - - let ciphersuite = Ciphersuite::resolve_ciphersuite( - KEM::X25519, - HashFunction::Blake3, - SignatureScheme::Ed25519, - None, - ) - .unwrap(); - - let (_session_secret, _context, request_data) = create_request( - ciphersuite, - initiator_ed25519_keypair.private_key(), - responder_x25519.public_key(), - ) - .unwrap(); - - // Gateway handles request but we provide WRONG verification key - let result = handle_request( - &request_data, - Some(wrong_keypair.public_key()), // Wrong key! - responder_ed25519_keypair.private_key(), - responder_x25519.private_key(), - &responder_kem_key, - ); - - // Should fail signature verification - assert!(result.is_err()); - if let Err(LpError::KKTError(_)) = result { - // Expected - } else { - panic!("Expected KKTError"); - } - } - - #[test] - fn test_hash_mismatch_rejected() { - let (init, resp) = mock_peers(); - let responder_kem_key = resp.encapsulate_kem_key().unwrap(); - - let ciphersuite = Ciphersuite::resolve_ciphersuite( - KEM::X25519, - HashFunction::Blake3, - SignatureScheme::Ed25519, - None, - ) - .unwrap(); - - // Use WRONG hash - let wrong_hash = [0u8; 32]; - - let (session_secret, context, request_data) = create_request( - ciphersuite, - init.ed25519.private_key(), - resp.x25519.public_key(), - ) - .unwrap(); - - let response_data = handle_request( - &request_data, - Some(init.ed25519.public_key()), - resp.ed25519.private_key(), - resp.x25519.private_key(), - &responder_kem_key, - ) - .unwrap(); - - // Client validates with WRONG hash - let result = process_response( - context, - &session_secret, - resp.ed25519.public_key(), - &wrong_hash, // Wrong! - &response_data, - ); - - // Should fail hash validation - assert!(result.is_err()); - if let Err(LpError::KKTError(_)) = result { - // Expected - } else { - panic!("Expected KKTError"); - } - } - - #[test] - fn test_malformed_request_rejected() { - let mut rng = rand09::rng(); - - let mut responder_secret = [0u8; 32]; - rng.fill_bytes(&mut responder_secret); - let responder_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(1)); - - let responder_x25519 = generate_keypair_x25519(&mut rng); - - let (_, responder_kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - let responder_kem_key = EncapsulationKey::X25519(responder_kem_pk); - - // Create malformed request data (invalid bytes) - let malformed_request = KKTRequestData(vec![0xFF; 100]); - - let result = handle_request( - &malformed_request, - None, - responder_ed25519_keypair.private_key(), - responder_x25519.private_key(), - &responder_kem_key, - ); - - // Should fail to parse - assert!(result.is_err()); - if let Err(LpError::KKTError(_)) = result { - // Expected - } else { - panic!("Expected KKTError"); - } - } - - #[test] - fn test_malformed_response_rejected() { - let mut rng = rand09::rng(); - - // Generate Ed25519 keypairs for both parties - let initiator_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(0)); - let responder_ed25519_keypair = generate_keypair_ed25519(&mut rng, Some(1)); - - let responder_x25519 = generate_keypair_x25519(&mut rng); - - let ciphersuite = Ciphersuite::resolve_ciphersuite( - KEM::X25519, - HashFunction::Blake3, - SignatureScheme::Ed25519, - None, - ) - .unwrap(); - - let (session_secret, context, _request_data) = create_request( - ciphersuite, - initiator_ed25519_keypair.private_key(), - responder_x25519.public_key(), - ) - .unwrap(); - - // Create malformed response data - let malformed_response = KKTResponseData(vec![0xFF; 100]); - let key_hash = [0u8; 32]; - - let result = process_response( - context, - &session_secret, - responder_ed25519_keypair.public_key(), - &key_hash, - &malformed_response, - ); - - // Should fail to parse - assert!(result.is_err()); - if let Err(LpError::KKTError(_)) = result { - // Expected - } else { - panic!("Expected KKTError"); - } - } -} diff --git a/common/nym-lp/src/lib.rs b/common/nym-lp/src/lib.rs index 75f1bf3916..ee4518b7ac 100644 --- a/common/nym-lp/src/lib.rs +++ b/common/nym-lp/src/lib.rs @@ -7,13 +7,9 @@ pub mod codec; // pub use config::LpConfig; pub mod error; -// georgio: no use for this -// pub mod kkt_orchestrator; pub mod message; -pub mod noise_protocol; pub mod packet; pub mod peer; -pub mod psk; pub mod psq; pub mod replay; pub mod session; @@ -22,17 +18,13 @@ pub mod session_manager; pub mod state_machine; pub use error::LpError; -pub use message::{ClientHelloData, LpMessage}; -pub use packet::{BOOTSTRAP_RECEIVER_IDX, LpPacket, OuterHeader}; +pub use message::LpMessage; +pub use packet::{LpPacket, OuterHeader}; pub use replay::{ReceivingKeyCounterValidator, ReplayError}; pub use session::LpSession; pub use session_manager::SessionManager; pub use state_machine::LpStateMachine; -// noiserm -pub const NOISE_PATTERN: &str = "Noise_XKpsk3_25519_ChaChaPoly_SHA256"; -pub const NOISE_PSK_INDEX: u8 = 3; - #[cfg(test)] pub fn kem_list() -> Vec { todo!() @@ -179,108 +171,109 @@ mod tests { use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, SignatureScheme}; // Import the new standalone functions - use crate::codec::{parse_lp_packet, serialize_lp_packet}; + use crate::codec::serialize_lp_packet; #[test] fn test_replay_protection_integration() { - for kem in kem_list() { - // Create session - let mut session = mock_session_for_test(); - - // === Packet 1 (Counter 0 - Should succeed) === - let packet1 = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, // Matches session's sending_index assumption for this test - counter: 0, - }, - message: LpMessage::Busy, - trailer: [0u8; TRAILER_LEN], - }; - - // Serialize packet - let mut buf1 = BytesMut::new(); - serialize_lp_packet(&packet1, &mut buf1, None).unwrap(); - - // Parse packet - let parsed_packet1 = parse_lp_packet(&buf1, None).unwrap(); - - // Perform replay check (should pass) - session - .receiving_counter_quick_check(parsed_packet1.header.counter) - .expect("Initial packet failed replay check"); - - // Mark received (simulating successful processing) - session - .receiving_counter_mark(parsed_packet1.header.counter) - .expect("Failed to mark initial packet received"); - - // === Packet 2 (Counter 0 - Replay, should fail check) === - let packet2 = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 0, // Same counter as before (replay) - }, - message: LpMessage::Busy, - trailer: [0u8; TRAILER_LEN], - }; - - // Serialize packet - let mut buf2 = BytesMut::new(); - serialize_lp_packet(&packet2, &mut buf2, None).unwrap(); - - // Parse packet - let parsed_packet2 = parse_lp_packet(&buf2, None).unwrap(); - - // Perform replay check (should fail) - let replay_result = - session.receiving_counter_quick_check(parsed_packet2.header.counter); - assert!(replay_result.is_err()); - match replay_result.unwrap_err() { - LpError::Replay(e) => { - assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter)); - } - e => panic!("Expected replay error, got {:?}", e), - } - // Do not mark received as it failed validation - - // === Packet 3 (Counter 1 - Should succeed) === - let packet3 = LpPacket { - header: LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 42, - counter: 1, // Incremented counter - }, - message: LpMessage::Busy, - trailer: [0u8; TRAILER_LEN], - }; - - // Serialize packet - let mut buf3 = BytesMut::new(); - serialize_lp_packet(&packet3, &mut buf3, None).unwrap(); - - // Parse packet - let parsed_packet3 = parse_lp_packet(&buf3, None).unwrap(); - - // Perform replay check (should pass) - session - .receiving_counter_quick_check(parsed_packet3.header.counter) - .expect("Packet 3 failed replay check"); - - // Mark received - session - .receiving_counter_mark(parsed_packet3.header.counter) - .expect("Failed to mark packet 3 received"); - - // Verify validator state directly on the session - let state = session.current_packet_cnt(); - assert_eq!(state.0, 2); // Next expected counter (correct - was 1, now expects 2) - assert_eq!(state.1, 2); // Total marked received (correct - packets 1 and 3) - } + todo!() + // for kem in kem_list() { + // // Create session + // let mut session = mock_session_for_test(); + // + // // === Packet 1 (Counter 0 - Should succeed) === + // let packet1 = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, // Matches session's sending_index assumption for this test + // counter: 0, + // }, + // message: LpMessage::Busy, + // trailer: [0u8; TRAILER_LEN], + // }; + // + // // Serialize packet + // let mut buf1 = BytesMut::new(); + // serialize_lp_packet(&packet1, &mut buf1, None).unwrap(); + // + // // Parse packet + // let parsed_packet1 = parse_lp_packet(&buf1, None).unwrap(); + // + // // Perform replay check (should pass) + // session + // .receiving_counter_quick_check(parsed_packet1.header.counter) + // .expect("Initial packet failed replay check"); + // + // // Mark received (simulating successful processing) + // session + // .receiving_counter_mark(parsed_packet1.header.counter) + // .expect("Failed to mark initial packet received"); + // + // // === Packet 2 (Counter 0 - Replay, should fail check) === + // let packet2 = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 0, // Same counter as before (replay) + // }, + // message: LpMessage::Busy, + // trailer: [0u8; TRAILER_LEN], + // }; + // + // // Serialize packet + // let mut buf2 = BytesMut::new(); + // serialize_lp_packet(&packet2, &mut buf2, None).unwrap(); + // + // // Parse packet + // let parsed_packet2 = parse_lp_packet(&buf2, None).unwrap(); + // + // // Perform replay check (should fail) + // let replay_result = + // session.receiving_counter_quick_check(parsed_packet2.header.counter); + // assert!(replay_result.is_err()); + // match replay_result.unwrap_err() { + // LpError::Replay(e) => { + // assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter)); + // } + // e => panic!("Expected replay error, got {:?}", e), + // } + // // Do not mark received as it failed validation + // + // // === Packet 3 (Counter 1 - Should succeed) === + // let packet3 = LpPacket { + // header: LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 42, + // counter: 1, // Incremented counter + // }, + // message: LpMessage::Busy, + // trailer: [0u8; TRAILER_LEN], + // }; + // + // // Serialize packet + // let mut buf3 = BytesMut::new(); + // serialize_lp_packet(&packet3, &mut buf3, None).unwrap(); + // + // // Parse packet + // let parsed_packet3 = parse_lp_packet(&buf3, None).unwrap(); + // + // // Perform replay check (should pass) + // session + // .receiving_counter_quick_check(parsed_packet3.header.counter) + // .expect("Packet 3 failed replay check"); + // + // // Mark received + // session + // .receiving_counter_mark(parsed_packet3.header.counter) + // .expect("Failed to mark packet 3 received"); + // + // // Verify validator state directly on the session + // let state = session.current_packet_cnt(); + // assert_eq!(state.0, 2); // Next expected counter (correct - was 1, now expects 2) + // assert_eq!(state.1, 2); // Total marked received (correct - packets 1 and 3) + // } } #[test] diff --git a/common/nym-lp/src/message.rs b/common/nym-lp/src/message.rs index afd881decb..7a504c0d0e 100644 --- a/common/nym-lp/src/message.rs +++ b/common/nym-lp/src/message.rs @@ -1,165 +1,31 @@ // Copyright 2025 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::packet::LpHeader; -use crate::peer::LpRemotePeer; -use crate::{BOOTSTRAP_RECEIVER_IDX, LpError, LpPacket}; +use crate::LpError; use bytes::{BufMut, BytesMut}; use num_enum::{IntoPrimitive, TryFromPrimitive}; -use nym_crypto::asymmetric::{ed25519, x25519}; -use rand::RngCore; +use nym_crypto::asymmetric::ed25519; use std::fmt::{self, Display}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; -/// Data structure for the ClientHello message -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct ClientHelloData { - /// Client-proposed receiver index for session identification (4 bytes) - /// Auto-generated randomly by the client - pub receiver_index: u32, - /// Client's LP x25519 public key (32 bytes) - derived from Ed25519 key - pub client_lp_public_key: x25519::PublicKey, - /// Client's Ed25519 public key (32 bytes) - for PSQ authentication - pub client_ed25519_public_key: ed25519::PublicKey, - // noiserm - /// Salt for PSK derivation (32 bytes: 8-byte timestamp + 24-byte nonce) - pub salt: [u8; 32], -} - -impl ClientHelloData { - // noiserm (remove 32 bytes for salt) - // 4 bytes for receiver index + 32 bytes for client lp key, 32 bytes for client ed25519 key + 32 bytes for salt - pub const LEN: usize = 100; - - pub fn into_lp_packet(self, protocol_version: u8) -> LpPacket { - LpPacket::new( - LpHeader::new( - BOOTSTRAP_RECEIVER_IDX, // session_id not yet established - 0, // counter starts at 0 - protocol_version, - ), - LpMessage::ClientHello(self), - ) - } - - fn len(&self) -> usize { - Self::LEN - } - - fn generate_receiver_index() -> u32 { - loop { - let candidate = rand::random(); - if candidate != BOOTSTRAP_RECEIVER_IDX { - return candidate; - } - } - } - - // noiserm - /// Generates a new ClientHelloData with fresh salt. - /// - /// Salt format: 8 bytes timestamp (u64 LE) + 24 bytes random nonce - /// - /// # Arguments - /// * `client_lp_public_key` - Client's x25519 public key (derived from Ed25519) - /// * `client_ed25519_public_key` - Client's Ed25519 public key (for PSQ authentication) - pub fn new_with_fresh_salt( - client_lp_public_key: x25519::PublicKey, - client_ed25519_public_key: ed25519::PublicKey, - timestamp: u64, - ) -> Self { - // Generate salt: timestamp + nonce - let mut salt = [0u8; 32]; - - // First 8 bytes: current timestamp as u64 little-endian - salt[..8].copy_from_slice(×tamp.to_le_bytes()); - - // Last 24 bytes: random nonce - rand::thread_rng().fill_bytes(&mut salt[8..]); - - Self { - receiver_index: Self::generate_receiver_index(), // Auto-generate random receiver index - client_lp_public_key, - client_ed25519_public_key, - salt, - } - } - // noiserm - /// Extracts the timestamp from the salt. - /// - /// # Returns - /// Unix timestamp in seconds - pub fn extract_timestamp(&self) -> u64 { - let mut timestamp_bytes = [0u8; 8]; - timestamp_bytes.copy_from_slice(&self.salt[..8]); - u64::from_le_bytes(timestamp_bytes) - } - - pub fn encode(&self, dst: &mut BytesMut) { - dst.put_u32_le(self.receiver_index); - dst.put_slice(self.client_lp_public_key.as_bytes()); - dst.put_slice(self.client_ed25519_public_key.as_bytes()); - // noiserm - dst.put_slice(&self.salt); - } - - pub fn decode(b: &[u8]) -> Result { - if b.len() != Self::LEN { - return Err(LpError::DeserializationError(format!( - "Expected {} bytes to deserialise ClientHelloData. got {}", - Self::LEN, - b.len() - ))); - } - - // SAFETY: we checked for valid byte lengths - #[allow(clippy::unwrap_used)] - let client_lp_public_key_bytes = b[4..36].try_into().unwrap(); - let client_ed25519_public_key_bytes = b[36..68].try_into().unwrap(); - - Ok(ClientHelloData { - receiver_index: u32::from_le_bytes([b[0], b[1], b[2], b[3]]), - client_lp_public_key: x25519::PublicKey::from_byte_array(client_lp_public_key_bytes), - client_ed25519_public_key: ed25519::PublicKey::from_byte_array( - client_ed25519_public_key_bytes, - )?, - // noiserm - salt: b[68..].try_into().unwrap(), - }) - } - - /// Attempt to construct remote peer information based on the data provided in this packet. - pub fn to_remote_peer(&self) -> LpRemotePeer { - LpRemotePeer::new(self.client_ed25519_public_key, self.client_lp_public_key) - } -} - #[derive(Debug, Copy, Clone, PartialEq, Eq, IntoPrimitive, TryFromPrimitive)] #[repr(u32)] pub enum MessageType { + /// The party is busy Busy = 0x0000, - Handshake = 0x0001, - EncryptedData = 0x0002, - ClientHello = 0x0003, - KKTRequest = 0x0004, - KKTResponse = 0x0005, - ForwardPacket = 0x0006, - /// Receiver index collision - client should retry with new index - Collision = 0x0007, - /// Acknowledgment - gateway confirms receipt of message - Ack = 0x0008, - /// Subsession request - client initiates subsession creation - SubsessionRequest = 0x0009, - // georgio: this should be the psq msg - /// Subsession KK1 - first message of Noise KK handshake - SubsessionKK1 = 0x000A, - /// Subsession KK2 - second message of Noise KK handshake - SubsessionKK2 = 0x000B, - /// Subsession ready - subsession established confirmation - SubsessionReady = 0x000C, - /// Subsession abort - race winner tells loser to become responder - SubsessionAbort = 0x000D, + /// Encrypted payload + EncryptedData = 0x0001, + + /// Receiver should forward this message via telescoping + ForwardPacket = 0x0002, + + /// Receiver index collision - client should retry with new index + Collision = 0x0003, + + /// Acknowledgment - gateway confirms receipt of message + Ack = 0x0004, + /// General error Error = 0x00FF, } @@ -175,31 +41,10 @@ impl MessageType { } #[derive(Debug, Clone, PartialEq, Eq)] -pub struct HandshakeData(pub Vec); +pub struct ApplicationData(pub Vec); -impl HandshakeData { - pub(crate) fn new(bytes: Vec) -> Self { - Self(bytes) - } - fn len(&self) -> usize { - self.0.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.0); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(HandshakeData(bytes.to_vec())) - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct EncryptedDataPayload(pub Vec); - -impl EncryptedDataPayload { - #[allow(dead_code)] - pub(crate) fn new(bytes: Vec) -> Self { +impl ApplicationData { + pub fn new(bytes: Vec) -> Self { Self(bytes) } @@ -212,51 +57,7 @@ impl EncryptedDataPayload { } fn decode(bytes: &[u8]) -> Result { - Ok(EncryptedDataPayload(bytes.to_vec())) - } -} - -/// KKT request frame data (serialized KKTFrame bytes) -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct KKTRequestData(pub Vec); - -impl KKTRequestData { - pub(crate) fn new(bytes: Vec) -> Self { - Self(bytes) - } - - fn len(&self) -> usize { - self.0.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.0); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(KKTRequestData(bytes.to_vec())) - } -} - -/// KKT response frame data (serialized KKTFrame bytes) -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct KKTResponseData(pub Vec); - -impl KKTResponseData { - pub(crate) fn new(bytes: Vec) -> Self { - Self(bytes) - } - - fn len(&self) -> usize { - self.0.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.0); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(KKTResponseData(bytes.to_vec())) + Ok(ApplicationData(bytes.to_vec())) } } @@ -306,42 +107,6 @@ impl ErrorPacketData { } } -/// PSQ request frame data (serialized bytes) -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct PSQRequestData(pub Vec); - -impl PSQRequestData { - fn len(&self) -> usize { - self.0.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.0); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(PSQRequestData(bytes.to_vec())) - } -} - -/// PSQ response frame data (serialized bytes) -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct PSQResponseData(pub Vec); - -impl PSQResponseData { - fn len(&self) -> usize { - self.0.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.0); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(PSQResponseData(bytes.to_vec())) - } -} - /// Packet forwarding request with embedded inner LP packet #[derive(Debug, Clone)] pub struct ForwardPacketData { @@ -467,146 +232,30 @@ impl ForwardPacketData { } } -// georgio: swap with psq -/// Subsession KK1 message - first message of Noise KK handshake -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SubsessionKK1Data { - /// Noise KK first message payload (ephemeral key + encrypted static) - pub payload: Vec, -} - -impl SubsessionKK1Data { - fn len(&self) -> usize { - self.payload.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.payload); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(SubsessionKK1Data { - payload: bytes.to_vec(), - }) - } -} - -/// Subsession KK2 message - second message of Noise KK handshake -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SubsessionKK2Data { - /// Noise KK second message payload (ephemeral key + encrypted response) - pub payload: Vec, -} - -impl SubsessionKK2Data { - fn len(&self) -> usize { - self.payload.len() - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_slice(&self.payload); - } - - fn decode(bytes: &[u8]) -> Result { - Ok(SubsessionKK2Data { - payload: bytes.to_vec(), - }) - } -} - -/// Subsession ready confirmation with new session index -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SubsessionReadyData { - /// New subsession's receiver index for routing - pub receiver_index: u32, -} - -impl SubsessionReadyData { - pub const LEN: usize = 4; - - fn len(&self) -> usize { - Self::LEN - } - - fn encode(&self, dst: &mut BytesMut) { - dst.put_u32_le(self.receiver_index); - } - - fn decode(bytes: &[u8]) -> Result { - if bytes.len() != 4 { - return Err(LpError::DeserializationError(format!( - "Expected 4 bytes to deserialise SubsessionReadyData. got {}", - bytes.len() - ))); - } - Ok(SubsessionReadyData { - receiver_index: u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]), - }) - } -} - #[derive(Debug, Clone)] pub enum LpMessage { + /// The party is busy Busy, - PSQRequest(PSQRequestData), - PSQResponse(PSQResponseData), - EncryptedData(EncryptedDataPayload), - ClientHello(ClientHelloData), - KKTRequest(KKTRequestData), - KKTResponse(KKTResponseData), + + /// Application payload is being sent + ApplicationData(ApplicationData), + + /// Receiver should forward this message via telescoping ForwardPacket(ForwardPacketData), + /// Receiver index collision - client should retry with new receiver_index Collision, + /// Acknowledgment - gateway confirms receipt of message Ack, - // georgio: this should become psq stuff - /// Subsession request - client initiates subsession creation (empty, signal only) - SubsessionRequest, - /// Subsession KK1 - first message of Noise KK handshake - SubsessionKK1(SubsessionKK1Data), - /// Subsession KK2 - second message of Noise KK handshake - SubsessionKK2(SubsessionKK2Data), - /// Subsession ready - subsession established confirmation - SubsessionReady(SubsessionReadyData), - /// Subsession abort - race winner tells loser to become responder (empty, signal only) - SubsessionAbort, + /// An error has occurred Error(ErrorPacketData), } -impl From for LpMessage { - fn from(value: PSQRequestData) -> Self { - LpMessage::PSQRequest(value) - } -} - -impl From for LpMessage { - fn from(value: PSQResponseData) -> Self { - LpMessage::PSQResponse(value) - } -} - -impl From for LpMessage { - fn from(value: EncryptedDataPayload) -> Self { - LpMessage::EncryptedData(value) - } -} - -impl From for LpMessage { - fn from(value: ClientHelloData) -> Self { - LpMessage::ClientHello(value) - } -} - -impl From for LpMessage { - fn from(value: KKTRequestData) -> Self { - LpMessage::KKTRequest(value) - } -} - -impl From for LpMessage { - fn from(value: KKTResponseData) -> Self { - LpMessage::KKTResponse(value) +impl From for LpMessage { + fn from(value: ApplicationData) -> Self { + LpMessage::ApplicationData(value) } } @@ -616,86 +265,40 @@ impl From for LpMessage { } } -impl From for LpMessage { - fn from(value: SubsessionKK1Data) -> Self { - LpMessage::SubsessionKK1(value) - } -} - -impl From for LpMessage { - fn from(value: SubsessionKK2Data) -> Self { - LpMessage::SubsessionKK2(value) - } -} - -impl From for LpMessage { - fn from(value: SubsessionReadyData) -> Self { - LpMessage::SubsessionReady(value) - } -} - impl Display for LpMessage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { LpMessage::Busy => write!(f, "Busy"), - LpMessage::EncryptedData(_) => write!(f, "EncryptedData"), - LpMessage::ClientHello(_) => write!(f, "ClientHello"), - LpMessage::KKTRequest(_) => write!(f, "KKTRequest"), - LpMessage::KKTResponse(_) => write!(f, "KKTResponse"), + LpMessage::ApplicationData(_) => write!(f, "EncryptedData"), LpMessage::ForwardPacket(_) => write!(f, "ForwardPacket"), LpMessage::Collision => write!(f, "Collision"), LpMessage::Ack => write!(f, "Ack"), - LpMessage::SubsessionRequest => write!(f, "SubsessionRequest"), - LpMessage::SubsessionKK1(_) => write!(f, "SubsessionKK1"), - LpMessage::SubsessionKK2(_) => write!(f, "SubsessionKK2"), - LpMessage::SubsessionReady(_) => write!(f, "SubsessionReady"), - LpMessage::SubsessionAbort => write!(f, "SubsessionAbort"), LpMessage::Error(_) => write!(f, "Error"), - LpMessage::PSQRequest(_) => write!(f, "PSQRequest"), - LpMessage::PSQResponse(_) => write!(f, "PSQResponse"), } } } impl LpMessage { + #[deprecated(note = "is it actually needed?")] pub fn payload(&self) -> &[u8] { match self { LpMessage::Busy => &[], - LpMessage::PSQRequest(payload) => payload.0.as_slice(), - LpMessage::PSQResponse(payload) => payload.0.as_slice(), - LpMessage::EncryptedData(payload) => payload.0.as_slice(), - LpMessage::ClientHello(_) => &[], // Structured data, serialized in encode_content - LpMessage::KKTRequest(payload) => payload.0.as_slice(), - LpMessage::KKTResponse(payload) => payload.0.as_slice(), + LpMessage::ApplicationData(payload) => payload.0.as_slice(), LpMessage::ForwardPacket(_) => &[], // Structured data, serialized in encode_content LpMessage::Collision => &[], LpMessage::Ack => &[], - LpMessage::SubsessionRequest => &[], - LpMessage::SubsessionKK1(_) => &[], // Structured data, serialized in encode_content - LpMessage::SubsessionKK2(_) => &[], // Structured data, serialized in encode_content - LpMessage::SubsessionReady(_) => &[], // Structured data, serialized in encode_content - LpMessage::SubsessionAbort => &[], LpMessage::Error(_) => &[], // Structured data, serialized in encode_content (?) } } + #[deprecated(note = "is it actually needed?")] pub fn is_empty(&self) -> bool { match self { LpMessage::Busy => true, - LpMessage::EncryptedData(payload) => payload.0.is_empty(), - LpMessage::ClientHello(_) => false, // Always has data - LpMessage::KKTRequest(payload) => payload.0.is_empty(), - LpMessage::KKTResponse(payload) => payload.0.is_empty(), + LpMessage::ApplicationData(payload) => payload.0.is_empty(), LpMessage::ForwardPacket(_) => false, // Always has data LpMessage::Collision => true, LpMessage::Ack => true, - LpMessage::SubsessionRequest => true, // Empty signal - LpMessage::SubsessionKK1(_) => false, // Always has payload - LpMessage::SubsessionKK2(_) => false, // Always has payload - LpMessage::SubsessionReady(_) => false, // Always has receiver_index - LpMessage::SubsessionAbort => true, // Empty signal - LpMessage::PSQRequest(_) => true, // Always had data (?) - LpMessage::PSQResponse(_) => true, // Always had data (?) LpMessage::Error(_) => false, } } @@ -703,20 +306,10 @@ impl LpMessage { pub fn len(&self) -> usize { match self { LpMessage::Busy => 0, - LpMessage::PSQRequest(payload) => payload.len(), - LpMessage::PSQResponse(payload) => payload.len(), - LpMessage::EncryptedData(payload) => payload.len(), - LpMessage::ClientHello(payload) => payload.len(), - LpMessage::KKTRequest(payload) => payload.len(), - LpMessage::KKTResponse(payload) => payload.len(), + LpMessage::ApplicationData(payload) => payload.len(), LpMessage::ForwardPacket(payload) => payload.len(), LpMessage::Collision => 0, LpMessage::Ack => 0, - LpMessage::SubsessionRequest => 0, - LpMessage::SubsessionKK1(payload) => payload.len(), - LpMessage::SubsessionKK2(payload) => payload.len(), - LpMessage::SubsessionReady(payload) => payload.len(), - LpMessage::SubsessionAbort => 0, LpMessage::Error(payload) => payload.len(), } } @@ -724,20 +317,10 @@ impl LpMessage { pub fn typ(&self) -> MessageType { match self { LpMessage::Busy => MessageType::Busy, - LpMessage::PSQRequest(_) => todo!(), - LpMessage::PSQResponse(_) => todo!(), - LpMessage::EncryptedData(_) => MessageType::EncryptedData, - LpMessage::ClientHello(_) => MessageType::ClientHello, - LpMessage::KKTRequest(_) => MessageType::KKTRequest, - LpMessage::KKTResponse(_) => MessageType::KKTResponse, + LpMessage::ApplicationData(_) => MessageType::EncryptedData, LpMessage::ForwardPacket(_) => MessageType::ForwardPacket, LpMessage::Collision => MessageType::Collision, LpMessage::Ack => MessageType::Ack, - LpMessage::SubsessionRequest => MessageType::SubsessionRequest, - LpMessage::SubsessionKK1(_) => MessageType::SubsessionKK1, - LpMessage::SubsessionKK2(_) => MessageType::SubsessionKK2, - LpMessage::SubsessionReady(_) => MessageType::SubsessionReady, - LpMessage::SubsessionAbort => MessageType::SubsessionAbort, LpMessage::Error(_) => MessageType::Error, } } @@ -745,20 +328,10 @@ impl LpMessage { pub fn encode_content(&self, dst: &mut BytesMut) { match self { LpMessage::Busy => { /* No content */ } - LpMessage::PSQRequest(payload) => payload.encode(dst), - LpMessage::PSQResponse(payload) => payload.encode(dst), - LpMessage::EncryptedData(payload) => payload.encode(dst), - LpMessage::ClientHello(data) => data.encode(dst), - LpMessage::KKTRequest(payload) => payload.encode(dst), - LpMessage::KKTResponse(payload) => payload.encode(dst), + LpMessage::ApplicationData(payload) => payload.encode(dst), LpMessage::ForwardPacket(data) => data.encode(dst), LpMessage::Collision => { /* No content */ } LpMessage::Ack => { /* No content */ } - LpMessage::SubsessionRequest => { /* No content - signal only */ } - LpMessage::SubsessionKK1(data) => data.encode(dst), - LpMessage::SubsessionKK2(data) => data.encode(dst), - LpMessage::SubsessionReady(data) => data.encode(dst), - LpMessage::SubsessionAbort => { /* No content - signal only */ } LpMessage::Error(data) => data.encode(dst), } } @@ -773,17 +346,9 @@ impl LpMessage { content.ensure_empty()?; Ok(LpMessage::Busy) } - MessageType::Handshake => todo!(), - MessageType::EncryptedData => Ok(LpMessage::EncryptedData( - EncryptedDataPayload::decode(content)?, - )), - MessageType::ClientHello => { - Ok(LpMessage::ClientHello(ClientHelloData::decode(content)?)) - } - MessageType::KKTRequest => Ok(LpMessage::KKTRequest(KKTRequestData::decode(content)?)), - MessageType::KKTResponse => { - Ok(LpMessage::KKTResponse(KKTResponseData::decode(content)?)) - } + MessageType::EncryptedData => Ok(LpMessage::ApplicationData(ApplicationData::decode( + content, + )?)), MessageType::ForwardPacket => Ok(LpMessage::ForwardPacket(ForwardPacketData::decode( content, )?)), @@ -795,23 +360,6 @@ impl LpMessage { content.ensure_empty()?; Ok(LpMessage::Ack) } - MessageType::SubsessionRequest => { - content.ensure_empty()?; - Ok(LpMessage::SubsessionRequest) - } - MessageType::SubsessionKK1 => Ok(LpMessage::SubsessionKK1(SubsessionKK1Data::decode( - content, - )?)), - MessageType::SubsessionKK2 => Ok(LpMessage::SubsessionKK2(SubsessionKK2Data::decode( - content, - )?)), - MessageType::SubsessionReady => Ok(LpMessage::SubsessionReady( - SubsessionReadyData::decode(content)?, - )), - MessageType::SubsessionAbort => { - content.ensure_empty()?; - Ok(LpMessage::SubsessionAbort) - } MessageType::Error => Ok(LpMessage::Error(ErrorPacketData::decode(content)?)), } } @@ -836,114 +384,40 @@ impl EnsureEmptyContent for &[u8] { #[cfg(test)] mod tests { - use std::time::{SystemTime, UNIX_EPOCH}; - use super::*; use crate::LpPacket; use crate::packet::{LpHeader, TRAILER_LEN}; #[test] fn encoding() { - let message = LpMessage::EncryptedData(EncryptedDataPayload(vec![11u8; 124])); - - let resp_header = LpHeader { - protocol_version: 1, - reserved: [0u8; 3], - receiver_idx: 0, - counter: 0, - }; - - let packet = LpPacket { - header: resp_header, - message, - trailer: [80; TRAILER_LEN], - }; - - // Just print packet for debug, will be captured in test output - println!("{packet:?}"); - - // Verify message type - assert!(matches!(packet.message.typ(), MessageType::EncryptedData)); - - // Verify correct data in message - match &packet.message { - LpMessage::EncryptedData(data) => { - assert_eq!(*data, EncryptedDataPayload(vec![11u8; 124])); - } - _ => panic!("Wrong message type"), - } - } - - #[test] - fn test_client_hello_salt_generation() { - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("System time before UNIX epoch") - .as_secs(); - - let mut rng = rand::thread_rng(); - let ed25519 = ed25519::KeyPair::new(&mut rng); - let x25519 = ed25519.to_x25519(); - - let client_key = *x25519.public_key(); - let client_ed25519_key = *ed25519.public_key(); - let hello1 = - ClientHelloData::new_with_fresh_salt(client_key, client_ed25519_key, timestamp); - let hello2 = - ClientHelloData::new_with_fresh_salt(client_key, client_ed25519_key, timestamp); - - // Different salts should be generated - assert_ne!(hello1.salt, hello2.salt); - - // But timestamps should be very close (within 1 second) - let ts1 = hello1.extract_timestamp(); - let ts2 = hello2.extract_timestamp(); - assert!((ts1 as i64 - ts2 as i64).abs() <= 1); - } - - #[test] - fn test_client_hello_timestamp_extraction() { - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("System time before UNIX epoch") - .as_secs(); - let mut rng = rand::thread_rng(); - let ed25519 = ed25519::KeyPair::new(&mut rng); - let x25519 = ed25519.to_x25519(); - - let client_key = *x25519.public_key(); - let client_ed25519_key = *ed25519.public_key(); - let hello = ClientHelloData::new_with_fresh_salt(client_key, client_ed25519_key, timestamp); - - let timestamp = hello.extract_timestamp(); - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - - // Timestamp should be within 1 second of now - assert!((timestamp as i64 - now as i64).abs() <= 1); - } - - #[test] - fn test_client_hello_salt_format() { - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("System time before UNIX epoch") - .as_secs(); - let mut rng = rand::thread_rng(); - let ed25519 = ed25519::KeyPair::new(&mut rng); - let x25519 = ed25519.to_x25519(); - - let client_key = *x25519.public_key(); - let client_ed25519_key = *ed25519.public_key(); - let hello = ClientHelloData::new_with_fresh_salt(client_key, client_ed25519_key, timestamp); - - // First 8 bytes should be non-zero timestamp - let timestamp_bytes = &hello.salt[..8]; - assert_ne!(timestamp_bytes, &[0u8; 8]); - - // Salt should be 32 bytes total - assert_eq!(hello.salt.len(), 32); + todo!() + // let message = LpMessage::EncryptedData(EncryptedDataPayload(vec![11u8; 124])); + // + // let resp_header = LpHeader { + // protocol_version: 1, + // reserved: [0u8; 3], + // receiver_idx: 0, + // counter: 0, + // }; + // + // let packet = LpPacket { + // header: resp_header, + // message, + // trailer: [80; TRAILER_LEN], + // }; + // + // // Just print packet for debug, will be captured in test output + // println!("{packet:?}"); + // + // // Verify message type + // assert!(matches!(packet.message.typ(), MessageType::EncryptedData)); + // + // // Verify correct data in message + // match &packet.message { + // LpMessage::EncryptedData(data) => { + // assert_eq!(*data, EncryptedDataPayload(vec![11u8; 124])); + // } + // _ => panic!("Wrong message type"), + // } } } diff --git a/common/nym-lp/src/noise_protocol.rs b/common/nym-lp/src/noise_protocol.rs deleted file mode 100644 index c56f7cd9b9..0000000000 --- a/common/nym-lp/src/noise_protocol.rs +++ /dev/null @@ -1,337 +0,0 @@ -// Copyright 2025 - Nym Technologies SA -// SPDX-License-Identifier: Apache-2.0 - -//! Sans-IO Noise protocol state machine, adapted from noise-psq. - -use snow::{TransportState, params::NoiseParams}; -use thiserror::Error; - -// --- Error Definition --- - -/// Errors related to the Noise protocol state machine. -#[derive(Error, Debug)] -pub enum NoiseError { - #[error("encountered a Noise decryption error")] - DecryptionError, - - #[error("encountered a Noise Protocol error - {0}")] - ProtocolError(snow::Error), - - #[error("operation is invalid in the current protocol state")] - IncorrectStateError, - - #[error("attempted transport mode operation without real PSK injection")] - PskNotInjected, - - #[error("Other Noise-related error: {0}")] - Other(String), - - #[error("session is read-only after demotion")] - SessionReadOnly, -} - -impl From for NoiseError { - fn from(err: snow::Error) -> Self { - match err { - snow::Error::Decrypt => NoiseError::DecryptionError, - err => NoiseError::ProtocolError(err), - } - } -} - -// --- Protocol State and Structs --- - -/// Represents the possible states of the Noise protocol machine. -#[derive(Debug)] -pub enum NoiseProtocolState { - /// The protocol is currently performing the handshake. - /// Contains the Snow handshake state. - Handshaking(Box), - - /// The handshake is complete, and the protocol is in transport mode. - /// Contains the Snow transport state. - Transport(TransportState), - - /// The protocol has encountered an unrecoverable error. - /// Stores the error description. - Failed(String), -} - -/// The core sans-io Noise protocol state machine. -#[derive(Debug)] -pub struct NoiseProtocol { - state: NoiseProtocolState, - // We might need buffers for incoming/outgoing data later if we add internal buffering - // read_buffer: Vec, - // write_buffer: Vec, -} - -/// Represents the outcome of processing received bytes via `read_message`. -#[derive(Debug, PartialEq)] -pub enum ReadResult { - /// A handshake or transport message was successfully processed, but yielded no application data - /// and did not complete the handshake. - NoOp, - /// A complete application data message was decrypted. - DecryptedData(Vec), - /// The handshake successfully completed during this read operation. - HandshakeComplete, - // NOTE: NeedMoreBytes variant removed as read_message expects full frames. -} - -// --- Implementation --- - -impl NoiseProtocol { - pub fn params() -> NoiseParams { - // SAFETY: the hardcoded pattern must be valid - // and if for some reason it was not, we MUST fail non-gracefully for there is no possible recovery - #[allow(clippy::unwrap_used)] - crate::NOISE_PATTERN.parse().unwrap() - } - - /// Creates a new `NoiseProtocol` instance in the Handshaking state. - /// - /// Takes an initialized `snow::HandshakeState` (e.g., from `snow::Builder`). - pub fn new(initial_state: snow::HandshakeState) -> Self { - NoiseProtocol { - state: NoiseProtocolState::Handshaking(Box::new(initial_state)), - } - } - - fn prepare_handshake_state<'a>( - local_private_key: &'a [u8], - remote_public_key: &'a [u8], - psk: &'a [u8], - ) -> snow::Builder<'a> { - let psk_index = crate::NOISE_PSK_INDEX; - let noise_params = NoiseProtocol::params(); - - snow::Builder::new(noise_params) - .local_private_key(local_private_key) - .remote_public_key(remote_public_key) - .psk(psk_index, psk) - } - - /// Builds a new `NoiseProtocol` initiator instance with the provided local private key, - /// remote public key and psk - pub fn build_new_initiator( - local_private_key: &[u8], - remote_public_key: &[u8], - psk: &[u8], - ) -> Result { - let handshake_state = - Self::prepare_handshake_state(local_private_key, remote_public_key, psk) - .build_initiator()?; - Ok(Self::new(handshake_state)) - } - - /// Builds a new `NoiseProtocol` responder instance with the provided local private key, - /// remote public key and psk - pub fn build_new_responder( - local_private_key: &[u8], - remote_public_key: &[u8], - psk: &[u8], - ) -> Result { - let handshake_state = - Self::prepare_handshake_state(local_private_key, remote_public_key, psk) - .build_responder()?; - Ok(Self::new(handshake_state)) - } - - /// Processes a single, complete incoming Noise message frame. - /// - /// Assumes the caller handles buffering and framing to provide one full message. - /// Returns the result of processing the message. - pub fn read_message(&mut self, input: &[u8]) -> Result { - // Allocate a buffer large enough for the maximum possible Noise message size. - // TODO: Consider reusing a buffer for efficiency. - let mut buffer = vec![0u8; 65535]; // Max Noise message size - - match &mut self.state { - NoiseProtocolState::Handshaking(handshake_state) => { - match handshake_state.read_message(input, &mut buffer) { - Ok(_) => { - if handshake_state.is_handshake_finished() { - // Transition to Transport state. - let current_state = std::mem::replace( - &mut self.state, - // Temporary placeholder needed for mem::replace - NoiseProtocolState::Failed( - NoiseError::IncorrectStateError.to_string(), - ), - ); - if let NoiseProtocolState::Handshaking(state_to_convert) = current_state - { - match state_to_convert.into_transport_mode() { - Ok(transport_state) => { - self.state = NoiseProtocolState::Transport(transport_state); - Ok(ReadResult::HandshakeComplete) - } - Err(e) => { - let err = NoiseError::from(e); - self.state = NoiseProtocolState::Failed(err.to_string()); - Err(err) - } - } - } else { - // Should be unreachable - let err = NoiseError::IncorrectStateError; - self.state = NoiseProtocolState::Failed(err.to_string()); - Err(err) - } - } else { - // Handshake continues - Ok(ReadResult::NoOp) - } - } - Err(e) => { - let err = NoiseError::from(e); - self.state = NoiseProtocolState::Failed(err.to_string()); - Err(err) - } - } - } - NoiseProtocolState::Transport(transport_state) => { - match transport_state.read_message(input, &mut buffer) { - Ok(len) => Ok(ReadResult::DecryptedData(buffer[..len].to_vec())), - Err(e) => { - let err = NoiseError::from(e); - self.state = NoiseProtocolState::Failed(err.to_string()); - Err(err) - } - } - } - NoiseProtocolState::Failed(_) => Err(NoiseError::IncorrectStateError), - } - } - - /// Checks if there are pending handshake messages to send. - /// - /// If in Handshaking state and it's our turn, generates the message. - /// Transitions state to Transport if the handshake completes after this message. - /// Returns `None` if not in Handshaking state or not our turn. - pub fn get_bytes_to_send(&mut self) -> Option, NoiseError>> { - match &mut self.state { - NoiseProtocolState::Handshaking(handshake_state) => { - if handshake_state.is_my_turn() { - let mut buffer = vec![0u8; 65535]; - match handshake_state.write_message(&[], &mut buffer) { - // Empty payload for handshake msg - Ok(len) => { - if handshake_state.is_handshake_finished() { - // Transition to Transport state. - let current_state = std::mem::replace( - &mut self.state, - NoiseProtocolState::Failed( - NoiseError::IncorrectStateError.to_string(), - ), - ); - if let NoiseProtocolState::Handshaking(state_to_convert) = - current_state - { - match state_to_convert.into_transport_mode() { - Ok(transport_state) => { - self.state = - NoiseProtocolState::Transport(transport_state); - Some(Ok(buffer[..len].to_vec())) // Return final handshake msg - } - Err(e) => { - let err = NoiseError::from(e); - self.state = - NoiseProtocolState::Failed(err.to_string()); - Some(Err(err)) - } - } - } else { - // Should be unreachable - let err = NoiseError::IncorrectStateError; - self.state = NoiseProtocolState::Failed(err.to_string()); - Some(Err(err)) - } - } else { - // Handshake continues - Some(Ok(buffer[..len].to_vec())) - } - } - Err(e) => { - let err = NoiseError::from(e); - self.state = NoiseProtocolState::Failed(err.to_string()); - Some(Err(err)) - } - } - } else { - // Not our turn - None - } - } - NoiseProtocolState::Transport(_) | NoiseProtocolState::Failed(_) => { - // No handshake messages to send in these states - None - } - } - } - - /// Encrypts an application data payload for sending during the Transport phase. - /// - /// Returns the ciphertext (payload + 16-byte tag). - /// Errors if not in Transport state or encryption fails. - pub fn write_message(&mut self, payload: &[u8]) -> Result, NoiseError> { - match &mut self.state { - NoiseProtocolState::Transport(transport_state) => { - let mut buffer = vec![0u8; payload.len() + 16]; // Payload + tag - match transport_state.write_message(payload, &mut buffer) { - Ok(len) => Ok(buffer[..len].to_vec()), - Err(e) => { - let err = NoiseError::from(e); - self.state = NoiseProtocolState::Failed(err.to_string()); - Err(err) - } - } - } - NoiseProtocolState::Handshaking(_) | NoiseProtocolState::Failed(_) => { - Err(NoiseError::IncorrectStateError) - } - } - } - - /// Returns true if the protocol is in the transport phase (handshake complete). - pub fn is_transport(&self) -> bool { - matches!(self.state, NoiseProtocolState::Transport(_)) - } - - /// Returns true if the protocol has failed. - pub fn is_failed(&self) -> bool { - matches!(self.state, NoiseProtocolState::Failed(_)) - } - - /// Check if the handshake has finished and the protocol is in transport mode. - pub fn is_handshake_finished(&self) -> bool { - matches!(self.state, NoiseProtocolState::Transport(_)) - } - - /// Inject a PSK into the Noise HandshakeState. - /// - /// This allows dynamic PSK injection after HandshakeState construction, - /// which is required for PSQ (Post-Quantum Secure PSK) integration where - /// the PSK is derived during the handshake process. - /// - /// # Arguments - /// * `index` - PSK index (typically 3 for XKpsk3 pattern) - /// * `psk` - The pre-shared key bytes to inject - /// - /// # Errors - /// Returns an error if: - /// - Not in handshake state - /// - The underlying snow library rejects the PSK - pub fn set_psk(&mut self, index: u8, psk: &[u8]) -> Result<(), NoiseError> { - match &mut self.state { - NoiseProtocolState::Handshaking(handshake_state) => { - handshake_state - .set_psk(index as usize, psk) - .map_err(NoiseError::ProtocolError)?; - Ok(()) - } - _ => Err(NoiseError::IncorrectStateError), - } - } -} diff --git a/common/nym-lp/src/packet.rs b/common/nym-lp/src/packet.rs index db2644b5fa..2081dcf8b3 100644 --- a/common/nym-lp/src/packet.rs +++ b/common/nym-lp/src/packet.rs @@ -6,10 +6,7 @@ use crate::message::{LpMessage, MessageType}; use crate::replay::ReceivingKeyCounterValidator; use bytes::{BufMut, BytesMut}; use nym_lp_common::format_debug_bytes; -use parking_lot::Mutex; -use std::fmt::Write; use std::fmt::{Debug, Formatter}; -use std::sync::Arc; use tracing::warn; #[allow(dead_code)] @@ -31,11 +28,38 @@ pub mod version { pub const CURRENT: u8 = 1; } +#[derive(Clone)] +pub struct EncryptedLpPacket { + // The outer header that's sent in plaintext + pub(crate) outer_header: OuterHeader, + + // The ciphertext containing the inner header and the payload + pub(crate) ciphertext: Vec, +} + +impl Debug for EncryptedLpPacket { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", format_debug_bytes(&self.debug_bytes())?) + } +} + +impl EncryptedLpPacket { + pub(crate) fn debug_bytes(&self) -> Vec { + let mut bytes = BytesMut::new(); + self.encode(&mut bytes); + bytes.freeze().to_vec() + } + + pub fn encode(&self, dst: &mut BytesMut) { + self.outer_header.encode(dst); + dst.put_slice(&self.ciphertext) + } +} + #[derive(Clone)] pub struct LpPacket { pub(crate) header: LpHeader, pub(crate) message: LpMessage, - pub(crate) trailer: [u8; TRAILER_LEN], } impl Debug for LpPacket { @@ -46,43 +70,13 @@ impl Debug for LpPacket { impl LpPacket { pub fn new(header: LpHeader, message: LpMessage) -> Self { - Self { - header, - message, - trailer: [0; TRAILER_LEN], - } + Self { header, message } } pub fn typ(&self) -> MessageType { self.message.typ() } - /// Compute a hash of the message payload - /// - /// This can be used for message integrity verification or deduplication - pub fn hash_payload(&self) -> [u8; 32] { - use sha2::{Digest, Sha256}; - - let mut hasher = Sha256::new(); - let mut buffer = BytesMut::new(); - - // Include message type and content in the hash - buffer.put_slice(&(self.message.typ() as u16).to_le_bytes()); - self.message.encode_content(&mut buffer); - - hasher.update(&buffer); - hasher.finalize().into() - } - - pub fn hash_payload_hex(&self) -> String { - let hash = self.hash_payload(); - hash.iter() - .fold(String::with_capacity(hash.len() * 2), |mut acc, byte| { - let _ = write!(acc, "{:02x}", byte); - acc - }) - } - pub fn message(&self) -> &LpMessage { &self.message } @@ -93,17 +87,15 @@ impl LpPacket { pub(crate) fn debug_bytes(&self) -> Vec { let mut bytes = BytesMut::new(); - self.encode(&mut bytes); + self.dbg_encode(&mut bytes); bytes.freeze().to_vec() } - pub(crate) fn encode(&self, dst: &mut BytesMut) { - self.header.encode(dst); + pub(crate) fn dbg_encode(&self, dst: &mut BytesMut) { + self.header.dbg_encode(dst); dst.put_slice(&(self.message.typ() as u16).to_le_bytes()); self.message.encode_content(dst); - - dst.put_slice(&self.trailer) } /// Validate packet counter against a replay protection validator @@ -112,10 +104,9 @@ impl LpPacket { /// any expensive processing is done. pub fn validate_counter( &self, - validator: &Arc>, + validator: &ReceivingKeyCounterValidator, ) -> Result<(), LpError> { - let guard = validator.lock(); - guard.will_accept_branchless(self.header.counter)?; + validator.will_accept_branchless(self.header.outer.counter)?; Ok(()) } @@ -124,22 +115,13 @@ impl LpPacket { /// This should be called after a packet has been successfully processed. pub fn mark_received( &self, - validator: &Arc>, + validator: &mut ReceivingKeyCounterValidator, ) -> Result<(), LpError> { - let mut guard = validator.lock(); - guard.mark_did_receive_branchless(self.header.counter)?; + validator.mark_did_receive_branchless(self.header.outer.counter)?; Ok(()) } } -/// Session ID used for ClientHello bootstrap packets before session is established. -/// -/// When a client first connects, it sends a ClientHello packet with receiver_idx=0 -/// because neither side can compute the deterministic session ID yet (requires -/// both parties' X25519 keys). After ClientHello is processed, both sides derive -/// the same session ID from their keys, and all subsequent packets use that ID. -pub const BOOTSTRAP_RECEIVER_IDX: u32 = 0; - /// Outer header (12 bytes) - always cleartext, used for routing. /// /// This is the first 12 bytes of every LP packet, containing only the fields @@ -171,66 +153,36 @@ impl OuterHeader { }) } - pub fn encode(&self) -> [u8; Self::SIZE] { - let mut buf = [0u8; Self::SIZE]; - buf[0..4].copy_from_slice(&self.receiver_idx.to_le_bytes()); - buf[4..12].copy_from_slice(&self.counter.to_le_bytes()); - buf - } - /// Encode directly into a BytesMut buffer - pub fn encode_into(&self, dst: &mut BytesMut) { + pub fn encode(&self, dst: &mut BytesMut) { dst.put_slice(&self.receiver_idx.to_le_bytes()); dst.put_slice(&self.counter.to_le_bytes()); } } -/// Internal LP header representation containing all logical header fields. -/// -/// **Note**: This struct represents the LOGICAL header, not the wire format. -/// On the wire, packets use the unified format where: -/// - `OuterHeader` (receiver_idx + counter) always comes first (12 bytes, cleartext) -/// - Inner content (version + reserved + payload) follows (cleartext or encrypted) -/// -/// The `LpHeader::encode()` method outputs the old logical format for debug purposes only. -/// Use `serialize_lp_packet()` in codec.rs for actual wire serialization. -#[derive(Debug, Clone)] -pub struct LpHeader { +/// InnerHeader header (8 bytes) - encrypted, used for message parsing +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct InnerHeader { pub protocol_version: u8, pub reserved: [u8; 3], - pub receiver_idx: u32, - pub counter: u64, + pub message_type: MessageType, } -impl LpHeader { - pub const SIZE: usize = 16; -} +impl InnerHeader { + pub const SIZE: usize = 8; // protocol_version(1) + reserved(3) + message_type(4) -impl LpHeader { - pub fn new(receiver_idx: u32, counter: u64, protocol_version: u8) -> Self { - Self { - protocol_version, - reserved: [0u8; 3], - receiver_idx, - counter, - } - } - - pub fn encode(&self, dst: &mut BytesMut) { + pub(crate) fn encode(&self, dst: &mut BytesMut) { // protocol version dst.put_u8(self.protocol_version); // reserved dst.put_slice(&self.reserved); - // sender index - dst.put_slice(&self.receiver_idx.to_le_bytes()); - - // counter - dst.put_slice(&self.counter.to_le_bytes()); + // message type + dst.put_slice(&(self.message_type as u32).to_le_bytes()); } - pub fn parse(src: &[u8]) -> Result { + pub(crate) fn parse(src: &[u8]) -> Result { if src.len() < Self::SIZE { return Err(LpError::InsufficientBufferSize); } @@ -259,30 +211,63 @@ impl LpHeader { warn!("received non-zero reserved bytes. got: {reserved:?}"); } - let mut receiver_idx_bytes = [0u8; 4]; - receiver_idx_bytes.copy_from_slice(&src[4..8]); - let receiver_idx = u32::from_le_bytes(receiver_idx_bytes); + let msg_type_raw = u32::from_le_bytes([src[4], src[5], src[6], src[7]]); + let message_type = MessageType::from_u32(msg_type_raw) + .ok_or_else(|| LpError::invalid_message_type(msg_type_raw))?; - let mut counter_bytes = [0u8; 8]; - counter_bytes.copy_from_slice(&src[8..16]); - let counter = u64::from_le_bytes(counter_bytes); - - Ok(LpHeader { + Ok(InnerHeader { protocol_version, reserved, - receiver_idx, - counter, + message_type, }) } +} + +/// Internal LP header representation containing all logical header fields. +/// +/// **Note**: This struct represents the LOGICAL header, not the wire format. +/// On the wire, packets use the unified format where: +/// - `OuterHeader` (receiver_idx + counter) always comes first (12 bytes, cleartext) +/// - Inner content (version + reserved + payload) follows (cleartext or encrypted) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LpHeader { + pub outer: OuterHeader, + pub inner: InnerHeader, +} + +impl LpHeader { + pub fn new( + receiver_idx: u32, + counter: u64, + protocol_version: u8, + message_type: MessageType, + ) -> Self { + Self { + outer: OuterHeader { + receiver_idx, + counter, + }, + inner: InnerHeader { + protocol_version, + reserved: [0u8; 3], + message_type, + }, + } + } + + pub(crate) fn dbg_encode(&self, dst: &mut BytesMut) { + self.outer.encode(dst); + self.inner.encode(dst); + } /// Get the counter value from the header pub fn counter(&self) -> u64 { - self.counter + self.outer.counter } /// Get the sender index from the header pub fn receiver_idx(&self) -> u32 { - self.receiver_idx + self.outer.receiver_idx } } diff --git a/common/nym-lp/src/peer.rs b/common/nym-lp/src/peer.rs index 85a0701884..ce0c01d1e3 100644 --- a/common/nym-lp/src/peer.rs +++ b/common/nym-lp/src/peer.rs @@ -1,12 +1,9 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::{ClientHelloData, LpError}; +use crate::LpError; use libcrux_psq::handshake::types::{DHKeyPair, DHPublicKey}; use nym_crypto::asymmetric::{ed25519, x25519}; -use nym_kkt::key_utils::{ - generate_keypair_mceliece, generate_keypair_mlkem, generate_keypair_x25519, -}; use nym_kkt::keys::KEMKeys; use nym_kkt_ciphersuite::{Ciphersuite, KEM, KEMKeyDigests, SignatureScheme, SigningKeyDigests}; use std::collections::HashMap; @@ -43,15 +40,6 @@ impl LpLocalPeer { } } - pub fn build_client_hello_data(&self, timestamp: u64) -> ClientHelloData { - todo!() - // ClientHelloData::new_with_fresh_salt( - // *self.x25519().public_key(), - // *self.ed25519().public_key(), - // timestamp, - // ) - } - #[must_use] pub fn with_kem_keys(mut self, kem_keys: Arc) -> Self { self.kem_keypairs = Some(kem_keys); @@ -66,14 +54,6 @@ impl LpLocalPeer { &self.x25519 } - // /// Returns the reference to the KEM Public key of the peer (if available). - // pub fn get_kem_key_handle(&self) -> Result<&x25519::PublicKey, LpError> { - // self.kem_psq - // .as_ref() - // .map(|kp| kp.public_key()) - // .ok_or(LpError::ResponderWithMissingKEMKey) - // } - /// Convert this `LpLocalPeer` into a valid `LpRemotePeer` that can be used within tests #[doc(hidden)] pub fn as_remote(&self) -> LpRemotePeer { @@ -202,15 +182,15 @@ pub fn random_peer<'a, R: rand::CryptoRng + rand::RngCore>(rng: &mut R) -> LpLoc let mut rng09 = nym_test_utils::helpers::seeded_rng_09(seed); - let x25519 = Arc::new(generate_keypair_x25519(&mut rng09)); + let x25519 = Arc::new(nym_kkt::key_utils::generate_keypair_x25519(&mut rng09)); LpLocalPeer { ciphersuite: Ciphersuite::default(), ed25519, x25519, kem_keypairs: Some(Arc::new(KEMKeys::new( - generate_keypair_mceliece(&mut rng09), - generate_keypair_mlkem(&mut rng09), + nym_kkt::key_utils::generate_keypair_mceliece(&mut rng09), + nym_kkt::key_utils::generate_keypair_mlkem(&mut rng09), ))), } } diff --git a/common/nym-lp/src/psk.rs b/common/nym-lp/src/psk.rs deleted file mode 100644 index 6cc40b3b35..0000000000 --- a/common/nym-lp/src/psk.rs +++ /dev/null @@ -1,788 +0,0 @@ -// Copyright 2025 - Nym Technologies SA -// SPDX-License-Identifier: Apache-2.0 - -//! PSK (Pre-Shared Key) derivation for LP sessions using Blake3 KDF. -//! -//! This module implements identity-bound PSK derivation where both client and gateway -//! derive the same PSK from their LP keypairs. -//! -//! PSQ is embedded in Noise (not separate protocol) because: -//! 1. Single round-trip: PSQ ciphertext piggybacks on Noise handshake messages -//! 2. PSK binding: Noise XKpsk3 pattern authenticates both ECDH and PSQ-derived PSK -//! 3. Simpler state machine: No separate PSQ negotiation phase needed -//! 4. Atomic security: Session establishment either succeeds fully or fails completely -//! -//! Two approaches are supported: -//! - **Legacy ECDH-only** (`derive_psk`) - Simple but no post-quantum security -//! - **PSQ-enhanced** (`derive_psk_with_psq_*`) - Combines ECDH with post-quantum KEM -//! -//! ## Error Handling Strategy -//! -//! **PSQ failures always abort the handshake cleanly with no retry or fallback.** -//! -//! ### Rationale -//! -//! PSQ errors indicate: -//! - **Authentication failures** (CredError) - Potential attack or misconfiguration -//! - **Timing failures** (TimestampElapsed) - Replay attacks or clock skew -//! - **Crypto failures** (CryptoError) - Library bugs or hardware faults -//! - **Serialization failures** (Serialization) - Protocol violations or corruption -//! -//! None of these are transient errors that benefit from retry. Falling back to -//! ECDH-only PSK would silently degrade post-quantum security. -//! -//! ### Error Recovery Behavior -//! -//! On any PSQ error: -//! 1. Function returns `Err(LpError)` immediately -//! 2. Session state remains unchanged (dummy PSK, clean Noise state) -//! 3. Handshake aborts - caller must start fresh connection -//! 4. Error is logged with diagnostic context -//! -//! ### State Guarantees on Error -//! -//! - **`psq_state`**: Remains in `NotStarted` (initiator) or `ResponderWaiting` (responder) -//! - **Noise `HandshakeState`**: PSK slot 3 = dummy `[0u8; 32]` (not modified on error) -//! - **No partial data**: All allocations are stack-local to failed function -//! - **No cleanup needed**: No state was mutated - -use crate::LpError; -use libcrux_psq::handshake::types::{DHPrivateKey, DHPublicKey}; -use libcrux_psq::v1::cred::{Authenticator, Ed25519}; -use libcrux_psq::v1::impls::X25519 as PsqX25519; -use libcrux_psq::v1::psk_registration::{Initiator, InitiatorMsg, Responder}; -use libcrux_psq::v1::traits::{Ciphertext as PsqCiphertext, PSQ}; -use nym_crypto::asymmetric::ed25519; -use std::time::Duration; -use tls_codec::{Deserialize as TlsDeserializeTrait, Serialize as TlsSerializeTrait}; - -/// Context string for Blake3 KDF domain separation (PSQ-enhanced). -const PSK_PSQ_CONTEXT: &str = "nym-lp-psk-psq-v1"; - -/// Session context for PSQ protocol. -const PSQ_SESSION_CONTEXT: &[u8] = b"nym-lp-psq-session"; - -/// Context string for subsession PSK derivation. -const SUBSESSION_PSK_CONTEXT: &str = "lp-subsession-psk-v1"; - -/// Result from PSQ initiator message creation. -/// -/// Contains all outputs needed for session establishment: -/// - `psk`: Final derived PSK for Noise handshake (ECDH || K_pq || salt → Blake3) -/// - `payload`: Serialized PSQ message to send to responder -/// - `pq_shared_secret`: Raw K_pq from KEM encapsulation (for subsession derivation) -#[derive(Debug)] -pub struct PsqInitiatorResult { - /// Final PSK for Noise XKpsk3 handshake - pub psk: [u8; 32], - /// Serialized PSQ payload to embed in handshake message - pub payload: Vec, - /// Raw PQ shared secret (K_pq) before KDF combination. - /// Used for deriving subsession PSKs to preserve PQ protection. - pub pq_shared_secret: [u8; 32], -} - -/// Result from PSQ responder message processing. -/// -/// Contains all outputs needed for session establishment: -/// - `psk`: Final derived PSK for Noise handshake (matches initiator's) -/// - `psk_handle`: Encrypted PSK handle (ctxt_B) to send back to initiator -/// - `pq_shared_secret`: Raw K_pq from KEM decapsulation (for subsession derivation) -#[derive(Debug)] -pub struct PsqResponderResult { - /// Final PSK for Noise XKpsk3 handshake - pub psk: [u8; 32], - /// Encrypted PSK handle (ctxt_B) from PSQ responder message - pub psk_handle: Vec, - /// Raw PQ shared secret (K_pq) before KDF combination. - /// Used for deriving subsession PSKs to preserve PQ protection. - pub pq_shared_secret: [u8; 32], -} - -/// Derives a PSK using PSQ (Post-Quantum Secure PSK) protocol - Initiator side. -/// -/// This function combines classical ECDH with post-quantum KEM to provide forward secrecy -/// and HNDL (Harvest-Now, Decrypt-Later) resistance. -/// -/// # Formula -/// ```text -/// ecdh_secret = ECDH(local_x25519_private, remote_x25519_public) -/// (psq_psk, ct) = PSQ_Encapsulate(remote_kem_public, session_context) -/// psk = Blake3_derive_key( -/// context="nym-lp-psk-psq-v1", -/// input=ecdh_secret || psq_psk || salt -/// ) -/// ``` -/// -/// # Arguments -/// * `local_x25519_private` - Initiator's X25519 private key (for Noise) -/// * `remote_x25519_public` - Responder's X25519 public key (for Noise) -/// * `remote_kem_public` - Responder's KEM public key (obtained via KKT) -/// * `salt` - 32-byte salt for session binding -/// -/// # Returns -/// * `Ok((psk, ciphertext))` - PSK and ciphertext to send to responder -/// * `Err(LpError)` - If PSQ encapsulation fails -/// -/// # Example -/// ```ignore -/// // Client side (after KKT exchange) -/// let (psk, ciphertext) = derive_psk_with_psq_initiator( -/// client_x25519_private, -/// gateway_x25519_public, -/// &gateway_kem_key, // from KKT -/// &salt -/// )?; -/// // Send ciphertext to gateway -/// ``` -pub fn derive_psk_with_psq_initiator( - local_x25519_private: &DHPrivateKey, - remote_x25519_public: &DHPublicKey, - // remote_kem_public: &EncapsulationKey, - remote_kem_public: (), - salt: &[u8; 32], -) -> Result<([u8; 32], Vec), LpError> { - // Step 1: Classical ECDH for baseline security - // let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); - - let ecdh_secret: [u8; 32] = unimplemented!("unexposed by libcrux"); - - todo!() - // - // // Step 2: PSQ encapsulation for post-quantum security - // // KEM algorithm migration path: - // // - X25519: Current default for testing/compatibility (no HNDL resistance) - // // - MlKem768: Future production default (NIST PQ Level 3, HNDL resistant) - // // - XWing: Maximum security option (hybrid X25519 + ML-KEM) - // // Migration: Update LpConfig.kem_algorithm, no protocol changes needed. - // // KKT protocol adapts automatically to different KEM key sizes. - // let kem_pk = match remote_kem_public { - // EncapsulationKey::X25519(pk) => pk, - // _ => { - // return Err(LpError::KKTError( - // "Only X25519 KEM is currently supported for PSQ".to_string(), - // )); - // } - // }; - // - // let mut rng = rand09::rng(); - // let (psq_psk, ciphertext) = - // PsqX25519::encapsulate_psq(kem_pk, PSQ_SESSION_CONTEXT, &mut rng) - // .map_err(|e| LpError::Internal(format!("PSQ encapsulation failed: {:?}", e)))?; - // - // // Step 3: Combine ECDH + PSQ via Blake3 KDF - // let mut combined = Vec::with_capacity(64 + psq_psk.len()); - // combined.extend_from_slice(&ecdh_secret); - // combined.extend_from_slice(&psq_psk); // psq_psk is [u8; 32], need & - // combined.extend_from_slice(salt); - // - // let final_psk = nym_crypto::hkdf::blake3::derive_key_blake3(PSK_PSQ_CONTEXT, &combined, &[]); - // - // // Serialize ciphertext using TLS encoding for transport - // let ct_bytes = ciphertext - // .tls_serialize_detached() - // .map_err(|e| LpError::Internal(format!("Ciphertext serialization failed: {:?}", e)))?; - // - // Ok((final_psk, ct_bytes)) -} - -/// Derives a PSK using PSQ (Post-Quantum Secure PSK) protocol - Responder side. -/// -/// This function decapsulates the ciphertext from the initiator and combines it with -/// ECDH to derive the same PSK. -/// -/// # Formula -/// ```text -/// ecdh_secret = ECDH(local_x25519_private, remote_x25519_public) -/// psq_psk = PSQ_Decapsulate(local_kem_keypair, ciphertext, session_context) -/// psk = Blake3_derive_key( -/// context="nym-lp-psk-psq-v1", -/// input=ecdh_secret || psq_psk || salt -/// ) -/// ``` -/// -/// # Arguments -/// * `local_x25519_private` - Responder's X25519 private key (for Noise) -/// * `remote_x25519_public` - Initiator's X25519 public key (for Noise) -/// * `local_kem_keypair` - Responder's KEM keypair (decapsulation key, public key) -/// * `ciphertext` - PSQ ciphertext from initiator -/// * `salt` - 32-byte salt for session binding -/// -/// # Returns -/// * `Ok(psk)` - Derived PSK -/// * `Err(LpError)` - If PSQ decapsulation fails -/// -/// # Example -/// ```ignore -/// // Gateway side (after receiving ciphertext) -/// let psk = derive_psk_with_psq_responder( -/// gateway_x25519_private, -/// client_x25519_public, -/// (&gateway_kem_sk, &gateway_kem_pk), -/// &ciphertext, // from client -/// &salt -/// )?; -/// ``` -pub fn derive_psk_with_psq_responder( - local_x25519_private: &DHPrivateKey, - remote_x25519_public: &DHPublicKey, - // local_kem_keypair: (&DecapsulationKey, &EncapsulationKey), - local_kem_keypair: ((), ()), - ciphertext: &[u8], - salt: &[u8; 32], -) -> Result<[u8; 32], LpError> { - // Step 1: Classical ECDH for baseline security - // let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); - - todo!() - // let ecdh_secret: [u8; 32] = unimplemented!("unexposed by libcrux"); - // - // // Step 2: Extract X25519 keypair from DecapsulationKey/EncapsulationKey - // let (kem_sk, kem_pk) = match (local_kem_keypair.0, local_kem_keypair.1) { - // (DecapsulationKey::X25519(sk), EncapsulationKey::X25519(pk)) => (sk, pk), - // _ => { - // return Err(LpError::KKTError( - // "Only X25519 KEM is currently supported for PSQ".to_string(), - // )); - // } - // }; - // - // // Step 3: Deserialize ciphertext using TLS decoding - // let ct = PsqCiphertext::::tls_deserialize(&mut &ciphertext[..]) - // .map_err(|e| LpError::Internal(format!("Ciphertext deserialization failed: {:?}", e)))?; - // - // // Step 4: PSQ decapsulation for post-quantum security - // let psq_psk = PsqX25519::decapsulate_psq(kem_sk, kem_pk, &ct, PSQ_SESSION_CONTEXT) - // .map_err(|e| LpError::Internal(format!("PSQ decapsulation failed: {:?}", e)))?; - // - // // Step 5: Combine ECDH + PSQ via Blake3 KDF (same formula as initiator) - // let mut combined = Vec::with_capacity(64 + psq_psk.len()); - // combined.extend_from_slice(&ecdh_secret); - // combined.extend_from_slice(&psq_psk); // psq_psk is [u8; 32], need & - // combined.extend_from_slice(salt); - // - // let final_psk = nym_crypto::hkdf::blake3::derive_key_blake3(PSK_PSQ_CONTEXT, &combined, &[]); - // - // Ok(final_psk) -} - -/// PSQ protocol wrapper for initiator (client) side. -/// -/// Creates a PSQ initiator message with Ed25519 authentication, following the protocol: -/// 1. Encapsulate PSK using responder's KEM key -/// 2. Derive PSK and AEAD keys from K_pq -/// 3. Sign the encapsulation with Ed25519 -/// 4. AEAD encrypt (timestamp || signature || public_key) -/// -/// Returns (PSK, serialized_payload) where payload includes enc_pq and encrypted auth data. -/// -/// # Arguments -/// * `local_x25519_private` - Client's X25519 private key (for hybrid ECDH) -/// * `remote_x25519_public` - Gateway's X25519 public key (for hybrid ECDH) -/// * `remote_kem_public` - Gateway's PQ KEM public key (from KKT) -/// * `client_ed25519_sk` - Client's Ed25519 signing key -/// * `client_ed25519_pk` - Client's Ed25519 public key (credential) -/// * `salt` - Session salt -/// * `session_context` - Context bytes for PSQ (e.g., b"nym-lp-psq-session") -/// -/// # Returns -/// `PsqInitiatorResult` containing PSK, payload, and raw PQ shared secret -pub fn psq_initiator_create_message( - local_x25519_private: &DHPrivateKey, - remote_x25519_public: &DHPublicKey, - // remote_kem_public: &EncapsulationKey, - remote_kem_public: (), - client_ed25519_sk: &ed25519::PrivateKey, - client_ed25519_pk: &ed25519::PublicKey, - salt: &[u8; 32], - session_context: &[u8], -) -> Result { - // Step 1: Classical ECDH for baseline security - // let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); - - let ecdh_secret: [u8; 32] = unimplemented!("unexposed by libcrux"); - todo!() - // - // // Step 2: PSQ v1 with Ed25519 authentication - // // Extract X25519 KEM key from EncapsulationKey - // let kem_pk = match remote_kem_public { - // EncapsulationKey::X25519(pk) => pk, - // _ => { - // return Err(LpError::KKTError( - // "Only X25519 KEM is currently supported for PSQ".to_string(), - // )); - // } - // }; - // - // // Convert nym Ed25519 keys to libcrux format - // type Ed25519VerificationKey = ::VerificationKey; - // let ed25519_sk_bytes = client_ed25519_sk.to_bytes(); - // let ed25519_pk_bytes = client_ed25519_pk.to_bytes(); - // let ed25519_verification_key = Ed25519VerificationKey::from_bytes(ed25519_pk_bytes); - // - // // Use PSQ v1 API with Ed25519 authentication - // let mut rng = rand09::rng(); - // let (state, initiator_msg) = Initiator::send_initial_message::( - // session_context, - // Duration::from_secs(3600), // 1 hour expiry - // kem_pk, - // &ed25519_sk_bytes, - // &ed25519_verification_key, - // &mut rng, - // ) - // .map_err(|e| { - // tracing::error!( - // "PSQ initiator failed - KEM encapsulation or signing error: {:?}", - // e - // ); - // LpError::Internal(format!("PSQ v1 send_initial_message failed: {:?}", e)) - // })?; - // - // // Extract PSQ shared secret (unregistered PSK) - this is K_pq - // let psq_psk = state.unregistered_psk(); - // - // // pq_shared_secret is the raw K_pq from KEM encapsulation. - // // Store it for subsession derivation before it's combined with ECDH. - // let pq_shared_secret: [u8; 32] = *psq_psk; - // - // // Step 3: Combine ECDH + PSQ via Blake3 KDF - // let mut combined = Vec::with_capacity(64 + psq_psk.len()); - // combined.extend_from_slice(&ecdh_secret); - // combined.extend_from_slice(psq_psk); // psq_psk is already a &[u8; 32] - // combined.extend_from_slice(salt); - // - // let final_psk = nym_crypto::hkdf::blake3::derive_key_blake3(PSK_PSQ_CONTEXT, &combined, &[]); - // - // // Serialize InitiatorMsg with TLS encoding for transport - // let msg_bytes = initiator_msg - // .tls_serialize_detached() - // .map_err(|e| LpError::Internal(format!("InitiatorMsg serialization failed: {:?}", e)))?; - // - // Ok(PsqInitiatorResult { - // psk: final_psk, - // payload: msg_bytes, - // pq_shared_secret, - // }) -} - -/// PSQ protocol wrapper for responder (gateway) side. -/// -/// Processes a PSQ initiator message, verifies authentication, and derives PSK. -/// Follows the protocol: -/// 1. Decapsulate to get K_pq -/// 2. Derive AEAD keys and verify encrypted auth data -/// 3. Verify Ed25519 signature -/// 4. Check timestamp validity -/// 5. Derive PSK -/// -/// # Arguments -/// * `local_x25519_private` - Gateway's X25519 private key (for hybrid ECDH) -/// * `remote_x25519_public` - Client's X25519 public key (for hybrid ECDH) -/// * `local_kem_keypair` - Gateway's PQ KEM keypair -/// * `initiator_ed25519_pk` - Client's Ed25519 public key (for signature verification) -/// * `psq_payload` - Serialized PSQ payload from initiator -/// * `salt` - Session salt (must match initiator's) -/// * `session_context` - Context bytes for PSQ -/// -/// # Returns -/// `PsqResponderResult` containing PSK, PSK handle, and raw PQ shared secret -pub fn psq_responder_process_message( - local_x25519_private: &DHPrivateKey, - remote_x25519_public: &DHPublicKey, - // local_kem_keypair: (&DecapsulationKey, &EncapsulationKey), - local_kem_keypair: ((), ()), - initiator_ed25519_pk: &ed25519::PublicKey, - psq_payload: &[u8], - salt: &[u8; 32], - session_context: &[u8], -) -> Result { - // Step 1: Classical ECDH for baseline security - // let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); - let ecdh_secret: [u8; 32] = unimplemented!("unexposed by libcrux"); - - todo!() - // - // // Step 2: Extract X25519 keypair from DecapsulationKey/EncapsulationKey - // let (kem_sk, kem_pk) = match (local_kem_keypair.0, local_kem_keypair.1) { - // (DecapsulationKey::X25519(sk), EncapsulationKey::X25519(pk)) => (sk, pk), - // _ => { - // return Err(LpError::KKTError( - // "Only X25519 KEM is currently supported for PSQ".to_string(), - // )); - // } - // }; - // - // // Step 3: Deserialize InitiatorMsg using TLS decoding - // let initiator_msg = InitiatorMsg::::tls_deserialize(&mut &psq_payload[..]) - // .map_err(|e| LpError::Internal(format!("InitiatorMsg deserialization failed: {:?}", e)))?; - // - // // Step 4: Convert nym Ed25519 public key to libcrux VerificationKey format - // type Ed25519VerificationKey = ::VerificationKey; - // let initiator_ed25519_pk_bytes = initiator_ed25519_pk.to_bytes(); - // let initiator_verification_key = Ed25519VerificationKey::from_bytes(initiator_ed25519_pk_bytes); - // - // // Step 5: PSQ v1 responder processing with Ed25519 verification - // let (registered_psk, responder_msg) = Responder::send::( - // b"nym-lp-handle", // PSK storage handle - // Duration::from_secs(3600), // 1 hour expiry (must match initiator) - // session_context, // Must match initiator's session_context - // kem_pk, // Responder's public key - // kem_sk, // Responder's secret key - // &initiator_verification_key, // Initiator's Ed25519 public key for verification - // &initiator_msg, // InitiatorMsg to verify and process - // ) - // .map_err(|e| { - // use libcrux_psq::v1::Error as PsqError; - // match e { - // PsqError::CredError => { - // tracing::warn!( - // "PSQ responder auth failure - invalid Ed25519 signature (potential attack)" - // ); - // } - // PsqError::TimestampElapsed | PsqError::RegistrationError => { - // tracing::warn!( - // "PSQ responder timing failure - TTL expired (potential replay attack)" - // ); - // } - // _ => { - // tracing::error!("PSQ responder failed - {:?}", e); - // } - // } - // LpError::Internal(format!("PSQ v1 responder send failed: {:?}", e)) - // })?; - // - // // Extract the PSQ PSK from the registered PSK - this is K_pq - // let psq_psk = registered_psk.psk; - // - // // pq_shared_secret is the raw K_pq from KEM decapsulation. - // // Store it for subsession derivation before it's combined with ECDH. - // let pq_shared_secret: [u8; 32] = psq_psk; - // - // // Step 6: Combine ECDH + PSQ via Blake3 KDF (same formula as initiator) - // let mut combined = Vec::with_capacity(64 + psq_psk.len()); - // combined.extend_from_slice(&ecdh_secret); - // combined.extend_from_slice(&psq_psk); // psq_psk is [u8; 32], need & - // combined.extend_from_slice(salt); - // - // let final_psk = nym_crypto::hkdf::blake3::derive_key_blake3(PSK_PSQ_CONTEXT, &combined, &[]); - // - // // Step 7: Serialize ResponderMsg (contains ctxt_B - encrypted PSK handle) - // use tls_codec::Serialize; - // let responder_msg_bytes = responder_msg - // .tls_serialize_detached() - // .map_err(|e| LpError::Internal(format!("ResponderMsg serialization failed: {:?}", e)))?; - // - // Ok(PsqResponderResult { - // psk: final_psk, - // psk_handle: responder_msg_bytes, - // pq_shared_secret, - // }) -} - -/// Derive subsession PSK from parent's PQ shared secret. -/// -/// Uses Blake3 KDF with domain separation to derive unique PSK for each subsession. -/// This preserves PQ protection: subsession keys inherit quantum resistance from -/// parent's KEM shared secret (K_pq). -/// -/// # Security Model -/// -/// Subsessions use Noise KKpsk0 pattern where: -/// - Both parties already know each other's static X25519 keys (from parent session) -/// - PSK provides PQ protection by deriving from parent's K_pq -/// - Each subsession gets unique PSK via index parameter (prevents key reuse) -/// -/// # Arguments -/// * `pq_shared_secret` - Parent session's K_pq (32 bytes from KEM) -/// * `subsession_index` - Monotonic index for this subsession (prevents reuse) -/// -/// # Returns -/// 32-byte PSK for Noise KKpsk0 handshake -pub fn derive_subsession_psk(pq_shared_secret: &[u8; 32], subsession_index: u64) -> [u8; 32] { - nym_crypto::hkdf::blake3::derive_key_blake3( - SUBSESSION_PSK_CONTEXT, - pq_shared_secret, - &subsession_index.to_le_bytes(), - ) -} - -#[cfg(test)] -mod tests { - use super::*; - use libcrux_psq::handshake::types::DHKeyPair; - - fn generate_x25519_keypair() -> DHKeyPair { - DHKeyPair::new(&mut rand09::rng()) - } - - #[test] - fn test_psk_derivation_is_symmetric() { - todo!() - // let keypair_1 = generate_x25519_keypair(); - // let keypair_2 = generate_x25519_keypair(); - // let salt = [2u8; 32]; - // - // let mut rng = &mut rand09::rng(); - // let (_kem_sk, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // let dec_key = DecapsulationKey::X25519(_kem_sk); - // - // // Client derives PSK - // let (client_psk, ciphertext) = - // derive_psk_with_psq_initiator(keypair_1.sk(), &keypair_2.pk, &enc_key, &salt).unwrap(); - // - // // Gateway derives PSK from their perspective - // let gateway_psk = derive_psk_with_psq_responder( - // keypair_2.sk(), - // &keypair_1.pk, - // (&dec_key, &enc_key), - // &ciphertext, - // &salt, - // ) - // .unwrap(); - // - // assert_eq!( - // client_psk, gateway_psk, - // "Both sides should derive identical PSK" - // ); - } - - #[test] - fn test_different_salts_produce_different_psks() { - todo!() - // let keypair_1 = generate_x25519_keypair(); - // let keypair_2 = generate_x25519_keypair(); - // - // let salt1 = [1u8; 32]; - // let salt2 = [2u8; 32]; - // let mut rng = &mut rand09::rng(); - // let (_kem_sk, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // - // let psk1 = - // derive_psk_with_psq_initiator(keypair_1.sk(), &keypair_2.pk, &enc_key, &salt1).unwrap(); - // let psk2 = - // derive_psk_with_psq_initiator(keypair_1.sk(), &keypair_2.pk, &enc_key, &salt2).unwrap(); - // - // assert_ne!(psk1, psk2, "Different salts should produce different PSKs"); - } - - #[test] - fn test_different_keys_produce_different_psks() { - todo!() - // let keypair_1 = generate_x25519_keypair(); - // let keypair_2 = generate_x25519_keypair(); - // let keypair_3 = generate_x25519_keypair(); - // let salt = [3u8; 32]; - // - // let mut rng = &mut rand09::rng(); - // let (_kem_sk, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // - // let psk1 = - // derive_psk_with_psq_initiator(keypair_1.sk(), &keypair_2.pk, &enc_key, &salt).unwrap(); - // let psk2 = - // derive_psk_with_psq_initiator(keypair_1.sk(), &keypair_3.pk, &enc_key, &salt).unwrap(); - // - // assert_ne!( - // psk1, psk2, - // "Different remote keys should produce different PSKs" - // ); - } - - #[test] - fn test_psq_derivation_deterministic() { - todo!() - // let mut rng = rand09::rng(); - // - // // Generate X25519 keypairs for Noise - // let client_keypair = generate_x25519_keypair(); - // let gateway_keypair = generate_x25519_keypair(); - // - // // Generate KEM keypair for PSQ - // let (kem_sk, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // let dec_key = DecapsulationKey::X25519(kem_sk); - // - // let salt = [1u8; 32]; - // - // // Derive PSK twice with same inputs (initiator side) - // let (_psk1, ct1) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt, - // ) - // .unwrap(); - // - // let (_psk2, _ct2) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt, - // ) - // .unwrap(); - // - // // PSKs will be different due to randomness in PSQ, but ciphertexts too - // // This test verifies the function is deterministic given the SAME ciphertext - // let psk_responder1 = derive_psk_with_psq_responder( - // gateway_keypair.sk(), - // &client_keypair.pk, - // (&dec_key, &enc_key), - // &ct1, - // &salt, - // ) - // .unwrap(); - // - // let psk_responder2 = derive_psk_with_psq_responder( - // gateway_keypair.sk(), - // &client_keypair.pk, - // (&dec_key, &enc_key), - // &ct1, // Same ciphertext - // &salt, - // ) - // .unwrap(); - // - // assert_eq!( - // psk_responder1, psk_responder2, - // "Same ciphertext should produce same PSK" - // ); - } - - #[test] - fn test_psq_derivation_symmetric() { - todo!() - // let mut rng = rand09::rng(); - // - // // Generate X25519 keypairs for Noise - // let client_keypair = generate_x25519_keypair(); - // let gateway_keypair = generate_x25519_keypair(); - // - // // Generate KEM keypair for PSQ - // let (kem_sk, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // let dec_key = DecapsulationKey::X25519(kem_sk); - // - // let salt = [2u8; 32]; - // - // // Client derives PSK (initiator) - // let (client_psk, ciphertext) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt, - // ) - // .unwrap(); - // - // // Gateway derives PSK from ciphertext (responder) - // let gateway_psk = derive_psk_with_psq_responder( - // gateway_keypair.sk(), - // &client_keypair.pk, - // (&dec_key, &enc_key), - // &ciphertext, - // &salt, - // ) - // .unwrap(); - // - // assert_eq!( - // client_psk, gateway_psk, - // "Both sides should derive identical PSK via PSQ" - // ); - } - - #[test] - fn test_different_kem_keys_different_psk() { - todo!() - // let mut rng = rand09::rng(); - // - // let client_keypair = generate_x25519_keypair(); - // let gateway_keypair = generate_x25519_keypair(); - // - // // Two different KEM keypairs - // let (_, kem_pk1) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let (_, kem_pk2) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // - // let enc_key1 = EncapsulationKey::X25519(kem_pk1); - // let enc_key2 = EncapsulationKey::X25519(kem_pk2); - // - // let salt = [3u8; 32]; - // - // let (psk1, _) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key1, - // &salt, - // ) - // .unwrap(); - // - // let (psk2, _) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key2, - // &salt, - // ) - // .unwrap(); - // - // assert_ne!( - // psk1, psk2, - // "Different KEM keys should produce different PSKs" - // ); - } - - #[test] - fn test_psq_psk_output_length() { - todo!() - // let mut rng = rand09::rng(); - // - // let client_keypair = generate_x25519_keypair(); - // let gateway_keypair = generate_x25519_keypair(); - // - // let (_, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // - // let salt = [4u8; 32]; - // - // let (psk, _) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt, - // ) - // .unwrap(); - // - // assert_eq!(psk.len(), 32, "PSQ PSK should be exactly 32 bytes"); - } - - #[test] - fn test_psq_different_salts_different_psks() { - todo!() - // let mut rng = rand09::rng(); - // - // let client_keypair = generate_x25519_keypair(); - // let gateway_keypair = generate_x25519_keypair(); - // - // let (_, kem_pk) = generate_keypair_libcrux(&mut rng, KEM::X25519).unwrap(); - // let enc_key = EncapsulationKey::X25519(kem_pk); - // - // let salt1 = [1u8; 32]; - // let salt2 = [2u8; 32]; - // - // let (psk1, _) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt1, - // ) - // .unwrap(); - // - // let (psk2, _) = derive_psk_with_psq_initiator( - // client_keypair.sk(), - // &gateway_keypair.pk, - // &enc_key, - // &salt2, - // ) - // .unwrap(); - // - // assert_ne!(psk1, psk2, "Different salts should produce different PSKs"); - } -} diff --git a/common/nym-lp/src/psq/helpers.rs b/common/nym-lp/src/psq/helpers.rs index 149fce2a2c..ab4edfa7cb 100644 --- a/common/nym-lp/src/psq/helpers.rs +++ b/common/nym-lp/src/psq/helpers.rs @@ -1,17 +1,8 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::codec::{OuterAeadKey, parse_lp_packet, serialize_lp_packet}; -use crate::{LpError, LpPacket}; -use bytes::BytesMut; -use nym_lp_transport::traits::LpTransport; - use libcrux_psq::handshake::ciphersuites::CiphersuiteName; -#[cfg(test)] -use mock_instant::thread_local::{SystemTime, UNIX_EPOCH}; use nym_kkt_ciphersuite::KEM; -#[cfg(not(test))] -use std::time::{SystemTime, UNIX_EPOCH}; pub(crate) fn kem_to_ciphersuite(kem: KEM) -> CiphersuiteName { match kem { @@ -19,43 +10,3 @@ pub(crate) fn kem_to_ciphersuite(kem: KEM) -> CiphersuiteName { KEM::McEliece => CiphersuiteName::X25519_CLASSICMCELIECE_X25519_AESGCM128_HKDFSHA256, } } - -pub(crate) fn current_timestamp() -> Result { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_err(|_| LpError::Internal("System time before UNIX epoch".into())) - .map(|d| d.as_secs()) -} - -// only used in internal code (and tests) -#[allow(async_fn_in_trait)] -pub trait LpTransportHandshakeExt: LpTransport { - // the outer key is temporary until the algorithm is changed with psqv2 - async fn receive_packet( - &mut self, - outer_key: Option<&OuterAeadKey>, - ) -> Result - where - Self: Unpin, - { - let raw = self.receive_raw_packet().await?; - parse_lp_packet(&raw, outer_key) - } - - async fn send_packet( - &mut self, - packet: LpPacket, - outer_key: Option<&OuterAeadKey>, - ) -> Result<(), LpError> - where - Self: Unpin, - { - let mut packet_buf = BytesMut::new(); - - serialize_lp_packet(&packet, &mut packet_buf, outer_key)?; - self.send_serialised_packet(&packet_buf).await?; - Ok(()) - } -} - -impl LpTransportHandshakeExt for T where T: LpTransport {} diff --git a/common/nym-lp/src/psq/initiator.rs b/common/nym-lp/src/psq/initiator.rs index 2b2fdf4e84..f5eca5350f 100644 --- a/common/nym-lp/src/psq/initiator.rs +++ b/common/nym-lp/src/psq/initiator.rs @@ -1,34 +1,27 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::codec::OuterAeadKey; -use crate::message::{HandshakeData, KKTRequestData, MessageType}; -use crate::noise_protocol::NoiseProtocol; use crate::peer::{LpLocalPeer, LpRemotePeer}; -use crate::psk::psq_initiator_create_message; -use crate::psq::helpers::{LpTransportHandshakeExt, current_timestamp, kem_to_ciphersuite}; +use crate::psq::helpers::kem_to_ciphersuite; use crate::psq::{ - AAD_INITIATOR_INNER_V1, AAD_INITIATOR_OUTER_V1, InitiatorData, MinimalSession, - PSQHandshakeState, SESSION_CONTEXT_V1, initiator, + AAD_INITIATOR_INNER_V1, AAD_INITIATOR_OUTER_V1, InitiatorData, PSQHandshakeState, + SESSION_CONTEXT_V1, }; -use crate::session::PqSharedSecret; -use crate::{ClientHelloData, LpError, LpMessage, LpSession}; +use crate::session::PersistentSessionBinding; +use crate::{LpError, LpSession}; use libcrux_psq::handshake::RegistrationInitiator; use libcrux_psq::handshake::builders::{ CiphersuiteBuilder, InitiatorCiphersuite, PrincipalBuilder, }; -use libcrux_psq::handshake::ciphersuites::CiphersuiteName; +use libcrux_psq::handshake::types::Authenticator; use libcrux_psq::{Channel, IntoSession}; -use nym_kkt::context::KKTContext; use nym_kkt::initiator::KKTInitiator; use nym_kkt::keys::EncapsulationKey; use nym_kkt::message::{KKTRequest, KKTResponse}; -use nym_kkt_ciphersuite::KEM; use nym_lp_transport::traits::LpTransport; -use rand09::rng; use tracing::debug; -pub(crate) struct PSQHandshakeStateInitiator<'a, S> { +pub struct PSQHandshakeStateInitiator<'a, S> { pub(super) inner_state: PSQHandshakeState<'a, S>, pub(super) initiator_data: InitiatorData, } @@ -81,23 +74,6 @@ impl<'a, S> PSQHandshakeStateInitiator<'a, S> where S: LpTransport + Unpin, { - fn build_psq_initiator_principal<'b>( - &'b self, - encapsulation_key: &'b EncapsulationKey, - ) -> Result, LpError> { - let initiator_ciphersuite = build_psq_ciphersuite( - &self.inner_state.local_peer, - &self.initiator_data.remote_peer, - &encapsulation_key, - )?; - let initiator = build_psq_principal( - rng(), - self.initiator_data.protocol_version, - initiator_ciphersuite, - )?; - Ok(initiator) - } - /// Attempt to send KKT request to begin the handshake async fn send_kkt_request(&mut self, request: KKTRequest) -> Result<(), LpError> { // TODO: extra header @@ -114,7 +90,7 @@ where Ok(KKTResponse::from_bytes(data)) } - pub async fn complete_handshake(mut self, rng: &mut R) -> Result + pub async fn complete_handshake(mut self, rng: &mut R) -> Result where S: LpTransport + Unpin, R: rand09::CryptoRng, @@ -143,7 +119,10 @@ where // 4. generate and send PSQ request let protocol = self.initiator_data.protocol_version; - let mut conn = self.inner_state.connection; + let conn = self.inner_state.connection; + + // note: the clone is cheap due to internal Arcs + let encapsulation_key = response.encapsulation_key.clone(); // build the PSQ initiator let initiator_ciphersuite = build_psq_ciphersuite( @@ -172,12 +151,14 @@ where )); } - let session = psq_initiator.into_session()?; - Ok(MinimalSession { - session, - encapsulation_key: Some(response.encapsulation_key), - init_authenticator: None, - }) + let binding = PersistentSessionBinding { + initiator_authenticator: Authenticator::Dh(self.inner_state.local_peer.x25519().pk), + responder_ecdh_pk: self.initiator_data.remote_peer.x25519_public, + responder_pq_pk: Some(encapsulation_key), + }; + + let psq_session = psq_initiator.into_session()?; + LpSession::new(psq_session, binding, protocol) } } @@ -186,10 +167,8 @@ mod tests { use super::*; use crate::peer::mock_peers; use crate::psq::responder; - use libcrux_psq::handshake::types::Authenticator; - use libcrux_psq::session::{Session, SessionBinding}; use nym_kkt::responder::KKTResponder; - use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, SignatureScheme}; + use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, KEM, SignatureScheme}; use nym_test_utils::helpers::{DeterministicRng09Send, u64_seeded_rng_09}; use nym_test_utils::mocks::async_read_write::MockIOStream; use nym_test_utils::traits::{Leak, Timeboxed}; @@ -277,10 +256,8 @@ mod tests { assert!(responder.is_handshake_finished()); - let session_init = init_fut.await???; + let mut session_init = init_fut.await???; - let encapsulation_key = session_init.encapsulation_key.unwrap(); - let mut i_transport = session_init.session; let mut r_transport = responder.into_session().unwrap(); // test serialization, deserialization @@ -288,7 +265,7 @@ mod tests { let mut payload_buf_responder = vec![0u8; 4096]; let mut payload_buf_initiator = vec![0u8; 4096]; - let mut channel_i = i_transport.transport_channel().unwrap(); + let mut channel_i = session_init.active_transport(); let mut channel_r = r_transport.transport_channel().unwrap(); assert_eq!(channel_i.identifier(), channel_r.identifier()); diff --git a/common/nym-lp/src/psq/mod.rs b/common/nym-lp/src/psq/mod.rs index d6618743a1..e269687ac0 100644 --- a/common/nym-lp/src/psq/mod.rs +++ b/common/nym-lp/src/psq/mod.rs @@ -1,42 +1,27 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::LpMessage; use crate::packet::version; use crate::peer::{LpLocalPeer, LpRemotePeer}; -use crate::psq::helpers::LpTransportHandshakeExt; -use crate::psq::initiator::PSQHandshakeStateInitiator; -use crate::psq::responder::PSQHandshakeStateResponder; -use libcrux_psq::handshake::types::Authenticator; -use libcrux_psq::session::Session; -use nym_kkt::keys::EncapsulationKey; use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, IntoEnumIterator, SignatureScheme}; use nym_lp_transport::traits::LpTransport; mod helpers; -mod initiator; -mod responder; +pub mod initiator; +pub mod responder; + +pub use initiator::PSQHandshakeStateInitiator; +pub use responder::PSQHandshakeStateResponder; pub(crate) const AAD_INITIATOR_OUTER_V1: &[u8] = b"NYM-PQ-AAD-INIT-OUTER-V1"; pub(crate) const AAD_INITIATOR_INNER_V1: &[u8] = b"NYM-PQ-AAD-INIT-INNER-V1"; pub(crate) const AAD_RESPONDER_V1: &[u8] = b"NYM-PQ-AAD-RESP-V1"; pub(crate) const SESSION_CONTEXT_V1: &[u8] = b"NYM-PQ-SESSION-CONTEXT-V1"; -pub struct MinimalSession { - session: Session, - encapsulation_key: Option, - init_authenticator: Option, -} - pub struct PSQHandshakeState<'a, S> { /// The underlying connection established for the handshake connection: &'a mut S, - /// Protocol version used for the exchange. - /// either known implicitly through the directory (initiator) - /// or established through KKTRequest (responder) - protocol_version: Option, - /// Ciphersuite selected for the KKT/PSQ exchange ciphersuite: Ciphersuite, @@ -94,7 +79,6 @@ where pub fn new(connection: &'a mut S, ciphersuite: Ciphersuite, local_peer: LpLocalPeer) -> Self { PSQHandshakeState { connection, - protocol_version: None, ciphersuite, local_peer, } @@ -119,7 +103,6 @@ where mod tests { use super::*; use crate::peer::mock_peers; - use crate::psq::helpers::LpTransportHandshakeExt; use libcrux_psq::handshake::types::Authenticator; use libcrux_psq::session::{Session, SessionBinding}; use libcrux_psq::{Channel, IntoSession}; @@ -133,15 +116,6 @@ mod tests { use nym_test_utils::traits::{Leak, TimeboxedSpawnable}; use tokio::join; - #[allow(dead_code)] - async fn extract_error(conn: &mut MockIOStream) -> String { - let packet = conn.receive_packet(None).await.unwrap(); - match packet.message { - LpMessage::Error(error) => error.message, - _ => panic!("non error packet"), - } - } - #[tokio::test] async fn e2e_psq_handshake() -> anyhow::Result<()> { let conn_init = MockIOStream::default(); @@ -183,8 +157,8 @@ mod tests { let session_resp = session_resp???; assert_eq!( - session_init.session.identifier(), - session_resp.session.identifier() + session_init.session_identifier(), + session_resp.session_identifier() ); Ok(()) diff --git a/common/nym-lp/src/psq/responder.rs b/common/nym-lp/src/psq/responder.rs index 31a3b3379c..dee898384f 100644 --- a/common/nym-lp/src/psq/responder.rs +++ b/common/nym-lp/src/psq/responder.rs @@ -1,25 +1,27 @@ // Copyright 2026 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::peer::{LpLocalPeer, LpRemotePeer}; +use crate::peer::LpLocalPeer; use crate::psq::helpers::kem_to_ciphersuite; -use crate::psq::{ - AAD_RESPONDER_V1, MinimalSession, PSQHandshakeState, ResponderData, SESSION_CONTEXT_V1, -}; -use crate::session::PqSharedSecret; -use crate::{ClientHelloData, LpError, LpSession}; +use crate::psq::{AAD_RESPONDER_V1, PSQHandshakeState, ResponderData, SESSION_CONTEXT_V1}; +use crate::session::PersistentSessionBinding; +use crate::{LpError, LpSession}; use libcrux_psq::handshake::Responder; use libcrux_psq::handshake::builders::{ CiphersuiteBuilder, PrincipalBuilder, ResponderCiphersuite, }; use libcrux_psq::{Channel, IntoSession}; -use nym_kkt::context::KKTContext; use nym_kkt::message::{KKTRequest, KKTResponse, ProcessedKKTRequest}; use nym_kkt::responder::KKTResponder; use nym_kkt_ciphersuite::KEM; use nym_lp_transport::traits::LpTransport; use tracing::debug; +pub struct PSQHandshakeStateResponder<'a, S> { + pub(super) inner_state: PSQHandshakeState<'a, S>, + pub(super) responder_data: ResponderData, +} + pub(crate) fn build_psq_principal( rng: R, version: u8, @@ -44,7 +46,7 @@ where pub(crate) fn build_psq_ciphersuite( peer: &LpLocalPeer, kem: KEM, -) -> Result { +) -> Result, LpError> { let Some(kem_keys) = peer.kem_keypairs.as_ref() else { return Err(LpError::ResponderWithMissingKEMKeys); }; @@ -64,11 +66,6 @@ pub(crate) fn build_psq_ciphersuite( .map_err(|inner| LpError::PSQResponderBuilderFailure { inner }) } -pub(crate) struct PSQHandshakeStateResponder<'a, S> { - pub(super) inner_state: PSQHandshakeState<'a, S>, - pub(super) responder_data: ResponderData, -} - impl<'a, S> PSQHandshakeStateResponder<'a, S> where S: LpTransport + Unpin, @@ -113,7 +110,7 @@ where Ok(self.inner_state.connection.receive_raw_packet().await?) } - pub async fn complete_handshake(mut self, rng: &mut R) -> Result + pub async fn complete_handshake(mut self, rng: &mut R) -> Result where S: LpTransport + Unpin, R: rand09::CryptoRng, @@ -144,7 +141,7 @@ where .ok_or(LpError::MissingInitiatorAuthenticator)?; // 4. send PSQ response - let mut conn = self.inner_state.connection; + let conn = self.inner_state.connection; let mut buf = [0u8; 128]; let n = psq_responder.write_message(&[], &mut buf)?; @@ -157,12 +154,25 @@ where )); } - let session = psq_responder.into_session()?; - Ok(MinimalSession { - session, - encapsulation_key: processed_req.remote_encapsulation_key, - init_authenticator: Some(initiator_authenticator), - }) + // SAFETY: we have completed the exchange so this key MUST HAVE been present + #[allow(clippy::unwrap_used)] + let kem_key = self + .inner_state + .local_peer + .kem_keypairs + .as_ref() + .unwrap() + .encapsulation_key(kem) + .unwrap(); + + let binding = PersistentSessionBinding { + initiator_authenticator, + responder_ecdh_pk: self.inner_state.local_peer.x25519().pk, + responder_pq_pk: Some(kem_key), + }; + + let psq_session = psq_responder.into_session()?; + LpSession::new(psq_session, binding, processed_req.outer_protocol_version) } } @@ -171,8 +181,6 @@ mod tests { use super::*; use crate::peer::mock_peers; use crate::psq::initiator; - use libcrux_psq::handshake::types::Authenticator; - use libcrux_psq::session::{Session, SessionBinding}; use nym_kkt::initiator::KKTInitiator; use nym_kkt_ciphersuite::Ciphersuite; use nym_test_utils::helpers::{ @@ -258,11 +266,9 @@ mod tests { assert!(initiator.is_handshake_finished()); - let session_resp = resp_fut.await???; - let init_auth = session_resp.init_authenticator.unwrap(); + let mut session_resp = resp_fut.await???; let mut i_transport = initiator.into_session().unwrap(); - let mut r_transport = session_resp.session; // test serialization, deserialization let mut msg_channel = vec![0u8; 2048]; @@ -270,7 +276,7 @@ mod tests { let mut payload_buf_initiator = vec![0u8; 4096]; let mut channel_i = i_transport.transport_channel().unwrap(); - let mut channel_r = r_transport.transport_channel().unwrap(); + let mut channel_r = session_resp.active_transport(); assert_eq!(channel_i.identifier(), channel_r.identifier()); diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 58c795a985..b914e21511 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -4,135 +4,121 @@ //! Session management for the Lewes Protocol. //! //! This module implements session management functionality, including replay protection -//! and Noise protocol state handling. -use crate::codec::OuterAeadKey; -use crate::message::EncryptedDataPayload; -// noiserm -use crate::noise_protocol::{NoiseError, NoiseProtocol, ReadResult}; -use crate::packet::LpHeader; +use crate::codec::{decrypt_lp_packet, encrypt_lp_packet}; +use crate::message::ApplicationData; +use crate::packet::{EncryptedLpPacket, LpHeader}; use crate::peer::{LpLocalPeer, LpRemotePeer}; -use crate::psk::derive_subsession_psk; -use crate::replay::ReceivingKeyCounterValidator; -use crate::{LpError, LpMessage, LpPacket}; -use libcrux_psq::v1::traits::PSQ; -use libcrux_psq::{Channel, IntoSession}; -use rand09::rngs::ThreadRng; - -use crate::psq::PSQHandshakeState; -use libcrux_psq::{ - handshake::{RegistrationInitiator, Responder, types::DHPublicKey}, - session::Session, +use crate::psq::{ + InitiatorData, PSQHandshakeState, PSQHandshakeStateInitiator, PSQHandshakeStateResponder, + ResponderData, }; -use nym_crypto::asymmetric::{ed25519, x25519}; -use nym_kkt::context::KKTContext; +use crate::{LpError, LpMessage, LpPacket, ReceivingKeyCounterValidator}; +use libcrux_psq::handshake::types::{Authenticator, DHPublicKey}; +use libcrux_psq::session::{Session, SessionBinding}; +use nym_kkt::keys::EncapsulationKey; use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, HashLength, KEM, SignatureScheme}; use nym_lp_transport::traits::LpTransport; -use parking_lot::Mutex; -use snow::Builder; -use zeroize::{Zeroize, ZeroizeOnDrop}; +use std::fmt::{Debug, Formatter}; -/// PQ shared secret wrapper with automatic memory zeroization. -/// Ensures K_pq is cleared from memory when dropped. -#[derive(Clone, Zeroize, ZeroizeOnDrop)] -pub struct PqSharedSecret([u8; 32]); - -impl PqSharedSecret { - pub fn new(secret: [u8; 32]) -> Self { - Self(secret) - } - - pub fn as_bytes(&self) -> &[u8; 32] { - &self.0 - } -} - -impl std::fmt::Debug for PqSharedSecret { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PqSharedSecret") - .field("secret", &"") - .finish() - } -} +pub type SessionId = [u8; 32]; /// A session in the Lewes Protocol, handling connection state with Noise. /// /// Sessions manage connection state, including LP replay protection. /// Each session has a unique receiving index and sending index for connection identification. -#[derive(Debug)] pub struct LpSession { - /// Id of the established session - session_id: u32, + /// The underlying established session + psq_session: Session, + + /// The public key material bound to the underlying session. Used for serialisation. + session_binding: PersistentSessionBinding, + + /// The current active transport channel + // In the future it might get split between UDP and TCP transports + active_transport: libcrux_psq::Transport, /// Negotiated protocol version from handshake. - /// Set during handshake completion from the ClientHello/ServerHello packet header. - /// Used for future version negotiation and compatibility checks. - version: u8, - - /// Representation of a local Lewes Protocol peer - /// encapsulating all the known information and keys. - local_peer: LpLocalPeer, - - /// Representation of a remote Lewes Protocol peer - /// encapsulating all the known information and keys. - remote_peer: LpRemotePeer, + protocol_version: u8, /// Counter for outgoing packets sending_counter: u64, /// Validator for incoming packet counters to prevent replay attacks receiving_counter: ReceivingKeyCounterValidator, +} - /// Monotonically increasing counter for subsession indices. - /// Each subsession gets a unique index to ensure unique PSK derivation. - /// Uses u64 to make overflow practically impossible (~585k years at 1M/sec). - subsession_counter: u64, +/// Wraps public key material that is bound to a session. +pub struct PersistentSessionBinding { + /// The initiator's authenticator value, i.e. a long-term DH public value or signature verification key. + pub initiator_authenticator: Authenticator, - /// True if this session has been demoted to read-only mode. - /// Demoted sessions can still receive/decrypt but cannot send/encrypt. - read_only: bool, + /// The responder's long term DH public value. + pub responder_ecdh_pk: DHPublicKey, - /// ID of the successor session that replaced this one. - /// Set when demote() is called. - successor_session_id: Option, + /// The responder's long term PQ-KEM public key (if any). + pub responder_pq_pk: Option, +} + +impl Debug for PersistentSessionBinding { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PersistentSessionBinding") + .field("initiator_authenticator", &"") + .field("responder_ecdh_pk", &self.responder_ecdh_pk) + .field("responder_pq_pk", &self.responder_pq_pk) + .finish() + } +} + +impl<'a> From<&'a PersistentSessionBinding> for SessionBinding<'a> { + fn from(value: &'a PersistentSessionBinding) -> Self { + SessionBinding { + initiator_authenticator: &value.initiator_authenticator, + responder_ecdh_pk: &value.responder_ecdh_pk, + responder_pq_pk: value + .responder_pq_pk + .as_ref() + .map(|k| k.as_pq_encapsulation_key()), + } + } +} + +impl Debug for LpSession { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LpSession") + .field("session_id", &self.psq_session.identifier()) + .field("session_binding", &self.session_binding) + .field("active_transport_id", &self.active_transport.identifier()) + .field("protocol_version", &self.protocol_version) + .field("sending_counter", &self.sending_counter) + .field("receiving_counter", &self.receiving_counter) + .finish() + } } impl LpSession { /// Creates a new session after completed KTT/PSQ exchange /// /// # Arguments - /// - /// * `session_id` - Session identifier - /// * `version` - Protocol version to attach in all `LpPacket`s - /// * `outer_aead_key` - Outer AEAD key for packet encryption - /// * `local_peer` - This side's LP peer's keys - /// * `remote_peer` - The remote's LP peer's keys - /// * `pq_shared_secret` - Raw PQ shared secret (K_pq) from PSQ KEM encapsulation/decapsulation. - /// * `noise_state` - Noise protocol state machine + pub fn new( - session_id: u32, - version: u8, - outer_aead_key: OuterAeadKey, - local_peer: LpLocalPeer, - remote_peer: LpRemotePeer, - pq_shared_secret: PqSharedSecret, - noise_state: NoiseProtocol, - ) -> Self { - todo!() - // LpSession { - // session_id, - // version, - // outer_aead_key, - // local_peer, - // remote_peer, - // pq_shared_secret, - // noise_state, - // sending_counter: 0, - // receiving_counter: Default::default(), - // subsession_counter: 0, - // read_only: false, - // successor_session_id: None, - // } + mut psq_session: Session, + session_binding: PersistentSessionBinding, + protocol_version: u8, + ) -> Result { + // attempt to derive initial transport + let transport = psq_session + .transport_channel() + .map_err(|inner| LpError::TransportDerivationFailure { inner })?; + + Ok(LpSession { + psq_session, + session_binding, + active_transport: transport, + protocol_version, + sending_counter: 0, + receiving_counter: Default::default(), + }) } /// Create an instance of `Ciphersuite` using hardcoded defaults. @@ -154,14 +140,12 @@ impl LpSession { local_peer: LpLocalPeer, remote_peer: LpRemotePeer, remote_protocol_version: u8, - ) -> PSQHandshakeState<'_, S> + ) -> PSQHandshakeStateInitiator<'_, S> where S: LpTransport + Unpin, { - todo!() - // PSQHandshakeState::new(connection, ciphersuite, local_peer) - // .with_protocol_version(remote_protocol_version) - // .with_remote_peer(remote_peer) + PSQHandshakeState::new(connection, ciphersuite, local_peer) + .as_initiator(InitiatorData::new(remote_protocol_version, remote_peer)) } /// Helper function to create `PSQHandshakeState` for the handshake responder @@ -169,79 +153,41 @@ impl LpSession { connection: &'_ mut S, ciphersuite: Ciphersuite, local_peer: LpLocalPeer, - ) -> PSQHandshakeState<'_, S> + ) -> PSQHandshakeStateResponder<'_, S> where S: LpTransport + Unpin, { - todo!() - // PSQHandshakeState::new(connection, ciphersuite, local_peer) + PSQHandshakeState::new(connection, ciphersuite, local_peer) + .as_responder(ResponderData::default()) + } + + pub fn session_binding(&self) -> &PersistentSessionBinding { + &self.session_binding + } + + pub(crate) fn active_transport(&mut self) -> &mut libcrux_psq::Transport { + &mut self.active_transport + } + + pub fn session_identifier(&self) -> &[u8; 32] { + self.psq_session.identifier() } pub fn id(&self) -> u32 { - self.session_id + todo!() + // self.session.identifier() } - // noiserm /// Returns the negotiated protocol version from the handshake. /// /// Set during `LpSession` creation after sending / receiving `ClientHelloData` pub fn negotiated_version(&self) -> u8 { - self.version - } - - /// Returns the local X25519 public key. - /// - /// This is used for KKT protocol when the responder needs to send their - /// KEM public key in the KKT response. - pub fn local_x25519_public(&self) -> x25519::PublicKey { - todo!() - // *self.local_peer.x25519.pk - } - - /// Returns the remote ed25519 public key. - pub fn remote_ed25519_public(&self) -> ed25519::PublicKey { - self.remote_peer.ed25519_public - } - - // pub fn local_kem_keys(&self) -> Result<&KemKeyPair, LpError> { - // todo!() - // // let kem = self.base.local_peer.ciphersuite.kem(); - // // - // // self.base - // // .local_peer - // // .kem_key(kem) - // // .ok_or_else(|| LpError::ResponderWithMissingKEMKey { kem }) - // // .map(|keys| keys.deref()) - // } - - /// Returns the remote X25519 public key. - /// - /// Used for tie-breaking in simultaneous subsession initiation. - /// Lower key loses and becomes responder. - pub fn remote_x25519_public(&self) -> &x25519::PublicKey { - todo!() - // &self.remote_peer.x25519_public - } - - // noiserm - /// Returns the outer AEAD key for packet encryption/decryption. - /// - /// Returns `None` before PSK is derived (during initial handshake), - /// `Some(&OuterAeadKey)` after PSK injection via PSQ. - /// - /// Callers should use `None` for packet encryption/decryption during - /// the handshake phase, and use the returned key for transport phase. - /// - /// Note: For sending packets during handshake, use `outer_aead_key_for_sending()` - /// which checks PSQ state to avoid encrypting before the responder can decrypt. - pub fn outer_aead_key(&self) -> &OuterAeadKey { - todo!() - // &self.outer_aead_key + self.protocol_version } pub fn next_packet(&mut self, message: LpMessage) -> Result { let counter = self.next_counter(); - let header = LpHeader::new(self.id(), counter, self.version); + let header = LpHeader::new(self.id(), counter, self.protocol_version, message.typ()); let packet = LpPacket::new(header, message); Ok(packet) } @@ -304,299 +250,44 @@ impl LpSession { self.receiving_counter.current_packet_cnt() } - /// Gets the next subsession index and increments the counter. - /// - /// Each subsession requires a unique index to ensure unique PSK derivation. - /// The index is monotonically increasing per session. - pub fn next_subsession_index(&mut self) -> u64 { - let next = self.subsession_counter; - self.subsession_counter += 1; - next - } - - /// Attempt to retrieve expected KEM key hash of the remote - /// for [KEM] key type and [HashFunction] specified by own [KKTContext] - fn expected_kem_key_hash(&self, context: KKTContext) -> Result<&Vec, LpError> { - todo!() - // let kem = context.ciphersuite().kem(); - // let hash_function = context.ciphersuite().hash_function(); - // - // let digests = self - // .base - // .remote_peer - // .expected_kem_key_digests - // .get(&kem) - // .ok_or(LpError::NoKnownKEMKeyDigests { kem, hash_function })?; - // - // digests - // .get(&hash_function) - // .ok_or(LpError::NoKnownKEMKeyDigests { kem, hash_function }) - } - - /// Returns true if this session is in read-only mode. - /// - /// Read-only sessions have been demoted after a subsession was promoted. - /// They can still decrypt incoming messages but cannot encrypt outgoing ones. - pub fn is_read_only(&self) -> bool { - self.read_only - } - - /// Demotes this session to read-only mode after a subsession replaces it. - /// - /// After demotion: - /// - `encrypt_data()` will return `NoiseError::SessionReadOnly` - /// - `decrypt_data()` still works (to drain in-flight messages) - /// - Session should be cleaned up after TTL expires - /// - /// # Arguments - /// * `successor_idx` - The receiver index of the session that replaced this one - pub fn demote(&mut self, successor_idx: u32) { - self.successor_session_id = Some(successor_idx); - self.read_only = true; - } - - /// Returns the successor session ID if this session was demoted. - pub fn successor_session_id(&self) -> Option { - self.successor_session_id - } - - /// Encrypts application data payload using the established Noise transport session. + /// Encrypts a produced application using the established transport session + /// and produce an `EncryptedLpPacket` /// /// # Arguments /// - /// * `payload` - The application data to encrypt. + /// * `data` - plaintext data to encrypt /// /// # Returns /// - /// * `Ok(Vec)` containing the encrypted Noise message ciphertext. - /// * `Err(NoiseError)` if the session is not in transport mode or encryption fails. - pub fn encrypt_data(&mut self, payload: &[u8]) -> Result { - // // Check if session is read-only (demoted) - // if self.read_only { - // return Err(NoiseError::SessionReadOnly); - // } - // - // let payload = self.noise_state.write_message(payload)?; - // Ok(LpMessage::EncryptedData(EncryptedDataPayload(payload))) - todo!() + /// * `Ok(EncryptedLpPacket)` containing the encrypted message ciphertext. + /// * `Err(LpError)` if the session is not in transport mode or encryption fails. + pub fn encrypt_application_data( + &mut self, + data: Vec, + ) -> Result { + let packet = self.next_packet(LpMessage::ApplicationData(ApplicationData::new(data)))?; + encrypt_lp_packet(packet, &mut self.active_transport) } - /// Decrypts an incoming Noise message containing application data. + /// Decrypts an incoming LpPacket /// /// # Arguments /// - /// * `noise_ciphertext` - The encrypted Noise message received from the peer. + /// * `ciphertext` - The encrypted packet /// /// # Returns /// - /// * `Ok(Vec)` containing the decrypted application data payload. - /// * `Err(NoiseError)` if the session is not in transport mode, decryption fails, or the message is not data. - pub fn decrypt_data(&mut self, noise_ciphertext: &LpMessage) -> Result, NoiseError> { - let payload = noise_ciphertext.payload(); - - todo!() - // match self.noise_state.read_message(payload)? { - // ReadResult::DecryptedData(data) => Ok(data), - // _ => Err(NoiseError::IncorrectStateError), - // } - } - - /// Creates a new subsession using Noise KKpsk0 pattern. - /// - /// KKpsk0 reuses parent's static X25519 keys (both parties know each other from parent session). - /// PSK is derived from parent's PQ shared secret, preserving quantum resistance. - /// - /// # Arguments - /// * `subsession_index` - Unique index for this subsession (use `next_subsession_index()`) - /// * `is_initiator` - True if this side initiates the subsession handshake - /// - /// # Returns - /// `SubsessionHandshake` ready for KK1/KK2 message exchange - /// - /// # Errors - /// * Returns error if parent handshake not complete - /// * Returns error if PQ shared secret not available - pub fn create_subsession( - &self, - subsession_index: u64, - is_initiator: bool, - ) -> Result { - todo!() - // // Get PQ shared secret - // let pq_secret = self.pq_shared_secret(); - // - // // Derive subsession PSK from parent's PQ shared secret - // let subsession_psk = derive_subsession_psk(pq_secret.as_bytes(), subsession_index); - // - // // Build KKpsk0 handshake - // // Pattern: Noise_KKpsk0_25519_ChaChaPoly_SHA256 - // // Both parties already know each other's static keys from parent session - // let pattern_name = "Noise_KKpsk0_25519_ChaChaPoly_SHA256"; - // let params = pattern_name.parse()?; - // - // let local_key_bytes = self.local_peer.x25519.private_key().to_bytes(); - // let remote_key_bytes = self.remote_x25519_public().to_bytes(); - // - // let builder = Builder::new(params) - // .local_private_key(&local_key_bytes) - // .remote_public_key(&remote_key_bytes) - // .psk(0, &subsession_psk); // PSK at position 0 for KKpsk0 - // - // let handshake_state = if is_initiator { - // builder.build_initiator().map_err(LpError::SnowKeyError)? - // } else { - // builder.build_responder().map_err(LpError::SnowKeyError)? - // }; - // - // Ok(SubsessionHandshake { - // index: subsession_index, - // noise_state: Mutex::new(NoiseProtocol::new(handshake_state)), - // is_initiator, - // local_peer: self.local_peer.clone(), - // remote_peer: self.remote_peer.clone(), - // pq_shared_secret: self.pq_shared_secret.clone(), - // subsession_psk, - // negotiated_version: self.version, - // }) - } -} - -/// Subsession created via Noise KKpsk0 handshake tunneled through parent session. -/// -/// Subsessions provide fresh session keys while inheriting PQ protection from parent's -/// ML-KEM shared secret. After handshake completes, the subsession can be promoted -/// to replace the parent session. -/// -/// # Lifecycle -/// 1. Parent calls `create_subsession()` to get `SubsessionHandshake` -/// 2. Initiator calls `prepare_message()` to get KK1 -/// 3. KK1 sent through parent session (encrypted tunnel) -/// 4. Responder calls `process_message(kk1)` to process KK1 -/// 5. Responder calls `prepare_message()` to get KK2 -/// 6. KK2 sent through parent session -/// 7. Initiator calls `process_message(kk2)` to complete handshake -/// 8. Both call `is_complete()` to verify -#[derive(Debug)] -pub struct SubsessionHandshake { - /// Subsession index (unique per parent session) - pub index: u64, - /// Noise KKpsk0 handshake state - // georgio: replace with psq - noise_state: Mutex, - /// Is this side the initiator? - is_initiator: bool, - - // Key material inherited from parent session for into_session() conversion - /// Representation of a local Lewes Protocol peer - /// encapsulating all the known information and keys. - local_peer: LpLocalPeer, - - /// Representation of a remote Lewes Protocol peer - /// encapsulating all the known information and keys. - remote_peer: LpRemotePeer, - - /// PQ shared secret inherited from parent (for creating further subsessions) - pq_shared_secret: PqSharedSecret, - - /// Subsession PSK (for deriving outer AEAD key) - subsession_psk: [u8; 32], - - /// Negotiated protocol version from handshake. - negotiated_version: u8, -} - -impl SubsessionHandshake { - /// Prepares the next KK handshake message (KK1 or KK2 depending on role/state). - /// - /// # Returns - /// Noise handshake message bytes to send through parent session tunnel. - pub fn prepare_message(&self) -> Result, LpError> { - let mut noise_state = self.noise_state.lock(); - noise_state - .get_bytes_to_send() - .ok_or_else(|| LpError::Internal("Not our turn to send".into()))? - .map_err(LpError::NoiseError) - } - - /// Processes a received KK handshake message (KK1 or KK2). - /// - /// # Arguments - /// * `message` - Noise handshake message received through parent session tunnel. - /// - /// # Returns - /// Any payload embedded in the handshake message (usually empty for KK). - pub fn process_message(&self, message: &[u8]) -> Result, LpError> { - let mut noise_state = self.noise_state.lock(); - let result = noise_state - .read_message(message) - .map_err(LpError::NoiseError)?; - match result { - ReadResult::HandshakeComplete | ReadResult::NoOp => Ok(vec![]), - ReadResult::DecryptedData(data) => Ok(data), - } - } - - /// Checks if the handshake is complete (ready for transport mode). - pub fn is_complete(&self) -> bool { - self.noise_state.lock().is_handshake_finished() - } - - /// Returns whether this side is the initiator. - pub fn is_initiator(&self) -> bool { - self.is_initiator - } - - /// Returns the subsession index. - pub fn subsession_index(&self) -> u64 { - self.index - } - - /// Convert completed subsession handshake into a full LpSession. - /// - /// This consumes the SubsessionHandshake and creates a new LpSession - /// that can be used as a replacement for the parent session. - /// - /// # Arguments - /// * `receiver_index` - New receiver index for the promoted session - /// - /// # Errors - /// Returns error if handshake is not complete - pub fn into_session(self, receiver_index: u32) -> Result { - if !self.is_complete() { - return Err(LpError::Internal( - "Cannot convert incomplete subsession to session".to_string(), - )); - } - - // Extract the noise state (now in transport mode) - let noise_state = self.noise_state.into_inner(); - - // Derive outer AEAD key from the subsession PSK - let outer_key = OuterAeadKey::from_psk(&self.subsession_psk); - - todo!() - // Ok(LpSession { - // // noiserm - // session_id: receiver_index, - // noise_state, - // sending_counter: 0, - // receiving_counter: ReceivingKeyCounterValidator::new(0), - // local_peer: self.local_peer, - // remote_peer: self.remote_peer, - // outer_aead_key: outer_key, - // pq_shared_secret: self.pq_shared_secret, - // subsession_counter: 0, - // read_only: false, - // successor_session_id: None, - // version: self.negotiated_version, - // }) + /// * `Ok(LpPacket)` containing the decrypted application data payload. + /// * `Err(LpError)` if the session is not in transport mode, decryption fails, or the message is not data. + pub fn decrypt_packet(&mut self, packet: EncryptedLpPacket) -> Result { + decrypt_lp_packet(packet, &mut self.active_transport) } } #[cfg(test)] mod tests { use super::*; - use crate::{SessionsMock, kem_list, replay::ReplayError, sessions_for_tests}; + use nym_crypto::asymmetric::x25519; use rand::thread_rng; // Helper function to generate keypairs for tests @@ -621,21 +312,6 @@ mod tests { // } } - // NOTE: These tests are obsolete after removing optional KEM parameters. - // PSQ now always runs using X25519 keys internally converted to KEM format. - // The new tests at the end of this file (test_psq_*) cover PSQ integration. - /* - #[test] - fn test_session_creation_with_psq_state_initiator() { - // OLD API - REMOVED - } - - #[test] - fn test_session_creation_with_psq_state_responder() { - // OLD API - REMOVED - } - */ - #[test] fn test_replay_protection_sequential() { todo!() @@ -705,138 +381,4 @@ mod tests { // assert_eq!(received, 2); // } } - - /* - // These tests remain commented as they rely on the old mock crypto functions - #[test] - fn test_mock_crypto() { - let mut session = create_test_session(true); - let data = [1, 2, 3, 4, 5]; - let mut encrypted = [0; 5]; - let mut decrypted = [0; 5]; - - // Mock encrypt should copy the data - // let encrypted_len = session.encrypt_packet(&data, &mut encrypted).unwrap(); // Removed method - // assert_eq!(encrypted_len, 5); - // assert_eq!(encrypted, data); - - // Mock decrypt should copy the data - // let decrypted_len = session.decrypt_packet(&encrypted, &mut decrypted).unwrap(); // Removed method - // assert_eq!(decrypted_len, 5); - // assert_eq!(decrypted, data); - } - - #[test] - fn test_mock_crypto_buffer_too_small() { - let mut session = create_test_session(true); - let data = [1, 2, 3, 4, 5]; - let mut too_small = [0; 3]; - - // Should fail with buffer too small - // let result = session.encrypt_packet(&data, &mut too_small); // Removed method - // assert!(result.is_err()); - // match result.unwrap_err() { - // LpError::InsufficientBufferSize => {} // Error type might change - // _ => panic!("Expected InsufficientBufferSize error"), - // } - } - */ - - /// Test that X25519 keys are correctly converted to KEM format - #[test] - fn test_x25519_to_kem_conversion() { - todo!() - // - // let initiator_keys = generate_x25519_keypair(); - // let responder_keys = generate_x25519_keypair(); - // - // // Verify we can convert X25519 public key to KEM format (as done in session.rs) - // let x25519_public_bytes = responder_keys.public_key().as_bytes(); - // let libcrux_public_key = - // libcrux_kem::PublicKey::decode(libcrux_kem::Algorithm::X25519, x25519_public_bytes) - // .expect("X25519 public key should convert to libcrux PublicKey"); - // - // let _kem_key = EncapsulationKey::X25519(libcrux_public_key); - // - // // Verify we can convert X25519 private key to KEM format - // let x25519_private_bytes = initiator_keys.private_key().to_bytes(); - // let _libcrux_private_key = - // libcrux_kem::PrivateKey::decode(libcrux_kem::Algorithm::X25519, &x25519_private_bytes) - // .expect("X25519 private key should convert to libcrux PrivateKey"); - // - // // Successful conversion is sufficient - actual encapsulation is tested in psk.rs - // // (libcrux_kem::PrivateKey is an enum with no len() method, conversion success is enough) - } - - #[test] - fn test_demote_sets_read_only() { - todo!() - // let sessions = SessionsMock::mock_post_handshake(12345); - // let mut session = sessions.initiator; - // for kem in kem_list() { - // let session = create_handshake_test_session(kem, 12345u32, true); - // - // // Initially not read-only - // assert!(!session.is_read_only()); - // assert!(session.successor_session_id().is_none()); - // - // // Demote the session - // session.demote(99999); - // - // // Now read-only with successor - // assert!(session.is_read_only()); - // assert_eq!(session.successor_session_id(), Some(99999)); - // } - } - - #[test] - fn test_encrypt_fails_after_demotion() { - let TODO = " for kem in kem_list() {"; - let receiver_index = 12345; - let sessions = SessionsMock::mock_post_handshake(receiver_index); - let mut initiator_session = sessions.initiator; - - // Encryption works before demotion - let plaintext = b"Hello before demotion"; - assert!(initiator_session.encrypt_data(plaintext).is_ok()); - - // Demote the session - initiator_session.demote(99999); - - // Encryption fails after demotion - let result = initiator_session.encrypt_data(plaintext); - assert!(result.is_err()); - match result.unwrap_err() { - NoiseError::SessionReadOnly => { - // Expected - } - e => panic!("Expected SessionReadOnly error, got: {:?}", e), - } - } - - #[test] - fn test_decrypt_works_after_demotion() { - // --- Setup Handshake --- - let TODO = " for kem in kem_list() {"; - let receiver_index = 12345; - let sessions = SessionsMock::mock_post_handshake(receiver_index); - let mut initiator_session = sessions.initiator; - let mut responder_session = sessions.responder; - - // Responder encrypts a message - let plaintext = b"Message to demoted initiator"; - let ciphertext = responder_session - .encrypt_data(plaintext) - .expect("Encryption failed"); - - // Demote the initiator session - initiator_session.demote(99999); - assert!(initiator_session.is_read_only()); - - // Decryption still works on demoted session (drain in-flight) - let decrypted = initiator_session - .decrypt_data(&ciphertext) - .expect("Decryption should work on demoted session"); - assert_eq!(decrypted, plaintext); - } } diff --git a/common/nym-lp/src/session_integration/mod.rs b/common/nym-lp/src/session_integration/mod.rs index 4fbc9a171f..8c945f4b81 100644 --- a/common/nym-lp/src/session_integration/mod.rs +++ b/common/nym-lp/src/session_integration/mod.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use crate::codec::{parse_lp_packet, serialize_lp_packet}; + use crate::codec::serialize_lp_packet; use crate::{ LpError, SessionsMock, kem_list, message::LpMessage, @@ -17,22 +17,10 @@ mod tests { message: LpMessage, ) -> LpPacket { // Create the header - let header = LpHeader { - protocol_version, - reserved: [0u8; 3], // reserved - receiver_idx, - counter, - }; - - // Create the trailer (zeros for now, in a real implementation this might be a MAC) - let trailer = [0u8; TRAILER_LEN]; + let header = LpHeader::new(receiver_idx, counter, protocol_version, message.typ()); // Create and return the packet directly - LpPacket { - header, - message, - trailer, - } + LpPacket { header, message } } /// Tests the complete session flow including: @@ -43,696 +31,698 @@ mod tests { /// - Session removal and cleanup #[test] fn test_full_session_flow() { - // 1. Initialize session manager - let mut session_manager_1 = SessionManager::new(); - let mut session_manager_2 = SessionManager::new(); + todo!() - let TODO = " for kem in kem_list() {"; - - let receiver_index = 12345; - let sessions = SessionsMock::mock_post_handshake(receiver_index); - - // 2. Create sessions using the pre-built Noise states - let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); - let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); - - // Verify session count - assert_eq!(session_manager_1.session_count(), 1); - assert_eq!(session_manager_2.session_count(), 1); - - // 3. Simulate Data Transfer (Post-Handshake) - println!("Starting data transfer simulation..."); - let plaintext_a_to_b = b"Hello from A!"; - - // A encrypts data - let ciphertext_a_to_b = session_manager_1 - .encrypt_data(peer_a_sm, plaintext_a_to_b) - .expect("A encrypt failed"); - - // A prepares packet - let counter_a = session_manager_1.next_counter(peer_a_sm).unwrap(); - let message_a_to_b = create_test_packet(1, receiver_index, counter_a, ciphertext_a_to_b); - let mut encoded_data_a_to_b = BytesMut::new(); - serialize_lp_packet(&message_a_to_b, &mut encoded_data_a_to_b, None) - .expect("A serialize data failed"); - - // B parses packet and checks replay - let decoded_packet_b = - parse_lp_packet(&encoded_data_a_to_b, None).expect("B parse data failed"); - assert_eq!(decoded_packet_b.header.counter, counter_a); - - // Check replay before decrypting - session_manager_2 - .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) - .expect("B data replay check failed (A->B)"); - - // B decrypts data - let decrypted_payload = session_manager_2 - .decrypt_data(peer_b_sm, &decoded_packet_b.message) - .expect("B decrypt failed"); - assert_eq!(decrypted_payload, plaintext_a_to_b); - // Mark counter only after successful decryption - session_manager_2 - .receiving_counter_mark(peer_b_sm, decoded_packet_b.header.counter) - .expect("B mark data counter failed"); - println!( - " A->B: Decrypted successfully: {:?}", - String::from_utf8_lossy(&decrypted_payload) - ); - - // B sends data to A - let plaintext_b_to_a = b"Hello from B!"; - let ciphertext_b_to_a = session_manager_2 - .encrypt_data(peer_b_sm, plaintext_b_to_a) - .expect("B encrypt failed"); - let counter_b = session_manager_2.next_counter(peer_b_sm).unwrap(); - let message_b_to_a = create_test_packet(1, receiver_index, counter_b, ciphertext_b_to_a); - let mut encoded_data_b_to_a = BytesMut::new(); - serialize_lp_packet(&message_b_to_a, &mut encoded_data_b_to_a, None) - .expect("B serialize data failed"); - - // A parses packet and checks replay - let decoded_packet_a = - parse_lp_packet(&encoded_data_b_to_a, None).expect("A parse data failed"); - assert_eq!(decoded_packet_a.header.counter, counter_b); - - // Check replay before decrypting - session_manager_1 - .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) - .expect("A data replay check failed (B->A)"); - - // A decrypts data - let decrypted_payload = session_manager_1 - .decrypt_data(peer_a_sm, &decoded_packet_a.message) - .expect("A decrypt failed"); - assert_eq!(decrypted_payload, plaintext_b_to_a); - // Mark counter only after successful decryption - session_manager_1 - .receiving_counter_mark(peer_a_sm, decoded_packet_a.header.counter) - .expect("A mark data counter failed"); - println!( - " B->A: Decrypted successfully: {:?}", - String::from_utf8_lossy(&decrypted_payload) - ); - - println!("Data transfer simulation completed."); - - // 4. Replay Protection Test (Data Packet) - println!("Testing data packet replay protection..."); - // Try to replay the last message from B to A - // Need to re-encode because decode consumes the buffer - let message_b_to_a_replay = create_test_packet( - 1, - receiver_index, - counter_b, - LpMessage::EncryptedData(crate::message::EncryptedDataPayload( - plaintext_b_to_a.to_vec(), - )), // Using plaintext here, but content doesn't matter for replay check - ); - let mut encoded_data_b_to_a_replay = BytesMut::new(); - serialize_lp_packet( - &message_b_to_a_replay, - &mut encoded_data_b_to_a_replay, - None, - ) - .expect("B serialize replay failed"); - - let parsed_replay_packet = - parse_lp_packet(&encoded_data_b_to_a_replay, None).expect("A parse replay failed"); - let replay_result = session_manager_1 - .receiving_counter_quick_check(peer_a_sm, parsed_replay_packet.header.counter); - assert!(replay_result.is_err(), "Data replay should be prevented"); - assert!( - matches!(replay_result.unwrap_err(), LpError::Replay(_)), - "Should be a replay protection error for data packet" - ); - println!("Data packet replay protection test passed."); - - // 5. Test out-of-order packet reception (send counter N+1 before counter N) - println!("Testing out-of-order data packet reception..."); - let counter_a_next = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 1 - let counter_a_skip = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 2 - - // Prepare data for counter_a_skip (N+1) - let plaintext_skip = b"Out of order message"; - let ciphertext_skip = session_manager_1 - .encrypt_data(peer_a_sm, plaintext_skip) - .expect("A encrypt skip failed"); - - let message_a_to_b_skip = create_test_packet( - 1, // protocol version - receiver_index, - counter_a_skip, // Send N+1 first - ciphertext_skip, - ); - - // Encode the skip message - let mut encoded_skip = BytesMut::new(); - serialize_lp_packet(&message_a_to_b_skip, &mut encoded_skip, None) - .expect("Failed to serialize skip message"); - - // B parses skip message and checks replay - let decoded_packet_skip = - parse_lp_packet(&encoded_skip, None).expect("B parse skip failed"); - session_manager_2 - .receiving_counter_quick_check(peer_b_sm, decoded_packet_skip.header.counter) - .expect("B replay check skip failed"); - assert_eq!(decoded_packet_skip.header.counter, counter_a_skip); - - // B decrypts skip message - let decrypted_payload = session_manager_2 - .decrypt_data(peer_b_sm, &decoded_packet_skip.message) - .expect("B decrypt skip failed"); - assert_eq!(decrypted_payload, plaintext_skip); - // Mark counter N+1 - session_manager_2 - .receiving_counter_mark(peer_b_sm, decoded_packet_skip.header.counter) - .expect("B mark skip counter failed"); - println!( - " A->B (Counter {}): Decrypted successfully: {:?}", - counter_a_skip, - String::from_utf8_lossy(&decrypted_payload) - ); - - // 6. Now send the skipped counter N message (should still work) - println!("Testing delayed data packet reception..."); - // Prepare data for counter_a_next (N) - let plaintext_delayed = b"Delayed message"; - let ciphertext_delayed = session_manager_1 - .encrypt_data(peer_a_sm, plaintext_delayed) - .expect("A encrypt delayed failed"); - - let message_a_to_b_delayed = create_test_packet( - 1, // protocol version - receiver_index, - counter_a_next, // counter N (delayed packet) - ciphertext_delayed, - ); - - // Encode the delayed message - let mut encoded_delayed = BytesMut::new(); - serialize_lp_packet(&message_a_to_b_delayed, &mut encoded_delayed, None) - .expect("Failed to serialize delayed message"); - - // Make a copy for replay test later - let encoded_delayed_copy = encoded_delayed.clone(); - - // B parses delayed message and checks replay - let decoded_packet_delayed = - parse_lp_packet(&encoded_delayed, None).expect("B parse delayed failed"); - session_manager_2 - .receiving_counter_quick_check(peer_b_sm, decoded_packet_delayed.header.counter) - .expect("B replay check delayed failed"); - assert_eq!(decoded_packet_delayed.header.counter, counter_a_next); - - // B decrypts delayed message - let decrypted_payload = session_manager_2 - .decrypt_data(peer_b_sm, &decoded_packet_delayed.message) - .expect("B decrypt delayed failed"); - assert_eq!(decrypted_payload, plaintext_delayed); - // Mark counter N - session_manager_2 - .receiving_counter_mark(peer_b_sm, decoded_packet_delayed.header.counter) - .expect("B mark delayed counter failed"); - println!( - " A->B (Counter {}): Decrypted successfully: {:?}", - counter_a_next, - String::from_utf8_lossy(&decrypted_payload) - ); - - println!("Delayed data packet reception test passed."); - - // 7. Try to replay message with counter N (should fail) - println!("Testing replay of delayed packet..."); - let parsed_delayed_replay = - parse_lp_packet(&encoded_delayed_copy, None).expect("Parse delayed replay failed"); - let result = session_manager_2 - .receiving_counter_quick_check(peer_b_sm, parsed_delayed_replay.header.counter); - assert!(result.is_err(), "Replay attack should be prevented"); - assert!( - matches!(result, Err(LpError::Replay(_))), - "Should be a replay protection error" - ); - - // 8. Session removal - assert!(session_manager_1.remove_state_machine(receiver_index)); - assert_eq!(session_manager_1.session_count(), 0); - - // Verify the session is gone - let session = session_manager_1.state_machine_exists(receiver_index); - assert!(!session, "Session should be removed"); - - // But the other session still exists - let session = session_manager_2.state_machine_exists(receiver_index); - assert!(session, "Session still exists in the other manager"); - } - - /// Tests simultaneous bidirectional communication between sessions - #[test] - fn test_bidirectional_communication() { - // 1. Initialize session manager - let mut session_manager_1 = SessionManager::new(); - let mut session_manager_2 = SessionManager::new(); - - let TODO = " for kem in kem_list() {"; - - let receiver_index = 12345; - let sessions = SessionsMock::mock_post_handshake(receiver_index); - - // 2. Create sessions using the pre-built Noise states - let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); - let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); - - // Counters after handshake - let mut counter_a = 0; // Next counter for A to send - let mut counter_b = 0; // Next counter for B to send - - // 3. Send multiple encrypted messages both ways - const NUM_MESSAGES: u64 = 5; - for i in 0..NUM_MESSAGES { - println!("Bidirectional test: Round {}", i); - // --- A sends to B --- - let plaintext_a = format!("A->B Message {}", i).into_bytes(); - let ciphertext_a = session_manager_1 - .encrypt_data(peer_a_sm, &plaintext_a) - .expect("A encrypt failed"); - let current_counter_a = counter_a; - counter_a += 1; - - let message_a = create_test_packet(1, receiver_index, current_counter_a, ciphertext_a); - let mut encoded_a = BytesMut::new(); - serialize_lp_packet(&message_a, &mut encoded_a, None).expect("A serialize failed"); - - // B parses and checks replay - let decoded_packet_b = parse_lp_packet(&encoded_a, None).expect("B parse failed"); - session_manager_2 - .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) - .expect("B replay check failed (A->B)"); - assert_eq!(decoded_packet_b.header.counter, current_counter_a); - let decrypted_payload = session_manager_2 - .decrypt_data(peer_b_sm, &decoded_packet_b.message) - .expect("B decrypt failed"); - assert_eq!(decrypted_payload, plaintext_a); - session_manager_2 - .receiving_counter_mark(peer_b_sm, current_counter_a) - .expect("B mark counter failed"); - - // --- B sends to A --- - let plaintext_b = format!("B->A Message {}", i).into_bytes(); - let ciphertext_b = session_manager_2 - .encrypt_data(peer_b_sm, &plaintext_b) - .expect("B encrypt failed"); - let current_counter_b = counter_b; - counter_b += 1; - - let message_b = create_test_packet(1, receiver_index, current_counter_b, ciphertext_b); - let mut encoded_b = BytesMut::new(); - serialize_lp_packet(&message_b, &mut encoded_b, None).expect("B serialize failed"); - - // A parses and checks replay - let decoded_packet_a = parse_lp_packet(&encoded_b, None).expect("A parse failed"); - session_manager_1 - .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) - .expect("A replay check failed (B->A)"); - assert_eq!(decoded_packet_a.header.counter, current_counter_b); - let decrypted_payload = session_manager_1 - .decrypt_data(peer_a_sm, &decoded_packet_a.message) - .expect("A decrypt failed"); - assert_eq!(decrypted_payload, plaintext_b); - session_manager_1 - .receiving_counter_mark(peer_a_sm, current_counter_b) - .expect("A mark counter failed"); - } - - // 5. Verify counter stats - // Note: current_packet_cnt() returns (next_expected_receive_counter, total_received) - let (next_recv_a, total_recv_a) = session_manager_1.current_packet_cnt(peer_a_sm).unwrap(); - let (next_recv_b, total_recv_b) = session_manager_2.current_packet_cnt(peer_b_sm).unwrap(); - - // Peer A sent handshake(0), handshake(1) + 5 data packets = 7 total. Next send counter = 7. - // Peer A received handshake(0) + 5 data packets = 6 total. Next expected recv counter = 6. - assert_eq!( - counter_a, NUM_MESSAGES, - "Peer A final send counter mismatch" - ); - assert_eq!( - total_recv_a, NUM_MESSAGES, - "Peer A total received count mismatch" - ); // Received 5 data - assert_eq!( - next_recv_a, NUM_MESSAGES, - "Peer A next expected receive counter mismatch" - ); // Expected counter for msg from B - - // Peer B sent handshake(0) + 5 data packets = 6 total. Next send counter = 6. - // Peer B received handshake(0), handshake(1) + 5 data packets = 7 total. Next expected recv counter = 7. - assert_eq!( - counter_b, NUM_MESSAGES, - "Peer B final send counter mismatch" - ); - assert_eq!( - total_recv_b, NUM_MESSAGES, - "Peer B total received count mismatch" - ); // Received 5 data - assert_eq!( - next_recv_b, NUM_MESSAGES, - "Peer B next expected receive counter mismatch" - ); // Expected counter for msg from A - - println!("Bidirectional test completed."); - } - - /// Tests error handling in session flow - #[test] - fn test_session_error_handling() { - // 1. Initialize session manager - let mut session_manager = SessionManager::new(); - - let TODO = " for kem in kem_list() {"; - let receiver_index = 123; - let session1 = SessionsMock::mock_post_handshake(receiver_index).initiator; - let session2 = SessionsMock::mock_post_handshake(124).initiator; - - // 2. Create a session (using real noise state) - let _session = session_manager.create_session_state_machine(session1); - - // 3. Try to get a non-existent session - let result = session_manager.state_machine_exists(999); - assert!(!result, "Non-existent session should return None"); - - // 4. Try to remove a non-existent session - let result = session_manager.remove_state_machine(999); - assert!( - !result, - "Remove session should not remove a non-existent session" - ); - - // 5. Create and immediately remove a session - let _temp_session = session_manager.create_session_state_machine(session2); - - assert!( - session_manager.remove_state_machine(124), - "Should remove the session" - ); - - // 6. Create a codec and test error cases - // let mut codec = LPCodec::new(session); - - // 7. Create an invalid message type packet - let mut buf = BytesMut::new(); - - // Add header - buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - buf.extend_from_slice(&receiver_index.to_le_bytes()); // Sender index - buf.extend_from_slice(&0u64.to_le_bytes()); // Counter - - // Add invalid message type - buf.extend_from_slice(&0xFFFFu16.to_le_bytes()); - - // Add some dummy data - buf.extend_from_slice(&[0u8; 80]); - - // Add trailer - buf.extend_from_slice(&[0u8; TRAILER_LEN]); - - // Try to parse the invalid message type - let result = parse_lp_packet(&buf, None); - assert!(result.is_err(), "Decoding invalid message type should fail"); - - // Add assertion for the specific error type - assert!(matches!( - result.unwrap_err(), - LpError::InvalidMessageType(0xFFFF) - )); - - // 8. Test partial packet decoding - let partial_packet = &buf[0..10]; // Too short to be a valid packet - let partial_bytes = BytesMut::from(partial_packet); - - let result = parse_lp_packet(&partial_bytes, None); - assert!(result.is_err(), "Parsing partial packet should fail"); - assert!(matches!( - result.unwrap_err(), - LpError::InsufficientBufferSize - )); - } - // Remove unused imports if SessionManager methods are no longer direct dependencies - // use crate::noise_protocol::{create_noise_state, create_noise_state_responder}; - use crate::state_machine::LpData; - use crate::state_machine::{LpAction, LpInput, LpStateBare}; - // Use Bytes for SendData input - - // Keep helper function for creating test packets if needed, - // but LpAction::SendPacket should provide the packets now. - // fn create_test_packet(...) -> LpPacket { ... } - - /// Tests the complete session flow using ONLY the process_input interface: - /// - Creation of sessions through session manager - /// - Handshake driven by StartHandshake, ReceivePacket inputs - /// - Data transfer driven by SendData, ReceivePacket inputs - /// - Actions like SendPacket, DeliverData handled from output - /// - Implicit replay protection via state machine logic - /// - Closing driven by Close input - #[test] - fn test_full_session_flow_with_process_input() { - // 1. Initialize session managers - let mut session_manager_1 = SessionManager::new(); - let mut session_manager_2 = SessionManager::new(); - - let TODO = " for kem in kem_list() {"; - - let receiver_index = 12345; - let sessions = SessionsMock::mock_post_handshake(receiver_index); - - // 2. Create sessions state machines - session_manager_1.create_session_state_machine(sessions.initiator); - session_manager_2.create_session_state_machine(sessions.responder); - - assert_eq!(session_manager_1.session_count(), 1); - assert_eq!(session_manager_2.session_count(), 1); - assert!(session_manager_1.state_machine_exists(receiver_index)); - assert!(session_manager_2.state_machine_exists(receiver_index)); - - // Verify initial states are Transport - assert_eq!( - session_manager_1.get_state(receiver_index).unwrap(), - LpStateBare::Transport - ); - assert_eq!( - session_manager_2.get_state(receiver_index).unwrap(), - LpStateBare::Transport - ); - - // --- 3. Simulate Data Transfer via process_input --- - println!("Starting data transfer simulation via process_input..."); - let plaintext_a_to_b = LpData::new_opaque(b"Hello from A via process_input!".to_vec()); - let plaintext_b_to_a = LpData::new_opaque(b"Hello from B via process_input!".to_vec()); - - // --- A sends to B --- - println!(" A sends to B"); - let action_a_send = session_manager_1 - .process_input(receiver_index, LpInput::SendData(plaintext_a_to_b.clone())) - .expect("A SendData should produce action") - .expect("A SendData failed"); - - let data_packet_a = if let LpAction::SendPacket(packet) = action_a_send { - packet - } else { - panic!("A SendData did not produce SendPacket"); - }; - - // Simulate network - let mut buf_data_a = BytesMut::new(); - serialize_lp_packet(&data_packet_a, &mut buf_data_a, None).unwrap(); - let parsed_data_a = parse_lp_packet(&buf_data_a, None).unwrap(); - - // B receives - println!(" B receives from A"); - let action_b_recv = session_manager_2 - .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_a)) - .expect("B ReceivePacket (data) should produce action") - .expect("B ReceivePacket (data) failed"); - - if let LpAction::DeliverData(data) = action_b_recv { - assert_eq!(data, plaintext_a_to_b, "Decrypted data mismatch A->B"); - println!( - " B successfully decrypted: {:?}", - String::from_utf8_lossy(&data.content) - ); - } else { - panic!("B ReceivePacket did not produce DeliverData"); - } - - // --- B sends to A --- - println!(" B sends to A"); - let action_b_send = session_manager_2 - .process_input(receiver_index, LpInput::SendData(plaintext_b_to_a.clone())) - .expect("B SendData should produce action") - .expect("B SendData failed"); - - let data_packet_b = if let LpAction::SendPacket(packet) = action_b_send { - packet - } else { - panic!("B SendData did not produce SendPacket"); - }; - // Keep a copy for replay test - let data_packet_b_replay = data_packet_b.clone(); - - // Simulate network - let mut buf_data_b = BytesMut::new(); - serialize_lp_packet(&data_packet_b, &mut buf_data_b, None).unwrap(); - let parsed_data_b = parse_lp_packet(&buf_data_b, None).unwrap(); - - // A receives - println!(" A receives from B"); - let action_a_recv = session_manager_1 - .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_b)) - .expect("A ReceivePacket (data) should produce action") - .expect("A ReceivePacket (data) failed"); - - if let LpAction::DeliverData(data) = action_a_recv { - assert_eq!(data, plaintext_b_to_a, "Decrypted data mismatch B->A"); - println!( - " A successfully decrypted: {:?}", - String::from_utf8_lossy(&data.content) - ); - } else { - panic!("A ReceivePacket did not produce DeliverData"); - } - println!("Data transfer simulation completed."); - - // --- 4. Replay Protection Test --- - println!("Testing data packet replay protection via process_input..."); - let replay_result = session_manager_1 - .process_input(receiver_index, LpInput::ReceivePacket(data_packet_b_replay)); // Use cloned packet - - assert!(replay_result.is_err(), "Replay should produce Err(...)"); - let error = replay_result.err().unwrap(); - assert!( - matches!(error, LpError::Replay(_)), - "Expected Replay error, got {:?}", - error - ); - println!("Data packet replay protection test passed."); - - // --- 5. Out-of-Order Test --- - println!("Testing out-of-order reception via process_input..."); - - // A prepares N+1 then N - let data_n_plus_1 = LpData::new_opaque(b"Message N+1".to_vec()); - let data_n = LpData::new_opaque(b"Message N".to_vec()); - - let action_send_n1 = session_manager_1 - .process_input(receiver_index, LpInput::SendData(data_n_plus_1.clone())) - .unwrap() - .unwrap(); - let packet_n1 = match action_send_n1 { - LpAction::SendPacket(p) => p, - _ => panic!("Expected SendPacket"), - }; - - let action_send_n = session_manager_1 - .process_input(receiver_index, LpInput::SendData(data_n.clone())) - .unwrap() - .unwrap(); - let packet_n = match action_send_n { - LpAction::SendPacket(p) => p, - _ => panic!("Expected SendPacket"), - }; - let packet_n_replay = packet_n.clone(); // For replay test - - // B receives N+1 first - println!(" B receives N+1"); - let action_recv_n1 = session_manager_2 - .process_input(receiver_index, LpInput::ReceivePacket(packet_n1)) - .unwrap() - .unwrap(); - match action_recv_n1 { - LpAction::DeliverData(d) => assert_eq!(d, data_n_plus_1, "Data N+1 mismatch"), - _ => panic!("Expected DeliverData for N+1"), - } - - // B receives N second (should work) - println!(" B receives N"); - let action_recv_n = session_manager_2 - .process_input(receiver_index, LpInput::ReceivePacket(packet_n)) - .unwrap() - .unwrap(); - match action_recv_n { - LpAction::DeliverData(d) => assert_eq!(d, data_n, "Data N mismatch"), - _ => panic!("Expected DeliverData for N"), - } - - // B tries to replay N (should fail) - println!(" B tries to replay N"); - let replay_n_result = session_manager_2 - .process_input(receiver_index, LpInput::ReceivePacket(packet_n_replay)); - assert!(replay_n_result.is_err(), "Replay N should produce Err"); - assert!( - matches!(replay_n_result.err().unwrap(), LpError::Replay(_)), - "Expected Replay error for N" - ); - println!("Out-of-order test passed."); - - // --- 6. Close Test --- - println!("Testing close via process_input..."); - - // A closes - let action_a_close = session_manager_1 - .process_input(receiver_index, LpInput::Close) - .expect("A Close should produce action") - .expect("A Close failed"); - assert!(matches!(action_a_close, LpAction::ConnectionClosed)); - assert_eq!( - session_manager_1.get_state(receiver_index).unwrap(), - LpStateBare::Closed - ); - - // Further actions on A fail - let send_after_close_a = session_manager_1.process_input( - receiver_index, - LpInput::SendData(LpData::new_opaque(b"fail".to_vec())), - ); - assert!(send_after_close_a.is_err()); - assert!(matches!( - send_after_close_a.err().unwrap(), - LpError::LpSessionClosed - )); - - // B closes - let action_b_close = session_manager_2 - .process_input(receiver_index, LpInput::Close) - .expect("B Close should produce action") - .expect("B Close failed"); - assert!(matches!(action_b_close, LpAction::ConnectionClosed)); - assert_eq!( - session_manager_2.get_state(receiver_index).unwrap(), - LpStateBare::Closed - ); - - // Further actions on B fail - let send_after_close_b = session_manager_2.process_input( - receiver_index, - LpInput::SendData(LpData::new_opaque(b"fail".to_vec())), - ); - assert!(send_after_close_b.is_err()); - assert!(matches!( - send_after_close_b.err().unwrap(), - LpError::LpSessionClosed - )); - println!("Close test passed."); - - // --- 7. Session Removal --- - assert!(session_manager_1.remove_state_machine(receiver_index)); - assert_eq!(session_manager_1.session_count(), 0); - assert!(!session_manager_1.state_machine_exists(receiver_index)); - - // B's session manager still has it until removed - assert!(session_manager_2.state_machine_exists(receiver_index)); - assert!(session_manager_2.remove_state_machine(receiver_index)); - assert_eq!(session_manager_2.session_count(), 0); - assert!(!session_manager_2.state_machine_exists(receiver_index)); - println!("Session removal test passed."); + // // 1. Initialize session manager + // let mut session_manager_1 = SessionManager::new(); + // let mut session_manager_2 = SessionManager::new(); + // + // let TODO = " for kem in kem_list() {"; + // + // let receiver_index = 12345; + // let sessions = SessionsMock::mock_post_handshake(receiver_index); + // + // // 2. Create sessions using the pre-built Noise states + // let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); + // let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); + // + // // Verify session count + // assert_eq!(session_manager_1.session_count(), 1); + // assert_eq!(session_manager_2.session_count(), 1); + // + // // 3. Simulate Data Transfer (Post-Handshake) + // println!("Starting data transfer simulation..."); + // let plaintext_a_to_b = b"Hello from A!"; + // + // // A encrypts data + // let ciphertext_a_to_b = session_manager_1 + // .encrypt_data(peer_a_sm, plaintext_a_to_b) + // .expect("A encrypt failed"); + // + // // A prepares packet + // let counter_a = session_manager_1.next_counter(peer_a_sm).unwrap(); + // let message_a_to_b = create_test_packet(1, receiver_index, counter_a, ciphertext_a_to_b); + // let mut encoded_data_a_to_b = BytesMut::new(); + // serialize_lp_packet(&message_a_to_b, &mut encoded_data_a_to_b, None) + // .expect("A serialize data failed"); + // + // // B parses packet and checks replay + // let decoded_packet_b = + // parse_lp_packet(&encoded_data_a_to_b, None).expect("B parse data failed"); + // assert_eq!(decoded_packet_b.header.counter, counter_a); + // + // // Check replay before decrypting + // session_manager_2 + // .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) + // .expect("B data replay check failed (A->B)"); + // + // // B decrypts data + // let decrypted_payload = session_manager_2 + // .decrypt_data(peer_b_sm, &decoded_packet_b.message) + // .expect("B decrypt failed"); + // assert_eq!(decrypted_payload, plaintext_a_to_b); + // // Mark counter only after successful decryption + // session_manager_2 + // .receiving_counter_mark(peer_b_sm, decoded_packet_b.header.counter) + // .expect("B mark data counter failed"); + // println!( + // " A->B: Decrypted successfully: {:?}", + // String::from_utf8_lossy(&decrypted_payload) + // ); + // + // // B sends data to A + // let plaintext_b_to_a = b"Hello from B!"; + // let ciphertext_b_to_a = session_manager_2 + // .encrypt_data(peer_b_sm, plaintext_b_to_a) + // .expect("B encrypt failed"); + // let counter_b = session_manager_2.next_counter(peer_b_sm).unwrap(); + // let message_b_to_a = create_test_packet(1, receiver_index, counter_b, ciphertext_b_to_a); + // let mut encoded_data_b_to_a = BytesMut::new(); + // serialize_lp_packet(&message_b_to_a, &mut encoded_data_b_to_a, None) + // .expect("B serialize data failed"); + // + // // A parses packet and checks replay + // let decoded_packet_a = + // parse_lp_packet(&encoded_data_b_to_a, None).expect("A parse data failed"); + // assert_eq!(decoded_packet_a.header.counter, counter_b); + // + // // Check replay before decrypting + // session_manager_1 + // .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) + // .expect("A data replay check failed (B->A)"); + // + // // A decrypts data + // let decrypted_payload = session_manager_1 + // .decrypt_data(peer_a_sm, &decoded_packet_a.message) + // .expect("A decrypt failed"); + // assert_eq!(decrypted_payload, plaintext_b_to_a); + // // Mark counter only after successful decryption + // session_manager_1 + // .receiving_counter_mark(peer_a_sm, decoded_packet_a.header.counter) + // .expect("A mark data counter failed"); + // println!( + // " B->A: Decrypted successfully: {:?}", + // String::from_utf8_lossy(&decrypted_payload) + // ); + // + // println!("Data transfer simulation completed."); + // + // // 4. Replay Protection Test (Data Packet) + // println!("Testing data packet replay protection..."); + // // Try to replay the last message from B to A + // // Need to re-encode because decode consumes the buffer + // let message_b_to_a_replay = create_test_packet( + // 1, + // receiver_index, + // counter_b, + // LpMessage::EncryptedData(crate::message::EncryptedDataPayload( + // plaintext_b_to_a.to_vec(), + // )), // Using plaintext here, but content doesn't matter for replay check + // ); + // let mut encoded_data_b_to_a_replay = BytesMut::new(); + // serialize_lp_packet( + // &message_b_to_a_replay, + // &mut encoded_data_b_to_a_replay, + // None, + // ) + // .expect("B serialize replay failed"); + // + // let parsed_replay_packet = + // parse_lp_packet(&encoded_data_b_to_a_replay, None).expect("A parse replay failed"); + // let replay_result = session_manager_1 + // .receiving_counter_quick_check(peer_a_sm, parsed_replay_packet.header.counter); + // assert!(replay_result.is_err(), "Data replay should be prevented"); + // assert!( + // matches!(replay_result.unwrap_err(), LpError::Replay(_)), + // "Should be a replay protection error for data packet" + // ); + // println!("Data packet replay protection test passed."); + // + // // 5. Test out-of-order packet reception (send counter N+1 before counter N) + // println!("Testing out-of-order data packet reception..."); + // let counter_a_next = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 1 + // let counter_a_skip = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 2 + // + // // Prepare data for counter_a_skip (N+1) + // let plaintext_skip = b"Out of order message"; + // let ciphertext_skip = session_manager_1 + // .encrypt_data(peer_a_sm, plaintext_skip) + // .expect("A encrypt skip failed"); + // + // let message_a_to_b_skip = create_test_packet( + // 1, // protocol version + // receiver_index, + // counter_a_skip, // Send N+1 first + // ciphertext_skip, + // ); + // + // // Encode the skip message + // let mut encoded_skip = BytesMut::new(); + // serialize_lp_packet(&message_a_to_b_skip, &mut encoded_skip, None) + // .expect("Failed to serialize skip message"); + // + // // B parses skip message and checks replay + // let decoded_packet_skip = + // parse_lp_packet(&encoded_skip, None).expect("B parse skip failed"); + // session_manager_2 + // .receiving_counter_quick_check(peer_b_sm, decoded_packet_skip.header.counter) + // .expect("B replay check skip failed"); + // assert_eq!(decoded_packet_skip.header.counter, counter_a_skip); + // + // // B decrypts skip message + // let decrypted_payload = session_manager_2 + // .decrypt_data(peer_b_sm, &decoded_packet_skip.message) + // .expect("B decrypt skip failed"); + // assert_eq!(decrypted_payload, plaintext_skip); + // // Mark counter N+1 + // session_manager_2 + // .receiving_counter_mark(peer_b_sm, decoded_packet_skip.header.counter) + // .expect("B mark skip counter failed"); + // println!( + // " A->B (Counter {}): Decrypted successfully: {:?}", + // counter_a_skip, + // String::from_utf8_lossy(&decrypted_payload) + // ); + // + // // 6. Now send the skipped counter N message (should still work) + // println!("Testing delayed data packet reception..."); + // // Prepare data for counter_a_next (N) + // let plaintext_delayed = b"Delayed message"; + // let ciphertext_delayed = session_manager_1 + // .encrypt_data(peer_a_sm, plaintext_delayed) + // .expect("A encrypt delayed failed"); + // + // let message_a_to_b_delayed = create_test_packet( + // 1, // protocol version + // receiver_index, + // counter_a_next, // counter N (delayed packet) + // ciphertext_delayed, + // ); + // + // // Encode the delayed message + // let mut encoded_delayed = BytesMut::new(); + // serialize_lp_packet(&message_a_to_b_delayed, &mut encoded_delayed, None) + // .expect("Failed to serialize delayed message"); + // + // // Make a copy for replay test later + // let encoded_delayed_copy = encoded_delayed.clone(); + // + // // B parses delayed message and checks replay + // let decoded_packet_delayed = + // parse_lp_packet(&encoded_delayed, None).expect("B parse delayed failed"); + // session_manager_2 + // .receiving_counter_quick_check(peer_b_sm, decoded_packet_delayed.header.counter) + // .expect("B replay check delayed failed"); + // assert_eq!(decoded_packet_delayed.header.counter, counter_a_next); + // + // // B decrypts delayed message + // let decrypted_payload = session_manager_2 + // .decrypt_data(peer_b_sm, &decoded_packet_delayed.message) + // .expect("B decrypt delayed failed"); + // assert_eq!(decrypted_payload, plaintext_delayed); + // // Mark counter N + // session_manager_2 + // .receiving_counter_mark(peer_b_sm, decoded_packet_delayed.header.counter) + // .expect("B mark delayed counter failed"); + // println!( + // " A->B (Counter {}): Decrypted successfully: {:?}", + // counter_a_next, + // String::from_utf8_lossy(&decrypted_payload) + // ); + // + // println!("Delayed data packet reception test passed."); + // + // // 7. Try to replay message with counter N (should fail) + // println!("Testing replay of delayed packet..."); + // let parsed_delayed_replay = + // parse_lp_packet(&encoded_delayed_copy, None).expect("Parse delayed replay failed"); + // let result = session_manager_2 + // .receiving_counter_quick_check(peer_b_sm, parsed_delayed_replay.header.counter); + // assert!(result.is_err(), "Replay attack should be prevented"); + // assert!( + // matches!(result, Err(LpError::Replay(_))), + // "Should be a replay protection error" + // ); + // + // // 8. Session removal + // assert!(session_manager_1.remove_state_machine(receiver_index)); + // assert_eq!(session_manager_1.session_count(), 0); + // + // // Verify the session is gone + // let session = session_manager_1.state_machine_exists(receiver_index); + // assert!(!session, "Session should be removed"); + // + // // But the other session still exists + // let session = session_manager_2.state_machine_exists(receiver_index); + // assert!(session, "Session still exists in the other manager"); } + // + // /// Tests simultaneous bidirectional communication between sessions + // #[test] + // fn test_bidirectional_communication() { + // // 1. Initialize session manager + // let mut session_manager_1 = SessionManager::new(); + // let mut session_manager_2 = SessionManager::new(); + // + // let TODO = " for kem in kem_list() {"; + // + // let receiver_index = 12345; + // let sessions = SessionsMock::mock_post_handshake(receiver_index); + // + // // 2. Create sessions using the pre-built Noise states + // let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); + // let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); + // + // // Counters after handshake + // let mut counter_a = 0; // Next counter for A to send + // let mut counter_b = 0; // Next counter for B to send + // + // // 3. Send multiple encrypted messages both ways + // const NUM_MESSAGES: u64 = 5; + // for i in 0..NUM_MESSAGES { + // println!("Bidirectional test: Round {}", i); + // // --- A sends to B --- + // let plaintext_a = format!("A->B Message {}", i).into_bytes(); + // let ciphertext_a = session_manager_1 + // .encrypt_data(peer_a_sm, &plaintext_a) + // .expect("A encrypt failed"); + // let current_counter_a = counter_a; + // counter_a += 1; + // + // let message_a = create_test_packet(1, receiver_index, current_counter_a, ciphertext_a); + // let mut encoded_a = BytesMut::new(); + // serialize_lp_packet(&message_a, &mut encoded_a, None).expect("A serialize failed"); + // + // // B parses and checks replay + // let decoded_packet_b = parse_lp_packet(&encoded_a, None).expect("B parse failed"); + // session_manager_2 + // .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) + // .expect("B replay check failed (A->B)"); + // assert_eq!(decoded_packet_b.header.counter, current_counter_a); + // let decrypted_payload = session_manager_2 + // .decrypt_data(peer_b_sm, &decoded_packet_b.message) + // .expect("B decrypt failed"); + // assert_eq!(decrypted_payload, plaintext_a); + // session_manager_2 + // .receiving_counter_mark(peer_b_sm, current_counter_a) + // .expect("B mark counter failed"); + // + // // --- B sends to A --- + // let plaintext_b = format!("B->A Message {}", i).into_bytes(); + // let ciphertext_b = session_manager_2 + // .encrypt_data(peer_b_sm, &plaintext_b) + // .expect("B encrypt failed"); + // let current_counter_b = counter_b; + // counter_b += 1; + // + // let message_b = create_test_packet(1, receiver_index, current_counter_b, ciphertext_b); + // let mut encoded_b = BytesMut::new(); + // serialize_lp_packet(&message_b, &mut encoded_b, None).expect("B serialize failed"); + // + // // A parses and checks replay + // let decoded_packet_a = parse_lp_packet(&encoded_b, None).expect("A parse failed"); + // session_manager_1 + // .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) + // .expect("A replay check failed (B->A)"); + // assert_eq!(decoded_packet_a.header.counter, current_counter_b); + // let decrypted_payload = session_manager_1 + // .decrypt_data(peer_a_sm, &decoded_packet_a.message) + // .expect("A decrypt failed"); + // assert_eq!(decrypted_payload, plaintext_b); + // session_manager_1 + // .receiving_counter_mark(peer_a_sm, current_counter_b) + // .expect("A mark counter failed"); + // } + // + // // 5. Verify counter stats + // // Note: current_packet_cnt() returns (next_expected_receive_counter, total_received) + // let (next_recv_a, total_recv_a) = session_manager_1.current_packet_cnt(peer_a_sm).unwrap(); + // let (next_recv_b, total_recv_b) = session_manager_2.current_packet_cnt(peer_b_sm).unwrap(); + // + // // Peer A sent handshake(0), handshake(1) + 5 data packets = 7 total. Next send counter = 7. + // // Peer A received handshake(0) + 5 data packets = 6 total. Next expected recv counter = 6. + // assert_eq!( + // counter_a, NUM_MESSAGES, + // "Peer A final send counter mismatch" + // ); + // assert_eq!( + // total_recv_a, NUM_MESSAGES, + // "Peer A total received count mismatch" + // ); // Received 5 data + // assert_eq!( + // next_recv_a, NUM_MESSAGES, + // "Peer A next expected receive counter mismatch" + // ); // Expected counter for msg from B + // + // // Peer B sent handshake(0) + 5 data packets = 6 total. Next send counter = 6. + // // Peer B received handshake(0), handshake(1) + 5 data packets = 7 total. Next expected recv counter = 7. + // assert_eq!( + // counter_b, NUM_MESSAGES, + // "Peer B final send counter mismatch" + // ); + // assert_eq!( + // total_recv_b, NUM_MESSAGES, + // "Peer B total received count mismatch" + // ); // Received 5 data + // assert_eq!( + // next_recv_b, NUM_MESSAGES, + // "Peer B next expected receive counter mismatch" + // ); // Expected counter for msg from A + // + // println!("Bidirectional test completed."); + // } + // + // /// Tests error handling in session flow + // #[test] + // fn test_session_error_handling() { + // // 1. Initialize session manager + // let mut session_manager = SessionManager::new(); + // + // let TODO = " for kem in kem_list() {"; + // let receiver_index = 123; + // let session1 = SessionsMock::mock_post_handshake(receiver_index).initiator; + // let session2 = SessionsMock::mock_post_handshake(124).initiator; + // + // // 2. Create a session (using real noise state) + // let _session = session_manager.create_session_state_machine(session1); + // + // // 3. Try to get a non-existent session + // let result = session_manager.state_machine_exists(999); + // assert!(!result, "Non-existent session should return None"); + // + // // 4. Try to remove a non-existent session + // let result = session_manager.remove_state_machine(999); + // assert!( + // !result, + // "Remove session should not remove a non-existent session" + // ); + // + // // 5. Create and immediately remove a session + // let _temp_session = session_manager.create_session_state_machine(session2); + // + // assert!( + // session_manager.remove_state_machine(124), + // "Should remove the session" + // ); + // + // // 6. Create a codec and test error cases + // // let mut codec = LPCodec::new(session); + // + // // 7. Create an invalid message type packet + // let mut buf = BytesMut::new(); + // + // // Add header + // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved + // buf.extend_from_slice(&receiver_index.to_le_bytes()); // Sender index + // buf.extend_from_slice(&0u64.to_le_bytes()); // Counter + // + // // Add invalid message type + // buf.extend_from_slice(&0xFFFFu16.to_le_bytes()); + // + // // Add some dummy data + // buf.extend_from_slice(&[0u8; 80]); + // + // // Add trailer + // buf.extend_from_slice(&[0u8; TRAILER_LEN]); + // + // // Try to parse the invalid message type + // let result = parse_lp_packet(&buf, None); + // assert!(result.is_err(), "Decoding invalid message type should fail"); + // + // // Add assertion for the specific error type + // assert!(matches!( + // result.unwrap_err(), + // LpError::InvalidMessageType(0xFFFF) + // )); + // + // // 8. Test partial packet decoding + // let partial_packet = &buf[0..10]; // Too short to be a valid packet + // let partial_bytes = BytesMut::from(partial_packet); + // + // let result = parse_lp_packet(&partial_bytes, None); + // assert!(result.is_err(), "Parsing partial packet should fail"); + // assert!(matches!( + // result.unwrap_err(), + // LpError::InsufficientBufferSize + // )); + // } + // // Remove unused imports if SessionManager methods are no longer direct dependencies + // // use crate::noise_protocol::{create_noise_state, create_noise_state_responder}; + // use crate::state_machine::LpData; + // use crate::state_machine::{LpAction, LpInput, LpStateBare}; + // // Use Bytes for SendData input + // + // // Keep helper function for creating test packets if needed, + // // but LpAction::SendPacket should provide the packets now. + // // fn create_test_packet(...) -> LpPacket { ... } + // + // /// Tests the complete session flow using ONLY the process_input interface: + // /// - Creation of sessions through session manager + // /// - Handshake driven by StartHandshake, ReceivePacket inputs + // /// - Data transfer driven by SendData, ReceivePacket inputs + // /// - Actions like SendPacket, DeliverData handled from output + // /// - Implicit replay protection via state machine logic + // /// - Closing driven by Close input + // #[test] + // fn test_full_session_flow_with_process_input() { + // // 1. Initialize session managers + // let mut session_manager_1 = SessionManager::new(); + // let mut session_manager_2 = SessionManager::new(); + // + // let TODO = " for kem in kem_list() {"; + // + // let receiver_index = 12345; + // let sessions = SessionsMock::mock_post_handshake(receiver_index); + // + // // 2. Create sessions state machines + // session_manager_1.create_session_state_machine(sessions.initiator); + // session_manager_2.create_session_state_machine(sessions.responder); + // + // assert_eq!(session_manager_1.session_count(), 1); + // assert_eq!(session_manager_2.session_count(), 1); + // assert!(session_manager_1.state_machine_exists(receiver_index)); + // assert!(session_manager_2.state_machine_exists(receiver_index)); + // + // // Verify initial states are Transport + // assert_eq!( + // session_manager_1.get_state(receiver_index).unwrap(), + // LpStateBare::Transport + // ); + // assert_eq!( + // session_manager_2.get_state(receiver_index).unwrap(), + // LpStateBare::Transport + // ); + // + // // --- 3. Simulate Data Transfer via process_input --- + // println!("Starting data transfer simulation via process_input..."); + // let plaintext_a_to_b = LpData::new_opaque(b"Hello from A via process_input!".to_vec()); + // let plaintext_b_to_a = LpData::new_opaque(b"Hello from B via process_input!".to_vec()); + // + // // --- A sends to B --- + // println!(" A sends to B"); + // let action_a_send = session_manager_1 + // .process_input(receiver_index, LpInput::SendData(plaintext_a_to_b.clone())) + // .expect("A SendData should produce action") + // .expect("A SendData failed"); + // + // let data_packet_a = if let LpAction::SendPacket(packet) = action_a_send { + // packet + // } else { + // panic!("A SendData did not produce SendPacket"); + // }; + // + // // Simulate network + // let mut buf_data_a = BytesMut::new(); + // serialize_lp_packet(&data_packet_a, &mut buf_data_a, None).unwrap(); + // let parsed_data_a = parse_lp_packet(&buf_data_a, None).unwrap(); + // + // // B receives + // println!(" B receives from A"); + // let action_b_recv = session_manager_2 + // .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_a)) + // .expect("B ReceivePacket (data) should produce action") + // .expect("B ReceivePacket (data) failed"); + // + // if let LpAction::DeliverData(data) = action_b_recv { + // assert_eq!(data, plaintext_a_to_b, "Decrypted data mismatch A->B"); + // println!( + // " B successfully decrypted: {:?}", + // String::from_utf8_lossy(&data.content) + // ); + // } else { + // panic!("B ReceivePacket did not produce DeliverData"); + // } + // + // // --- B sends to A --- + // println!(" B sends to A"); + // let action_b_send = session_manager_2 + // .process_input(receiver_index, LpInput::SendData(plaintext_b_to_a.clone())) + // .expect("B SendData should produce action") + // .expect("B SendData failed"); + // + // let data_packet_b = if let LpAction::SendPacket(packet) = action_b_send { + // packet + // } else { + // panic!("B SendData did not produce SendPacket"); + // }; + // // Keep a copy for replay test + // let data_packet_b_replay = data_packet_b.clone(); + // + // // Simulate network + // let mut buf_data_b = BytesMut::new(); + // serialize_lp_packet(&data_packet_b, &mut buf_data_b, None).unwrap(); + // let parsed_data_b = parse_lp_packet(&buf_data_b, None).unwrap(); + // + // // A receives + // println!(" A receives from B"); + // let action_a_recv = session_manager_1 + // .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_b)) + // .expect("A ReceivePacket (data) should produce action") + // .expect("A ReceivePacket (data) failed"); + // + // if let LpAction::DeliverData(data) = action_a_recv { + // assert_eq!(data, plaintext_b_to_a, "Decrypted data mismatch B->A"); + // println!( + // " A successfully decrypted: {:?}", + // String::from_utf8_lossy(&data.content) + // ); + // } else { + // panic!("A ReceivePacket did not produce DeliverData"); + // } + // println!("Data transfer simulation completed."); + // + // // --- 4. Replay Protection Test --- + // println!("Testing data packet replay protection via process_input..."); + // let replay_result = session_manager_1 + // .process_input(receiver_index, LpInput::ReceivePacket(data_packet_b_replay)); // Use cloned packet + // + // assert!(replay_result.is_err(), "Replay should produce Err(...)"); + // let error = replay_result.err().unwrap(); + // assert!( + // matches!(error, LpError::Replay(_)), + // "Expected Replay error, got {:?}", + // error + // ); + // println!("Data packet replay protection test passed."); + // + // // --- 5. Out-of-Order Test --- + // println!("Testing out-of-order reception via process_input..."); + // + // // A prepares N+1 then N + // let data_n_plus_1 = LpData::new_opaque(b"Message N+1".to_vec()); + // let data_n = LpData::new_opaque(b"Message N".to_vec()); + // + // let action_send_n1 = session_manager_1 + // .process_input(receiver_index, LpInput::SendData(data_n_plus_1.clone())) + // .unwrap() + // .unwrap(); + // let packet_n1 = match action_send_n1 { + // LpAction::SendPacket(p) => p, + // _ => panic!("Expected SendPacket"), + // }; + // + // let action_send_n = session_manager_1 + // .process_input(receiver_index, LpInput::SendData(data_n.clone())) + // .unwrap() + // .unwrap(); + // let packet_n = match action_send_n { + // LpAction::SendPacket(p) => p, + // _ => panic!("Expected SendPacket"), + // }; + // let packet_n_replay = packet_n.clone(); // For replay test + // + // // B receives N+1 first + // println!(" B receives N+1"); + // let action_recv_n1 = session_manager_2 + // .process_input(receiver_index, LpInput::ReceivePacket(packet_n1)) + // .unwrap() + // .unwrap(); + // match action_recv_n1 { + // LpAction::DeliverData(d) => assert_eq!(d, data_n_plus_1, "Data N+1 mismatch"), + // _ => panic!("Expected DeliverData for N+1"), + // } + // + // // B receives N second (should work) + // println!(" B receives N"); + // let action_recv_n = session_manager_2 + // .process_input(receiver_index, LpInput::ReceivePacket(packet_n)) + // .unwrap() + // .unwrap(); + // match action_recv_n { + // LpAction::DeliverData(d) => assert_eq!(d, data_n, "Data N mismatch"), + // _ => panic!("Expected DeliverData for N"), + // } + // + // // B tries to replay N (should fail) + // println!(" B tries to replay N"); + // let replay_n_result = session_manager_2 + // .process_input(receiver_index, LpInput::ReceivePacket(packet_n_replay)); + // assert!(replay_n_result.is_err(), "Replay N should produce Err"); + // assert!( + // matches!(replay_n_result.err().unwrap(), LpError::Replay(_)), + // "Expected Replay error for N" + // ); + // println!("Out-of-order test passed."); + // + // // --- 6. Close Test --- + // println!("Testing close via process_input..."); + // + // // A closes + // let action_a_close = session_manager_1 + // .process_input(receiver_index, LpInput::Close) + // .expect("A Close should produce action") + // .expect("A Close failed"); + // assert!(matches!(action_a_close, LpAction::ConnectionClosed)); + // assert_eq!( + // session_manager_1.get_state(receiver_index).unwrap(), + // LpStateBare::Closed + // ); + // + // // Further actions on A fail + // let send_after_close_a = session_manager_1.process_input( + // receiver_index, + // LpInput::SendData(LpData::new_opaque(b"fail".to_vec())), + // ); + // assert!(send_after_close_a.is_err()); + // assert!(matches!( + // send_after_close_a.err().unwrap(), + // LpError::LpSessionClosed + // )); + // + // // B closes + // let action_b_close = session_manager_2 + // .process_input(receiver_index, LpInput::Close) + // .expect("B Close should produce action") + // .expect("B Close failed"); + // assert!(matches!(action_b_close, LpAction::ConnectionClosed)); + // assert_eq!( + // session_manager_2.get_state(receiver_index).unwrap(), + // LpStateBare::Closed + // ); + // + // // Further actions on B fail + // let send_after_close_b = session_manager_2.process_input( + // receiver_index, + // LpInput::SendData(LpData::new_opaque(b"fail".to_vec())), + // ); + // assert!(send_after_close_b.is_err()); + // assert!(matches!( + // send_after_close_b.err().unwrap(), + // LpError::LpSessionClosed + // )); + // println!("Close test passed."); + // + // // --- 7. Session Removal --- + // assert!(session_manager_1.remove_state_machine(receiver_index)); + // assert_eq!(session_manager_1.session_count(), 0); + // assert!(!session_manager_1.state_machine_exists(receiver_index)); + // + // // B's session manager still has it until removed + // assert!(session_manager_2.state_machine_exists(receiver_index)); + // assert!(session_manager_2.remove_state_machine(receiver_index)); + // assert_eq!(session_manager_2.session_count(), 0); + // assert!(!session_manager_2.state_machine_exists(receiver_index)); + // println!("Session removal test passed."); + // } // ... other tests ... } diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index 2bb18f956a..4b9e89f3ef 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -6,20 +6,17 @@ //! This module implements session lifecycle management functionality, handling //! creation, retrieval, and storage of sessions. -use crate::peer::{LpLocalPeer, LpRemotePeer}; -use crate::state_machine::{LpAction, LpInput, LpState, LpStateBare}; -use crate::{LpError, LpMessage, LpSession, LpStateMachine}; +use crate::session::SessionId; +use crate::state_machine::{LpAction, LpInput, LpStateBare}; +use crate::{LpError, LpSession, LpStateMachine}; use std::collections::HashMap; -#[cfg(test)] -use libcrux_psq::handshake::types::DHPublicKey; - /// Manages the lifecycle of Lewes Protocol sessions. /// /// The SessionManager is responsible for creating, storing, and retrieving sessions pub struct SessionManager { /// Manages state machines directly, keyed by lp_id - state_machines: HashMap, + state_machines: HashMap, } impl Default for SessionManager { @@ -38,62 +35,54 @@ impl SessionManager { pub fn process_input( &mut self, - lp_id: u32, + lp_id: SessionId, input: LpInput, ) -> Result, LpError> { self.with_state_machine_mut(lp_id, |sm| sm.process_input(input).transpose())? } - pub fn closed(&self, lp_id: u32) -> Result { + pub fn closed(&self, lp_id: SessionId) -> Result { Ok(self.get_state(lp_id)? == LpStateBare::Closed) } - pub fn transport(&self, lp_id: u32) -> Result { + pub fn transport(&self, lp_id: SessionId) -> Result { Ok(self.get_state(lp_id)? == LpStateBare::Transport) } #[cfg(test)] - fn get_state_machine_id(&self, lp_id: u32) -> Result { + fn get_state_machine_id(&self, lp_id: SessionId) -> Result { self.with_state_machine(lp_id, |sm| sm.id())? } - pub fn get_state(&self, lp_id: u32) -> Result { + pub fn get_state(&self, lp_id: SessionId) -> Result { self.with_state_machine(lp_id, |sm| Ok(sm.bare_state()))? } - pub fn receiving_counter_quick_check(&self, lp_id: u32, counter: u64) -> Result<(), LpError> { + pub fn receiving_counter_quick_check( + &self, + lp_id: SessionId, + counter: u64, + ) -> Result<(), LpError> { self.with_state_machine(lp_id, |sm| { sm.session()?.receiving_counter_quick_check(counter) })? } - pub fn receiving_counter_mark(&mut self, lp_id: u32, counter: u64) -> Result<(), LpError> { + pub fn receiving_counter_mark( + &mut self, + lp_id: SessionId, + counter: u64, + ) -> Result<(), LpError> { self.with_state_machine_mut(lp_id, |sm| { sm.session_mut()?.receiving_counter_mark(counter) })? } - pub fn next_counter(&mut self, lp_id: u32) -> Result { + pub fn next_counter(&mut self, lp_id: SessionId) -> Result { self.with_state_machine_mut(lp_id, |sm| Ok(sm.session_mut()?.next_counter()))? } - pub fn decrypt_data(&mut self, lp_id: u32, message: &LpMessage) -> Result, LpError> { - self.with_state_machine_mut(lp_id, |sm| { - sm.session_mut()? - .decrypt_data(message) - .map_err(LpError::NoiseError) - })? - } - - pub fn encrypt_data(&mut self, lp_id: u32, message: &[u8]) -> Result { - self.with_state_machine_mut(lp_id, |sm| { - sm.session_mut()? - .encrypt_data(message) - .map_err(LpError::NoiseError) - })? - } - - pub fn current_packet_cnt(&self, lp_id: u32) -> Result<(u64, u64), LpError> { + pub fn current_packet_cnt(&self, lp_id: SessionId) -> Result<(u64, u64), LpError> { self.with_state_machine(lp_id, |sm| Ok(sm.session()?.current_packet_cnt()))? } @@ -101,11 +90,11 @@ impl SessionManager { self.state_machines.len() } - pub fn state_machine_exists(&self, lp_id: u32) -> bool { + pub fn state_machine_exists(&self, lp_id: SessionId) -> bool { self.state_machines.contains_key(&lp_id) } - pub fn with_state_machine(&self, lp_id: u32, f: F) -> Result + pub fn with_state_machine(&self, lp_id: SessionId, f: F) -> Result where F: FnOnce(&LpStateMachine) -> R, { @@ -117,7 +106,7 @@ impl SessionManager { } // For mutable access (like running process_input) - pub fn with_state_machine_mut(&mut self, lp_id: u32, f: F) -> Result + pub fn with_state_machine_mut(&mut self, lp_id: SessionId, f: F) -> Result where F: FnOnce(&mut LpStateMachine) -> R, // Closure takes mutable ref { @@ -128,15 +117,15 @@ impl SessionManager { } } - pub fn create_session_state_machine(&mut self, lp_session: LpSession) -> u32 { - let receiver_index = lp_session.id(); + pub fn create_session_state_machine(&mut self, lp_session: LpSession) -> SessionId { + let session_id = *lp_session.session_identifier(); let sm = LpStateMachine::new(lp_session); - self.state_machines.insert(receiver_index, sm); - receiver_index + self.state_machines.insert(session_id, sm); + session_id } /// Method to remove a state machine - pub fn remove_state_machine(&mut self, lp_id: u32) -> bool { + pub fn remove_state_machine(&mut self, lp_id: SessionId) -> bool { let removed = self.state_machines.remove(&lp_id); removed.is_some() diff --git a/common/nym-lp/src/state_machine.rs b/common/nym-lp/src/state_machine.rs index e7a40d3466..b31fd3456d 100644 --- a/common/nym-lp/src/state_machine.rs +++ b/common/nym-lp/src/state_machine.rs @@ -13,37 +13,27 @@ //! State machine ensures protocol steps execute in correct order. Invalid transitions //! return LpError, preventing protocol violations. -use crate::{ - LpError, - message::{LpMessage, SubsessionKK1Data, SubsessionKK2Data, SubsessionReadyData}, - noise_protocol::NoiseError, - packet::LpPacket, - session::{LpSession, SubsessionHandshake}, -}; +use crate::packet::EncryptedLpPacket; +use crate::{LpError, message::LpMessage, packet::LpPacket, session::LpSession}; use bytes::{Buf, Bytes}; use num_enum::{IntoPrimitive, TryFromPrimitive}; use std::mem; -use tracing::debug; + +#[derive(Debug)] +pub struct LpTransportState { + /// The underlying session in the transport state + session: Box, +} /// Represents the possible states of the Lewes Protocol connection. #[derive(Debug, Default)] pub enum LpState { /// Handshake complete, ready for data transport. - Transport { session: Box }, - - /// Performing subsession KK handshake while parent remains active. - /// Parent can still send/receive; subsession messages tunneled through parent. - SubsessionHandshaking { - session: Box, - subsession: Box, - }, - - /// Parent session demoted after subsession promoted. - /// Can only receive (drain in-flight), cannot send. - ReadOnlyTransport { session: Box }, + Transport(LpTransportState), /// An error occurred, or the connection was intentionally closed. Closed { reason: String }, + /// Processing an input event. #[default] Processing, @@ -52,8 +42,6 @@ pub enum LpState { #[derive(Debug, Clone, PartialEq, Eq)] pub enum LpStateBare { Transport, - SubsessionHandshaking, - ReadOnlyTransport, Closed, Processing, } @@ -62,8 +50,6 @@ impl From<&LpState> for LpStateBare { fn from(state: &LpState) -> Self { match state { LpState::Transport { .. } => LpStateBare::Transport, - LpState::SubsessionHandshaking { .. } => LpStateBare::SubsessionHandshaking, - LpState::ReadOnlyTransport { .. } => LpStateBare::ReadOnlyTransport, LpState::Closed { .. } => LpStateBare::Closed, LpState::Processing => LpStateBare::Processing, } @@ -76,38 +62,25 @@ impl From<&LpState> for LpStateBare { pub enum LpInput { /// Received an LP Packet from the network. ReceivePacket(LpPacket), + /// Application wants to send data (only valid in Transport state). SendData(LpData), + /// Close the connection. Close, - /// Initiate a subsession handshake (only valid in Transport state). - /// Creates SubsessionHandshake and sends KK1 message. - InitiateSubsession, } /// Represents actions the state machine requests the environment to perform. #[derive(Debug)] pub enum LpAction { /// Send an LP Packet over the network. - SendPacket(LpPacket), + SendPacket(EncryptedLpPacket), + /// Deliver decrypted application data received from the peer. DeliverData(LpData), + /// Inform the environment that the connection is closed. ConnectionClosed, - /// Subsession KK handshake initiated by this side. - /// Contains the KK1 packet to send and the subsession index for tracking. - SubsessionInitiated { - packet: LpPacket, - subsession_index: u64, - }, - /// Subsession handshake complete, ready for promotion. - /// Contains the packet to send (Some for initiator with SubsessionReady, None for responder), - /// the completed SubsessionHandshake for into_session(), and the new receiver_index. - SubsessionComplete { - packet: Option, - subsession: Box, - new_receiver_index: u32, - }, } /// Represent application data being sent in Transport mode @@ -191,9 +164,7 @@ impl LpStateMachine { pub fn session_mut(&mut self) -> Result<&mut LpSession, LpError> { match &mut self.state { - LpState::Transport { session } - | LpState::SubsessionHandshaking { session, .. } - | LpState::ReadOnlyTransport { session } => Ok(session), + LpState::Transport(transport) => Ok(&mut transport.session), LpState::Closed { .. } => Err(LpError::LpSessionClosed), LpState::Processing => Err(LpError::LpSessionProcessing), } @@ -201,9 +172,7 @@ impl LpStateMachine { pub fn session(&self) -> Result<&LpSession, LpError> { match &self.state { - LpState::Transport { session } - | LpState::SubsessionHandshaking { session, .. } - | LpState::ReadOnlyTransport { session } => Ok(session), + LpState::Transport(transport) => Ok(&transport.session), LpState::Closed { .. } => Err(LpError::LpSessionClosed), LpState::Processing => Err(LpError::LpSessionProcessing), } @@ -214,50 +183,103 @@ impl LpStateMachine { /// ownership of the session to the caller. pub fn into_session(self) -> Result { match self.state { - LpState::Transport { session } - | LpState::SubsessionHandshaking { session, .. } - | LpState::ReadOnlyTransport { session } => Ok(*session), + LpState::Transport(transport) => Ok(*transport.session), LpState::Closed { .. } => Err(LpError::LpSessionClosed), LpState::Processing => Err(LpError::LpSessionProcessing), } } pub fn id(&self) -> Result { - Ok(self.session()?.id()) + todo!() + // Ok(self.session()?.id()) } /// Creates a new state machine in `Transport` state post-KKT/PSQ handshake pub fn new(session: LpSession) -> Self { LpStateMachine { - state: LpState::Transport { + state: LpState::Transport(LpTransportState { session: Box::new(session), - }, + }), } } - /// Creates a state machine in Transport state from a completed subsession handshake. - /// - /// This is used when a subsession (rekeying) completes and we need a new state machine - /// for the promoted session that can handle further subsession initiations (chained rekeying). - /// - /// # Arguments - /// - /// * `subsession` - The completed subsession handshake - /// * `receiver_index` - The new session's receiver index - /// - /// # Errors - /// - /// Returns error if the subsession handshake is not complete. - pub fn from_subsession( - subsession: SubsessionHandshake, - receiver_index: u32, - ) -> Result { - let session = subsession.into_session(receiver_index)?; - Ok(LpStateMachine { - state: LpState::Transport { - session: Box::new(session), - }, - }) + fn process_input_transport( + &mut self, + mut state: LpTransportState, + input: LpInput, + ) -> (LpState, Option>) { + let mut session = &mut state.session; + match input { + LpInput::ReceivePacket(packet) => { + // Check if packet lp_id matches our session + if packet.header.receiver_idx() != session.id() { + let result_action = + Some(Err(LpError::UnknownSessionId(packet.header.receiver_idx()))); + return (LpState::Transport(state), result_action); + } + + let ctr = packet.header.counter(); + // Check message type + match packet.message { + // Normal encrypted data + LpMessage::ApplicationData(payload) => { + // 1. Check replay protection + if let Err(e) = session.receiving_counter_quick_check(ctr) { + return (LpState::Transport(state), Some(Err(e))); + } + + // 2. Mark counter as received + if let Err(e) = session.receiving_counter_mark(ctr) { + return (LpState::Transport(state), Some(Err(e))); + } + + // 3. Deliver data + match payload.0.try_into() { + Ok(data) => { + let result_action = Some(Ok(LpAction::DeliverData(data))); + (LpState::Transport(state), result_action) + } + Err(e) => { + let reason = e.to_string(); + (LpState::Closed { reason }, Some(Err(e))) + } + } + } + other => { + // Unexpected message type in Transport state + let err = LpError::InvalidStateTransition { + state: "Transport".to_string(), + input: format!("Unexpected message type: {other}"), + }; + (LpState::Transport(state), Some(Err(err))) + } + } + } + LpInput::SendData(data) => { + // Encrypt and send application data + let result_action = match self.prepare_data_packet(&mut session, data) { + Ok(packet) => Some(Ok(LpAction::SendPacket(packet))), + Err(e) => { + // If prepare fails, should we close? Let's report error and stay Transport for now. + // Alternative: transition to Closed state. + Some(Err(e.into())) + } + }; + // Remain in transport state + (LpState::Transport(state), result_action) + } + + // --- Close Transition --- + LpInput::Close => { + // Transition to Closed state + ( + LpState::Closed { + reason: "Closed by user".to_string(), + }, + Some(Ok(LpAction::ConnectionClosed)), + ) + } + } } /// Processes an input event and returns a list of actions to perform. @@ -270,606 +292,10 @@ impl LpStateMachine { // 2. Match on the owned current_state. Each arm calculates and returns the NEXT state. let next_state = match (current_state, input) { // --- Transport State --- - (LpState::Transport { mut session }, LpInput::ReceivePacket(packet)) => { - // Check if packet lp_id matches our session - if packet.header.receiver_idx() != session.id() { - result_action = - Some(Err(LpError::UnknownSessionId(packet.header.receiver_idx()))); - LpState::Transport { session } - } else { - // Check message type - handle subsession initiation from peer - match &packet.message { - // Peer initiated subsession - we become responder - LpMessage::SubsessionKK1(kk1_data) => { - // Create subsession as responder - let subsession_index = session.next_subsession_index(); - match session.create_subsession(subsession_index, false) { - Ok(subsession) => { - // Process KK1 - match subsession.process_message(&kk1_data.payload) { - Ok(_) => { - // Prepare KK2 response - match subsession.prepare_message() { - Ok(kk2_payload) => { - let kk2_msg = LpMessage::SubsessionKK2( - SubsessionKK2Data { - payload: kk2_payload, - }, - ); - match session.next_packet(kk2_msg) { - Ok(response_packet) => { - result_action = - Some(Ok(LpAction::SendPacket( - response_packet, - ))); - // Stay in SubsessionHandshaking, wait for SubsessionReady - LpState::SubsessionHandshaking { - session, - subsession: Box::new(subsession), - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - // Normal encrypted data - LpMessage::EncryptedData(_) => { - // 1. Check replay protection - if let Err(e) = - session.receiving_counter_quick_check(packet.header.counter) - { - result_action = Some(Err(e)); - LpState::Transport { session } - } else { - // 2. Decrypt data - match session.decrypt_data(&packet.message) { - Ok(plaintext) => { - // 3. Mark counter as received - if let Err(e) = - session.receiving_counter_mark(packet.header.counter) - { - result_action = Some(Err(e)); - LpState::Transport { session } - } else { - // 4. Deliver data - match plaintext.try_into() { - Ok(data) => { - result_action = - Some(Ok(LpAction::DeliverData(data))); - LpState::Transport { session } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e.into())); - LpState::Closed { reason } - } - } - } - } - // Stale abort in Transport state - race already resolved. - // This can happen if abort arrives after loser already returned to Transport - // via KK1 processing (loser detected local < remote and became responder). - // The winner's abort message arrived late. Silently ignore. - LpMessage::SubsessionAbort => { - debug!("Ignoring stale SubsessionAbort in Transport state"); - result_action = None; - LpState::Transport { session } - } - _ => { - // Unexpected message type in Transport state - let err = LpError::InvalidStateTransition { - state: "Transport".to_string(), - input: format!("Unexpected message type: {}", packet.message), - }; - result_action = Some(Err(err)); - LpState::Transport { session } - } - } - } - } - (LpState::Transport { mut session }, LpInput::SendData(data)) => { - // Encrypt and send application data - match self.prepare_data_packet(&mut session, data) { - Ok(packet) => result_action = Some(Ok(LpAction::SendPacket(packet))), - Err(e) => { - // If prepare fails, should we close? Let's report error and stay Transport for now. - // Alternative: transition to Closed state. - result_action = Some(Err(e.into())); - } - } - // Remain in transport state - LpState::Transport { session } - } - - // --- Transport + InitiateSubsession → SubsessionHandshaking --- - (LpState::Transport { mut session }, LpInput::InitiateSubsession) => { - // Get next subsession index - let subsession_index = session.next_subsession_index(); - - // Create subsession handshake (this side is initiator) - match session.create_subsession(subsession_index, true) { - Ok(subsession) => { - // Prepare KK1 message - match subsession.prepare_message() { - Ok(kk1_payload) => { - let kk1_msg = LpMessage::SubsessionKK1(SubsessionKK1Data { - payload: kk1_payload, - }); - match session.next_packet(kk1_msg) { - Ok(packet) => { - // Emit SubsessionInitiated with packet and index - result_action = Some(Ok(LpAction::SubsessionInitiated { - packet, - subsession_index, - })); - LpState::SubsessionHandshaking { - session, - subsession: Box::new(subsession), - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - - // --- SubsessionHandshaking State --- - ( - LpState::SubsessionHandshaking { - mut session, - subsession, - }, - LpInput::ReceivePacket(packet), - ) => { - // Check if packet receiver_idx matches our session - if packet.header.receiver_idx() != session.id() { - result_action = - Some(Err(LpError::UnknownSessionId(packet.header.receiver_idx()))); - LpState::SubsessionHandshaking { - session, - subsession, - } - } else { - match &packet.message { - LpMessage::SubsessionKK1(kk1_data) if !subsession.is_initiator() => { - // Responder processes KK1, prepares KK2 - // Responder stays in SubsessionHandshaking after sending KK2, - // waiting for SubsessionReady from initiator before completing - match subsession.process_message(&kk1_data.payload) { - Ok(_) => { - match subsession.prepare_message() { - Ok(kk2_payload) => { - let kk2_msg = - LpMessage::SubsessionKK2(SubsessionKK2Data { - payload: kk2_payload, - }); - match session.next_packet(kk2_msg) { - Ok(response_packet) => { - result_action = Some(Ok(LpAction::SendPacket( - response_packet, - ))); - // Stay in SubsessionHandshaking, wait for SubsessionReady - LpState::SubsessionHandshaking { - session, - subsession, - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - LpMessage::SubsessionKK1(kk1_data) if subsession.is_initiator() => { - // Simultaneous initiation race detected. - // Both sides called InitiateSubsession and sent KK1 to each other. - // Use X25519 public key comparison as deterministic tie-breaker. - // Lower key loses and becomes responder. - let local_key = session.local_x25519_public(); - let remote_key = session.remote_x25519_public(); - - if local_key.as_ref() < remote_key.as_ref() { - // We LOSE - become responder - // Use the same index as our initiator subsession, which should - // match the winner's index if subsession counters are in sync. - // This works because both sides independently picked the same index when - // they initiated simultaneously (both counters were at the same value). - let subsession_index = subsession.index; - match session.create_subsession(subsession_index, false) { - Ok(new_subsession) => { - match new_subsession.process_message(&kk1_data.payload) { - Ok(_) => { - match new_subsession.prepare_message() { - Ok(kk2_payload) => { - let kk2_msg = LpMessage::SubsessionKK2( - SubsessionKK2Data { - payload: kk2_payload, - }, - ); - match session.next_packet(kk2_msg) { - Ok(response_packet) => { - result_action = - Some(Ok(LpAction::SendPacket( - response_packet, - ))); - // Replace old initiator subsession with new responder subsession - LpState::SubsessionHandshaking { - session, - subsession: Box::new( - new_subsession, - ), - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } else { - // We WIN - stay initiator, notify peer they lost - // Send SubsessionAbort to explicitly tell peer to become responder - let abort_msg = LpMessage::SubsessionAbort; - match session.next_packet(abort_msg) { - Ok(abort_packet) => { - result_action = - Some(Ok(LpAction::SendPacket(abort_packet))); - LpState::SubsessionHandshaking { - session, - subsession, - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - } - LpMessage::SubsessionKK2(kk2_data) if subsession.is_initiator() => { - // Initiator processes KK2, completes handshake - // Initiator emits SubsessionComplete with SubsessionReady packet - // and the subsession for caller to promote via into_session() - match subsession.process_message(&kk2_data.payload) { - Ok(_) if subsession.is_complete() => { - // Generate new receiver_index for subsession - let new_receiver_index: u32 = rand::random(); - session.demote(new_receiver_index); - - // Send SubsessionReady with new index - let ready_msg = - LpMessage::SubsessionReady(SubsessionReadyData { - receiver_index: new_receiver_index, - }); - match session.next_packet(ready_msg) { - Ok(ready_packet) => { - result_action = - Some(Ok(LpAction::SubsessionComplete { - packet: Some(ready_packet), - subsession, - new_receiver_index, - })); - LpState::ReadOnlyTransport { session } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Ok(_) => { - // Handshake not complete yet, shouldn't happen for KK - let err = LpError::Internal( - "Subsession handshake incomplete after KK2".to_string(), - ); - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - LpMessage::EncryptedData(_) => { - // Parent still processes normal traffic during subsession handshake - // Same as Transport state handling - if let Err(e) = - session.receiving_counter_quick_check(packet.header.counter) - { - result_action = Some(Err(e)); - LpState::SubsessionHandshaking { - session, - subsession, - } - } else { - match session.decrypt_data(&packet.message) { - Ok(plaintext) => { - if let Err(e) = - session.receiving_counter_mark(packet.header.counter) - { - result_action = Some(Err(e)); - LpState::SubsessionHandshaking { - session, - subsession, - } - } else { - match plaintext.try_into() { - Ok(data) => { - result_action = - Some(Ok(LpAction::DeliverData(data))); - LpState::SubsessionHandshaking { - session, - subsession, - } - } - Err(err) => { - result_action = Some(Err(err)); - LpState::SubsessionHandshaking { - session, - subsession, - } - } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e.into())); - LpState::Closed { reason } - } - } - } - } - LpMessage::SubsessionReady(ready_data) if !subsession.is_initiator() => { - // Responder receives SubsessionReady from initiator - // Responder completes handshake here, uses initiator's receiver_index - // The subsession handshake should already be complete (after KK2) - if subsession.is_complete() { - let new_receiver_index = ready_data.receiver_index; - session.demote(new_receiver_index); - result_action = Some(Ok(LpAction::SubsessionComplete { - packet: None, // Responder has no packet to send - subsession, - new_receiver_index, - })); - LpState::ReadOnlyTransport { session } - } else { - // Shouldn't happen - handshake should be complete after KK2 - let err = LpError::Internal( - "Received SubsessionReady but handshake not complete" - .to_string(), - ); - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - } - LpMessage::SubsessionAbort if subsession.is_initiator() => { - // We received abort from peer - we lost the simultaneous initiation race. - // Peer has higher X25519 key and is staying as initiator. - // Discard our initiator subsession and return to Transport to receive peer's KK1. - // Peer's KK1 should already be in flight or queued. - result_action = None; - LpState::Transport { session } - } - LpMessage::SubsessionAbort if !subsession.is_initiator() => { - // Race was already resolved via KK1 - this abort is stale. - // We already became responder when we received KK1 and detected local < remote. - // The winner's abort message arrived after we processed their KK1. - // Silently ignore it - we're in the correct state. - result_action = None; - LpState::SubsessionHandshaking { - session, - subsession, - } - } - _ => { - // Wrong message type for subsession handshake - let err = LpError::InvalidStateTransition { - state: "SubsessionHandshaking".to_string(), - input: format!("Unexpected message type: {:?}", packet.message), - }; - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - } - } - } - - // Parent can still send data during subsession handshake - ( - LpState::SubsessionHandshaking { - mut session, - subsession, - }, - LpInput::SendData(data), - ) => { - match self.prepare_data_packet(&mut session, data) { - Ok(packet) => result_action = Some(Ok(LpAction::SendPacket(packet))), - Err(e) => { - result_action = Some(Err(e.into())); - } - } - LpState::SubsessionHandshaking { - session, - subsession, - } - } - - // Reject other inputs during subsession handshake - ( - LpState::SubsessionHandshaking { - session, - subsession, - }, - LpInput::InitiateSubsession, - ) => { - result_action = Some(Err(LpError::InvalidStateTransition { - state: "SubsessionHandshaking".to_string(), - input: "InitiateSubsession".to_string(), - })); - LpState::SubsessionHandshaking { - session, - subsession, - } - } - - // --- ReadOnlyTransport State --- - (LpState::ReadOnlyTransport { mut session }, LpInput::ReceivePacket(packet)) => { - // Can still receive and decrypt, but state stays ReadOnlyTransport - if packet.header.receiver_idx() != session.id() { - result_action = - Some(Err(LpError::UnknownSessionId(packet.header.receiver_idx()))); - LpState::ReadOnlyTransport { session } - } else if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) - { - result_action = Some(Err(e)); - LpState::ReadOnlyTransport { session } - } else { - match session.decrypt_data(&packet.message) { - Ok(plaintext) => { - if let Err(e) = session.receiving_counter_mark(packet.header.counter) { - result_action = Some(Err(e)); - LpState::ReadOnlyTransport { session } - } else { - match plaintext.try_into() { - Ok(data) => { - result_action = Some(Ok(LpAction::DeliverData(data))); - LpState::ReadOnlyTransport { session } - } - Err(err) => { - result_action = Some(Err(err)); - LpState::ReadOnlyTransport { session } - } - } - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e.into())); - LpState::Closed { reason } - } - } - } - } - - // Reject SendData in read-only mode - (LpState::ReadOnlyTransport { session }, LpInput::SendData(_)) => { - result_action = Some(Err(LpError::NoiseError(NoiseError::SessionReadOnly))); - LpState::ReadOnlyTransport { session } - } - - // Reject other inputs in read-only mode - (LpState::ReadOnlyTransport { session }, LpInput::InitiateSubsession) => { - result_action = Some(Err(LpError::InvalidStateTransition { - state: "ReadOnlyTransport".to_string(), - input: "InitiateSubsession".to_string(), - })); - LpState::ReadOnlyTransport { session } - } - - // --- Close Transition (applies to ReadyToHandshake, KKTExchange, Handshaking, Transport, SubsessionHandshaking, ReadOnlyTransport) --- - ( - LpState::Transport { .. } - | LpState::SubsessionHandshaking { .. } - | LpState::ReadOnlyTransport { .. }, - LpInput::Close, - ) => { - result_action = Some(Ok(LpAction::ConnectionClosed)); - // Transition to Closed state - LpState::Closed { - reason: "Closed by user".to_string(), - } + (LpState::Transport(transport), input) => { + let (next_state, action) = self.process_input_transport(transport, input); + result_action = action; + next_state } // Ignore Close if already Closed (closed_state @ LpState::Closed { .. }, LpInput::Close) => { @@ -877,11 +303,6 @@ impl LpStateMachine { // Return the original closed state closed_state } - // Ignore StartHandshake if Closed - // (closed_state @ LpState::Closed { .. }, LpInput::StartHandshake) => { - // result_action = Some(Err(LpError::LpSessionClosed)); - // closed_state - // } // Ignore ReceivePacket if Closed (closed_state @ LpState::Closed { .. }, LpInput::ReceivePacket(_)) => { result_action = Some(Err(LpError::LpSessionClosed)); @@ -900,19 +321,6 @@ impl LpStateMachine { result_action = Some(Err(err)); LpState::Closed { reason } } - - // --- Default: Invalid input for current state (if any combinations missed) --- - // Consider if this should transition to Closed state. For now, just report error - // and transition to Closed as a safety measure. - (invalid_state, input) => { - let err = LpError::InvalidStateTransition { - state: format!("{:?}", invalid_state), // Use owned state for debug info - input: format!("{:?}", input), - }; - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } }; // 3. Put the calculated next state back into the machine. @@ -927,11 +335,8 @@ impl LpStateMachine { &self, session: &mut LpSession, data: LpData, - ) -> Result { - let encrypted_message = session.encrypt_data(Vec::::from(data).as_ref())?; - session - .next_packet(encrypted_message) - .map_err(|e| NoiseError::Other(e.to_string())) // Improve error conversion? + ) -> Result { + session.encrypt_application_data(data.to_vec()) } } @@ -961,71 +366,72 @@ mod tests { #[test] fn test_state_machine_simplified_flow() { - let TODO = " for kem in kem_list() {"; - - let receiver_index: u32 = 123; - let mock_sessions = SessionsMock::mock_post_handshake(123); - - // Create state machines (already in Transport) - let mut initiator = LpStateMachine::new(mock_sessions.initiator); - let mut responder = LpStateMachine::new(mock_sessions.responder); - - assert_eq!(initiator.id().unwrap(), responder.id().unwrap()); - - // --- Transport Phase --- - println!("--- Step 1: Initiator sends data ---"); - let data_to_send_1 = LpData::new_opaque(b"hello responder".to_vec()); - let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); - let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_4 { - packet.clone() - } else { - panic!("Initiator should send data packet"); - }; - assert_eq!(data_packet_1.header.receiver_idx(), receiver_index); - - println!("--- Step 2: Responder receives data ---"); - let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); - let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_5 { - data - } else { - panic!("Responder should deliver data"); - }; - assert_eq!(resp_data_1, data_to_send_1); - - println!("--- Step 3: Responder sends data ---"); - let data_to_send_2 = LpData::new_opaque(b"hello initiator".to_vec()); - let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); - let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_6 { - packet.clone() - } else { - panic!("Responder should send data packet"); - }; - assert_eq!(data_packet_2.header.receiver_idx(), receiver_index); - - println!("--- Step 4: Initiator receives data ---"); - let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); - if let Some(Ok(LpAction::DeliverData(data))) = init_actions_5 { - assert_eq!(data, data_to_send_2); - } else { - panic!("Initiator should deliver data"); - } - - // --- Close --- - println!("--- Step 5: Initiator closes ---"); - let init_actions_6 = initiator.process_input(LpInput::Close); - assert!(matches!( - init_actions_6, - Some(Ok(LpAction::ConnectionClosed)) - )); - assert!(matches!(initiator.state, LpState::Closed { .. })); - - println!("--- Step 6: Responder closes ---"); - let resp_actions_7 = responder.process_input(LpInput::Close); - assert!(matches!( - resp_actions_7, - Some(Ok(LpAction::ConnectionClosed)) - )); - assert!(matches!(responder.state, LpState::Closed { .. })); + todo!() + // let TODO = " for kem in kem_list() {"; + // + // let receiver_index: u32 = 123; + // let mock_sessions = SessionsMock::mock_post_handshake(123); + // + // // Create state machines (already in Transport) + // let mut initiator = LpStateMachine::new(mock_sessions.initiator); + // let mut responder = LpStateMachine::new(mock_sessions.responder); + // + // assert_eq!(initiator.id().unwrap(), responder.id().unwrap()); + // + // // --- Transport Phase --- + // println!("--- Step 1: Initiator sends data ---"); + // let data_to_send_1 = LpData::new_opaque(b"hello responder".to_vec()); + // let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); + // let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_4 { + // packet.clone() + // } else { + // panic!("Initiator should send data packet"); + // }; + // assert_eq!(data_packet_1.header.receiver_idx(), receiver_index); + // + // println!("--- Step 2: Responder receives data ---"); + // let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); + // let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_5 { + // data + // } else { + // panic!("Responder should deliver data"); + // }; + // assert_eq!(resp_data_1, data_to_send_1); + // + // println!("--- Step 3: Responder sends data ---"); + // let data_to_send_2 = LpData::new_opaque(b"hello initiator".to_vec()); + // let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); + // let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_6 { + // packet.clone() + // } else { + // panic!("Responder should send data packet"); + // }; + // assert_eq!(data_packet_2.header.receiver_idx(), receiver_index); + // + // println!("--- Step 4: Initiator receives data ---"); + // let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); + // if let Some(Ok(LpAction::DeliverData(data))) = init_actions_5 { + // assert_eq!(data, data_to_send_2); + // } else { + // panic!("Initiator should deliver data"); + // } + // + // // --- Close --- + // println!("--- Step 5: Initiator closes ---"); + // let init_actions_6 = initiator.process_input(LpInput::Close); + // assert!(matches!( + // init_actions_6, + // Some(Ok(LpAction::ConnectionClosed)) + // )); + // assert!(matches!(initiator.state, LpState::Closed { .. })); + // + // println!("--- Step 6: Responder closes ---"); + // let resp_actions_7 = responder.process_input(LpInput::Close); + // assert!(matches!( + // resp_actions_7, + // Some(Ok(LpAction::ConnectionClosed)) + // )); + // assert!(matches!(responder.state, LpState::Closed { .. })); } /// Helper function to complete a full handshake between initiator and responder, @@ -1037,147 +443,4 @@ mod tests { LpStateMachine::new(sessions.responder), ) } - - #[test] - fn test_simultaneous_subsession_initiation() { - // Test for simultaneous subsession initiation race condition. - // Both sides call InitiateSubsession at the same time, sending KK1 to each other. - // The tie-breaker uses X25519 public key comparison: lower key becomes responder. - for kem in kem_list() { - let (mut alice, mut bob) = setup_transport_sessions(kem); - - // Get X25519 public keys to determine expected winner - let alice_x25519 = alice.session().unwrap().local_x25519_public(); - let bob_x25519 = bob.session().unwrap().local_x25519_public(); - - // Determine who should win (higher key stays initiator) - let alice_wins = alice_x25519.as_ref() > bob_x25519.as_ref(); - - // --- Both sides initiate subsession simultaneously --- - // Alice initiates subsession - let alice_kk1_packet = - if let Some(Ok(LpAction::SubsessionInitiated { packet, .. })) = - alice.process_input(LpInput::InitiateSubsession) - { - packet - } else { - panic!("Alice should initiate subsession with KK1"); - }; - assert!(matches!(alice.state, LpState::SubsessionHandshaking { .. })); - - // Bob initiates subsession (simultaneously) - let bob_kk1_packet = if let Some(Ok(LpAction::SubsessionInitiated { packet, .. })) = - bob.process_input(LpInput::InitiateSubsession) - { - packet - } else { - panic!("Bob should initiate subsession with KK1"); - }; - assert!(matches!(bob.state, LpState::SubsessionHandshaking { .. })); - - // --- Cross-delivery of KK1 packets (race resolution) --- - // Alice receives Bob's KK1 - let alice_response = alice.process_input(LpInput::ReceivePacket(bob_kk1_packet)); - - // Bob receives Alice's KK1 - let bob_response = bob.process_input(LpInput::ReceivePacket(alice_kk1_packet)); - - // --- Verify tie-breaker worked correctly --- - if alice_wins { - // Alice has higher key - she stays initiator, sends SubsessionAbort - assert!( - matches!(alice_response, Some(Ok(LpAction::SendPacket(_)))), - "Alice (winner) should send SubsessionAbort" - ); - assert!( - matches!(alice.state, LpState::SubsessionHandshaking { .. }), - "Alice should still be SubsessionHandshaking as initiator" - ); - - // Bob has lower key - he becomes responder, sends KK2 - let bob_kk2_packet = if let Some(Ok(LpAction::SendPacket(p))) = bob_response { - p - } else { - panic!("Bob (loser) should send KK2 as new responder"); - }; - assert!( - matches!(bob.state, LpState::SubsessionHandshaking { .. }), - "Bob should be SubsessionHandshaking as responder" - ); - - // Complete the handshake: Alice receives KK2 - let alice_completion = alice.process_input(LpInput::ReceivePacket(bob_kk2_packet)); - match alice_completion { - Some(Ok(LpAction::SubsessionComplete { - packet: Some(ready_packet), - .. - })) => { - assert!( - matches!(alice.state, LpState::ReadOnlyTransport { .. }), - "Alice should be ReadOnlyTransport after SubsessionComplete" - ); - - // Bob receives SubsessionReady - let bob_final = bob.process_input(LpInput::ReceivePacket(ready_packet)); - assert!( - matches!(bob_final, Some(Ok(LpAction::SubsessionComplete { .. }))), - "Bob should complete with SubsessionComplete" - ); - assert!( - matches!(bob.state, LpState::ReadOnlyTransport { .. }), - "Bob should be ReadOnlyTransport" - ); - } - other => panic!("Alice should complete subsession, got: {:?}", other), - } - } else { - // Bob has higher key - he stays initiator, sends SubsessionAbort - assert!( - matches!(bob_response, Some(Ok(LpAction::SendPacket(_)))), - "Bob (winner) should send SubsessionAbort" - ); - assert!( - matches!(bob.state, LpState::SubsessionHandshaking { .. }), - "Bob should still be SubsessionHandshaking as initiator" - ); - - // Alice has lower key - she becomes responder, sends KK2 - let alice_kk2_packet = if let Some(Ok(LpAction::SendPacket(p))) = alice_response { - p - } else { - panic!("Alice (loser) should send KK2 as new responder"); - }; - assert!( - matches!(alice.state, LpState::SubsessionHandshaking { .. }), - "Alice should be SubsessionHandshaking as responder" - ); - - // Complete the handshake: Bob receives KK2 - let bob_completion = bob.process_input(LpInput::ReceivePacket(alice_kk2_packet)); - match bob_completion { - Some(Ok(LpAction::SubsessionComplete { - packet: Some(ready_packet), - .. - })) => { - assert!( - matches!(bob.state, LpState::ReadOnlyTransport { .. }), - "Bob should be ReadOnlyTransport after SubsessionComplete" - ); - - // Alice receives SubsessionReady - let alice_final = alice.process_input(LpInput::ReceivePacket(ready_packet)); - assert!( - matches!(alice_final, Some(Ok(LpAction::SubsessionComplete { .. }))), - "Alice should complete with SubsessionComplete" - ); - assert!( - matches!(alice.state, LpState::ReadOnlyTransport { .. }), - "Alice should be ReadOnlyTransport" - ); - } - other => panic!("Bob should complete subsession, got: {:?}", other), - } - } - } - } } diff --git a/gateway/src/node/lp_listener/handler.rs b/gateway/src/node/lp_listener/handler.rs index c65efc8443..c20c98abd3 100644 --- a/gateway/src/node/lp_listener/handler.rs +++ b/gateway/src/node/lp_listener/handler.rs @@ -6,6 +6,7 @@ use crate::error::GatewayError; use bytes::BytesMut; use nym_crypto::asymmetric::{ed25519, x25519}; use nym_lp::codec::serialize_lp_packet; +use nym_lp::message::ApplicationData; use nym_lp::state_machine::{LpAction, LpData, LpDataKind, LpInput}; use nym_lp::{ message::ForwardPacketData, LpMessage, LpPacket, LpSession, LpStateMachine, OuterHeader, @@ -512,23 +513,14 @@ where let wrapped_lp_data = LpData::new(response_kind, serialised_response); let data_bytes = wrapped_lp_data.to_vec(); - let encrypted_message = session.encrypt_data(&data_bytes).map_err(|e| { + let encrypted_message = session.encrypt_application_data(data_bytes).map_err(|e| { GatewayError::LpProtocolError(format!("Failed to encrypt response: {e}")) })?; - let response_packet = session.next_packet(encrypted_message).map_err(|e| { - GatewayError::LpProtocolError(format!("Failed to create response packet: {e}")) - })?; - - let outer_key = session.outer_aead_key(); - // make sure to drop the entry before the .await call // Serialize the packet (encrypted if outer_key provided) let mut packet_buf = BytesMut::new(); - serialize_lp_packet(&response_packet, &mut packet_buf, Some(outer_key)).map_err(|e| { - GatewayError::LpProtocolError(format!("Failed to serialize packet: {}", e)) - })?; - drop(session_entry); + encrypted_message.encode(&mut packet_buf); // Send response (encrypted with outer AEAD) self.send_serialised_packet(&packet_buf).await?; @@ -1038,7 +1030,7 @@ mod tests { use crate::node::ActiveClientsStore; use bytes::BytesMut; use nym_lp::codec::{parse_lp_packet, serialize_lp_packet, OuterAeadKey}; - use nym_lp::message::{EncryptedDataPayload, LpMessage}; + use nym_lp::message::{ApplicationData, LpMessage}; use nym_lp::packet::{LpHeader, LpPacket}; use nym_lp::peer::LpLocalPeer; use nym_lp::SessionsMock; @@ -1404,7 +1396,7 @@ mod tests { receiver_idx, counter: 20, }, - LpMessage::EncryptedData(EncryptedDataPayload(encrypted_payload)), + LpMessage::ApplicationData(ApplicationData(encrypted_payload)), ); handler.send_lp_packet(packet).await }); @@ -1418,8 +1410,8 @@ mod tests { assert_eq!(received.header().receiver_idx, 200); assert_eq!(received.header().counter, 20); match received.message() { - LpMessage::EncryptedData(data) => { - assert_eq!(data, &EncryptedDataPayload(expected_payload)) + LpMessage::ApplicationData(data) => { + assert_eq!(data, &ApplicationData(expected_payload)) } _ => panic!("Expected EncryptedData message"), }