From 25ef5f52b4de92aff007ee2e70a654a075e96ba3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Mon, 23 Feb 2026 12:16:26 +0000 Subject: [PATCH] wip --- common/nym-kkt-ciphersuite/src/lib.rs | 2 + common/nym-lp/src/lib.rs | 199 ++++++++----------- common/nym-lp/src/psq/initiator.rs | 5 +- common/nym-lp/src/psq/mod.rs | 6 +- common/nym-lp/src/psq/responder.rs | 2 +- common/nym-lp/src/session.rs | 147 +++++++------- common/nym-lp/src/session_integration/mod.rs | 3 - common/nym-lp/src/session_manager.rs | 122 +++++------- common/nym-lp/src/state_machine.rs | 165 +++++++-------- 9 files changed, 293 insertions(+), 358 deletions(-) diff --git a/common/nym-kkt-ciphersuite/src/lib.rs b/common/nym-kkt-ciphersuite/src/lib.rs index fcd8d02b13..52ef33ed3b 100644 --- a/common/nym-kkt-ciphersuite/src/lib.rs +++ b/common/nym-kkt-ciphersuite/src/lib.rs @@ -208,6 +208,7 @@ impl SignatureScheme { EnumIter, EnumString, Display, + Default, )] #[strum(ascii_case_insensitive)] #[strum(serialize_all = "lowercase")] @@ -215,6 +216,7 @@ impl SignatureScheme { pub enum KEM { // unsupported // XWing = 0, + #[default] MlKem768 = 1, McEliece = 2, } diff --git a/common/nym-lp/src/lib.rs b/common/nym-lp/src/lib.rs index b11ed47a44..a3b1fa3ab4 100644 --- a/common/nym-lp/src/lib.rs +++ b/common/nym-lp/src/lib.rs @@ -2,10 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 pub mod codec; -// georgio: config does not seem to be used anywhere -// pub mod config; -// pub use config::LpConfig; - pub mod error; pub mod packet; pub mod peer; @@ -16,6 +12,10 @@ mod session_integration; pub mod session_manager; pub mod state_machine; +use crate::psq::{PSQ_MSG2_SIZE, initiator, psq_msg1_size, responder}; +use crate::session::PersistentSessionBinding; +use libcrux_psq::{Channel, IntoSession}; + pub use error::LpError; pub use nym_kkt_ciphersuite::{ Ciphersuite, HashFunction, HashLength, KEM, KEMKeyDigests, SignatureScheme, @@ -30,12 +30,8 @@ pub use session::LpSession; pub use session_manager::SessionManager; pub use state_machine::LpStateMachine; -#[cfg(test)] -pub fn kem_list() -> Vec { - todo!() - // use nym_kkt::ciphersuite::KEM; - // vec![KEM::MlKem768, KEM::McEliece, KEM::X25519] -} +#[cfg(any(feature = "mock", test))] +use nym_test_utils::helpers::u64_seeded_rng_09; #[cfg(any(feature = "mock", test))] pub struct SessionsMock { @@ -45,119 +41,88 @@ pub struct SessionsMock { #[cfg(any(feature = "mock", test))] impl SessionsMock { - pub fn mock_post_handshake(session_id: u32) -> SessionsMock { - todo!() - // use crate::peer::mock_peers; - // use nym_kkt::ciphersuite::{DecapsulationKey, EncapsulationKey}; - // - // let (mut init, mut resp) = mock_peers(); - // let resp_remote = resp.as_remote(); - // let init_remote = init.as_remote(); - // let salt = [42u8; 32]; - // let session_id_bytes = session_id.to_le_bytes(); - // - // // skip KKT by just deriving the kem key locally - // let kem_keys = resp.kem_psq.as_ref().unwrap(); - // - // let libcrux_private_key = libcrux_kem::PrivateKey::decode( - // libcrux_kem::Algorithm::X25519, - // kem_keys.private_key().as_bytes(), - // ) - // .unwrap(); - // let decapsulation_key = DecapsulationKey::X25519(libcrux_private_key); - // - // let libcrux_public_key = libcrux_kem::PublicKey::decode( - // libcrux_kem::Algorithm::X25519, - // kem_keys.public_key().as_bytes(), - // ) - // .unwrap(); - // let encapsulation_key = EncapsulationKey::X25519(libcrux_public_key); - // - // // INIT -> RESP: PSQ MSG1 - // let psq_initiator = crate::psk::psq_initiator_create_message( - // init.x25519.private_key(), - // &resp_remote.x25519_public, - // &encapsulation_key, - // init.ed25519.private_key(), - // init.ed25519.public_key(), - // &salt, - // &session_id_bytes, - // ) - // .unwrap(); - // - // let psk = psq_initiator.psk; - // let psq_payload = psq_initiator.payload; - // let outer_aead_key = crate::codec::OuterAeadKey::from_psk(&psk); - // - // let noise_state_init = snow::Builder::new(crate::noise_protocol::NoiseProtocol::params()) - // .local_private_key(init.x25519().private_key().as_bytes()) - // .remote_public_key(resp_remote.x25519_public.as_bytes()) - // .psk(crate::NOISE_PSK_INDEX, &psk) - // .build_initiator() - // .unwrap(); - // let mut noise_protocol_init = crate::noise_protocol::NoiseProtocol::new(noise_state_init); - // let noise_msg1 = noise_protocol_init.get_bytes_to_send().unwrap().unwrap(); - // - // let psq_responder = crate::psk::psq_responder_process_message( - // resp.x25519.private_key(), - // &init_remote.x25519_public, - // (&decapsulation_key, &encapsulation_key), - // &init_remote.ed25519_public, - // &psq_payload, - // &salt, - // &session_id_bytes, - // ) - // .unwrap(); - // - // let noise_state_resp = snow::Builder::new(crate::noise_protocol::NoiseProtocol::params()) - // .local_private_key(resp.x25519().private_key().as_bytes()) - // .remote_public_key(init_remote.x25519_public.as_bytes()) - // .psk(crate::NOISE_PSK_INDEX, &psk) - // .build_responder() - // .unwrap(); - // let mut noise_protocol_resp = crate::noise_protocol::NoiseProtocol::new(noise_state_resp); - // noise_protocol_resp.read_message(&noise_msg1).unwrap(); - // - // let noise_msg2 = noise_protocol_resp.get_bytes_to_send().unwrap().unwrap(); - // noise_protocol_init.read_message(&noise_msg2).unwrap(); - // let noise_msg3 = noise_protocol_init.get_bytes_to_send().unwrap().unwrap(); - // - // assert!(noise_protocol_init.is_handshake_finished()); - // - // noise_protocol_resp.read_message(&noise_msg3).unwrap(); - // assert!(noise_protocol_resp.is_handshake_finished()); - // - // SessionsMock { - // initiator: LpSession::new( - // session_id, - // 1, - // outer_aead_key.clone(), - // init, - // resp_remote, - // crate::session::PqSharedSecret::new(psq_initiator.pq_shared_secret), - // noise_protocol_init, - // ), - // responder: LpSession::new( - // session_id, - // 1, - // outer_aead_key, - // resp, - // init_remote, - // crate::session::PqSharedSecret::new(psq_responder.pq_shared_secret), - // noise_protocol_resp, - // ), - // } + pub fn mock_seeded_post_handshake(seed: u64, kem: KEM) -> SessionsMock { + use crate::peer::mock_peers; + + let (init, resp) = mock_peers(); + let resp_remote = resp.as_remote(); + + let init_rng = u64_seeded_rng_09(seed); + let resp_rng = u64_seeded_rng_09(seed + 1); + + let kem_keys = resp.kem_keypairs.as_ref().unwrap(); + + // skip KKT by just deriving the kem key locally + let encapsulation_key = kem_keys.encapsulation_key(kem).unwrap(); + let enc_key = encapsulation_key.clone(); + + let initiator_ciphersuite = + initiator::build_psq_ciphersuite(&init, &resp_remote, &enc_key).unwrap(); + let mut initiator = + initiator::build_psq_principal(init_rng, 1, initiator_ciphersuite).unwrap(); + + let responder_ciphersuite = responder::build_psq_ciphersuite(&resp, kem).unwrap(); + let mut responder = + responder::build_psq_principal(resp_rng, 1, responder_ciphersuite).unwrap(); + + // run PSQ + let mut payload_buf_responder = vec![0u8; 4096]; + let mut payload_buf_initiator = vec![0u8; 4096]; + + // Send first message + let mut buf = vec![0u8; psq_msg1_size(kem)]; + let len_i = initiator.write_message(&[], &mut buf).unwrap(); + assert_eq!(len_i, buf.len()); + + // Read first message + let (_, _) = responder + .read_message(&buf, &mut payload_buf_responder) + .unwrap(); + + // Get the authenticator out here, so we can deserialize the session later. + let Some(initiator_authenticator) = responder.initiator_authenticator() else { + panic!("No initiator authenticator found") + }; + + // Respond + let mut buf = [0u8; PSQ_MSG2_SIZE]; + let len_r = responder.write_message(&[], &mut buf).unwrap(); + assert_eq!(len_r, buf.len()); + + // Finalize on registration initiator + let (_, _) = initiator + .read_message(&buf, &mut payload_buf_initiator) + .unwrap(); + + assert!(initiator.is_handshake_finished()); + assert!(responder.is_handshake_finished()); + + let binding = PersistentSessionBinding { + initiator_authenticator, + responder_ecdh_pk: resp_remote.x25519_public, + responder_pq_pk: Some(encapsulation_key), + }; + + SessionsMock { + initiator: LpSession::new(initiator.into_session().unwrap(), binding.clone(), 1) + .unwrap(), + responder: LpSession::new(responder.into_session().unwrap(), binding, 1).unwrap(), + } + } + + pub fn mock_post_handshake(kem: KEM) -> SessionsMock { + Self::mock_seeded_post_handshake(1, kem) } // we just need a dummy 'valid' session for simpler tests pub fn mock_initiator() -> LpSession { - Self::mock_post_handshake(1234).initiator + Self::mock_post_handshake(KEM::default()).initiator } } #[cfg(any(feature = "mock", test))] pub fn sessions_for_tests() -> (LpSession, LpSession) { - let sessions = SessionsMock::mock_post_handshake(69); + let sessions = SessionsMock::mock_post_handshake(KEM::default()); (sessions.initiator, sessions.responder) } @@ -169,9 +134,9 @@ pub fn mock_session_for_test() -> LpSession { #[cfg(test)] mod tests { use crate::session_manager::SessionManager; - use crate::{LpError, SessionsMock, kem_list, mock_session_for_test}; + use crate::{LpError, SessionsMock, mock_session_for_test}; use bytes::BytesMut; - use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, SignatureScheme}; + use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, IntoEnumIterator, KEM, SignatureScheme}; // Import the new standalone functions use crate::codec::serialize_lp_packet; @@ -285,7 +250,7 @@ mod tests { let mut local_manager = SessionManager::new(); let mut remote_manager = SessionManager::new(); - for kem in kem_list() { + for kem in KEM::iter() { todo!() // // Generate Ed25519 keypairs for PSQ authentication // let (init, resp) = mock_peers(kem); diff --git a/common/nym-lp/src/psq/initiator.rs b/common/nym-lp/src/psq/initiator.rs index f4ad3eb0ae..05b81a854d 100644 --- a/common/nym-lp/src/psq/initiator.rs +++ b/common/nym-lp/src/psq/initiator.rs @@ -263,7 +263,10 @@ mod tests { // 2. process let processed_req = kkt_responder.process_request(req)?; conn_resp - .send_handshake_message(processed_req.response.into(), kem) + .send_handshake_message::( + processed_req.response.into(), + kem, + ) .timeboxed() .await??; diff --git a/common/nym-lp/src/psq/mod.rs b/common/nym-lp/src/psq/mod.rs index 497e75485a..2dac8c1566 100644 --- a/common/nym-lp/src/psq/mod.rs +++ b/common/nym-lp/src/psq/mod.rs @@ -117,7 +117,7 @@ mod tests { use libcrux_psq::{Channel, IntoSession}; use nym_kkt::initiator::KKTInitiator; use nym_kkt::responder::KKTResponder; - use nym_kkt_ciphersuite::{HashFunction, KEM, SignatureScheme}; + use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, KEM, SignatureScheme}; use nym_test_utils::helpers::{ DeterministicRng09Send, deterministic_rng_09, u64_seeded_rng_09, }; @@ -269,7 +269,7 @@ mod tests { assert_eq!(len_i, buf.len()); // Read first message - let (len_r_deserialized, len_r_payload) = responder + let (_, _) = responder .read_message(&buf, &mut payload_buf_responder) .unwrap(); @@ -284,7 +284,7 @@ mod tests { assert_eq!(len_r, buf.len()); // Finalize on registration initiator - let (len_i_deserialized, len_i_payload) = initiator + let (len_i_deserialized, _) = initiator .read_message(&buf, &mut payload_buf_initiator) .unwrap(); diff --git a/common/nym-lp/src/psq/responder.rs b/common/nym-lp/src/psq/responder.rs index 06c7b8c9b4..3b44f555d1 100644 --- a/common/nym-lp/src/psq/responder.rs +++ b/common/nym-lp/src/psq/responder.rs @@ -268,7 +268,7 @@ mod tests { // 1. send kkt request conn_init - .send_handshake_message(request.into(), kem) + .send_handshake_message::(request.into(), kem) .timeboxed() .await??; diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 34564e78a5..f2703a597a 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -48,6 +48,7 @@ pub struct LpSession { } /// Wraps public key material that is bound to a session. +#[derive(Clone)] pub struct PersistentSessionBinding { /// The initiator's authenticator value, i.e. a long-term DH public value or signature verification key. pub initiator_authenticator: Authenticator, @@ -276,98 +277,88 @@ impl LpSession { #[cfg(test)] mod tests { use super::*; - use nym_crypto::asymmetric::x25519; - // use rand::thread_rng; - - // // Helper function to generate keypairs for tests - // fn generate_x25519_keypair() -> x25519::KeyPair { - // x25519::KeyPair::new(&mut thread_rng()) - // } + use crate::{ReplayError, SessionsMock, sessions_for_tests}; + use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; #[test] fn test_session_creation() { - todo!() - // let mut session = sessions_for_tests().0; - // for kem in kem_list() { - // let session = sessions_for_tests(kem).0; - // - // // Initial counter should be zero - // let counter = session.next_counter(); - // assert_eq!(counter, 0); - // - // // Counter should increment - // let counter = session.next_counter(); - // assert_eq!(counter, 1); - // } + for kem in KEM::iter() { + let mut session = SessionsMock::mock_post_handshake(kem).responder; + + // Initial counter should be zero + let counter = session.next_counter(); + assert_eq!(counter, 0); + + // Counter should increment + let counter = session.next_counter(); + assert_eq!(counter, 1); + } } #[test] fn test_replay_protection_sequential() { - todo!() - // for kem in kem_list() { - // let mut session = sessions_for_tests(kem).1; - // - // // Sequential counters should be accepted - // assert!(session.receiving_counter_quick_check(0).is_ok()); - // assert!(session.receiving_counter_mark(0).is_ok()); - // - // assert!(session.receiving_counter_quick_check(1).is_ok()); - // assert!(session.receiving_counter_mark(1).is_ok()); - // - // // Duplicates should be rejected - // assert!(session.receiving_counter_quick_check(0).is_err()); - // let err = session.receiving_counter_mark(0).unwrap_err(); - // match err { - // LpError::Replay(replay_error) => { - // assert!(matches!(replay_error, ReplayError::DuplicateCounter)); - // } - // _ => panic!("Expected replay error"), - // } - // } + for kem in KEM::iter() { + let mut session = SessionsMock::mock_post_handshake(kem).responder; + + // Sequential counters should be accepted + assert!(session.receiving_counter_quick_check(0).is_ok()); + assert!(session.receiving_counter_mark(0).is_ok()); + + assert!(session.receiving_counter_quick_check(1).is_ok()); + assert!(session.receiving_counter_mark(1).is_ok()); + + // Duplicates should be rejected + assert!(session.receiving_counter_quick_check(0).is_err()); + let err = session.receiving_counter_mark(0).unwrap_err(); + match err { + LpError::Replay(replay_error) => { + assert!(matches!(replay_error, ReplayError::DuplicateCounter)); + } + _ => panic!("Expected replay error"), + } + } } #[test] fn test_replay_protection_out_of_order() { - todo!() - // for kem in kem_list() { - // let mut session = sessions_for_tests(kem).1; - // - // // Receive packets in order - // assert!(session.receiving_counter_mark(0).is_ok()); - // assert!(session.receiving_counter_mark(1).is_ok()); - // assert!(session.receiving_counter_mark(2).is_ok()); - // - // // Skip ahead - // assert!(session.receiving_counter_mark(10).is_ok()); - // - // // Can still receive out-of-order packets within window - // assert!(session.receiving_counter_quick_check(5).is_ok()); - // assert!(session.receiving_counter_mark(5).is_ok()); - // - // // But duplicates are still rejected - // assert!(session.receiving_counter_quick_check(5).is_err()); - // assert!(session.receiving_counter_mark(5).is_err()); - // } + for kem in KEM::iter() { + let mut session = SessionsMock::mock_post_handshake(kem).responder; + + // Receive packets in order + assert!(session.receiving_counter_mark(0).is_ok()); + assert!(session.receiving_counter_mark(1).is_ok()); + assert!(session.receiving_counter_mark(2).is_ok()); + + // Skip ahead + assert!(session.receiving_counter_mark(10).is_ok()); + + // Can still receive out-of-order packets within window + assert!(session.receiving_counter_quick_check(5).is_ok()); + assert!(session.receiving_counter_mark(5).is_ok()); + + // But duplicates are still rejected + assert!(session.receiving_counter_quick_check(5).is_err()); + assert!(session.receiving_counter_mark(5).is_err()); + } } #[test] fn test_packet_stats() { - todo!() - // for kem in kem_list() { - // let mut session = sessions_for_tests(kem).1; - // - // // Initial stats - // let (next, received) = session.current_packet_cnt(); - // assert_eq!(next, 0); - // assert_eq!(received, 0); - // - // // After receiving packets - // assert!(session.receiving_counter_mark(0).is_ok()); - // assert!(session.receiving_counter_mark(1).is_ok()); - // - // let (next, received) = session.current_packet_cnt(); - // assert_eq!(next, 2); - // assert_eq!(received, 2); - // } + for kem in KEM::iter() { + let mut session = SessionsMock::mock_post_handshake(kem).responder; + + // Initial stats + let (next, received) = session.current_packet_cnt(); + assert_eq!(next, 0); + assert_eq!(received, 0); + + // After receiving packets + assert!(session.receiving_counter_mark(0).is_ok()); + assert!(session.receiving_counter_mark(1).is_ok()); + + let (next, received) = session.current_packet_cnt(); + assert_eq!(next, 2); + assert_eq!(received, 2); + } } } diff --git a/common/nym-lp/src/session_integration/mod.rs b/common/nym-lp/src/session_integration/mod.rs index d22e46a0cd..8934ec25a9 100644 --- a/common/nym-lp/src/session_integration/mod.rs +++ b/common/nym-lp/src/session_integration/mod.rs @@ -1,8 +1,5 @@ #[cfg(test)] mod tests { - use crate::codec::serialize_lp_packet; - use crate::{LpError, SessionsMock, kem_list, session_manager::SessionManager}; - use bytes::BytesMut; use nym_lp_packet::{LpHeader, LpMessage, LpPacket}; // Function to create a test packet - similar to how it's done in codec.rs tests diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index b7e42f0b11..b3d60b2efc 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -90,8 +90,8 @@ impl SessionManager { self.state_machines.len() } - pub fn state_machine_exists(&self, lp_id: SessionId) -> bool { - self.state_machines.contains_key(&lp_id) + pub fn state_machine_exists(&self, lp_id: &SessionId) -> bool { + self.state_machines.contains_key(lp_id) } pub fn with_state_machine(&self, lp_id: SessionId, f: F) -> Result @@ -136,85 +136,73 @@ impl SessionManager { mod tests { use super::*; use crate::{SessionsMock, mock_session_for_test}; + use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; #[test] fn test_session_manager_get() { - todo!() - // let mut manager = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let local_session = mock_session_for_test(); - // let id = local_session.id(); - // - // let sm_1_id = manager.create_session_state_machine(local_session); - // assert_eq!(sm_1_id, id); - // - // let retrieved = manager.state_machine_exists(id); - // assert!(retrieved); - // - // let not_found = manager.state_machine_exists(99); - // assert!(!not_found); + let mut manager = SessionManager::new(); + let local_session = mock_session_for_test(); + let id = *local_session.session_identifier(); + + let sm_1_id = manager.create_session_state_machine(local_session); + assert_eq!(sm_1_id, id); + + let retrieved = manager.state_machine_exists(&id); + assert!(retrieved); + + let not_found = manager.state_machine_exists(&[99u8; 32]); + assert!(!not_found); } #[test] fn test_session_manager_remove() { - todo!() - // let mut manager = SessionManager::new(); - // let local_session = mock_session_for_test(); - // - // let TODO = " for kem in kem_list() {"; - // - // let sm_1_id = manager.create_session_state_machine(local_session); - // - // let removed = manager.remove_state_machine(sm_1_id); - // assert!(removed); - // assert_eq!(manager.session_count(), 0); - // - // let removed_again = manager.remove_state_machine(sm_1_id); - // assert!(!removed_again); + let mut manager = SessionManager::new(); + let local_session = mock_session_for_test(); + let sm_1_id = manager.create_session_state_machine(local_session); + + let removed = manager.remove_state_machine(sm_1_id); + assert!(removed); + assert_eq!(manager.session_count(), 0); + + let removed_again = manager.remove_state_machine(sm_1_id); + assert!(!removed_again); } #[test] fn test_multiple_sessions() { - todo!() - // let mut manager = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let session1 = SessionsMock::mock_post_handshake(123).initiator; - // let session2 = SessionsMock::mock_post_handshake(124).initiator; - // let session3 = SessionsMock::mock_post_handshake(125).initiator; - // - // let sm_1 = manager.create_session_state_machine(session1); - // let sm_2 = manager.create_session_state_machine(session2); - // let sm_3 = manager.create_session_state_machine(session3); - // - // 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(); - // - // assert_eq!(retrieved1, sm_1); - // assert_eq!(retrieved2, sm_2); - // assert_eq!(retrieved3, sm_3); + for kem in KEM::iter() { + let mut manager = SessionManager::new(); + let session1 = SessionsMock::mock_seeded_post_handshake(123, kem).initiator; + 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); + let sm_2 = manager.create_session_state_machine(session2); + let sm_3 = manager.create_session_state_machine(session3); + + 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(); + + assert_eq!(retrieved1, sm_1); + assert_eq!(retrieved2, sm_2); + assert_eq!(retrieved3, sm_3); + } } #[test] fn test_session_manager_create_session() { - todo!() - // let mut manager = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let sesion = mock_session_for_test(); - // - // let sm = manager.create_session_state_machine(sesion); - // assert_eq!(manager.session_count(), 1); - // - // let retrieved = manager.get_state_machine_id(sm); - // assert!(retrieved.is_ok()); - // assert_eq!(retrieved.unwrap(), sm); + let mut manager = SessionManager::new(); + + let sesion = mock_session_for_test(); + + let sm = manager.create_session_state_machine(sesion); + assert_eq!(manager.session_count(), 1); + + let retrieved = manager.get_state_machine_id(sm); + assert!(retrieved.is_ok()); + assert_eq!(retrieved.unwrap(), sm); } } diff --git a/common/nym-lp/src/state_machine.rs b/common/nym-lp/src/state_machine.rs index a2eccd16cb..e5a6086260 100644 --- a/common/nym-lp/src/state_machine.rs +++ b/common/nym-lp/src/state_machine.rs @@ -61,7 +61,7 @@ impl From<&LpState> for LpStateBare { #[allow(clippy::large_enum_variant)] #[derive(Debug)] pub enum LpInput { - /// Received an LP Packet from the network. + /// Received an encrypted LP Packet from the network. ReceivePacket(LpPacket), /// Application wants to send data (only valid in Transport state). @@ -350,103 +350,92 @@ impl LpStateMachine { mod tests { use super::*; use crate::SessionsMock; - use nym_kkt_ciphersuite::KEM; + use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; #[test] fn test_state_machine_init() { - let TODO = " for kem in kem_list() {"; + for kem in KEM::iter() { + let mock_sessions = SessionsMock::mock_post_handshake(kem); - let mock_sessions = SessionsMock::mock_post_handshake(123); + let initiator_sm = LpStateMachine::new(mock_sessions.initiator); + assert!(matches!(initiator_sm.state, LpState::Transport { .. })); + let init_session = initiator_sm.session().unwrap(); - 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(); - 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.id(), resp_session.id()); + // Check both state machines use the same receiver_index + assert_eq!(init_session.id(), resp_session.id()); + } } #[test] fn test_state_machine_simplified_flow() { - todo!() - // let TODO = " for kem in kem_list() {"; - // - // let receiver_index: u32 = 123; - // let mock_sessions = SessionsMock::mock_post_handshake(123); - // - // // Create state machines (already in Transport) - // let mut initiator = LpStateMachine::new(mock_sessions.initiator); - // let mut responder = LpStateMachine::new(mock_sessions.responder); - // - // assert_eq!(initiator.id().unwrap(), responder.id().unwrap()); - // - // // --- Transport Phase --- - // println!("--- Step 1: Initiator sends data ---"); - // let data_to_send_1 = LpData::new_opaque(b"hello responder".to_vec()); - // let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); - // let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_4 { - // packet.clone() - // } else { - // panic!("Initiator should send data packet"); - // }; - // assert_eq!(data_packet_1.header.receiver_idx(), receiver_index); - // - // println!("--- Step 2: Responder receives data ---"); - // let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); - // let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_5 { - // data - // } else { - // panic!("Responder should deliver data"); - // }; - // assert_eq!(resp_data_1, data_to_send_1); - // - // println!("--- Step 3: Responder sends data ---"); - // let data_to_send_2 = LpData::new_opaque(b"hello initiator".to_vec()); - // let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); - // let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_6 { - // packet.clone() - // } else { - // panic!("Responder should send data packet"); - // }; - // assert_eq!(data_packet_2.header.receiver_idx(), receiver_index); - // - // println!("--- Step 4: Initiator receives data ---"); - // let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); - // if let Some(Ok(LpAction::DeliverData(data))) = init_actions_5 { - // assert_eq!(data, data_to_send_2); - // } else { - // panic!("Initiator should deliver data"); - // } - // - // // --- Close --- - // println!("--- Step 5: Initiator closes ---"); - // let init_actions_6 = initiator.process_input(LpInput::Close); - // assert!(matches!( - // init_actions_6, - // Some(Ok(LpAction::ConnectionClosed)) - // )); - // assert!(matches!(initiator.state, LpState::Closed { .. })); - // - // println!("--- Step 6: Responder closes ---"); - // let resp_actions_7 = responder.process_input(LpInput::Close); - // assert!(matches!( - // resp_actions_7, - // Some(Ok(LpAction::ConnectionClosed)) - // )); - // assert!(matches!(responder.state, LpState::Closed { .. })); - } + for kem in KEM::iter() { + let receiver_index: u32 = 123; + let mock_sessions = SessionsMock::mock_post_handshake(kem); - /// Helper function to complete a full handshake between initiator and responder, - /// returning both in Transport state ready for subsession testing. - fn setup_transport_sessions(kem: KEM) -> (LpStateMachine, LpStateMachine) { - let sessions = SessionsMock::mock_post_handshake(12345); - ( - LpStateMachine::new(sessions.initiator), - LpStateMachine::new(sessions.responder), - ) + // Create state machines (already in Transport) + let mut initiator = LpStateMachine::new(mock_sessions.initiator); + let mut responder = LpStateMachine::new(mock_sessions.responder); + + assert_eq!(initiator.id().unwrap(), responder.id().unwrap()); + + // --- Transport Phase --- + println!("--- Step 1: Initiator sends data ---"); + let data_to_send_1 = LpData::new_opaque(b"hello responder".to_vec()); + let init_actions_4 = initiator.process_input(LpInput::SendData(data_to_send_1.clone())); + let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_4 { + packet.clone() + } else { + panic!("Initiator should send data packet"); + }; + assert_eq!(data_packet_1.header.receiver_idx(), receiver_index); + + println!("--- Step 2: Responder receives data ---"); + let resp_actions_5 = responder.process_input(LpInput::ReceivePacket(data_packet_1)); + let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_5 { + data + } else { + panic!("Responder should deliver data"); + }; + assert_eq!(resp_data_1, data_to_send_1); + + println!("--- Step 3: Responder sends data ---"); + let data_to_send_2 = LpData::new_opaque(b"hello initiator".to_vec()); + let resp_actions_6 = responder.process_input(LpInput::SendData(data_to_send_2.clone())); + let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_6 { + packet.clone() + } else { + panic!("Responder should send data packet"); + }; + assert_eq!(data_packet_2.header.receiver_idx(), receiver_index); + + println!("--- Step 4: Initiator receives data ---"); + let init_actions_5 = initiator.process_input(LpInput::ReceivePacket(data_packet_2)); + if let Some(Ok(LpAction::DeliverData(data))) = init_actions_5 { + assert_eq!(data, data_to_send_2); + } else { + panic!("Initiator should deliver data"); + } + + // --- Close --- + println!("--- Step 5: Initiator closes ---"); + let init_actions_6 = initiator.process_input(LpInput::Close); + assert!(matches!( + init_actions_6, + Some(Ok(LpAction::ConnectionClosed)) + )); + assert!(matches!(initiator.state, LpState::Closed { .. })); + + println!("--- Step 6: Responder closes ---"); + let resp_actions_7 = responder.process_input(LpInput::Close); + assert!(matches!( + resp_actions_7, + Some(Ok(LpAction::ConnectionClosed)) + )); + assert!(matches!(responder.state, LpState::Closed { .. })); + } } }