diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 319db635d4..95fbb93d04 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -44,6 +44,7 @@ use parking_lot::Mutex; use rand::RngCore; use snow::Builder; +use crate::state_machine::LpData; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use zeroize::{Zeroize, ZeroizeOnDrop}; @@ -115,18 +116,6 @@ pub enum KKTState { }, } -impl KKTState { - pub fn carrier_mut(&mut self) -> Result<&mut Carrier, LpError> { - match &mut self { - KKTState::NotStarted | KKTState::InitiatorWaiting { .. } => { - Err(LpError::KKTError("incomplete KKT exchange".to_string())) - } - KKTState::Completed { carrier, .. } => Ok(carrier), - KKTState::ResponderProcessed { carrier } => Ok(carrier), - } - } -} - impl std::fmt::Debug for KKTState { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -737,7 +726,7 @@ impl<'a> LpSession<'a> { &self.local_peer.ciphersuite, &self.id, &self.local_peer.x25519, - self.remote_x25519_public(), + &self.remote_peer.x25519_public, remote_kem, ); @@ -899,7 +888,7 @@ impl<'a> LpSession<'a> { } /// Checks if the Noise handshake phase is complete. - pub fn is_handshake_complete(&'a self) -> bool { + pub fn is_handshake_complete(&self) -> bool { todo!() // self.noise_state.lock().is_handshake_finished() } @@ -908,7 +897,7 @@ impl<'a> LpSession<'a> { /// /// 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(&'a self) -> Option<[u8; 32]> { + pub fn pq_shared_secret(&self) -> Option<[u8; 32]> { todo!() // self.pq_shared_secret.lock().as_ref().map(|s| *s.as_bytes()) } @@ -980,6 +969,14 @@ impl<'a> LpSession<'a> { // Ok(LpMessage::EncryptedData(EncryptedDataPayload(payload))) } + // Helper to prepare an outgoing data packet + // Kept as it doesn't mutate self.state + pub fn prepare_data_packet(&self, data: LpData) -> Result { + let encrypted_message = self.encrypt_data(Vec::::from(data).as_ref())?; + self.next_packet(encrypted_message) + .map_err(|e| NoiseError::Other(e.to_string())) // Improve error conversion? + } + /// Decrypts an incoming Noise message containing application data. /// /// This should only be called after the handshake is complete (`is_handshake_complete` returns true) @@ -1094,6 +1091,8 @@ impl<'a> LpSession<'a> { subsession_psk, }) } + + // } /// Subsession created via Noise KKpsk0 handshake tunneled through parent session. diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index 26e5c2207d..adc3756502 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -10,19 +10,18 @@ use crate::noise_protocol::ReadResult; use crate::peer::{LpLocalPeer, LpRemotePeer}; use crate::state_machine::{LpAction, LpInput, LpState, LpStateBare}; use crate::{LpError, LpMessage, LpSession, LpStateMachine}; -use dashmap::DashMap; -use dashmap::mapref::one::{Ref, RefMut}; +use std::collections::HashMap; + #[cfg(test)] use libcrux_psq::handshake::types::DHPublicKey; use nym_kkt::ciphersuite::Ciphersuite; /// Manages the lifecycle of Lewes Protocol sessions. /// -/// The SessionManager is responsible for creating, storing, and retrieving sessions, -/// ensuring proper thread-safety for concurrent access. +/// The SessionManager is responsible for creating, storing, and retrieving sessions pub struct SessionManager<'a> { /// Manages state machines directly, keyed by lp_id - state_machines: DashMap>, + state_machines: HashMap>, } impl<'a> Default for SessionManager<'a> { @@ -35,12 +34,12 @@ impl<'a> SessionManager<'a> { /// Creates a new session manager with empty session storage. pub fn new() -> Self { Self { - state_machines: DashMap::new(), + state_machines: HashMap::new(), } } pub fn process_input( - &'a self, + &'a mut self, lp_id: u32, input: LpInput, ) -> Result, LpError> { @@ -49,7 +48,7 @@ impl<'a> SessionManager<'a> { .transpose() } - pub fn add(&self, session: LpSession<'a>) -> Result<(), LpError> { + pub fn add(&mut self, session: LpSession<'a>) -> Result<(), LpError> { let sm = LpStateMachine { state: LpState::ReadyToHandshake { session: Box::new(session), @@ -89,7 +88,7 @@ impl<'a> SessionManager<'a> { } pub fn receiving_counter_quick_check( - &'a self, + &mut self, lp_id: u32, counter: u64, ) -> Result<(), LpError> { @@ -120,7 +119,7 @@ impl<'a> SessionManager<'a> { .is_handshake_complete()) } - pub fn next_counter(&'a self, lp_id: u32) -> Result { + pub fn next_counter(&mut self, lp_id: u32) -> Result { Ok(self.state_machine_mut(lp_id)?.session()?.next_counter()) } @@ -160,23 +159,13 @@ impl<'a> SessionManager<'a> { self.state_machines.contains_key(&lp_id) } - // this method must NOT be used without some serious considerations, - // in particular the value MUST NOT be held across await points - // (even if the compiler is fine with it) - // as it could lead to a deadlock - #[doc(hidden)] - fn state_machine(&'a self, lp_id: u32) -> Result, LpError> { + fn state_machine(&'a self, lp_id: u32) -> Result<&LpStateMachine, LpError> { self.state_machines .get(&lp_id) .ok_or_else(|| LpError::StateMachineNotFound { lp_id }) } - // this method must NOT be used without some serious considerations, - // in particular the value MUST NOT be held across await points - // (even if the compiler is fine with it) - // as it could lead to a deadlock - #[doc(hidden)] - fn state_machine_mut(&'a self, lp_id: u32) -> Result, LpError> { + fn state_machine_mut(&mut self, lp_id: u32) -> Result<&mut LpStateMachine<'a>, LpError> { self.state_machines .get_mut(&lp_id) .ok_or_else(|| LpError::StateMachineNotFound { lp_id }) @@ -207,7 +196,7 @@ impl<'a> SessionManager<'a> { // } pub fn create_session_state_machine( - &self, + &mut self, receiver_index: u32, is_initiator: bool, ciphersuite: Ciphersuite, @@ -229,7 +218,7 @@ impl<'a> SessionManager<'a> { } /// Method to remove a state machine - pub fn remove_state_machine(&self, lp_id: u32) -> bool { + pub fn remove_state_machine(&mut self, lp_id: u32) -> bool { let removed = self.state_machines.remove(&lp_id); removed.is_some() @@ -239,7 +228,7 @@ impl<'a> SessionManager<'a> { /// This allows integration tests to bypass KKT exchange and directly test PSQ/handshake. #[cfg(test)] pub fn init_kkt_for_test( - &self, + &mut self, lp_id: u32, remote_x25519_pub: &DHPublicKey, ) -> Result<(), LpError> { @@ -261,7 +250,7 @@ mod tests { #[test] fn test_session_manager_get() { - let manager = SessionManager::new(); + let mut manager = SessionManager::new(); for kem in kem_list() { let (local, peer1) = mock_peers(kem); @@ -282,7 +271,7 @@ mod tests { } #[test] fn test_session_manager_remove() { - let manager = SessionManager::new(); + let mut manager = SessionManager::new(); for kem in kem_list() { let (local, peer1) = mock_peers(kem); @@ -304,7 +293,7 @@ mod tests { #[test] fn test_multiple_sessions() { - let manager = SessionManager::new(); + let mut manager = SessionManager::new(); for kem in kem_list() { let (local, peer1) = mock_peers(kem); let (peer2, peer3) = mock_peers(kem); @@ -337,7 +326,7 @@ mod tests { #[test] fn test_session_manager_create_session() { - let manager = SessionManager::new(); + let mut manager = SessionManager::new(); for kem in kem_list() { let (init, resp) = mock_peers(kem); diff --git a/common/nym-lp/src/state_machine.rs b/common/nym-lp/src/state_machine.rs index 03a07273a6..fbe431d186 100644 --- a/common/nym-lp/src/state_machine.rs +++ b/common/nym-lp/src/state_machine.rs @@ -58,11 +58,906 @@ pub enum LpState<'a> { /// An error occurred, or the connection was intentionally closed. Closed { reason: String }, + /// Processing an input event. #[default] Processing, } +impl<'a> LpState<'a> { + pub fn process_input(&'a mut self, input: LpInput) -> Option> { + let mut result_action = None; + let mut original_state = mem::take(self); + + match (original_state, input) { + // --- Handshaking State --- + (LpState::Handshaking { 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()))); + // Don't change state, return the original state variant + *self = LpState::Handshaking { session }; + return result_action; + } + // --- Inline handle_handshake_packet logic --- + // 1. Check replay protection *before* processing + if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) { + result_action = Some(Err(e)); + *self = LpState::Handshaking { session }; + return result_action; + } + + // 2. Process the handshake message + if let Err(err) = session.process_handshake_message(&packet.message) { + // Error from process_handshake_message + let reason = err.to_string(); + result_action = Some(Err(err)); + *self = LpState::Closed { reason }; + return result_action; + } + + // 3. Mark counter as received *after* successful processing + if let Err(e) = session.receiving_counter_mark(packet.header.counter) { + let _reason = e.to_string(); + result_action = Some(Err(e)); + *self = LpState::Handshaking { session }; + return result_action; + } + + // 4. First check if we need to send a handshake message (before checking completion) + let Some(handshake_message) = session.prepare_handshake_message() else { + // 5. No message to send - check if handshake is complete + return if session.is_handshake_complete() { + result_action = Some(Ok(LpAction::HandshakeComplete)); + // Transition to Transport + *self = LpState::Transport { session }; + result_action + } else { + // Handshake stalled unexpectedly + let err = LpError::NoiseError(NoiseError::Other( + "Handshake stalled unexpectedly".to_string(), + )); + let reason = err.to_string(); + result_action = Some(Err(err)); + *self = LpState::Closed { reason }; + result_action + }; + }; + + let message = match handshake_message { + Err(err) => { + let reason = err.to_string(); + result_action = Some(Err(err)); + *self = LpState::Closed { reason }; + return result_action; + } + Ok(message) => message, + }; + + let state = match session.next_packet(message) { + Ok(response_packet) => { + result_action = Some(Ok(LpAction::SendPacket(response_packet))); + // Check if handshake became complete after preparing message + if session.is_handshake_complete() { + LpState::Transport { session } // Transition to Transport + } else { + LpState::Handshaking { session } // Remain Handshaking + } + } + Err(e) => { + let reason = e.to_string(); + result_action = Some(Err(e)); + LpState::Closed { reason } + } + }; + + *self = state; + return result_action; + // return (state, result_action); + + // --- End inline handle_handshake_packet logic --- + } + + _ => todo!(), + }; + + // // 2. Match on the owned current_state. Each arm calculates and returns the NEXT state. + // let next_state = match (self, input) { + // // --- ReadyToHandshake State --- + // (LpState::ReadyToHandshake { mut session }, LpInput::StartHandshake) => { + // if session.is_initiator() { + // // Initiator starts by requesting KEM key via KKT + // match session.prepare_kkt_request() { + // Some(Ok(kkt_message)) => { + // match session.next_packet(kkt_message) { + // Ok(kkt_packet) => { + // result_action = Some(Ok(LpAction::SendPacket(kkt_packet))); + // LpState::KKTExchange { session } // Transition to KKTExchange + // } + // Err(e) => { + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // } + // } + // Some(Err(e)) => { + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // None => { + // // Should not happen for initiator + // let err = LpError::Internal( + // "prepare_kkt_request returned None for initiator".to_string(), + // ); + // let reason = err.to_string(); + // result_action = Some(Err(err)); + // LpState::Closed { reason } + // } + // } + // } else { + // // Responder waits for KKT request + // LpState::KKTExchange { session } + // // No action needed yet, result_action remains None. + // } + // } + // + // // --- KKTExchange State --- + // (LpState::KKTExchange { 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::KKTExchange { session } + // } else { + // use crate::message::LpMessage; + // + // // Packet message is already parsed, match on it directly + // match &packet.message { + // LpMessage::KKTRequest(kkt_request) if !session.is_initiator() => { + // match session.process_kkt_request(&kkt_request.0) { + // Ok(kkt_response_message) => { + // match session.next_packet(kkt_response_message) { + // Ok(response_packet) => { + // result_action = Some(Ok(LpAction::SendPacket(response_packet))); + // // After KKT exchange, move to Handshaking + // LpState::Handshaking { 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)); + // LpState::Closed { reason } + // } + // } + // } + // LpMessage::KKTResponse(kkt_response) if session.is_initiator() => { + // match session.process_kkt_response(&kkt_response.0) { + // Ok(()) => { + // result_action = Some(Ok(LpAction::KKTComplete)); + // // After successful KKT, move to Handshaking + // LpState::Handshaking { session } + // } + // Err(e) => { + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // } + // } + // _ => { + // // Wrong message type for KKT state + // let err = LpError::InvalidStateTransition { + // state: "KKTExchange".to_string(), + // input: format!("Unexpected message type: {:?}", packet.message), + // }; + // let reason = err.to_string(); + // result_action = Some(Err(err)); + // LpState::Closed { reason } + // } + // } + // } + // } + // + // // Reject SendData during KKT exchange + // (LpState::KKTExchange { session }, LpInput::SendData(_)) => { + // result_action = Some(Err(LpError::InvalidStateTransition { + // state: "KKTExchange".to_string(), + // input: "SendData".to_string(), + // })); + // LpState::KKTExchange { session } + // } + // + // // Reject StartHandshake if already in KKT exchange + // (LpState::KKTExchange { session }, LpInput::StartHandshake) => { + // result_action = Some(Err(LpError::InvalidStateTransition { + // state: "KKTExchange".to_string(), + // input: "StartHandshake".to_string(), + // })); + // LpState::KKTExchange { session } + // } + // + // // --- Handshaking State --- + // (LpState::Handshaking { 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()))); + // // Don't change state, return the original state variant + // LpState::Handshaking { session } + // } else { + // // --- Inline handle_handshake_packet logic --- + // // 1. Check replay protection *before* processing + // if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) { + // let _reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Handshaking { session } + // // LpState::Closed { reason } + // } else { + // // 2. Process the handshake message + // match session.process_handshake_message(&packet.message) { + // Ok(_) => { + // // 3. Mark counter as received *after* successful processing + // if let Err(e) = session.receiving_counter_mark(packet.header.counter) { + // let _reason = e.to_string(); + // result_action = Some(Err(e)); + // // LpState::Closed { reason } + // LpState::Handshaking { session } + // } else { + // // 4. First check if we need to send a handshake message (before checking completion) + // match session.prepare_handshake_message() { + // Some(Ok(message)) => { + // match session.next_packet(message) { + // Ok(response_packet) => { + // result_action = Some(Ok(LpAction::SendPacket(response_packet))); + // // Check if handshake became complete after preparing message + // if session.is_handshake_complete() { + // LpState::Transport { session } // Transition to Transport + // } else { + // LpState::Handshaking { session } // Remain Handshaking + // } + // } + // Err(e) => { + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // } + // } + // Some(Err(e)) => { + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // None => { + // // 5. No message to send - check if handshake is complete + // if session.is_handshake_complete() { + // result_action = Some(Ok(LpAction::HandshakeComplete)); + // LpState::Transport { session } // Transition to Transport + // } else { + // // Handshake stalled unexpectedly + // let err = LpError::NoiseError(NoiseError::Other( + // "Handshake stalled unexpectedly".to_string(), + // )); + // let reason = err.to_string(); + // result_action = Some(Err(err)); + // LpState::Closed { reason } + // } + // } + // } + // } + // } + // Err(e) => { // Error from process_handshake_message + // let reason = e.to_string(); + // result_action = Some(Err(e)); + // LpState::Closed { reason } + // } + // } + // } + // // --- End inline handle_handshake_packet logic --- + // } + // } + // // Reject SendData during handshake + // (LpState::Handshaking { session }, LpInput::SendData(_)) => { // Keep session if returning to this state + // result_action = Some(Err(LpError::InvalidStateTransition { + // state: "Handshaking".to_string(), + // input: "SendData".to_string(), + // })); + // // Invalid input, remain in Handshaking state + // LpState::Handshaking { session } + // } + // // Reject StartHandshake if already handshaking + // (LpState::Handshaking { session }, LpInput::StartHandshake) => { // Keep session + // result_action = Some(Err(LpError::InvalidStateTransition { + // state: "Handshaking".to_string(), + // input: "StartHandshake".to_string(), + // })); + // // Invalid input, remain in Handshaking state + // LpState::Handshaking { session } + // } + // + // // --- Transport State --- + // (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()))); + // 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 { session }, LpInput::SendData(data)) => { + // // Encrypt and send application data + // match session.prepare_data_packet(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 } + // } + // // Reject StartHandshake if already in transport + // (LpState::Transport { session }, LpInput::StartHandshake) => { // Keep session + // result_action = Some(Err(LpError::InvalidStateTransition { + // state: "Transport".to_string(), + // input: "StartHandshake".to_string(), + // })); + // // Invalid input, remain in Transport state + // LpState::Transport { session } + // } + // + // // --- 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: 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 { 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 { session, subsession }, LpInput::SendData(data)) => { + // match session.prepare_data_packet(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 { + // 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::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::SubsessionHandshaking { .. } + // | LpState::ReadOnlyTransport { .. }, + // LpInput::Close, + // ) => { + // result_action = Some(Ok(LpAction::ConnectionClosed)); + // // Transition to Closed state + // LpState::Closed { reason: "Closed by user".to_string() } + // } + // // Ignore Close if already Closed + // (closed_state @ LpState::Closed { .. }, LpInput::Close) => { + // // result_action remains None + // // 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)); + // closed_state + // } + // // Ignore SendData if Closed + // (closed_state @ LpState::Closed { .. }, LpInput::SendData(_)) => { + // result_action = Some(Err(LpError::LpSessionClosed)); + // closed_state + // } + // // Processing state should not be matched directly if using replace + // (LpState::Processing, _) => { + // // This case should ideally be unreachable if placeholder logic is correct + // let err = LpError::Internal("Reached Processing state unexpectedly".to_string()); + // let reason = err.to_string(); + // 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 } + // } + // }; + + // (next_state, result_action) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum LpStateBare { ReadyToHandshake, @@ -304,824 +1199,21 @@ impl<'a> LpStateMachine<'a> { } /// Processes an input event and returns a list of actions to perform. - pub fn process_input<'b: 'a>( - &'b mut self, - input: LpInput, - ) -> Option> { + pub fn process_input(&'a mut self, input: LpInput) -> Option> { // 1. Replace current state with a placeholder, taking ownership of the real current state. - let current_state = mem::take(&mut self.state); - let mut result_action: Option> = None; - - // 2. Match on the owned current_state. Each arm calculates and returns the NEXT state. - let next_state = match (current_state, input) { - // --- ReadyToHandshake State --- - (LpState::ReadyToHandshake { mut session }, LpInput::StartHandshake) => { - if session.is_initiator() { - // Initiator starts by requesting KEM key via KKT - match session.prepare_kkt_request() { - Some(Ok(kkt_message)) => { - match session.next_packet(kkt_message) { - Ok(kkt_packet) => { - result_action = Some(Ok(LpAction::SendPacket(kkt_packet))); - LpState::KKTExchange { session } // Transition to KKTExchange - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Some(Err(e)) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - None => { - // Should not happen for initiator - let err = LpError::Internal( - "prepare_kkt_request returned None for initiator".to_string(), - ); - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - } - } else { - // Responder waits for KKT request - LpState::KKTExchange { session } - // No action needed yet, result_action remains None. - } - } - - // --- KKTExchange State --- - (LpState::KKTExchange { 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::KKTExchange { session } - } else { - use crate::message::LpMessage; - - // Packet message is already parsed, match on it directly - match &packet.message { - LpMessage::KKTRequest(kkt_request) if !session.is_initiator() => { - match session.process_kkt_request(&kkt_request.0) { - Ok(kkt_response_message) => { - match session.next_packet(kkt_response_message) { - Ok(response_packet) => { - result_action = Some(Ok(LpAction::SendPacket(response_packet))); - // After KKT exchange, move to Handshaking - LpState::Handshaking { 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)); - LpState::Closed { reason } - } - } - } - LpMessage::KKTResponse(kkt_response) if session.is_initiator() => { - match session.process_kkt_response(&kkt_response.0) { - Ok(()) => { - result_action = Some(Ok(LpAction::KKTComplete)); - // After successful KKT, move to Handshaking - LpState::Handshaking { session } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - _ => { - // Wrong message type for KKT state - let err = LpError::InvalidStateTransition { - state: "KKTExchange".to_string(), - input: format!("Unexpected message type: {:?}", packet.message), - }; - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - } - } - } - - // Reject SendData during KKT exchange - (LpState::KKTExchange { session }, LpInput::SendData(_)) => { - result_action = Some(Err(LpError::InvalidStateTransition { - state: "KKTExchange".to_string(), - input: "SendData".to_string(), - })); - LpState::KKTExchange { session } - } - - // Reject StartHandshake if already in KKT exchange - (LpState::KKTExchange { session }, LpInput::StartHandshake) => { - result_action = Some(Err(LpError::InvalidStateTransition { - state: "KKTExchange".to_string(), - input: "StartHandshake".to_string(), - })); - LpState::KKTExchange { session } - } - - // --- Handshaking State --- - (LpState::Handshaking { 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()))); - // Don't change state, return the original state variant - LpState::Handshaking { session } - } else { - // --- Inline handle_handshake_packet logic --- - // 1. Check replay protection *before* processing - if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) { - let _reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Handshaking { session } - // LpState::Closed { reason } - } else { - // 2. Process the handshake message - match session.process_handshake_message(&packet.message) { - Ok(_) => { - // 3. Mark counter as received *after* successful processing - if let Err(e) = session.receiving_counter_mark(packet.header.counter) { - let _reason = e.to_string(); - result_action = Some(Err(e)); - // LpState::Closed { reason } - LpState::Handshaking { session } - } else { - // 4. First check if we need to send a handshake message (before checking completion) - match session.prepare_handshake_message() { - Some(Ok(message)) => { - match session.next_packet(message) { - Ok(response_packet) => { - result_action = Some(Ok(LpAction::SendPacket(response_packet))); - // Check if handshake became complete after preparing message - if session.is_handshake_complete() { - LpState::Transport { session } // Transition to Transport - } else { - LpState::Handshaking { session } // Remain Handshaking - } - } - Err(e) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - Some(Err(e)) => { - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - None => { - // 5. No message to send - check if handshake is complete - if session.is_handshake_complete() { - result_action = Some(Ok(LpAction::HandshakeComplete)); - LpState::Transport { session } // Transition to Transport - } else { - // Handshake stalled unexpectedly - let err = LpError::NoiseError(NoiseError::Other( - "Handshake stalled unexpectedly".to_string(), - )); - let reason = err.to_string(); - result_action = Some(Err(err)); - LpState::Closed { reason } - } - } - } - } - } - Err(e) => { // Error from process_handshake_message - let reason = e.to_string(); - result_action = Some(Err(e)); - LpState::Closed { reason } - } - } - } - // --- End inline handle_handshake_packet logic --- - } - } - // Reject SendData during handshake - (LpState::Handshaking { session }, LpInput::SendData(_)) => { // Keep session if returning to this state - result_action = Some(Err(LpError::InvalidStateTransition { - state: "Handshaking".to_string(), - input: "SendData".to_string(), - })); - // Invalid input, remain in Handshaking state - LpState::Handshaking { session } - } - // Reject StartHandshake if already handshaking - (LpState::Handshaking { session }, LpInput::StartHandshake) => { // Keep session - result_action = Some(Err(LpError::InvalidStateTransition { - state: "Handshaking".to_string(), - input: "StartHandshake".to_string(), - })); - // Invalid input, remain in Handshaking state - LpState::Handshaking { session } - } - - // --- Transport State --- - (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()))); - 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 { session }, LpInput::SendData(data)) => { - // Encrypt and send application data - match self.prepare_data_packet(&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 } - } - // Reject StartHandshake if already in transport - (LpState::Transport { session }, LpInput::StartHandshake) => { // Keep session - result_action = Some(Err(LpError::InvalidStateTransition { - state: "Transport".to_string(), - input: "StartHandshake".to_string(), - })); - // Invalid input, remain in Transport state - LpState::Transport { session } - } - - // --- 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: 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 { 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 { 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 { - 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::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::SubsessionHandshaking { .. } - | LpState::ReadOnlyTransport { .. }, - LpInput::Close, - ) => { - result_action = Some(Ok(LpAction::ConnectionClosed)); - // Transition to Closed state - LpState::Closed { reason: "Closed by user".to_string() } - } - // Ignore Close if already Closed - (closed_state @ LpState::Closed { .. }, LpInput::Close) => { - // result_action remains None - // 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)); - closed_state - } - // Ignore SendData if Closed - (closed_state @ LpState::Closed { .. }, LpInput::SendData(_)) => { - result_action = Some(Err(LpError::LpSessionClosed)); - closed_state - } - // Processing state should not be matched directly if using replace - (LpState::Processing, _) => { - // This case should ideally be unreachable if placeholder logic is correct - let err = LpError::Internal("Reached Processing state unexpectedly".to_string()); - let reason = err.to_string(); - 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. - self.state = next_state; - - result_action // Return the determined action (or None) - } - - // Helper to prepare an outgoing data packet - // Kept as it doesn't mutate self.state - fn prepare_data_packet( - &self, - session: &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? + self.state.process_input(input) + // let mut current_state = mem::take(&mut self.state); + // + // let result_action = current_state.process_input(input); + // self.state = current_state; + // + // result_action + // // let (next_state, result_action) = current_state.process_input(input); + // // + // // self.state = next_state; + // // + // // result_action // Return the determined action (or None) } }