diff --git a/common/nym-lp/src/codec.rs b/common/nym-lp/src/codec.rs index a4e4100e30..29af2fc788 100644 --- a/common/nym-lp/src/codec.rs +++ b/common/nym-lp/src/codec.rs @@ -4,7 +4,8 @@ use crate::LpError; use crate::message::{ ClientHelloData, EncryptedDataPayload, ForwardPacketData, HandshakeData, KKTRequestData, - KKTResponseData, LpMessage, MessageType, + KKTResponseData, LpMessage, MessageType, SubsessionKK1Data, SubsessionKK2Data, + SubsessionReadyData, }; use crate::packet::{LpHeader, LpPacket, OuterHeader, TRAILER_LEN}; use bytes::{BufMut, BytesMut}; @@ -138,6 +139,39 @@ fn parse_message_from_type_and_content( } Ok(LpMessage::Ack) } + MessageType::SubsessionRequest => { + if !content.is_empty() { + return Err(LpError::InvalidPayloadSize { + expected: 0, + actual: content.len(), + }); + } + Ok(LpMessage::SubsessionRequest) + } + MessageType::SubsessionKK1 => { + let data: SubsessionKK1Data = bincode::deserialize(content) + .map_err(|e| LpError::DeserializationError(e.to_string()))?; + Ok(LpMessage::SubsessionKK1(data)) + } + MessageType::SubsessionKK2 => { + let data: SubsessionKK2Data = bincode::deserialize(content) + .map_err(|e| LpError::DeserializationError(e.to_string()))?; + Ok(LpMessage::SubsessionKK2(data)) + } + MessageType::SubsessionReady => { + let data: SubsessionReadyData = bincode::deserialize(content) + .map_err(|e| LpError::DeserializationError(e.to_string()))?; + Ok(LpMessage::SubsessionReady(data)) + } + MessageType::SubsessionAbort => { + // Empty signal message - no content to deserialize + if !content.is_empty() { + return Err(LpError::DeserializationError( + "SubsessionAbort should have no payload".to_string(), + )); + } + Ok(LpMessage::SubsessionAbort) + } } } @@ -364,7 +398,7 @@ mod tests { use crate::packet::{LpHeader, LpPacket, TRAILER_LEN}; use bytes::BytesMut; - // AIDEV-NOTE: With unified format, outer header (receiver_idx + counter) is always first + // 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 @@ -1133,4 +1167,181 @@ mod tests { _ => 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: 0, + 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: 0, + 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: 0, + 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: 0, + 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_u16().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: 0, + 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"), + } + } } diff --git a/common/nym-lp/src/message.rs b/common/nym-lp/src/message.rs index 2b23d43d3f..006a8ca22e 100644 --- a/common/nym-lp/src/message.rs +++ b/common/nym-lp/src/message.rs @@ -81,6 +81,16 @@ pub enum MessageType { Collision = 0x0007, /// Acknowledgment - gateway confirms receipt of message Ack = 0x0008, + /// Subsession request - client initiates subsession creation + SubsessionRequest = 0x0009, + /// 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, } impl MessageType { @@ -121,6 +131,27 @@ pub struct ForwardPacketData { pub inner_packet_bytes: Vec, } +/// Subsession KK1 message - first message of Noise KK handshake +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SubsessionKK1Data { + /// Noise KK first message payload (ephemeral key + encrypted static) + pub payload: Vec, +} + +/// Subsession KK2 message - second message of Noise KK handshake +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SubsessionKK2Data { + /// Noise KK second message payload (ephemeral key + encrypted response) + pub payload: Vec, +} + +/// Subsession ready confirmation with new session index +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SubsessionReadyData { + /// New subsession's receiver index for routing + pub receiver_index: u32, +} + #[derive(Debug, Clone)] pub enum LpMessage { Busy, @@ -134,6 +165,16 @@ pub enum LpMessage { Collision, /// Acknowledgment - gateway confirms receipt of message Ack, + /// 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, } impl Display for LpMessage { @@ -148,6 +189,11 @@ impl Display for LpMessage { 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"), } } } @@ -164,6 +210,11 @@ impl LpMessage { 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 => &[], } } @@ -178,6 +229,11 @@ impl LpMessage { 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 } } @@ -195,6 +251,13 @@ impl LpMessage { } LpMessage::Collision => 0, LpMessage::Ack => 0, + LpMessage::SubsessionRequest => 0, + // Variable length: bincode overhead (~8 bytes for Vec length) + payload + LpMessage::SubsessionKK1(data) => 8 + data.payload.len(), + LpMessage::SubsessionKK2(data) => 8 + data.payload.len(), + // 4 bytes u32 + bincode overhead (~4 bytes) + LpMessage::SubsessionReady(_) => 8, + LpMessage::SubsessionAbort => 0, } } @@ -209,6 +272,11 @@ impl LpMessage { 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, } } @@ -240,6 +308,23 @@ impl LpMessage { } LpMessage::Collision => { /* No content */ } LpMessage::Ack => { /* No content */ } + LpMessage::SubsessionRequest => { /* No content - signal only */ } + LpMessage::SubsessionKK1(data) => { + let serialized = + bincode::serialize(data).expect("Failed to serialize SubsessionKK1Data"); + dst.put_slice(&serialized); + } + LpMessage::SubsessionKK2(data) => { + let serialized = + bincode::serialize(data).expect("Failed to serialize SubsessionKK2Data"); + dst.put_slice(&serialized); + } + LpMessage::SubsessionReady(data) => { + let serialized = + bincode::serialize(data).expect("Failed to serialize SubsessionReadyData"); + dst.put_slice(&serialized); + } + LpMessage::SubsessionAbort => { /* No content - signal only */ } } } } diff --git a/common/nym-lp/src/noise_protocol.rs b/common/nym-lp/src/noise_protocol.rs index 42b5e0308f..41e601494b 100644 --- a/common/nym-lp/src/noise_protocol.rs +++ b/common/nym-lp/src/noise_protocol.rs @@ -25,6 +25,9 @@ pub enum NoiseError { #[error("Other Noise-related error: {0}")] Other(String), + + #[error("session is read-only after demotion")] + SessionReadOnly, } impl From for NoiseError { diff --git a/common/nym-lp/src/psk.rs b/common/nym-lp/src/psk.rs index f4f416b3a7..bf529e859b 100644 --- a/common/nym-lp/src/psk.rs +++ b/common/nym-lp/src/psk.rs @@ -57,6 +57,43 @@ 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 @@ -230,7 +267,7 @@ pub fn derive_psk_with_psq_responder( /// * `session_context` - Context bytes for PSQ (e.g., b"nym-lp-psq-session") /// /// # Returns -/// `(psk, psq_payload_bytes)` - PSK for Noise and serialized PSQ payload to embed +/// `PsqInitiatorResult` containing PSK, payload, and raw PQ shared secret pub fn psq_initiator_create_message( local_x25519_private: &PrivateKey, remote_x25519_public: &PublicKey, @@ -239,7 +276,7 @@ pub fn psq_initiator_create_message( client_ed25519_pk: &ed25519::PublicKey, salt: &[u8; 32], session_context: &[u8], -) -> Result<([u8; 32], Vec), LpError> { +) -> Result { // Step 1: Classical ECDH for baseline security let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); @@ -278,9 +315,13 @@ pub fn psq_initiator_create_message( LpError::Internal(format!("PSQ v1 send_initial_message failed: {:?}", e)) })?; - // Extract PSQ shared secret (unregistered PSK) + // 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.as_bytes()); @@ -294,7 +335,11 @@ pub fn psq_initiator_create_message( .tls_serialize_detached() .map_err(|e| LpError::Internal(format!("InitiatorMsg serialization failed: {:?}", e)))?; - Ok((final_psk, msg_bytes)) + Ok(PsqInitiatorResult { + psk: final_psk, + payload: msg_bytes, + pq_shared_secret, + }) } /// PSQ protocol wrapper for responder (gateway) side. @@ -317,11 +362,7 @@ pub fn psq_initiator_create_message( /// * `session_context` - Context bytes for PSQ /// /// # Returns -/// `psk` - Derived PSK for Noise -/// Processes a PSQ initiator message and generates a PSK with encrypted handle. -/// -/// Returns a tuple of (derived_psk, responder_msg_bytes) where responder_msg_bytes -/// contains the encrypted PSK handle (ctxt_B) that should be sent to the initiator. +/// `PsqResponderResult` containing PSK, PSK handle, and raw PQ shared secret pub fn psq_responder_process_message( local_x25519_private: &PrivateKey, remote_x25519_public: &PublicKey, @@ -330,7 +371,7 @@ pub fn psq_responder_process_message( psq_payload: &[u8], salt: &[u8; 32], session_context: &[u8], -) -> Result<([u8; 32], Vec), LpError> { +) -> Result { // Step 1: Classical ECDH for baseline security let ecdh_secret = local_x25519_private.diffie_hellman(remote_x25519_public); @@ -383,9 +424,13 @@ pub fn psq_responder_process_message( LpError::Internal(format!("PSQ v1 responder send failed: {:?}", e)) })?; - // Extract the PSQ PSK from the registered PSK + // 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.as_bytes()); @@ -400,7 +445,38 @@ pub fn psq_responder_process_message( .tls_serialize_detached() .map_err(|e| LpError::Internal(format!("ResponderMsg serialization failed: {:?}", e)))?; - Ok((final_psk, responder_msg_bytes)) + 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::kdf::derive_key_blake3( + SUBSESSION_PSK_CONTEXT, + pq_shared_secret, + &subsession_index.to_le_bytes(), + ) } #[cfg(test)] diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 2d1476f940..316849b211 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -11,7 +11,7 @@ use crate::keypair::{PrivateKey, PublicKey}; use crate::message::{EncryptedDataPayload, HandshakeData}; use crate::noise_protocol::{NoiseError, NoiseProtocol, ReadResult}; use crate::packet::LpHeader; -use crate::psk::{psq_initiator_create_message, psq_responder_process_message}; +use crate::psk::{derive_subsession_psk, psq_initiator_create_message, psq_responder_process_message}; use crate::replay::ReceivingKeyCounterValidator; use crate::{LpError, LpMessage, LpPacket}; use nym_crypto::asymmetric::ed25519; @@ -19,6 +19,30 @@ use nym_kkt::ciphersuite::{DecapsulationKey, EncapsulationKey}; use parking_lot::Mutex; use snow::Builder; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use zeroize::{Zeroize, ZeroizeOnDrop}; + +/// 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() + } +} /// KKT (KEM Key Transfer) exchange state. /// @@ -170,6 +194,25 @@ pub struct LpSession { /// Outer AEAD key for packet encryption (derived from PSK after PSQ handshake). /// None before PSK is available, Some after PSK injection. outer_aead_key: Mutex>, + + /// Raw PQ shared secret (K_pq) from PSQ KEM encapsulation/decapsulation. + /// Stored after PSQ handshake completes for subsession PSK derivation. + /// This preserves PQ protection when creating subsessions via KKpsk0. + /// Wrapped in PqSharedSecret for automatic memory zeroization on drop. + pq_shared_secret: Mutex>, + + /// 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: AtomicU64, + + /// True if this session has been demoted to read-only mode. + /// Demoted sessions can still receive/decrypt but cannot send/encrypt. + read_only: AtomicBool, + + /// ID of the successor session that replaced this one. + /// Set when demote() is called. + successor_session_id: Mutex>, } /// Generates a fresh salt for PSK derivation. @@ -222,6 +265,14 @@ impl LpSession { self.local_x25519_private.public_key() } + /// 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) -> &PublicKey { + &self.remote_x25519_public + } + /// Returns the outer AEAD key for packet encryption/decryption. /// /// Returns `None` before PSK is derived (during initial handshake), @@ -318,6 +369,10 @@ impl LpSession { remote_x25519_public: remote_x25519_key.clone(), salt: *salt, outer_aead_key: Mutex::new(None), + pq_shared_secret: Mutex::new(None), + subsession_counter: AtomicU64::new(0), + read_only: AtomicBool::new(false), + successor_session_id: Mutex::new(None), }) } @@ -632,7 +687,7 @@ impl LpSession { // Generate PSQ payload and PSK using KKT-authenticated KEM key let session_context = self.id.to_le_bytes(); - let (psk, psq_payload) = match psq_initiator_create_message( + let psq_result = match psq_initiator_create_message( &self.local_x25519_private, &self.remote_x25519_public, remote_kem, @@ -647,6 +702,11 @@ impl LpSession { return Some(Err(e)); } }; + let psk = psq_result.psk; + let psq_payload = psq_result.payload; + + // Store PQ shared secret for subsession PSK derivation + *self.pq_shared_secret.lock() = Some(PqSharedSecret::new(psq_result.pq_shared_secret)); // Inject PSK into Noise HandshakeState if let Err(e) = noise_state.set_psk(3, &psk) { @@ -797,7 +857,7 @@ impl LpSession { // Decapsulate PSK from PSQ payload using X25519 as DHKEM let session_context = self.id.to_le_bytes(); - let (psk, responder_msg_bytes) = match psq_responder_process_message( + let psq_result = match psq_responder_process_message( &self.local_x25519_private, &self.remote_x25519_public, (&dec_key, &enc_key), @@ -812,11 +872,15 @@ impl LpSession { return Err(e); } }; + let psk = psq_result.psk; + + // Store PQ shared secret for subsession PSK derivation + *self.pq_shared_secret.lock() = Some(PqSharedSecret::new(psq_result.pq_shared_secret)); // Store the PSK handle (ctxt_B) for transmission in next message { let mut psk_handle = self.psk_handle.lock(); - *psk_handle = Some(responder_msg_bytes); + *psk_handle = Some(psq_result.psk_handle); } // Inject PSK into Noise HandshakeState @@ -887,6 +951,49 @@ impl LpSession { self.noise_state.lock().is_handshake_finished() } + /// Returns the PQ shared secret (K_pq) if available. + /// + /// This is the raw KEM output from PSQ before Blake3 KDF combination. + /// Used for deriving subsession PSKs to maintain PQ protection. + pub fn pq_shared_secret(&self) -> Option<[u8; 32]> { + self.pq_shared_secret.lock().as_ref().map(|s| *s.as_bytes()) + } + + /// 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(&self) -> u64 { + self.subsession_counter.fetch_add(1, Ordering::Relaxed) + } + + /// 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.load(Ordering::Acquire) + } + + /// 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(&self, successor_idx: u32) { + *self.successor_session_id.lock() = Some(successor_idx); + self.read_only.store(true, Ordering::Release); + } + + /// Returns the successor session ID if this session was demoted. + pub fn successor_session_id(&self) -> Option { + *self.successor_session_id.lock() + } + /// Encrypts application data payload using the established Noise transport session. /// /// This should only be called after the handshake is complete (`is_handshake_complete` returns true). @@ -900,6 +1007,11 @@ impl LpSession { /// * `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(&self, payload: &[u8]) -> Result { + // Check if session is read-only (demoted) + if self.read_only.load(Ordering::Acquire) { + return Err(NoiseError::SessionReadOnly); + } + let mut noise_state = self.noise_state.lock(); // Safety: Prevent transport mode with dummy PSK if !self.psk_injected.load(Ordering::Acquire) { @@ -961,6 +1073,220 @@ impl LpSession { kem_pk: Box::new(kem_pk), }; } + + /// 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 { + // Verify parent handshake is complete + if !self.is_handshake_complete() { + return Err(LpError::Internal( + "Parent handshake not complete".into(), + )); + } + + // Get PQ shared secret + let pq_secret = self + .pq_shared_secret() + .ok_or_else(|| LpError::Internal("PQ shared secret not available".into()))?; + + // Derive subsession PSK from parent's PQ shared secret + let subsession_psk = derive_subsession_psk(&pq_secret, 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_x25519_private.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, + // Copy key material from parent for into_session() conversion + local_ed25519_private: ed25519::PrivateKey::from_bytes( + &self.local_ed25519_private.to_bytes(), + ).expect("Valid Ed25519 private key from parent"), + local_ed25519_public: ed25519::PublicKey::from_bytes(&self.local_ed25519_public.to_bytes()) + .expect("Valid Ed25519 public key from parent"), + remote_ed25519_public: ed25519::PublicKey::from_bytes(&self.remote_ed25519_public.to_bytes()) + .expect("Valid Ed25519 public key from parent"), + local_x25519_private: self.local_x25519_private.clone(), + remote_x25519_public: self.remote_x25519_public.clone(), + pq_shared_secret: PqSharedSecret::new(pq_secret), + subsession_psk, + }) + } +} + +/// 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 + noise_state: Mutex, + /// Is this side the initiator? + is_initiator: bool, + + // Key material inherited from parent session for into_session() conversion + /// Local Ed25519 private key (for PSQ auth if needed) + local_ed25519_private: ed25519::PrivateKey, + /// Local Ed25519 public key + local_ed25519_public: ed25519::PublicKey, + /// Remote Ed25519 public key + remote_ed25519_public: ed25519::PublicKey, + /// Local X25519 private key (Noise static key) + local_x25519_private: PrivateKey, + /// Remote X25519 public key (Noise static key) + remote_x25519_public: PublicKey, + /// 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], +} + +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(); + + // Generate fresh salt for the new session + let salt = generate_fresh_salt(); + + // Derive outer AEAD key from the subsession PSK + let outer_key = OuterAeadKey::from_psk(&self.subsession_psk); + + Ok(LpSession { + id: receiver_index, + is_initiator: self.is_initiator, + noise_state: Mutex::new(noise_state), + // KKT: subsession inherits from parent, mark as processed + kkt_state: Mutex::new(KKTState::ResponderProcessed), + // PSQ: subsession uses PSK derived from parent's PQ secret + psq_state: Mutex::new(PSQState::Completed { psk: self.subsession_psk }), + psk_handle: Mutex::new(None), // Subsession doesn't have its own handle + sending_counter: AtomicU64::new(0), + receiving_counter: Mutex::new(ReceivingKeyCounterValidator::new(0)), + psk_injected: AtomicBool::new(true), // PSK was in KKpsk0 + local_ed25519_private: self.local_ed25519_private, + local_ed25519_public: self.local_ed25519_public, + remote_ed25519_public: self.remote_ed25519_public, + local_x25519_private: self.local_x25519_private, + remote_x25519_public: self.remote_x25519_public, + salt, + outer_aead_key: Mutex::new(Some(outer_key)), + pq_shared_secret: Mutex::new(Some(self.pq_shared_secret)), + subsession_counter: AtomicU64::new(0), + read_only: AtomicBool::new(false), + successor_session_id: Mutex::new(None), + }) + } } #[cfg(test)] @@ -1925,4 +2251,120 @@ mod tests { e => panic!("Expected PskNotInjected error, got: {:?}", e), } } + + #[test] + fn test_demote_sets_read_only() { + let initiator_keys = generate_keypair(); + let responder_keys = generate_keypair(); + + let session = + create_handshake_test_session(12345u32, true, &initiator_keys, responder_keys.public_key()); + + // 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() { + // --- Setup Handshake --- + let initiator_keys = generate_keypair(); + let responder_keys = generate_keypair(); + + let initiator_session = + create_handshake_test_session(12345u32, true, &initiator_keys, responder_keys.public_key()); + let responder_session = + create_handshake_test_session(12345u32, false, &responder_keys, initiator_keys.public_key()); + + // Drive handshake to completion + let i_msg = initiator_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + responder_session.process_handshake_message(&i_msg).unwrap(); + let r_msg = responder_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + initiator_session.process_handshake_message(&r_msg).unwrap(); + let i_msg = initiator_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + responder_session.process_handshake_message(&i_msg).unwrap(); + + assert!(initiator_session.is_handshake_complete()); + + // 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 initiator_keys = generate_keypair(); + let responder_keys = generate_keypair(); + + let initiator_session = + create_handshake_test_session(12345u32, true, &initiator_keys, responder_keys.public_key()); + let responder_session = + create_handshake_test_session(12345u32, false, &responder_keys, initiator_keys.public_key()); + + // Drive handshake to completion + let i_msg = initiator_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + responder_session.process_handshake_message(&i_msg).unwrap(); + let r_msg = responder_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + initiator_session.process_handshake_message(&r_msg).unwrap(); + let i_msg = initiator_session + .prepare_handshake_message() + .unwrap() + .unwrap(); + responder_session.process_handshake_message(&i_msg).unwrap(); + + assert!(initiator_session.is_handshake_complete()); + assert!(responder_session.is_handshake_complete()); + + // 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/state_machine.rs b/common/nym-lp/src/state_machine.rs index 11352c1493..c32e6c9f4e 100644 --- a/common/nym-lp/src/state_machine.rs +++ b/common/nym-lp/src/state_machine.rs @@ -6,9 +6,10 @@ use crate::{ LpError, keypair::{Keypair, PrivateKey as LpPrivateKey, PublicKey as LpPublicKey}, + message::{LpMessage, SubsessionKK1Data, SubsessionKK2Data, SubsessionReadyData}, noise_protocol::NoiseError, packet::LpPacket, - session::LpSession, + session::{LpSession, SubsessionHandshake}, }; use bytes::BytesMut; use nym_crypto::asymmetric::ed25519; @@ -31,6 +32,18 @@ pub enum LpState { /// Handshake complete, ready for data transport. Transport { session: LpSession }, + + /// Performing subsession KK handshake while parent remains active. + /// Parent can still send/receive; subsession messages tunneled through parent. + SubsessionHandshaking { + session: LpSession, + subsession: SubsessionHandshake, + }, + + /// Parent session demoted after subsession promoted. + /// Can only receive (drain in-flight), cannot send. + ReadOnlyTransport { session: LpSession }, + /// An error occurred, or the connection was intentionally closed. Closed { reason: String }, /// Processing an input event. @@ -44,6 +57,8 @@ pub enum LpStateBare { KKTExchange, Handshaking, Transport, + SubsessionHandshaking, + ReadOnlyTransport, Closed, Processing, } @@ -55,6 +70,8 @@ impl From<&LpState> for LpStateBare { LpState::KKTExchange { .. } => LpStateBare::KKTExchange, LpState::Handshaking { .. } => LpStateBare::Handshaking, LpState::Transport { .. } => LpStateBare::Transport, + LpState::SubsessionHandshaking { .. } => LpStateBare::SubsessionHandshaking, + LpState::ReadOnlyTransport { .. } => LpStateBare::ReadOnlyTransport, LpState::Closed { .. } => LpStateBare::Closed, LpState::Processing => LpStateBare::Processing, } @@ -72,6 +89,9 @@ pub enum LpInput { SendData(Vec), // Using Bytes for efficiency /// 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. @@ -87,6 +107,20 @@ pub enum LpAction { HandshakeComplete, /// 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: SubsessionHandshake, + new_receiver_index: u32, + }, } /// The Lewes Protocol State Machine. @@ -104,7 +138,9 @@ impl LpStateMachine { LpState::ReadyToHandshake { session } | LpState::KKTExchange { session } | LpState::Handshaking { session } - | LpState::Transport { session } => Ok(session), + | LpState::Transport { session } + | LpState::SubsessionHandshaking { session, .. } + | LpState::ReadOnlyTransport { session } => Ok(session), LpState::Closed { .. } => Err(LpError::LpSessionClosed), LpState::Processing => Err(LpError::LpSessionProcessing), } @@ -118,7 +154,9 @@ impl LpStateMachine { LpState::ReadyToHandshake { session } | LpState::KKTExchange { session } | LpState::Handshaking { session } - | LpState::Transport { session } => Ok(session), + | LpState::Transport { session } + | LpState::SubsessionHandshaking { session, .. } + | LpState::ReadOnlyTransport { session } => Ok(session), LpState::Closed { .. } => Err(LpError::LpSessionClosed), LpState::Processing => Err(LpError::LpSessionProcessing), } @@ -450,43 +488,99 @@ impl LpStateMachine { } // --- Transport State --- - (LpState::Transport { session }, LpInput::ReceivePacket(packet)) => { // Needs mut session for marking counter + (LpState::Transport { 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()))); - // Remain in transport state LpState::Transport { session } } else { - // --- Inline handle_data_packet logic --- - // 1. Check replay protection - if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) { - let _reason = e.to_string(); - 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) { - let _reason = e.to_string(); + // 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 } + } + 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::Transport{ session } - } else { - // 4. Deliver data - result_action = Some(Ok(LpAction::DeliverData(BytesMut::from(plaintext.as_slice())))); - // Remain in transport state - LpState::Transport { session } + LpState::Closed { reason } } } - Err(e) => { // Error decrypting data - let reason = e.to_string(); - result_action = Some(Err(e.into())); - 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 + result_action = Some(Ok(LpAction::DeliverData(BytesMut::from(plaintext.as_slice())))); + LpState::Transport { session } + } + } + Err(e) => { + let reason = e.to_string(); + result_action = Some(Err(e.into())); + LpState::Closed { reason } + } + } } } + _ => { + // 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 } + } } - // --- End inline handle_data_packet logic --- } } (LpState::Transport { session }, LpInput::SendData(data)) => { @@ -512,12 +606,377 @@ impl LpStateMachine { LpState::Transport { session } } - // --- Close Transition (applies to ReadyToHandshake, KKTExchange, Handshaking, Transport) --- + // --- Transport + InitiateSubsession → SubsessionHandshaking --- + (LpState::Transport { 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 } + } + 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 { 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_bytes() < remote_key.as_bytes() { + // 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: 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 { + result_action = Some(Ok(LpAction::DeliverData(BytesMut::from(plaintext.as_slice())))); + 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 { session, subsession }, LpInput::SendData(data)) => { + match self.prepare_data_packet(&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::StartHandshake) => { + result_action = Some(Err(LpError::InvalidStateTransition { + state: "SubsessionHandshaking".to_string(), + input: "StartHandshake".to_string(), + })); + LpState::SubsessionHandshaking { session, subsession } + } + + (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 { 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 { + result_action = Some(Ok(LpAction::DeliverData(BytesMut::from(plaintext.as_slice())))); + 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::StartHandshake) => { + result_action = Some(Err(LpError::InvalidStateTransition { + state: "ReadOnlyTransport".to_string(), + input: "StartHandshake".to_string(), + })); + LpState::ReadOnlyTransport { session } + } + + (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::ReadyToHandshake { .. } // We consume the session here | LpState::KKTExchange { .. } | LpState::Handshaking { .. } - | LpState::Transport { .. }, + | LpState::Transport { .. } + | LpState::SubsessionHandshaking { .. } + | LpState::ReadOnlyTransport { .. }, LpInput::Close, ) => { result_action = Some(Ok(LpAction::ConnectionClosed)); @@ -1056,4 +1515,240 @@ mod tests { )); assert!(matches!(initiator.state, LpState::KKTExchange { .. })); // Still in KKTExchange } + + /// Helper function to complete a full handshake between initiator and responder, + /// returning both in Transport state ready for subsession testing. + fn setup_transport_sessions() -> (LpStateMachine, LpStateMachine) { + // Use different seeds to get different X25519 keys. + // The tie-breaker compares X25519 public keys. + let ed25519_keypair_a = ed25519::KeyPair::from_secret([30u8; 32], 0); + let ed25519_keypair_b = ed25519::KeyPair::from_secret([31u8; 32], 1); + + let salt = [60u8; 32]; + let receiver_index: u32 = 111111; + + // Create state machines - Alice is initiator, Bob is responder + let mut alice = LpStateMachine::new( + receiver_index, + true, + ( + ed25519_keypair_a.private_key(), + ed25519_keypair_a.public_key(), + ), + ed25519_keypair_b.public_key(), + &salt, + ) + .unwrap(); + + let mut bob = LpStateMachine::new( + receiver_index, + false, + ( + ed25519_keypair_b.private_key(), + ed25519_keypair_b.public_key(), + ), + ed25519_keypair_a.public_key(), + &salt, + ) + .unwrap(); + + // --- Complete KKT Exchange --- + // Alice starts handshake + let kkt_request = if let Some(Ok(LpAction::SendPacket(p))) = + alice.process_input(LpInput::StartHandshake) + { + p + } else { + panic!("Alice should send KKT request"); + }; + + // Bob starts handshake + let _ = bob.process_input(LpInput::StartHandshake); + + // Bob receives KKT request, sends response + let kkt_response = if let Some(Ok(LpAction::SendPacket(p))) = + bob.process_input(LpInput::ReceivePacket(kkt_request)) + { + p + } else { + panic!("Bob should send KKT response"); + }; + + // Alice receives KKT response + let _ = alice.process_input(LpInput::ReceivePacket(kkt_response)); + + // --- Complete Noise Handshake --- + // Alice prepares first Noise message + let noise1_msg = alice.session().unwrap().prepare_handshake_message().unwrap().unwrap(); + let noise1_packet = alice.session().unwrap().next_packet(noise1_msg).unwrap(); + + // Bob receives noise1, sends noise2 + let noise2_packet = if let Some(Ok(LpAction::SendPacket(p))) = + bob.process_input(LpInput::ReceivePacket(noise1_packet)) + { + p + } else { + panic!("Bob should send Noise packet 2"); + }; + + // Alice receives noise2, sends noise3 + let noise3_packet = if let Some(Ok(LpAction::SendPacket(p))) = + alice.process_input(LpInput::ReceivePacket(noise2_packet)) + { + p + } else { + panic!("Alice should send Noise packet 3"); + }; + assert!(matches!(alice.state, LpState::Transport { .. })); + + // Bob receives noise3, completes handshake + let _ = bob.process_input(LpInput::ReceivePacket(noise3_packet)); + assert!(matches!(bob.state, LpState::Transport { .. })); + + (alice, bob) + } + + #[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. + + let (mut alice, mut bob) = setup_transport_sessions(); + + // 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_bytes() > bob_x25519.as_bytes(); + + // --- 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 e0c70df368..5760c00296 100644 --- a/gateway/src/node/lp_listener/handler.rs +++ b/gateway/src/node/lp_listener/handler.rs @@ -388,7 +388,7 @@ impl LpConnectionHandler { let session = &session_entry.value().state; - // AIDEV-NOTE: Validate counter BEFORE decryption to prevent replay DoS attacks. + // Validate counter BEFORE decryption to prevent replay DoS attacks. // Counter is from cleartext header but authenticated by AEAD AAD, so this is safe. session.receiving_counter_quick_check(counter).map_err(|e| { inc!("lp_errors_replay_check");