From 0d9d97e31ee0db39e5a48f4c1e0a6116382212a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Tue, 3 Mar 2026 16:20:27 +0000 Subject: [PATCH 1/2] remove redundant LP state machine in favour of in place processing --- common/nym-lp/src/lib.rs | 18 +- common/nym-lp/src/psq/initiator.rs | 11 +- common/nym-lp/src/psq/responder.rs | 11 +- common/nym-lp/src/session.rs | 126 ++++++- common/nym-lp/src/session_integration/mod.rs | 72 +--- common/nym-lp/src/session_manager.rs | 40 +- common/nym-lp/src/state_machine.rs | 352 ------------------ nym-node/src/node/lp/cleanup.rs | 6 +- nym-node/src/node/lp/control/handler.rs | 30 +- nym-node/src/node/lp/error.rs | 5 +- nym-node/src/node/lp/state.rs | 4 +- .../src/lp_client/client.rs | 38 +- .../src/lp_client/error.rs | 5 +- .../src/lp_client/helpers.rs | 2 +- .../lp_client/nested_session/connection.rs | 14 +- .../src/lp_client/nested_session/mod.rs | 19 +- .../src/lp_client/state_machine_helpers.rs | 16 +- 17 files changed, 213 insertions(+), 556 deletions(-) delete mode 100644 common/nym-lp/src/state_machine.rs diff --git a/common/nym-lp/src/lib.rs b/common/nym-lp/src/lib.rs index af1a273c2a..a1193bd153 100644 --- a/common/nym-lp/src/lib.rs +++ b/common/nym-lp/src/lib.rs @@ -11,7 +11,6 @@ pub mod replay; pub mod session; mod session_integration; pub mod session_manager; -pub mod state_machine; pub mod transport; pub use error::LpError; @@ -21,9 +20,8 @@ pub use nym_kkt_ciphersuite::{ #[cfg(any(feature = "mock", test))] pub use replay::{ReceivingKeyCounterValidator, ReplayError}; -pub use session::LpSession; +pub use session::LpTransportSession; pub use session_manager::SessionManager; -pub use state_machine::LpStateMachine; #[cfg(any(feature = "mock", test))] use nym_test_utils::helpers::u64_seeded_rng_09; @@ -39,8 +37,8 @@ use libcrux_psq::{Channel, IntoSession}; #[cfg(any(feature = "mock", test))] pub struct SessionsMock { - pub initiator: LpSession, - pub responder: LpSession, + pub initiator: LpTransportSession, + pub responder: LpTransportSession, } #[cfg(any(feature = "mock", test))] @@ -113,14 +111,14 @@ impl SessionsMock { }; SessionsMock { - initiator: LpSession::new( + initiator: LpTransportSession::new( initiator.into_session().unwrap(), binding.clone(), receiver_index, 1, ) .unwrap(), - responder: LpSession::new( + responder: LpTransportSession::new( responder.into_session().unwrap(), binding, receiver_index, @@ -135,18 +133,18 @@ impl SessionsMock { } // we just need a dummy 'valid' session for simpler tests - pub fn mock_initiator() -> LpSession { + pub fn mock_initiator() -> LpTransportSession { Self::mock_post_handshake(KEM::default()).initiator } } #[cfg(any(feature = "mock", test))] -pub fn sessions_for_tests() -> (LpSession, LpSession) { +pub fn sessions_for_tests() -> (LpTransportSession, LpTransportSession) { let sessions = SessionsMock::mock_post_handshake(KEM::default()); (sessions.initiator, sessions.responder) } #[cfg(any(feature = "mock", test))] -pub fn mock_session_for_test() -> LpSession { +pub fn mock_session_for_test() -> LpTransportSession { SessionsMock::mock_initiator() } diff --git a/common/nym-lp/src/psq/initiator.rs b/common/nym-lp/src/psq/initiator.rs index 7875e1196a..81fd571e01 100644 --- a/common/nym-lp/src/psq/initiator.rs +++ b/common/nym-lp/src/psq/initiator.rs @@ -11,7 +11,7 @@ use crate::psq::{ }; use crate::session::PersistentSessionBinding; use crate::transport::traits::LpHandshakeChannel; -use crate::{LpError, LpSession}; +use crate::{LpError, LpTransportSession}; use libcrux_psq::handshake::RegistrationInitiator; use libcrux_psq::handshake::builders::{ CiphersuiteBuilder, InitiatorCiphersuite, PrincipalBuilder, @@ -113,7 +113,7 @@ where Ok(resp.into()) } - pub async fn complete_handshake(self) -> Result + pub async fn complete_handshake(self) -> Result where S: LpHandshakeChannel + Unpin, { @@ -121,7 +121,10 @@ where self.complete_handshake_with_rng(&mut rng).await } - pub async fn complete_handshake_with_rng(mut self, rng: &mut R) -> Result + pub async fn complete_handshake_with_rng( + mut self, + rng: &mut R, + ) -> Result where S: LpHandshakeChannel + Unpin, R: rand09::CryptoRng, @@ -230,7 +233,7 @@ where }; let psq_session = psq_initiator.into_session()?; - LpSession::new(psq_session, binding, receiver_index, protocol) + LpTransportSession::new(psq_session, binding, receiver_index, protocol) } } diff --git a/common/nym-lp/src/psq/responder.rs b/common/nym-lp/src/psq/responder.rs index 59cf61046d..8e3f26c58f 100644 --- a/common/nym-lp/src/psq/responder.rs +++ b/common/nym-lp/src/psq/responder.rs @@ -11,7 +11,7 @@ use crate::psq::{ }; use crate::session::PersistentSessionBinding; use crate::transport::traits::{HandshakeMessage, LpHandshakeChannel}; -use crate::{LpError, LpSession}; +use crate::{LpError, LpTransportSession}; use libcrux_psq::handshake::Responder; use libcrux_psq::handshake::builders::{ CiphersuiteBuilder, PrincipalBuilder, ResponderCiphersuite, @@ -146,7 +146,7 @@ where Ok(msg.into_bytes()) } - pub async fn complete_handshake(self) -> Result + pub async fn complete_handshake(self) -> Result where S: LpHandshakeChannel + Unpin, { @@ -154,7 +154,10 @@ where self.complete_handshake_with_rng(&mut rng).await } - pub async fn complete_handshake_with_rng(mut self, rng: &mut R) -> Result + pub async fn complete_handshake_with_rng( + mut self, + rng: &mut R, + ) -> Result where S: LpHandshakeChannel + Unpin, R: rand09::CryptoRng, @@ -229,7 +232,7 @@ where }; let psq_session = psq_responder.into_session()?; - LpSession::new( + LpTransportSession::new( psq_session, binding, receiver_index, diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 9497648175..02d5effa23 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -21,13 +21,34 @@ use libcrux_psq::session::{Session, SessionBinding}; use nym_kkt::keys::EncapsulationKey; use std::fmt::{Debug, Formatter}; +/// Represents inputs that drive the state machine transitions. +#[allow(clippy::large_enum_variant)] +#[derive(Debug)] +pub enum LpInput { + /// Received an encrypted LP Packet from the network. + ReceivePacket(EncryptedLpPacket), + + /// Application wants to send data (only valid in Transport state). + SendData(LpMessage), +} + +/// Represents actions the state machine requests the environment to perform. +#[derive(Debug)] +pub enum LpAction { + /// Send an LP Packet over the network. + SendPacket(EncryptedLpPacket), + + /// Deliver decrypted application data received from the peer. + DeliverData(LpMessage), +} + pub type SessionId = [u8; 32]; /// A session in the Lewes Protocol, handling connection state with Noise. /// /// Sessions manage connection state, including LP replay protection. /// Each session has a unique receiving index and sending index for connection identification. -pub struct LpSession { +pub struct LpTransportSession { /// The underlying established session psq_session: Session, @@ -90,7 +111,7 @@ impl<'a> From<&'a PersistentSessionBinding> for SessionBinding<'a> { } } -impl Debug for LpSession { +impl Debug for LpTransportSession { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("LpSession") .field("session_id", &self.psq_session.identifier()) @@ -103,7 +124,7 @@ impl Debug for LpSession { } } -impl LpSession { +impl LpTransportSession { /// Creates a new session after completed KTT/PSQ exchange pub fn new( mut psq_session: Session, @@ -116,7 +137,7 @@ impl LpSession { .transport_channel() .map_err(|inner| LpError::TransportDerivationFailure { inner })?; - Ok(LpSession { + Ok(LpTransportSession { psq_session, session_binding, active_transport: transport, @@ -275,6 +296,48 @@ impl LpSession { ) -> Result { decrypt_lp_packet(packet, &mut self.active_transport) } + + /// Processes an input event and returns an action to perform. + pub fn process_input(&mut self, input: LpInput) -> Result { + match input { + LpInput::ReceivePacket(packet) => { + // Check if packet lp_id matches our session + if packet.outer_header().receiver_idx != self.receiver_index() { + return Err(LpError::UnknownSessionId( + packet.outer_header().receiver_idx, + )); + } + + let ctr = packet.outer_header().counter; + + // 1. Check replay protection + if let Err(e) = self.receiving_counter_quick_check(ctr) { + return Err(e); + } + + // 2. decrypt the packet and attempt to deliver data + let packet = match self.decrypt_packet(packet) { + Ok(packet) => packet, + Err(e) => return Err(e), + }; + + // 3. Mark counter as received + if let Err(e) = self.receiving_counter_mark(ctr) { + return Err(e); + } + + // 4. deliver the message + Ok(LpAction::DeliverData(packet.message)) + } + LpInput::SendData(data) => { + // Encrypt and send application data + match self.encrypt_application_data(data) { + Ok(packet) => Ok(LpAction::SendPacket(packet)), + Err(e) => Err(e), + } + } + } + } } #[cfg(test)] @@ -364,4 +427,59 @@ mod tests { assert_eq!(packet_count.received, 2); } } + + #[test] + fn test_state_machine_simplified_flow() { + for kem in KEM::iter() { + let mock_sessions = SessionsMock::mock_post_handshake(kem); + let receiver_index = mock_sessions.responder.receiver_index(); + + // Create state machines (already in Transport) + let mut initiator = mock_sessions.initiator; + let mut responder = mock_sessions.responder; + + assert_eq!( + initiator.session_identifier(), + responder.session_identifier() + ); + + // --- Transport Phase --- + println!("--- Step 1: Initiator sends data ---"); + let data_to_send_1 = LpMessage::new_opaque(b"hello responder".to_vec()); + let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); + let data_packet_1 = if let Ok(LpAction::SendPacket(packet)) = init_actions_4 { + packet.clone() + } else { + panic!("Initiator should send data packet"); + }; + assert_eq!(data_packet_1.outer_header().receiver_idx, receiver_index); + + println!("--- Step 2: Responder receives data ---"); + let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); + let resp_data_1 = if let Ok(LpAction::DeliverData(data)) = resp_actions_5 { + data + } else { + panic!("Responder should deliver data"); + }; + assert_eq!(resp_data_1, data_to_send_1); + + println!("--- Step 3: Responder sends data ---"); + let data_to_send_2 = LpMessage::new_opaque(b"hello initiator".to_vec()); + let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); + let data_packet_2 = if let Ok(LpAction::SendPacket(packet)) = resp_actions_6 { + packet.clone() + } else { + panic!("Responder should send data packet"); + }; + assert_eq!(data_packet_2.outer_header().receiver_idx, receiver_index); + + println!("--- Step 4: Initiator receives data ---"); + let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); + if let Ok(LpAction::DeliverData(data)) = init_actions_5 { + assert_eq!(data, data_to_send_2); + } else { + panic!("Initiator should deliver data"); + } + } + } } diff --git a/common/nym-lp/src/session_integration/mod.rs b/common/nym-lp/src/session_integration/mod.rs index ae2387a745..b628c4324b 100644 --- a/common/nym-lp/src/session_integration/mod.rs +++ b/common/nym-lp/src/session_integration/mod.rs @@ -1,7 +1,7 @@ #[cfg(test)] mod tests { use crate::packet::{EncryptedLpPacket, LpMessage}; - use crate::state_machine::{LpAction, LpInput, LpStateBare}; + use crate::session::{LpAction, LpInput}; use crate::{LpError, SessionManager, SessionsMock}; use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; @@ -62,7 +62,6 @@ mod tests { let decrypted_payload = session_manager_2 .receive_packet(peer_b_sm, ciphertext_a) .unwrap() - .unwrap() .data(); assert_eq!(decrypted_payload.content, plaintext_a); @@ -77,7 +76,6 @@ mod tests { let decrypted_payload = session_manager_1 .receive_packet(peer_a_sm, ciphertext_b) .unwrap() - .unwrap() .data(); assert_eq!(decrypted_payload.content, plaintext_b); } @@ -183,16 +181,6 @@ mod tests { assert!(session_manager_1.state_machine_exists(session_id)); assert!(session_manager_2.state_machine_exists(session_id)); - // Verify initial states are Transport - assert_eq!( - session_manager_1.get_state(session_id).unwrap(), - LpStateBare::Transport - ); - assert_eq!( - session_manager_2.get_state(session_id).unwrap(), - LpStateBare::Transport - ); - // --- 3. Simulate Data Transfer via process_input --- println!("Starting data transfer simulation via process_input..."); let plaintext_a_to_b = @@ -204,7 +192,6 @@ mod tests { println!(" A sends to B"); let action_a_send = session_manager_1 .process_input(session_id, LpInput::SendData(plaintext_a_to_b.clone())) - .expect("A SendData should produce action") .expect("A SendData failed"); let data_packet_a = action_a_send.ciphertext(); @@ -213,7 +200,6 @@ mod tests { println!(" B receives from A"); let action_b_recv = session_manager_2 .process_input(session_id, LpInput::ReceivePacket(data_packet_a)) - .expect("B ReceivePacket (data) should produce action") .expect("B ReceivePacket (data) failed"); if let LpAction::DeliverData(data) = action_b_recv { @@ -230,7 +216,6 @@ mod tests { println!(" B sends to A"); let action_b_send = session_manager_2 .process_input(session_id, LpInput::SendData(plaintext_b_to_a.clone())) - .expect("B SendData should produce action") .expect("B SendData failed"); let data_packet_b = action_b_send.ciphertext(); @@ -242,7 +227,6 @@ mod tests { println!(" A receives from B"); let action_a_recv = session_manager_1 .process_input(session_id, LpInput::ReceivePacket(data_packet_b)) - .expect("A ReceivePacket (data) should produce action") .expect("A ReceivePacket (data) failed"); if let LpAction::DeliverData(data) = action_a_recv { @@ -279,7 +263,6 @@ mod tests { let action_send_n1 = session_manager_1 .process_input(session_id, LpInput::SendData(data_n_plus_1.clone())) - .unwrap() .unwrap(); let packet_n1 = match action_send_n1 { LpAction::SendPacket(p) => p, @@ -288,7 +271,6 @@ mod tests { let action_send_n = session_manager_1 .process_input(session_id, LpInput::SendData(data_n.clone())) - .unwrap() .unwrap(); let packet_n = match action_send_n { LpAction::SendPacket(p) => p, @@ -300,7 +282,6 @@ mod tests { println!(" B receives N+1"); let action_recv_n1 = session_manager_2 .process_input(session_id, LpInput::ReceivePacket(packet_n1)) - .unwrap() .unwrap(); match action_recv_n1 { LpAction::DeliverData(d) => assert_eq!(d, data_n_plus_1, "Data N+1 mismatch"), @@ -311,7 +292,6 @@ mod tests { println!(" B receives N"); let action_recv_n = session_manager_2 .process_input(session_id, LpInput::ReceivePacket(packet_n)) - .unwrap() .unwrap(); match action_recv_n { LpAction::DeliverData(d) => assert_eq!(d, data_n, "Data N mismatch"), @@ -329,55 +309,7 @@ mod tests { ); println!("Out-of-order test passed."); - // --- 6. Close Test --- - println!("Testing close via process_input..."); - - // A closes - let action_a_close = session_manager_1 - .process_input(session_id, LpInput::Close) - .expect("A Close should produce action") - .expect("A Close failed"); - assert!(matches!(action_a_close, LpAction::ConnectionClosed)); - assert_eq!( - session_manager_1.get_state(session_id).unwrap(), - LpStateBare::Closed - ); - - // Further actions on A fail - let send_after_close_a = session_manager_1.process_input( - session_id, - LpInput::SendData(LpMessage::new_opaque(b"fail".to_vec())), - ); - assert!(send_after_close_a.is_err()); - assert!(matches!( - send_after_close_a.err().unwrap(), - LpError::LpSessionClosed - )); - - // B closes - let action_b_close = session_manager_2 - .process_input(session_id, LpInput::Close) - .expect("B Close should produce action") - .expect("B Close failed"); - assert!(matches!(action_b_close, LpAction::ConnectionClosed)); - assert_eq!( - session_manager_2.get_state(session_id).unwrap(), - LpStateBare::Closed - ); - - // Further actions on B fail - let send_after_close_b = session_manager_2.process_input( - session_id, - LpInput::SendData(LpMessage::new_opaque(b"fail".to_vec())), - ); - assert!(send_after_close_b.is_err()); - assert!(matches!( - send_after_close_b.err().unwrap(), - LpError::LpSessionClosed - )); - println!("Close test passed."); - - // --- 7. Session Removal --- + // --- 6. Session Removal --- assert!(session_manager_1.remove_state_machine(session_id)); assert_eq!(session_manager_1.session_count(), 0); assert!(!session_manager_1.state_machine_exists(session_id)); diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index 30e1cc7921..d944b478ca 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -8,18 +8,18 @@ use crate::packet::{EncryptedLpPacket, LpMessage}; use crate::peer_config::LpReceiverIndex; -use crate::state_machine::{LpAction, LpInput, LpStateBare}; -use crate::{LpError, LpSession, LpStateMachine}; +use crate::{LpError, LpTransportSession}; use std::collections::HashMap; pub use crate::replay::validator::PacketCount; +use crate::session::{LpAction, LpInput}; /// Manages the lifecycle of Lewes Protocol sessions. /// /// The SessionManager is responsible for creating, storing, and retrieving sessions pub struct SessionManager { /// Manages state machines directly, keyed by lp_id - state_machines: HashMap, + state_machines: HashMap, } impl Default for SessionManager { @@ -40,8 +40,8 @@ impl SessionManager { &mut self, lp_id: LpReceiverIndex, input: LpInput, - ) -> Result, LpError> { - self.with_state_machine_mut(lp_id, |sm| sm.process_input(input).transpose())? + ) -> Result { + self.with_state_machine_mut(lp_id, |sm| sm.process_input(input))? } pub fn send_data( @@ -49,37 +49,24 @@ impl SessionManager { lp_id: LpReceiverIndex, data: LpMessage, ) -> Result { - self.process_input(lp_id, LpInput::SendData(data))? - .ok_or(LpError::NotInTransport) + self.process_input(lp_id, LpInput::SendData(data)) } pub fn receive_packet( &mut self, lp_id: LpReceiverIndex, packet: EncryptedLpPacket, - ) -> Result, LpError> { + ) -> Result { self.process_input(lp_id, LpInput::ReceivePacket(packet)) } - pub fn closed(&self, lp_id: LpReceiverIndex) -> Result { - Ok(self.get_state(lp_id)? == LpStateBare::Closed) - } - - pub fn transport(&self, lp_id: LpReceiverIndex) -> Result { - Ok(self.get_state(lp_id)? == LpStateBare::Transport) - } - #[cfg(test)] fn get_state_machine_id(&self, lp_id: LpReceiverIndex) -> Result { - self.with_state_machine(lp_id, |sm| sm.receiver_index())? - } - - pub fn get_state(&self, lp_id: LpReceiverIndex) -> Result { - self.with_state_machine(lp_id, |sm| Ok(sm.bare_state()))? + self.with_state_machine(lp_id, |sm| sm.receiver_index()) } pub fn current_packet_cnt(&self, lp_id: LpReceiverIndex) -> Result { - self.with_state_machine(lp_id, |sm| Ok(sm.session()?.current_packet_cnt()))? + self.with_state_machine(lp_id, |sm| Ok(sm.current_packet_cnt()))? } pub fn session_count(&self) -> usize { @@ -92,7 +79,7 @@ impl SessionManager { pub fn with_state_machine(&self, lp_id: LpReceiverIndex, f: F) -> Result where - F: FnOnce(&LpStateMachine) -> R, + F: FnOnce(&LpTransportSession) -> R, { if let Some(sm) = self.state_machines.get(&lp_id) { Ok(f(sm)) @@ -108,7 +95,7 @@ impl SessionManager { f: F, ) -> Result where - F: FnOnce(&mut LpStateMachine) -> R, // Closure takes mutable ref + F: FnOnce(&mut LpTransportSession) -> R, // Closure takes mutable ref { if let Some(sm) = self.state_machines.get_mut(&lp_id) { Ok(f(sm)) @@ -119,7 +106,7 @@ impl SessionManager { pub fn create_session_state_machine( &mut self, - lp_session: LpSession, + lp_session: LpTransportSession, ) -> Result { let session_id = lp_session.receiver_index(); @@ -127,8 +114,7 @@ impl SessionManager { return Err(LpError::DuplicateSessionId(session_id)); } - let sm = LpStateMachine::new(lp_session); - self.state_machines.insert(session_id, sm); + self.state_machines.insert(session_id, lp_session); Ok(session_id) } diff --git a/common/nym-lp/src/state_machine.rs b/common/nym-lp/src/state_machine.rs deleted file mode 100644 index 442bff94b7..0000000000 --- a/common/nym-lp/src/state_machine.rs +++ /dev/null @@ -1,352 +0,0 @@ -// Copyright 2025 - Nym Technologies SA -// SPDX-License-Identifier: Apache-2.0 - -//! Lewes Protocol State Machine for managing connection lifecycle. -//! State machine ensures protocol steps execute in correct order. Invalid transitions -//! return LpError, preventing protocol violations. - -use crate::packet::EncryptedLpPacket; -use crate::packet::message::LpMessage; -use crate::peer_config::LpReceiverIndex; -use crate::session::SessionId; -use crate::{LpError, session::LpSession}; -use std::mem; - -#[derive(Debug)] -pub struct LpTransportState { - /// The underlying session in the transport state - session: Box, -} - -/// Represents the possible states of the Lewes Protocol connection. -#[derive(Debug, Default)] -pub enum LpState { - /// Handshake complete, ready for data transport. - Transport(LpTransportState), - - /// An error occurred, or the connection was intentionally closed. - Closed { reason: String }, - - /// Processing an input event. - #[default] - Processing, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum LpStateBare { - Transport, - Closed, - Processing, -} - -impl From<&LpState> for LpStateBare { - fn from(state: &LpState) -> Self { - match state { - LpState::Transport { .. } => LpStateBare::Transport, - LpState::Closed { .. } => LpStateBare::Closed, - LpState::Processing => LpStateBare::Processing, - } - } -} - -/// Represents inputs that drive the state machine transitions. -#[allow(clippy::large_enum_variant)] -#[derive(Debug)] -pub enum LpInput { - /// Received an encrypted LP Packet from the network. - ReceivePacket(EncryptedLpPacket), - - /// Application wants to send data (only valid in Transport state). - SendData(LpMessage), - - /// Close the connection. - Close, -} - -/// Represents actions the state machine requests the environment to perform. -#[derive(Debug)] -pub enum LpAction { - /// Send an LP Packet over the network. - SendPacket(EncryptedLpPacket), - - /// Deliver decrypted application data received from the peer. - DeliverData(LpMessage), - - /// Inform the environment that the connection is closed. - ConnectionClosed, -} - -/// The Lewes Protocol State Machine. -pub struct LpStateMachine { - pub state: LpState, -} - -impl LpStateMachine { - pub fn bare_state(&self) -> LpStateBare { - LpStateBare::from(&self.state) - } - - pub fn session_mut(&mut self) -> Result<&mut LpSession, LpError> { - match &mut self.state { - LpState::Transport(transport) => Ok(&mut transport.session), - LpState::Closed { .. } => Err(LpError::LpSessionClosed), - LpState::Processing => Err(LpError::LpSessionProcessing), - } - } - - pub fn session(&self) -> Result<&LpSession, LpError> { - match &self.state { - LpState::Transport(transport) => Ok(&transport.session), - LpState::Closed { .. } => Err(LpError::LpSessionClosed), - LpState::Processing => Err(LpError::LpSessionProcessing), - } - } - - /// Consume the state machine and return the session with ownership. - /// This is useful when the handshake is complete and you want to transfer - /// ownership of the session to the caller. - pub fn into_session(self) -> Result { - match self.state { - LpState::Transport(transport) => Ok(*transport.session), - LpState::Closed { .. } => Err(LpError::LpSessionClosed), - LpState::Processing => Err(LpError::LpSessionProcessing), - } - } - - pub fn session_identifier(&self) -> Result { - Ok(*self.session()?.session_identifier()) - } - - pub fn receiver_index(&self) -> Result { - Ok(self.session()?.receiver_index()) - } - - /// Creates a new state machine in `Transport` state post-KKT/PSQ handshake - pub fn new(session: LpSession) -> Self { - LpStateMachine { - state: LpState::Transport(LpTransportState { - session: Box::new(session), - }), - } - } - - fn process_input_transport( - &mut self, - mut state: LpTransportState, - input: LpInput, - ) -> (LpState, Option>) { - let session = &mut state.session; - match input { - LpInput::ReceivePacket(packet) => { - // Check if packet lp_id matches our session - if packet.outer_header().receiver_idx != session.receiver_index() { - let result_action = Some(Err(LpError::UnknownSessionId( - packet.outer_header().receiver_idx, - ))); - return (LpState::Transport(state), result_action); - } - - let ctr = packet.outer_header().counter; - - // 1. Check replay protection - if let Err(e) = session.receiving_counter_quick_check(ctr) { - return (LpState::Transport(state), Some(Err(e))); - } - - // 2. decrypt the packet and attempt to deliver data - let packet = match session.decrypt_packet(packet) { - Ok(packet) => packet, - Err(e) => return (LpState::Transport(state), Some(Err(e))), - }; - - // 3. Mark counter as received - if let Err(e) = session.receiving_counter_mark(ctr) { - return (LpState::Transport(state), Some(Err(e))); - } - - // 4. deliver the message - let message = packet.message; - let result_action = Some(Ok(LpAction::DeliverData(message))); - (LpState::Transport(state), result_action) - } - LpInput::SendData(data) => { - // Encrypt and send application data - let result_action = match self.prepare_data_packet(session, data) { - Ok(packet) => Some(Ok(LpAction::SendPacket(packet))), - Err(e) => { - // If prepare fails, should we close? Let's report error and stay Transport for now. - // Alternative: transition to Closed state. - Some(Err(e)) - } - }; - // Remain in transport state - (LpState::Transport(state), result_action) - } - - // --- Close Transition --- - LpInput::Close => { - // Transition to Closed state - ( - LpState::Closed { - reason: "Closed by user".to_string(), - }, - Some(Ok(LpAction::ConnectionClosed)), - ) - } - } - } - - /// Processes an input event and returns a list of actions to perform. - pub fn process_input(&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) { - // --- Transport State --- - (LpState::Transport(transport), input) => { - let (next_state, action) = self.process_input_transport(transport, input); - result_action = action; - next_state - } - // Ignore Close if already Closed - (closed_state @ LpState::Closed { .. }, LpInput::Close) => { - // result_action remains None - // Return the original closed state - 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 } - } - }; - - // 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: &mut LpSession, - data: LpMessage, - ) -> Result { - session.encrypt_application_data(data) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::SessionsMock; - use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; - - #[test] - fn test_state_machine_init() { - for kem in KEM::iter() { - let mock_sessions = SessionsMock::mock_post_handshake(kem); - - let initiator_sm = LpStateMachine::new(mock_sessions.initiator); - assert!(matches!(initiator_sm.state, LpState::Transport { .. })); - let init_session = initiator_sm.session().unwrap(); - - let responder_sm = LpStateMachine::new(mock_sessions.responder); - assert!(matches!(responder_sm.state, LpState::Transport { .. })); - let resp_session = responder_sm.session().unwrap(); - - // Check both state machines use the same receiver_index - assert_eq!(init_session.receiver_index(), resp_session.receiver_index()); - } - } - - #[test] - fn test_state_machine_simplified_flow() { - for kem in KEM::iter() { - let mock_sessions = SessionsMock::mock_post_handshake(kem); - let receiver_index = mock_sessions.responder.receiver_index(); - - // Create state machines (already in Transport) - let mut initiator = LpStateMachine::new(mock_sessions.initiator); - let mut responder = LpStateMachine::new(mock_sessions.responder); - - assert_eq!( - initiator.session_identifier().unwrap(), - responder.session_identifier().unwrap() - ); - - // --- Transport Phase --- - println!("--- Step 1: Initiator sends data ---"); - let data_to_send_1 = LpMessage::new_opaque(b"hello responder".to_vec()); - let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); - let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_4 { - packet.clone() - } else { - panic!("Initiator should send data packet"); - }; - assert_eq!(data_packet_1.outer_header().receiver_idx, receiver_index); - - println!("--- Step 2: Responder receives data ---"); - let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); - let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_5 { - data - } else { - panic!("Responder should deliver data"); - }; - assert_eq!(resp_data_1, data_to_send_1); - - println!("--- Step 3: Responder sends data ---"); - let data_to_send_2 = LpMessage::new_opaque(b"hello initiator".to_vec()); - let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); - let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_6 { - packet.clone() - } else { - panic!("Responder should send data packet"); - }; - assert_eq!(data_packet_2.outer_header().receiver_idx, receiver_index); - - println!("--- Step 4: Initiator receives data ---"); - let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); - if let Some(Ok(LpAction::DeliverData(data))) = init_actions_5 { - assert_eq!(data, data_to_send_2); - } else { - panic!("Initiator should deliver data"); - } - - // --- Close --- - println!("--- Step 5: Initiator closes ---"); - let init_actions_6 = initiator.process_input(LpInput::Close); - assert!(matches!( - init_actions_6, - Some(Ok(LpAction::ConnectionClosed)) - )); - assert!(matches!(initiator.state, LpState::Closed { .. })); - - println!("--- Step 6: Responder closes ---"); - let resp_actions_7 = responder.process_input(LpInput::Close); - assert!(matches!( - resp_actions_7, - Some(Ok(LpAction::ConnectionClosed)) - )); - assert!(matches!(responder.state, LpState::Closed { .. })); - } - } -} diff --git a/nym-node/src/node/lp/cleanup.rs b/nym-node/src/node/lp/cleanup.rs index 920469dca2..4b3ff23de6 100644 --- a/nym-node/src/node/lp/cleanup.rs +++ b/nym-node/src/node/lp/cleanup.rs @@ -3,7 +3,7 @@ use crate::config::LpDebug; use dashmap::DashMap; -use nym_lp::LpStateMachine; +use nym_lp::LpTransportSession; use nym_lp::peer_config::LpReceiverIndex; use nym_metrics::inc_by; use std::sync::Arc; @@ -74,14 +74,14 @@ impl TimestampedState { } pub(crate) struct CleanupTask { - session_states: Arc>>, + session_states: Arc>>, cfg: LpDebug, shutdown: nym_task::ShutdownToken, } impl CleanupTask { pub fn new( - session_states: Arc>>, + session_states: Arc>>, cfg: LpDebug, shutdown: nym_task::ShutdownToken, ) -> Self { diff --git a/nym-node/src/node/lp/control/handler.rs b/nym-node/src/node/lp/control/handler.rs index aaddf37e9d..b9bcd1d6cb 100644 --- a/nym-node/src/node/lp/control/handler.rs +++ b/nym-node/src/node/lp/control/handler.rs @@ -8,10 +8,10 @@ use dashmap::mapref::one::RefMut; use nym_lp::packet::message::LpMessageType; use nym_lp::packet::{EncryptedLpPacket, ForwardPacketData, LpMessage}; use nym_lp::peer_config::LpReceiverIndex; -use nym_lp::state_machine::{LpAction, LpInput}; +use nym_lp::session::{LpAction, LpInput}; use nym_lp::transport::LpHandshakeChannel; use nym_lp::transport::traits::LpTransportChannel; -use nym_lp::{LpSession, LpStateMachine, packet::message::ExpectedResponseSize}; +use nym_lp::{LpTransportSession, packet::message::ExpectedResponseSize}; use nym_metrics::{add_histogram_obs, inc}; use nym_registration_common::{LpRegistrationRequest, RegistrationStatus}; use std::net::SocketAddr; @@ -119,7 +119,8 @@ where /// It is vital it's never held across await points or this might lead to a deadlock. fn state_entry_mut( &self, - ) -> Result>, LpHandlerError> { + ) -> Result>, LpHandlerError> + { let receiver_index = self.bound_receiver_index()?; self.state .shared @@ -152,7 +153,7 @@ where let stream = &mut self.stream; let session = match tokio::time::timeout(timeout, async move { - LpSession::psq_handshake_responder(stream, local_peer) + LpTransportSession::psq_handshake_responder(stream, local_peer) .complete_handshake() .await }) @@ -179,11 +180,10 @@ where let receiver_idx = session.receiver_index(); // 2. insert the state machine into the shared state - let state_machine = LpStateMachine::new(session); self.state .shared .session_states - .insert(receiver_idx, TimestampedState::new(state_machine)); + .insert(receiver_idx, TimestampedState::new(session)); self.bound_receiver_idx = Some(receiver_idx); // 3. handle any new incoming packet @@ -277,28 +277,18 @@ where ); // Process packet through state machine - let action = state_machine - .process_input(LpInput::ReceivePacket(encrypted_packet)) - .ok_or(LpHandlerError::UnexpectedStateMachineHalt)??; + let action = state_machine.process_input(LpInput::ReceivePacket(encrypted_packet))?; drop(state_entry); match action { LpAction::SendPacket(response_packet) => { - self.send_serialised_packet(&response_packet).await?; - Ok(()) + self.send_serialised_packet(&response_packet).await } LpAction::DeliverData(data) => { // Decrypted application data - process as registration/forwarding self.handle_decrypted_payload(receiver_index, data).await } - other @ LpAction::ConnectionClosed => { - warn!( - "Unexpected action in transport from {}: {other:?}", - self.remote_addr - ); - Err(LpHandlerError::UnexpectedStateMachineAction { action: other }) - } } } @@ -357,9 +347,7 @@ where let wrapped_lp_data = LpMessage::new(response_kind, serialised_response); // Process packet through state machine - let action = state_machine - .process_input(LpInput::SendData(wrapped_lp_data)) - .ok_or(LpHandlerError::UnexpectedStateMachineHalt)??; + let action = state_machine.process_input(LpInput::SendData(wrapped_lp_data))?; let packet = match action { LpAction::SendPacket(packet) => packet, diff --git a/nym-node/src/node/lp/error.rs b/nym-node/src/node/lp/error.rs index 2f5ffa43a0..8952e9f680 100644 --- a/nym-node/src/node/lp/error.rs +++ b/nym-node/src/node/lp/error.rs @@ -3,7 +3,7 @@ use nym_lp::packet::message::LpMessageType; use nym_lp::peer_config::LpReceiverIndex; -use nym_lp::state_machine::LpAction; +use nym_lp::session::LpAction; use nym_lp::transport::LpTransportError; use nym_lp::{LpError, packet::MalformedLpPacketError}; use std::net::SocketAddr; @@ -32,9 +32,6 @@ pub enum LpHandlerError { received: LpReceiverIndex, }, - #[error("no action has been emitted from the LP State Machine")] - UnexpectedStateMachineHalt, - #[error("the state machine instructed an unexpected action: {action:?}")] UnexpectedStateMachineAction { action: LpAction }, diff --git a/nym-node/src/node/lp/state.rs b/nym-node/src/node/lp/state.rs index 2c8f1bb95a..e003dded6a 100644 --- a/nym-node/src/node/lp/state.rs +++ b/nym-node/src/node/lp/state.rs @@ -5,7 +5,7 @@ use crate::config::LpConfig; use crate::node::lp::cleanup::TimestampedState; use dashmap::DashMap; use nym_gateway::node::wireguard::PeerRegistrator; -use nym_lp::LpStateMachine; +use nym_lp::LpTransportSession; use nym_lp::peer::LpLocalPeer; use nym_lp::peer_config::LpReceiverIndex; use nym_mixnet_client::forwarder::MixForwardingSender; @@ -60,5 +60,5 @@ pub struct SharedLpState { /// Established sessions keyed by receiver index /// /// Wrapped in TimestampedState for TTL-based cleanup of inactive sessions. - pub session_states: Arc>>, + pub session_states: Arc>>, } diff --git a/nym-registration-client/src/lp_client/client.rs b/nym-registration-client/src/lp_client/client.rs index 81d6c5f67d..da6dc3bd57 100644 --- a/nym-registration-client/src/lp_client/client.rs +++ b/nym-registration-client/src/lp_client/client.rs @@ -11,10 +11,9 @@ use crate::lp_client::state_machine_helpers::{extract_forwarded_response, prepar use nym_bandwidth_controller::{BandwidthTicketProvider, DEFAULT_TICKETS_TO_SPEND}; use nym_credentials_interface::TicketType; use nym_crypto::asymmetric::{ed25519, x25519}; -use nym_lp::LpSession; +use nym_lp::LpTransportSession; use nym_lp::peer::{DHKeyPair, LpLocalPeer, LpRemotePeer}; use nym_lp::peer_config::LpReceiverIndex; -use nym_lp::state_machine::LpStateMachine; use nym_lp::transport::traits::LpTransportChannel; use nym_lp::transport::{LpHandshakeChannel, LpTransportError}; use nym_lp::{Ciphersuite, packet::EncryptedLpPacket, packet::version}; @@ -58,9 +57,9 @@ pub struct LpRegistrationClient { /// Included in case we have to downgrade our version. gateway_supported_lp_protocol_version: u8, - /// LP state machine for managing connection lifecycle. + /// LP transport session /// Created during handshake initiation. - state_machine: Option, + transport_session: Option, /// Configuration for timeouts and TCP parameters. pub(crate) config: LpRegistrationConfig, @@ -110,7 +109,7 @@ where gateway_lp_peer, gateway_lp_address, gateway_supported_lp_protocol_version: lp_protocol, - state_machine: None, + transport_session: None, config, stream: None, } @@ -153,14 +152,14 @@ where ) } - pub(crate) fn state_machine(&self) -> Result<&LpStateMachine> { - self.state_machine + pub(crate) fn transport_session(&self) -> Result<&LpTransportSession> { + self.transport_session .as_ref() .ok_or(LpClientError::IncompleteHandshake) } - pub(crate) fn state_machine_mut(&mut self) -> Result<&mut LpStateMachine> { - self.state_machine + pub(crate) fn transport_session_mut(&mut self) -> Result<&mut LpTransportSession> { + self.transport_session .as_mut() .ok_or(LpClientError::IncompleteHandshake) } @@ -171,7 +170,7 @@ where /// Returns whether the client has completed the handshake and is ready for registration. pub fn is_handshake_complete(&self) -> bool { - self.state_machine.is_some() + self.transport_session.is_some() } /// Returns the gateway LP address this client is configured for. @@ -393,7 +392,7 @@ where let connection = self.stream_mut()?; // TODO: - let session = LpSession::psq_handshake_initiator( + let session = LpTransportSession::psq_handshake_initiator( connection, local_peer, remote_peer, @@ -403,7 +402,7 @@ where .await?; // Store the state machine (with established session) for later use - self.state_machine = Some(LpStateMachine::new(session)); + self.transport_session = Some(session); Ok(()) } @@ -461,7 +460,7 @@ where let lp_data = request.to_lp_data()?; // 4. Encrypt and prepare packet via state machine - let state_machine = self.state_machine_mut()?; + let state_machine = self.transport_session_mut()?; let request_packet = prepare_send_packet(lp_data, state_machine)?; // 5. Send initial request and receive response on persistent connection with timeout @@ -473,7 +472,7 @@ where .await?; // 6. Decrypt via state machine (re-borrow) - let state_machine = self.state_machine_mut()?; + let state_machine = self.transport_session_mut()?; let received_data = extract_forwarded_response(response_packet, state_machine)?; // 7. Extract decrypted data and deserialise the response @@ -548,7 +547,7 @@ where let lp_data = request.to_lp_data()?; // 3. Encrypt and prepare packet via state machine - let state_machine = self.state_machine_mut()?; + let state_machine = self.transport_session_mut()?; let request_packet = prepare_send_packet(lp_data, state_machine)?; // 4. Send initial request and receive response on persistent connection with timeout @@ -560,7 +559,7 @@ where .await?; // 5. Decrypt via state machine (re-borrow) - let state_machine = self.state_machine_mut()?; + let state_machine = self.transport_session_mut()?; let received_data = extract_forwarded_response(response_packet, state_machine)?; // 6. Extract decrypted data and deserialise the response @@ -665,7 +664,7 @@ where if self.stream.is_none() || attempt > 0 { // Clear any stale state before re-handshaking self.close(); - self.state_machine = None; + self.transport_session = None; if let Err(e) = self.perform_handshake().await { tracing::warn!("Handshake failed on attempt {attempt_display}: {e}"); @@ -713,10 +712,7 @@ where /// # Errors /// Returns an error if handshake has not been completed. pub fn session_id(&self) -> Result { - self.state_machine()? - .session() - .map(|s| s.receiver_index()) - .map_err(Into::into) + Ok(self.transport_session()?.receiver_index()) } } diff --git a/nym-registration-client/src/lp_client/error.rs b/nym-registration-client/src/lp_client/error.rs index 16bd1d66c3..1bf138fabb 100644 --- a/nym-registration-client/src/lp_client/error.rs +++ b/nym-registration-client/src/lp_client/error.rs @@ -6,7 +6,7 @@ use nym_lp::LpError; use nym_lp::packet::MalformedLpPacketError; use nym_lp::packet::message::LpMessageType; -use nym_lp::state_machine::LpAction; +use nym_lp::session::LpAction; use nym_lp::transport::LpTransportError; use thiserror::Error; @@ -33,9 +33,6 @@ pub enum LpClientError { #[error(transparent)] LpProtocolError(#[from] LpError), - #[error("no action has been emitted from the LP State Machine")] - UnexpectedStateMachineHalt, - #[error("the state machine instructed an unexpected action: {action:?}")] UnexpectedStateMachineAction { action: LpAction }, diff --git a/nym-registration-client/src/lp_client/helpers.rs b/nym-registration-client/src/lp_client/helpers.rs index 2ed712f9a0..57ba4d4eb7 100644 --- a/nym-registration-client/src/lp_client/helpers.rs +++ b/nym-registration-client/src/lp_client/helpers.rs @@ -7,7 +7,7 @@ use crate::LpClientError; use nym_lp::packet::message::LpMessageType; use nym_lp::packet::{ForwardPacketData, LpMessage}; use nym_lp::peer::LpRemotePeer; -use nym_lp::state_machine::{LpAction, LpInput}; +use nym_lp::session::{LpAction, LpInput}; use nym_registration_common::{ LpRegistrationRequest, LpRegistrationResponse, NymNodeLPInformation, }; diff --git a/nym-registration-client/src/lp_client/nested_session/connection.rs b/nym-registration-client/src/lp_client/nested_session/connection.rs index 70432e9bb0..c4e54015ac 100644 --- a/nym-registration-client/src/lp_client/nested_session/connection.rs +++ b/nym-registration-client/src/lp_client/nested_session/connection.rs @@ -6,7 +6,7 @@ use crate::{LpClientError, LpRegistrationClient}; use bytes::{BufMut, BytesMut}; use nym_lp::KEM; use nym_lp::packet::{EncryptedLpPacket, ForwardPacketData, message::ExpectedResponseSize}; -use nym_lp::state_machine::{LpAction, LpInput}; +use nym_lp::session::{LpAction, LpInput}; use nym_lp::transport::traits::{HandshakeMessage, LpTransportChannel}; use nym_lp::transport::{LpHandshakeChannel, LpTransportError}; use std::io; @@ -67,11 +67,9 @@ impl<'a, S> NestedConnection<'a, S> { let input = convert_forward_data(data)?; // 2. Encrypt and prepare packet via state machine - let state_machine = self.outer_client.state_machine_mut()?; + let state_machine = self.outer_client.transport_session_mut()?; - let action = state_machine - .process_input(input) - .ok_or(LpClientError::UnexpectedStateMachineHalt)??; + let action = state_machine.process_input(input)?; let forward_packet = match action { LpAction::SendPacket(packet) => packet, @@ -102,10 +100,8 @@ impl<'a, S> NestedConnection<'a, S> { .map_err(|_| LpClientError::ConnectionTimeout)??; // 2. Decrypt via state machine (re-borrow) - let state_machine = self.outer_client.state_machine_mut()?; - let action = state_machine - .process_input(LpInput::ReceivePacket(response_packet)) - .ok_or(LpClientError::UnexpectedStateMachineHalt)??; + let state_machine = self.outer_client.transport_session_mut()?; + let action = state_machine.process_input(LpInput::ReceivePacket(response_packet))?; // 3. Extract decrypted response data let response_data = try_convert_forward_response(action)?; diff --git a/nym-registration-client/src/lp_client/nested_session/mod.rs b/nym-registration-client/src/lp_client/nested_session/mod.rs index f8150433fe..c86b0fdf62 100644 --- a/nym-registration-client/src/lp_client/nested_session/mod.rs +++ b/nym-registration-client/src/lp_client/nested_session/mod.rs @@ -28,10 +28,9 @@ use nym_crypto::asymmetric::{ed25519, x25519}; use nym_lp::packet::version; use nym_lp::packet::{EncryptedLpPacket, LpMessage}; use nym_lp::peer::{DHKeyPair, LpLocalPeer, LpRemotePeer}; -use nym_lp::state_machine::LpStateMachine; use nym_lp::transport::LpHandshakeChannel; use nym_lp::transport::traits::LpTransportChannel; -use nym_lp::{Ciphersuite, KEM, LpSession}; +use nym_lp::{Ciphersuite, KEM, LpTransportSession}; use nym_registration_common::dvpn::LpDvpnRegistrationResponseMessageContent; use nym_registration_common::{ LpRegistrationRequest, LpRegistrationResponse, WireguardConfiguration, @@ -79,8 +78,8 @@ pub struct NestedLpSession { /// Included in case we have to downgrade our version. gateway_supported_lp_protocol_version: u8, - /// LP state machine for exit gateway session (populated after handshake) - state_machine: Option, + /// LP transport session for exit gateway session (populated after handshake) + transport_session: Option, } impl NestedLpSession { @@ -116,12 +115,12 @@ impl NestedLpSession { lp_local_peer, gateway_lp_peer, gateway_supported_lp_protocol_version: lp_protocol, - state_machine: None, + transport_session: None, } } - fn state_machine_mut(&mut self) -> Result<&mut LpStateMachine> { - self.state_machine + fn state_machine_mut(&mut self) -> Result<&mut LpTransportSession> { + self.transport_session .as_mut() .ok_or(LpClientError::IncompleteHandshake) } @@ -181,7 +180,7 @@ impl NestedLpSession { let remote_peer = self.gateway_lp_peer.clone(); let protocol_version = self.gateway_supported_lp_protocol_version; - let session = LpSession::psq_handshake_initiator( + let session = LpTransportSession::psq_handshake_initiator( &mut nested_connection, local_peer, remote_peer, @@ -191,7 +190,7 @@ impl NestedLpSession { .await?; // Store the state machine (with established session) for later use - self.state_machine = Some(LpStateMachine::new(session)); + self.transport_session = Some(session); debug!("completed nested handshake"); Ok(()) } @@ -532,7 +531,7 @@ impl NestedLpSession { tokio::time::sleep(delay).await; // Clear state machine before retry - handshake needs fresh start - self.state_machine = None; + self.transport_session = None; } match self diff --git a/nym-registration-client/src/lp_client/state_machine_helpers.rs b/nym-registration-client/src/lp_client/state_machine_helpers.rs index 3b2a5e7f65..bbb15fd177 100644 --- a/nym-registration-client/src/lp_client/state_machine_helpers.rs +++ b/nym-registration-client/src/lp_client/state_machine_helpers.rs @@ -3,18 +3,16 @@ use crate::LpClientError; use nym_lp::packet::LpMessage; -use nym_lp::state_machine::{LpAction, LpInput}; -use nym_lp::{LpStateMachine, packet::EncryptedLpPacket}; +use nym_lp::session::{LpAction, LpInput}; +use nym_lp::{LpTransportSession, packet::EncryptedLpPacket}; /// Attempt to prepare the provided data for sending by wrapping it in appropriate `LpAction`, /// and attempting to extract `EncryptedLpPacket` from the provided state machine. pub(crate) fn prepare_send_packet( data: LpMessage, - state_machine: &mut LpStateMachine, + state_machine: &mut LpTransportSession, ) -> Result { - let action = state_machine - .process_input(LpInput::SendData(data)) - .ok_or(LpClientError::UnexpectedStateMachineHalt)??; + let action = state_machine.process_input(LpInput::SendData(data))?; match action { LpAction::SendPacket(packet) => Ok(packet), @@ -26,11 +24,9 @@ pub(crate) fn prepare_send_packet( /// using the provided state machine. pub(crate) fn extract_forwarded_response( response_packet: EncryptedLpPacket, - state_machine: &mut LpStateMachine, + state_machine: &mut LpTransportSession, ) -> Result { - let action = state_machine - .process_input(LpInput::ReceivePacket(response_packet)) - .ok_or(LpClientError::UnexpectedStateMachineHalt)??; + let action = state_machine.process_input(LpInput::ReceivePacket(response_packet))?; match action { LpAction::DeliverData(data) => Ok(data), From 4464d12103abd675a521611112e4baa603889ddc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Wed, 4 Mar 2026 09:08:56 +0000 Subject: [PATCH 2/2] clippy and review comments --- common/nym-lp/src/session.rs | 13 +--- common/nym-lp/src/session_integration/mod.rs | 32 ++++----- common/nym-lp/src/session_manager.rs | 68 +++++++++---------- nym-node/src/node/lp/control/handler.rs | 14 ++-- .../src/lp_client/client.rs | 2 +- nym-registration-client/src/lp_client/mod.rs | 2 +- .../src/lp_client/nested_session/mod.rs | 2 +- ..._machine_helpers.rs => session_helpers.rs} | 0 8 files changed, 60 insertions(+), 73 deletions(-) rename nym-registration-client/src/lp_client/{state_machine_helpers.rs => session_helpers.rs} (100%) diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 02d5effa23..d7b8d29999 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -311,20 +311,13 @@ impl LpTransportSession { let ctr = packet.outer_header().counter; // 1. Check replay protection - if let Err(e) = self.receiving_counter_quick_check(ctr) { - return Err(e); - } + self.receiving_counter_quick_check(ctr)?; // 2. decrypt the packet and attempt to deliver data - let packet = match self.decrypt_packet(packet) { - Ok(packet) => packet, - Err(e) => return Err(e), - }; + let packet = self.decrypt_packet(packet)?; // 3. Mark counter as received - if let Err(e) = self.receiving_counter_mark(ctr) { - return Err(e); - } + self.receiving_counter_mark(ctr)?; // 4. deliver the message Ok(LpAction::DeliverData(packet.message)) diff --git a/common/nym-lp/src/session_integration/mod.rs b/common/nym-lp/src/session_integration/mod.rs index b628c4324b..f60f8e5d13 100644 --- a/common/nym-lp/src/session_integration/mod.rs +++ b/common/nym-lp/src/session_integration/mod.rs @@ -41,10 +41,10 @@ mod tests { // 2. Create sessions using the pre-built Noise states let peer_a_sm = session_manager_1 - .create_session_state_machine(sessions.initiator) + .insert_session(sessions.initiator) .unwrap(); let peer_b_sm = session_manager_2 - .create_session_state_machine(sessions.responder) + .insert_session(sessions.responder) .unwrap(); // 3. Send multiple encrypted messages both ways @@ -129,24 +129,24 @@ mod tests { let session2 = sessions.responder; // 2. Create a session (using real noise state) - let _session = session_manager.create_session_state_machine(session1); + let _session = session_manager.insert_session(session1); // 3. Try to get a non-existent session - let result = session_manager.state_machine_exists(non_existent); + let result = session_manager.session_exists(non_existent); assert!(!result, "Non-existent session should return None"); // 4. Try to remove a non-existent session - let result = session_manager.remove_state_machine(non_existent); + let result = session_manager.remove_session(non_existent); assert!( !result, "Remove session should not remove a non-existent session" ); // 5. Create and immediately remove a session - let _temp_session = session_manager.create_session_state_machine(session2); + let _temp_session = session_manager.insert_session(session2); assert!( - session_manager.remove_state_machine(session_id), + session_manager.remove_session(session_id), "Should remove the session" ); } @@ -170,16 +170,16 @@ mod tests { // 2. Create sessions state machines session_manager_1 - .create_session_state_machine(sessions.initiator) + .insert_session(sessions.initiator) .unwrap(); session_manager_2 - .create_session_state_machine(sessions.responder) + .insert_session(sessions.responder) .unwrap(); assert_eq!(session_manager_1.session_count(), 1); assert_eq!(session_manager_2.session_count(), 1); - assert!(session_manager_1.state_machine_exists(session_id)); - assert!(session_manager_2.state_machine_exists(session_id)); + assert!(session_manager_1.session_exists(session_id)); + assert!(session_manager_2.session_exists(session_id)); // --- 3. Simulate Data Transfer via process_input --- println!("Starting data transfer simulation via process_input..."); @@ -310,15 +310,15 @@ mod tests { println!("Out-of-order test passed."); // --- 6. Session Removal --- - assert!(session_manager_1.remove_state_machine(session_id)); + assert!(session_manager_1.remove_session(session_id)); assert_eq!(session_manager_1.session_count(), 0); - assert!(!session_manager_1.state_machine_exists(session_id)); + assert!(!session_manager_1.session_exists(session_id)); // B's session manager still has it until removed - assert!(session_manager_2.state_machine_exists(session_id)); - assert!(session_manager_2.remove_state_machine(session_id)); + assert!(session_manager_2.session_exists(session_id)); + assert!(session_manager_2.remove_session(session_id)); assert_eq!(session_manager_2.session_count(), 0); - assert!(!session_manager_2.state_machine_exists(session_id)); + assert!(!session_manager_2.session_exists(session_id)); println!("Session removal test passed."); } } diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index d944b478ca..a005b30cdd 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -19,7 +19,7 @@ use crate::session::{LpAction, LpInput}; /// The SessionManager is responsible for creating, storing, and retrieving sessions pub struct SessionManager { /// Manages state machines directly, keyed by lp_id - state_machines: HashMap, + sessions: HashMap, } impl Default for SessionManager { @@ -32,7 +32,7 @@ impl SessionManager { /// Creates a new session manager with empty session storage. pub fn new() -> Self { Self { - state_machines: HashMap::new(), + sessions: HashMap::new(), } } @@ -41,7 +41,7 @@ impl SessionManager { lp_id: LpReceiverIndex, input: LpInput, ) -> Result { - self.with_state_machine_mut(lp_id, |sm| sm.process_input(input))? + self.with_session_mut(lp_id, |sm| sm.process_input(input))? } pub fn send_data( @@ -61,27 +61,27 @@ impl SessionManager { } #[cfg(test)] - fn get_state_machine_id(&self, lp_id: LpReceiverIndex) -> Result { - self.with_state_machine(lp_id, |sm| sm.receiver_index()) + fn get_session_id(&self, lp_id: LpReceiverIndex) -> Result { + self.with_session(lp_id, |sm| sm.receiver_index()) } pub fn current_packet_cnt(&self, lp_id: LpReceiverIndex) -> Result { - self.with_state_machine(lp_id, |sm| Ok(sm.current_packet_cnt()))? + self.with_session(lp_id, |sm| Ok(sm.current_packet_cnt()))? } pub fn session_count(&self) -> usize { - self.state_machines.len() + self.sessions.len() } - pub fn state_machine_exists(&self, lp_id: LpReceiverIndex) -> bool { - self.state_machines.contains_key(&lp_id) + pub fn session_exists(&self, lp_id: LpReceiverIndex) -> bool { + self.sessions.contains_key(&lp_id) } - pub fn with_state_machine(&self, lp_id: LpReceiverIndex, f: F) -> Result + pub fn with_session(&self, lp_id: LpReceiverIndex, f: F) -> Result where F: FnOnce(&LpTransportSession) -> R, { - if let Some(sm) = self.state_machines.get(&lp_id) { + if let Some(sm) = self.sessions.get(&lp_id) { Ok(f(sm)) } else { Err(LpError::StateMachineNotFound(lp_id)) @@ -89,38 +89,34 @@ impl SessionManager { } // For mutable access (like running process_input) - pub fn with_state_machine_mut( - &mut self, - lp_id: LpReceiverIndex, - f: F, - ) -> Result + pub fn with_session_mut(&mut self, lp_id: LpReceiverIndex, f: F) -> Result where F: FnOnce(&mut LpTransportSession) -> R, // Closure takes mutable ref { - if let Some(sm) = self.state_machines.get_mut(&lp_id) { + if let Some(sm) = self.sessions.get_mut(&lp_id) { Ok(f(sm)) } else { Err(LpError::StateMachineNotFound(lp_id)) } } - pub fn create_session_state_machine( + pub fn insert_session( &mut self, lp_session: LpTransportSession, ) -> Result { let session_id = lp_session.receiver_index(); - if self.state_machines.contains_key(&session_id) { + if self.sessions.contains_key(&session_id) { return Err(LpError::DuplicateSessionId(session_id)); } - self.state_machines.insert(session_id, lp_session); + self.sessions.insert(session_id, lp_session); Ok(session_id) } /// Method to remove a state machine - pub fn remove_state_machine(&mut self, lp_id: LpReceiverIndex) -> bool { - let removed = self.state_machines.remove(&lp_id); + pub fn remove_session(&mut self, lp_id: LpReceiverIndex) -> bool { + let removed = self.sessions.remove(&lp_id); removed.is_some() } @@ -138,13 +134,13 @@ mod tests { let local_session = mock_session_for_test(); let id = local_session.receiver_index(); - let sm_1_id = manager.create_session_state_machine(local_session).unwrap(); + let sm_1_id = manager.insert_session(local_session).unwrap(); assert_eq!(sm_1_id, id); - let retrieved = manager.state_machine_exists(id); + let retrieved = manager.session_exists(id); assert!(retrieved); - let not_found = manager.state_machine_exists(123); + let not_found = manager.session_exists(123); assert!(!not_found); } @@ -152,13 +148,13 @@ mod tests { fn test_session_manager_remove() { let mut manager = SessionManager::new(); let local_session = mock_session_for_test(); - let sm_1_id = manager.create_session_state_machine(local_session).unwrap(); + let sm_1_id = manager.insert_session(local_session).unwrap(); - let removed = manager.remove_state_machine(sm_1_id); + let removed = manager.remove_session(sm_1_id); assert!(removed); assert_eq!(manager.session_count(), 0); - let removed_again = manager.remove_state_machine(sm_1_id); + let removed_again = manager.remove_session(sm_1_id); assert!(!removed_again); } @@ -170,15 +166,15 @@ mod tests { let session2 = SessionsMock::mock_seeded_post_handshake(124, kem).initiator; let session3 = SessionsMock::mock_seeded_post_handshake(125, kem).initiator; - let sm_1 = manager.create_session_state_machine(session1).unwrap(); - let sm_2 = manager.create_session_state_machine(session2).unwrap(); - let sm_3 = manager.create_session_state_machine(session3).unwrap(); + let sm_1 = manager.insert_session(session1).unwrap(); + let sm_2 = manager.insert_session(session2).unwrap(); + let sm_3 = manager.insert_session(session3).unwrap(); assert_eq!(manager.session_count(), 3); - let retrieved1 = manager.get_state_machine_id(sm_1).unwrap(); - let retrieved2 = manager.get_state_machine_id(sm_2).unwrap(); - let retrieved3 = manager.get_state_machine_id(sm_3).unwrap(); + let retrieved1 = manager.get_session_id(sm_1).unwrap(); + let retrieved2 = manager.get_session_id(sm_2).unwrap(); + let retrieved3 = manager.get_session_id(sm_3).unwrap(); assert_eq!(retrieved1, sm_1); assert_eq!(retrieved2, sm_2); @@ -192,10 +188,10 @@ mod tests { let sesion = mock_session_for_test(); - let sm = manager.create_session_state_machine(sesion).unwrap(); + let sm = manager.insert_session(sesion).unwrap(); assert_eq!(manager.session_count(), 1); - let retrieved = manager.get_state_machine_id(sm); + let retrieved = manager.get_session_id(sm); assert!(retrieved.is_ok()); assert_eq!(retrieved.unwrap(), sm); } diff --git a/nym-node/src/node/lp/control/handler.rs b/nym-node/src/node/lp/control/handler.rs index b9bcd1d6cb..e84d300399 100644 --- a/nym-node/src/node/lp/control/handler.rs +++ b/nym-node/src/node/lp/control/handler.rs @@ -715,8 +715,8 @@ mod tests { let (init, resp) = sessions_for_tests(); let mut init_sm = SessionManager::new(); let mut resp_sm = SessionManager::new(); - resp_sm.create_session_state_machine(resp).unwrap(); - let id = init_sm.create_session_state_machine(init).unwrap(); + resp_sm.insert_session(resp).unwrap(); + let id = init_sm.insert_session(init).unwrap(); // Bind to localhost let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -731,8 +731,7 @@ mod tests { let packet = handler.receive_raw_packet().await?; let header = packet.outer_header(); assert_eq!(packet.outer_header().receiver_idx, id); - let Some(LpAction::DeliverData(data)) = resp_sm.receive_packet(id, packet).unwrap() - else { + let LpAction::DeliverData(data) = resp_sm.receive_packet(id, packet)? else { panic!("illegal state") }; Ok::<_, LpHandlerError>((header, data)) @@ -772,8 +771,8 @@ mod tests { let (init, resp) = sessions_for_tests(); let mut init_sm = SessionManager::new(); let mut resp_sm = SessionManager::new(); - resp_sm.create_session_state_machine(resp).unwrap(); - let id = init_sm.create_session_state_machine(init).unwrap(); + resp_sm.insert_session(resp).unwrap(); + let id = init_sm.insert_session(init).unwrap(); let server_task = tokio::spawn(async move { let (mut stream, _) = listener.accept().await.unwrap(); @@ -802,8 +801,7 @@ mod tests { .await .unwrap(); let header = received.outer_header(); - let Some(LpAction::DeliverData(data)) = init_sm.receive_packet(id, received).unwrap() - else { + let LpAction::DeliverData(data) = init_sm.receive_packet(id, received).unwrap() else { panic!("illegal state") }; diff --git a/nym-registration-client/src/lp_client/client.rs b/nym-registration-client/src/lp_client/client.rs index da6dc3bd57..2daa774f07 100644 --- a/nym-registration-client/src/lp_client/client.rs +++ b/nym-registration-client/src/lp_client/client.rs @@ -7,7 +7,7 @@ use super::config::LpRegistrationConfig; use super::error::{LpClientError, Result}; use crate::lp_client::helpers::{LpDataDeliverExt, LpDataSendExt}; use crate::lp_client::nested_session::connection::NestedConnection; -use crate::lp_client::state_machine_helpers::{extract_forwarded_response, prepare_send_packet}; +use crate::lp_client::session_helpers::{extract_forwarded_response, prepare_send_packet}; use nym_bandwidth_controller::{BandwidthTicketProvider, DEFAULT_TICKETS_TO_SPEND}; use nym_credentials_interface::TicketType; use nym_crypto::asymmetric::{ed25519, x25519}; diff --git a/nym-registration-client/src/lp_client/mod.rs b/nym-registration-client/src/lp_client/mod.rs index 492b01ed2e..3e7221c3b0 100644 --- a/nym-registration-client/src/lp_client/mod.rs +++ b/nym-registration-client/src/lp_client/mod.rs @@ -36,7 +36,7 @@ mod config; pub(crate) mod error; pub(crate) mod helpers; mod nested_session; -mod state_machine_helpers; +mod session_helpers; pub use client::LpRegistrationClient; pub use config::LpRegistrationConfig; diff --git a/nym-registration-client/src/lp_client/nested_session/mod.rs b/nym-registration-client/src/lp_client/nested_session/mod.rs index c86b0fdf62..a1e3002976 100644 --- a/nym-registration-client/src/lp_client/nested_session/mod.rs +++ b/nym-registration-client/src/lp_client/nested_session/mod.rs @@ -21,7 +21,7 @@ use super::client::LpRegistrationClient; use super::error::{LpClientError, Result}; use crate::lp_client::helpers::{LpDataDeliverExt, LpDataSendExt}; -use crate::lp_client::state_machine_helpers::{extract_forwarded_response, prepare_send_packet}; +use crate::lp_client::session_helpers::{extract_forwarded_response, prepare_send_packet}; use nym_bandwidth_controller::{BandwidthTicketProvider, DEFAULT_TICKETS_TO_SPEND}; use nym_credentials_interface::TicketType; use nym_crypto::asymmetric::{ed25519, x25519}; diff --git a/nym-registration-client/src/lp_client/state_machine_helpers.rs b/nym-registration-client/src/lp_client/session_helpers.rs similarity index 100% rename from nym-registration-client/src/lp_client/state_machine_helpers.rs rename to nym-registration-client/src/lp_client/session_helpers.rs