GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #6512 from nymtech/lp/remove-state-machine-states

Browse Source 
remove redundant LP state machine in favour of in place processing
This commit is contained in:
Jędrzej Stuczyński
2026-03-04 09:33:02 +00:00
committed by GitHub
parent a7705a5f2c 4464d12103
commit 8de574ec97
18 changed files with 258 additions and 614 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/nym-lp/src/lib.rs
+8 -10
View File
@@ -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()
}
common/nym-lp/src/psq/initiator.rs
+7 -4
View File
@@ -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<LpSession, LpError>
pub async fn complete_handshake(self) -> Result<LpTransportSession, LpError>
where
S: LpHandshakeChannel + Unpin,
{
@@ -121,7 +121,10 @@ where
self.complete_handshake_with_rng(&mut rng).await
}
pub async fn complete_handshake_with_rng<R>(mut self, rng: &mut R) -> Result<LpSession, LpError>
pub async fn complete_handshake_with_rng<R>(
mut self,
rng: &mut R,
) -> Result<LpTransportSession, LpError>
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)
}
}
common/nym-lp/src/psq/responder.rs
+7 -4
View File
@@ -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<LpSession, LpError>
pub async fn complete_handshake(self) -> Result<LpTransportSession, LpError>
where
S: LpHandshakeChannel + Unpin,
{
@@ -154,7 +154,10 @@ where
self.complete_handshake_with_rng(&mut rng).await
}
pub async fn complete_handshake_with_rng<R>(mut self, rng: &mut R) -> Result<LpSession, LpError>
pub async fn complete_handshake_with_rng<R>(
mut self,
rng: &mut R,
) -> Result<LpTransportSession, LpError>
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,
common/nym-lp/src/session.rs
+115 -4
View File
@@ -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,41 @@ impl LpSession {
) -> Result<LpPacket, LpError> {
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<LpAction, LpError> {
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
self.receiving_counter_quick_check(ctr)?;
// 2. decrypt the packet and attempt to deliver data
let packet = self.decrypt_packet(packet)?;
// 3. Mark counter as received
self.receiving_counter_mark(ctr)?;
// 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 +420,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");
}
}
}
}
common/nym-lp/src/session_integration/mod.rs
+18 -86
View File
@@ -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};
@@ -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
@@ -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);
}
@@ -131,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"
);
}
@@ -172,26 +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));
// 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
);
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...");
@@ -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,64 +309,16 @@ 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 ---
assert!(session_manager_1.remove_state_machine(session_id));
// --- 6. Session Removal ---
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.");
}
}
common/nym-lp/src/session_manager.rs
+40 -58
View File
@@ -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<LpReceiverIndex, LpStateMachine>,
sessions: HashMap<LpReceiverIndex, LpTransportSession>,
}
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(),
}
}
@@ -40,8 +40,8 @@ impl SessionManager {
&mut self,
lp_id: LpReceiverIndex,
input: LpInput,
) -> Result<Option<LpAction>, LpError> {
self.with_state_machine_mut(lp_id, |sm| sm.process_input(input).transpose())?
) -> Result<LpAction, LpError> {
self.with_session_mut(lp_id, |sm| sm.process_input(input))?
}
pub fn send_data(
@@ -49,52 +49,39 @@ impl SessionManager {
lp_id: LpReceiverIndex,
data: LpMessage,
) -> Result<LpAction, LpError> {
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<Option<LpAction>, LpError> {
) -> Result<LpAction, LpError> {
self.process_input(lp_id, LpInput::ReceivePacket(packet))
}
pub fn closed(&self, lp_id: LpReceiverIndex) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::Closed)
}
pub fn transport(&self, lp_id: LpReceiverIndex) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::Transport)
}
#[cfg(test)]
fn get_state_machine_id(&self, lp_id: LpReceiverIndex) -> Result<LpReceiverIndex, LpError> {
self.with_state_machine(lp_id, |sm| sm.receiver_index())?
}
pub fn get_state(&self, lp_id: LpReceiverIndex) -> Result<LpStateBare, LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.bare_state()))?
fn get_session_id(&self, lp_id: LpReceiverIndex) -> Result<LpReceiverIndex, LpError> {
self.with_session(lp_id, |sm| sm.receiver_index())
}
pub fn current_packet_cnt(&self, lp_id: LpReceiverIndex) -> Result<PacketCount, LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.session()?.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<F, R>(&self, lp_id: LpReceiverIndex, f: F) -> Result<R, LpError>
pub fn with_session<F, R>(&self, lp_id: LpReceiverIndex, f: F) -> Result<R, LpError>
where
F: FnOnce(&LpStateMachine) -> R,
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))
@@ -102,39 +89,34 @@ impl SessionManager {
}
// For mutable access (like running process_input)
pub fn with_state_machine_mut<F, R>(
&mut self,
lp_id: LpReceiverIndex,
f: F,
) -> Result<R, LpError>
pub fn with_session_mut<F, R>(&mut self, lp_id: LpReceiverIndex, f: F) -> Result<R, LpError>
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) {
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: LpSession,
lp_session: LpTransportSession,
) -> Result<LpReceiverIndex, LpError> {
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));
}
let sm = LpStateMachine::new(lp_session);
self.state_machines.insert(session_id, sm);
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()
}
@@ -152,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);
}
@@ -166,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);
}
@@ -184,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);
@@ -206,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);
}
common/nym-lp/src/state_machine.rs
-352
View File
@@ -1,352 +0,0 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// 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<LpSession>,
}
/// 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<LpSession, LpError> {
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<SessionId, LpError> {
Ok(*self.session()?.session_identifier())
}
pub fn receiver_index(&self) -> Result<LpReceiverIndex, LpError> {
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<Result<LpAction, LpError>>) {
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<Result<LpAction, LpError>> {
// 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<Result<LpAction, LpError>> = 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<EncryptedLpPacket, LpError> {
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 { .. }));
}
}
}
nym-node/src/node/lp/cleanup.rs
+3 -3
View File
@@ -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<T> TimestampedState<T> {
}
pub(crate) struct CleanupTask {
session_states: Arc<DashMap<LpReceiverIndex, TimestampedState<LpStateMachine>>>,
session_states: Arc<DashMap<LpReceiverIndex, TimestampedState<LpTransportSession>>>,
cfg: LpDebug,
shutdown: nym_task::ShutdownToken,
}
impl CleanupTask {
pub fn new(
session_states: Arc<DashMap<LpReceiverIndex, TimestampedState<LpStateMachine>>>,
session_states: Arc<DashMap<LpReceiverIndex, TimestampedState<LpTransportSession>>>,
cfg: LpDebug,
shutdown: nym_task::ShutdownToken,
) -> Self {
nym-node/src/node/lp/control/handler.rs
+15 -29
View File
@@ -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<RefMut<'_, LpReceiverIndex, TimestampedState<LpStateMachine>>, LpHandlerError> {
) -> Result<RefMut<'_, LpReceiverIndex, TimestampedState<LpTransportSession>>, 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,
@@ -727,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();
@@ -743,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))
@@ -784,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();
@@ -814,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")
};
nym-node/src/node/lp/error.rs
+1 -4
View File
@@ -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 },
nym-node/src/node/lp/state.rs
+2 -2
View File
@@ -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<DashMap<LpReceiverIndex, TimestampedState<LpStateMachine>>>,
pub session_states: Arc<DashMap<LpReceiverIndex, TimestampedState<LpTransportSession>>>,
}
nym-registration-client/src/lp_client/client.rs
+18 -22
View File
@@ -7,14 +7,13 @@ 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};
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<S = TcpStream> {
/// 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<LpStateMachine>,
transport_session: Option<LpTransportSession>,
/// 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<LpReceiverIndex> {
self.state_machine()?
.session()
.map(|s| s.receiver_index())
.map_err(Into::into)
Ok(self.transport_session()?.receiver_index())
}
}
nym-registration-client/src/lp_client/error.rs
+1 -4
View File
@@ -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 },
nym-registration-client/src/lp_client/helpers.rs
+1 -1
View File
@@ -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,
};
nym-registration-client/src/lp_client/mod.rs
+1 -1
View File
@@ -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;
nym-registration-client/src/lp_client/nested_session/connection.rs
+5 -9
View File
@@ -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)?;
nym-registration-client/src/lp_client/nested_session/mod.rs
+10 -11
View File
@@ -21,17 +21,16 @@
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};
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<LpStateMachine>,
/// LP transport session for exit gateway session (populated after handshake)
transport_session: Option<LpTransportSession>,
}
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
nym-registration-client/src/lp_client/state_machine_helpers.rs → nym-registration-client/src/lp_client/session_helpers.rs
+6 -10
View File
@@ -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<EncryptedLpPacket, LpClientError> {
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<LpMessage, LpClientError> {
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),