From 595d034f647e97cdc7b56f6c62adeaf5c230dea5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Tue, 24 Feb 2026 10:57:31 +0000 Subject: [PATCH] restoring nym-lp tests --- common/nym-lp-packet/src/message.rs | 2 +- common/nym-lp-packet/src/packet.rs | 2 +- common/nym-lp/benches/replay_protection.rs | 15 +- common/nym-lp/src/codec.rs | 1110 ++++------------- common/nym-lp/src/error.rs | 6 + common/nym-lp/src/lib.rs | 237 ---- common/nym-lp/src/packet.rs | 2 +- common/nym-lp/src/peer.rs | 2 +- common/nym-lp/src/psq/initiator.rs | 16 +- common/nym-lp/src/psq/mod.rs | 16 +- common/nym-lp/src/psq/responder.rs | 7 +- common/nym-lp/src/replay/validator.rs | 47 +- common/nym-lp/src/session.rs | 26 +- common/nym-lp/src/session_integration/mod.rs | 1088 ++++++---------- common/nym-lp/src/session_manager.rs | 68 +- common/nym-lp/src/state_machine.rs | 163 ++- common/registration/src/lib.rs | 2 +- gateway/src/node/lp_listener/handler.rs | 4 +- integration-tests/src/lp_registration.rs | 176 +-- nym-registration-client/src/clients/lp.rs | 3 - .../src/lp_client/client.rs | 21 +- .../src/lp_client/error.rs | 45 - .../src/lp_client/helpers.rs | 1 - .../lp_client/nested_session/connection.rs | 4 - .../src/lp_client/nested_session/mod.rs | 25 +- .../src/lp_client/state_machine_helpers.rs | 2 +- 26 files changed, 887 insertions(+), 2203 deletions(-) diff --git a/common/nym-lp-packet/src/message.rs b/common/nym-lp-packet/src/message.rs index 4c28d66e26..870c48c7fc 100644 --- a/common/nym-lp-packet/src/message.rs +++ b/common/nym-lp-packet/src/message.rs @@ -262,7 +262,7 @@ impl ForwardPacketData { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum LpMessage { /// The party is busy Busy, diff --git a/common/nym-lp-packet/src/packet.rs b/common/nym-lp-packet/src/packet.rs index da3c579d8d..daa45e77ef 100644 --- a/common/nym-lp-packet/src/packet.rs +++ b/common/nym-lp-packet/src/packet.rs @@ -53,7 +53,7 @@ impl EncryptedLpPacket { } } -#[derive(Clone)] +#[derive(Clone, PartialEq, Eq)] pub struct LpPacket { pub(crate) header: LpHeader, pub(crate) message: LpMessage, diff --git a/common/nym-lp/benches/replay_protection.rs b/common/nym-lp/benches/replay_protection.rs index 6c2894241b..63c1b615d9 100644 --- a/common/nym-lp/benches/replay_protection.rs +++ b/common/nym-lp/benches/replay_protection.rs @@ -1,9 +1,8 @@ use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use nym_lp::replay::ReceivingKeyCounterValidator; use nym_test_utils::helpers::deterministic_rng_09; -use parking_lot::Mutex; use rand09::Rng; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; fn bench_sequential_counters(c: &mut Criterion) { let mut group = c.benchmark_group("replay_sequential"); @@ -75,19 +74,15 @@ fn bench_thread_safety(c: &mut Criterion) { BenchmarkId::new("thread_safe_validator", size), &size, |b, &size| { - let validator = Arc::new(Mutex::new(ReceivingKeyCounterValidator::default())); + let mut validator = ReceivingKeyCounterValidator::default(); let counters: Vec = (0..size).collect(); b.iter(|| { for &counter in &counters { - let result = { - let guard = validator.lock(); - black_box(guard.will_accept_branchless(counter)) - }; + let result = { black_box(validator.will_accept_branchless(counter)) }; if result.is_ok() { - let mut guard = validator.lock(); - let _ = black_box(guard.mark_did_receive_branchless(counter)); + let _ = black_box(validator.mark_did_receive_branchless(counter)); } } }); @@ -202,7 +197,7 @@ fn bench_concurrency_scaling(c: &mut Criterion) { let mut success_count = 0; for i in 0..100 { let counter = t * 1000 + i; - let mut guard = validator_clone.lock(); + let mut guard = validator_clone.lock().unwrap(); if guard.mark_did_receive_branchless(counter as u64).is_ok() { success_count += 1; } diff --git a/common/nym-lp/src/codec.rs b/common/nym-lp/src/codec.rs index 488ce2e43a..cce6a13e55 100644 --- a/common/nym-lp/src/codec.rs +++ b/common/nym-lp/src/codec.rs @@ -4,39 +4,22 @@ use crate::{LpError, LpMessage}; use bytes::BytesMut; use libcrux_psq::Channel; -use nym_lp_packet::{EncryptedLpPacket, InnerHeader, LpHeader, LpPacket, OuterHeader}; -use tracing::error; +use nym_lp_packet::{EncryptedLpPacket, InnerHeader, LpHeader, LpPacket}; -// wtf, why are those different?! -pub(crate) const ENC_OVERHEAD: usize = 26; -pub(crate) const DEC_OVERHEAD: usize = 25; +// needs to be equal or above to the actual overhead +pub(crate) const SANE_ENC_OVERHEAD: usize = 32; -/// Parse only the outer header from raw packet bytes. -/// -/// Used for routing before session lookup. The outer header (receiver_idx + counter) -/// is always cleartext at bytes 0-12 in the unified packet format. -/// -/// # Arguments -/// * `src` - Raw packet bytes (at least OuterHeader::SIZE bytes) -/// -/// # Errors -/// * `LpError::InsufficientBufferSize` - Packet too small for outer header -pub(crate) fn parse_lp_header_only(src: &[u8]) -> Result { - Ok(OuterHeader::parse(src)?) -} +// needs to be equal or below the actual overhead +pub(crate) const SANE_DEC_OVERHEAD: usize = 24; pub(crate) fn encrypt_data( plaintext: &[u8], transport: &mut libcrux_psq::session::Transport, ) -> Result, LpError> { - let mut ciphertext = vec![0u8; plaintext.len() + ENC_OVERHEAD]; + let mut ciphertext = vec![0u8; plaintext.len() + SANE_ENC_OVERHEAD]; let n = transport.write_message(plaintext, &mut ciphertext)?; - if plaintext.len() + ENC_OVERHEAD != n { - error!( - "FIXME: inconsistent ciphertext overhead. expected: {ENC_OVERHEAD}, actual: {}", - n - plaintext.len() - ); + if plaintext.len() + SANE_ENC_OVERHEAD != n { ciphertext.truncate(n); } @@ -47,18 +30,13 @@ pub(crate) fn decrypt_data( ciphertext: &[u8], transport: &mut libcrux_psq::session::Transport, ) -> Result, LpError> { - if ciphertext.len() < DEC_OVERHEAD { + if ciphertext.len() < SANE_DEC_OVERHEAD { return Err(LpError::InsufficientBufferSize); } - let mut plaintext = vec![0u8; ciphertext.len() - DEC_OVERHEAD]; + let mut plaintext = vec![0u8; ciphertext.len() - SANE_DEC_OVERHEAD]; let (_, n) = transport.read_message(ciphertext, &mut plaintext)?; - if n != ciphertext.len() - DEC_OVERHEAD { - // TODO: check consistency - error!( - "FIXME: inconsistent ciphertext overhead. expected: {DEC_OVERHEAD}, actual: {}", - ciphertext.len() - n - ); + if n != ciphertext.len() - SANE_DEC_OVERHEAD { plaintext.truncate(n); } Ok(plaintext) @@ -81,7 +59,7 @@ pub(crate) fn decrypt_lp_packet( packet: EncryptedLpPacket, transport: &mut libcrux_psq::session::Transport, ) -> Result { - if packet.ciphertext().len() < InnerHeader::SIZE + ENC_OVERHEAD { + if packet.ciphertext().len() < InnerHeader::SIZE + SANE_DEC_OVERHEAD { return Err(LpError::InsufficientBufferSize); } @@ -100,877 +78,209 @@ pub(crate) fn decrypt_lp_packet( )) } -/// Serializes an LpPacket into the provided BytesMut buffer. -/// -/// ## Unified Packet Format -/// -/// Both cleartext and encrypted packets have the same structure: -/// - Outer header (12B): receiver_idx(4) + counter(8) - always cleartext -/// - Inner payload: proto(1) + reserved(3) + msg_type(4) + content - encrypted -/// -/// # Arguments -/// * `item` - Packet to serialize -/// * `dst` - Output buffer -/// * `transport` - AEAD encryption channel -pub(crate) fn serialize_lp_packet( - item: LpPacket, - dst: &mut BytesMut, - transport: &mut libcrux_psq::session::Transport, -) -> Result<(), LpError> { - // 1. encrypt the inner header and payload - let encrypted_packet = encrypt_lp_packet(item, transport)?; - - // 2. Write outer header (always cleartext) followed by the ciphertext - encrypted_packet.encode(dst); - - Ok(()) -} - #[cfg(test)] mod tests { + use crate::LpError; + use crate::codec::{decrypt_data, decrypt_lp_packet, encrypt_data, encrypt_lp_packet}; + use crate::peer::mock_peers; + use crate::psq::initiator::{build_psq_ciphersuite, build_psq_principal}; + use crate::psq::{PSQ_MSG2_SIZE, psq_msg1_size, responder}; + use libcrux_psq::{Channel, IntoSession}; + use nym_kkt_ciphersuite::KEM; + use nym_lp_packet::{EncryptedLpPacket, LpHeader, LpMessage, LpPacket, MessageType}; + use nym_test_utils::helpers::u64_seeded_rng_09; - // === Cleartext Encode/Decode Tests === + fn mock_transport() -> ( + libcrux_psq::session::Transport, + libcrux_psq::session::Transport, + ) { + let kem = KEM::MlKem768; + let rng1 = u64_seeded_rng_09(1); + let rng2 = u64_seeded_rng_09(2); + let (init, resp) = mock_peers(); + let remote_resp = resp.as_remote(); + let encapsulation_key = resp + .kem_keypairs + .as_ref() + .unwrap() + .encapsulation_key(kem) + .unwrap(); + + let initiator_ciphersuite = + build_psq_ciphersuite(&init, &remote_resp, &encapsulation_key).unwrap(); + let mut psq_initiator = build_psq_principal(rng1, 1, initiator_ciphersuite).unwrap(); + + let responder_ciphersuite = responder::build_psq_ciphersuite(&resp, kem).unwrap(); + let mut psq_responder = + responder::build_psq_principal(rng2, 1, responder_ciphersuite).unwrap(); + + // Send first message + let mut buf = vec![0u8; psq_msg1_size(kem)]; + + let mut payload_buf_responder = vec![0u8; 4096]; + let mut payload_buf_initiator = vec![0u8; 4096]; + + let len_i = psq_initiator.write_message(&[], &mut buf).unwrap(); + assert_eq!(len_i, buf.len()); + + // Read first message + let (_, _) = psq_responder + .read_message(&buf, &mut payload_buf_responder) + .unwrap(); + + // Respond + let mut buf = [0u8; PSQ_MSG2_SIZE]; + let len_r = psq_responder.write_message(&[], &mut buf).unwrap(); + assert_eq!(len_r, buf.len()); + + // Finalize on registration initiator + let (len_i_deserialized, _) = psq_initiator + .read_message(&buf, &mut payload_buf_initiator) + .unwrap(); + + // We read the same amount of data. + assert_eq!(len_r, len_i_deserialized); + + // Ready for transport mode + assert!(psq_initiator.is_handshake_finished()); + assert!(psq_responder.is_handshake_finished()); + + let transport_initiator = psq_initiator + .into_session() + .unwrap() + .transport_channel() + .unwrap(); + + let transport_responder = psq_responder + .into_session() + .unwrap() + .transport_channel() + .unwrap(); + + (transport_initiator, transport_responder) + } #[test] - fn restore_below_tests() { - todo!() + fn basic_plain_encryption_test() { + let (mut init_transport, mut resp_transport) = mock_transport(); + + for msg_size in [1usize, 10, 100, 1000, 10000, 65535] { + let message1 = vec![42u8; msg_size]; + + let mut ciphertext = vec![0u8; msg_size + 64]; + let written_init1 = init_transport + .write_message(&message1, &mut ciphertext) + .unwrap(); + + let init_ciphertext_overhead = written_init1 - msg_size; + let ciphertext_content = &ciphertext[..written_init1]; + + let mut plaintext = vec![0u8; msg_size + 64]; + let (read_resp1, written_resp1) = resp_transport + .read_message(ciphertext_content, &mut plaintext) + .unwrap(); + let resp_plaintext_overhead = ciphertext_content.len() - written_resp1; + + assert_eq!( + written_init1, read_resp1, + "should work for message {msg_size}" + ); + let message1_content = &plaintext[..written_resp1]; + assert_eq!( + message1_content, &message1, + "should work for message {msg_size}" + ); + + // reverse the communication + let message2 = vec![43u8; msg_size]; + + let mut ciphertext2 = vec![0u8; msg_size + 64]; + let written_resp2 = resp_transport + .write_message(&message2, &mut ciphertext2) + .unwrap(); + + let resp_ciphertext_overhead = written_resp2 - msg_size; + let ciphertext_content2 = &ciphertext2[..written_resp2]; + + let mut plaintext2 = vec![0u8; msg_size + 64]; + let (read_init2, written_init2) = init_transport + .read_message(ciphertext_content2, &mut plaintext2) + .unwrap(); + let init_plaintext_overhead = ciphertext_content2.len() - written_init2; + + assert_eq!( + written_resp2, read_init2, + "should work for message {msg_size}" + ); + let message2_content = &plaintext2[..written_init2]; + assert_eq!( + message2_content, &message2, + "should work for message {msg_size}" + ); + + // check consistent overheads + // enc/enc + assert_eq!(init_ciphertext_overhead, resp_ciphertext_overhead); + + // dec/dec + assert_eq!(resp_plaintext_overhead, init_plaintext_overhead); + + // enc/dec + assert_eq!(init_ciphertext_overhead, resp_plaintext_overhead); + } + } + + #[test] + fn basic_encryption() { + let (mut init_transport, mut resp_transport) = mock_transport(); + + // happy path + let msg = b"foomp".to_vec(); + let ciphertext = encrypt_data(&msg, &mut init_transport).unwrap(); + let plaintext = decrypt_data(&ciphertext, &mut resp_transport).unwrap(); + assert_eq!(msg, plaintext); + + // incomplete ciphertext + let msg2 = b"foomp".to_vec(); + let ciphertext2 = encrypt_data(&msg2, &mut init_transport).unwrap(); + let len = ciphertext2.len(); + let dec_err = decrypt_data(&ciphertext2[..len - 1], &mut resp_transport).unwrap_err(); + assert!(matches!(dec_err, LpError::PSQSessionFailure { .. })); + + // too small buffer + let msg3 = b"foomp".to_vec(); + let ciphertext3 = encrypt_data(&msg3, &mut resp_transport).unwrap(); + let dec_err = decrypt_data(&ciphertext3[..10], &mut init_transport).unwrap_err(); + assert!(matches!(dec_err, LpError::InsufficientBufferSize)); + } + + #[test] + fn basic_packet_encryption() { + let (mut init_transport, mut resp_transport) = mock_transport(); + + // happy path + let packet = LpPacket::new(LpHeader::new(123, 0, 1, MessageType::Busy), LpMessage::Busy); + + let ciphertext = encrypt_lp_packet(packet.clone(), &mut init_transport).unwrap(); + assert_eq!(packet.header().outer, ciphertext.outer_header()); + + let plaintext = decrypt_lp_packet(ciphertext, &mut resp_transport).unwrap(); + assert_eq!(packet, plaintext); + + // incomplete ciphertext + let packet = LpPacket::new(LpHeader::new(123, 1, 1, MessageType::Busy), LpMessage::Busy); + let ciphertext2 = encrypt_lp_packet(packet, &mut init_transport).unwrap(); + let l = ciphertext2.ciphertext().len(); + let malformed_content = ciphertext2.ciphertext()[..l - 1].to_vec(); + let malformed = EncryptedLpPacket::new(ciphertext2.outer_header(), malformed_content); + let dec_err = decrypt_lp_packet(malformed, &mut resp_transport).unwrap_err(); + assert!(matches!(dec_err, LpError::PSQSessionFailure { .. })); + + // too small buffer + let packet = LpPacket::new(LpHeader::new(123, 1, 1, MessageType::Busy), LpMessage::Busy); + let ciphertext3 = encrypt_lp_packet(packet, &mut resp_transport).unwrap(); + let malformed = EncryptedLpPacket::new(ciphertext3.outer_header(), vec![]); + let dec_err = decrypt_lp_packet(malformed, &mut init_transport).unwrap_err(); + assert!(matches!(dec_err, LpError::InsufficientBufferSize)); } - // - // #[test] - // fn test_serialize_parse_busy() { - // let mut dst = BytesMut::new(); - // - // // Create a Busy packet - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 123, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // // Serialize the packet (cleartext) - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // - // // Parse the packet (cleartext) - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // // Verify the packet fields - // assert_eq!(decoded.header.protocol_version, 1); - // assert_eq!(decoded.header.receiver_idx, 42); - // assert_eq!(decoded.header.counter, 123); - // assert!(matches!(decoded.message, LpMessage::Busy)); - // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - // } - // - // #[test] - // fn test_serialize_parse_handshake() { - // let mut dst = BytesMut::new(); - // - // // Create a Handshake message packet - // let payload = vec![42u8; 80]; // Example payload size - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 123, - // }, - // message: LpMessage::PSQRequest(PSQRequestData(payload.clone())), - // trailer: [0; TRAILER_LEN], - // }; - // - // // Serialize the packet (cleartext) - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // - // // Parse the packet (cleartext) - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // // Verify the packet fields - // assert_eq!(decoded.header.protocol_version, 1); - // assert_eq!(decoded.header.receiver_idx, 42); - // assert_eq!(decoded.header.counter, 123); - // - // // Verify message type and data - // match decoded.message { - // LpMessage::PSQRequest(decoded_payload) => { - // assert_eq!(decoded_payload, PSQRequestData(payload)); - // } - // _ => panic!("Expected Handshake message"), - // } - // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - // } - // - // #[test] - // fn test_serialize_parse_encrypted_data() { - // let mut dst = BytesMut::new(); - // - // // Create an EncryptedData message packet - // let payload = vec![43u8; 124]; // Example payload size - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 123, - // }, - // message: LpMessage::EncryptedData(EncryptedDataPayload(payload.clone())), - // trailer: [0; TRAILER_LEN], - // }; - // - // // Serialize the packet (cleartext) - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // - // // Parse the packet (cleartext) - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // // Verify the packet fields - // assert_eq!(decoded.header.protocol_version, 1); - // assert_eq!(decoded.header.receiver_idx, 42); - // assert_eq!(decoded.header.counter, 123); - // - // // Verify message type and data - // match decoded.message { - // LpMessage::EncryptedData(decoded_payload) => { - // assert_eq!(decoded_payload, EncryptedDataPayload(payload)); - // } - // _ => panic!("Expected EncryptedData message"), - // } - // assert_eq!(decoded.trailer, [0; TRAILER_LEN]); - // } - // - // // === Incomplete Data Tests === - // - // #[test] - // fn test_parse_incomplete_header() { - // // Create a buffer with incomplete header - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&[1, 0, 0, 0]); // Only 4 bytes, not enough for LpHeader::SIZE - // - // // Attempt to parse - expect error - // let result = parse_lp_packet(&buf, None); - // assert!(result.is_err()); - // assert!(matches!( - // result.unwrap_err(), - // LpError::InsufficientBufferSize - // )); - // } - // - // #[test] - // fn test_parse_incomplete_message_type() { - // // Create a buffer with complete header but incomplete message type - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf.extend_from_slice(&[0]); // Only 1 byte of message type (need 2) - // - // // Buffer length = 16 + 1 = 17. Min size = 16 + 2 + 16 = 34. - // let result = parse_lp_packet(&buf, None); - // assert!(result.is_err()); - // assert!(matches!( - // result.unwrap_err(), - // LpError::InsufficientBufferSize - // )); - // } - // - // #[test] - // fn test_parse_incomplete_message_data() { - // // Create a buffer simulating Handshake but missing trailer and maybe partial payload - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type - // buf.extend_from_slice(&[42; 40]); // 40 bytes of payload data - // - // // Buffer length = 16 + 2 + 40 = 58. Min size = 16 + 2 + 16 = 34. - // // Payload size calculated as 58 - 34 = 24. - // // Trailer expected at index 16 + 2 + 24 = 42. - // // Trailer read attempts src[42..58]. - // // This *should* parse successfully based on the logic, but the trailer is garbage. - // // Let's rethink: parse_lp_packet assumes the *entire slice* is the packet. - // // If the slice doesn't end exactly where the trailer should, it's an error. - // // In this case, total length is 58. OuterHdr(12) + InnerPrefix(4) + Type(2) + Trailer(16) = 34. Payload = 58-34=24. - // // Trailer starts at 16+2+24 = 42. Ends at 42+16=58. It fits exactly. - // // This test *still* doesn't test incompleteness correctly for the datagram parser. - // - // // Let's test a buffer that's *too short* even for header+type+trailer+min_payload - // // Note: Buffer order doesn't matter for this test since we fail on minimum size check - // let mut buf_too_short = BytesMut::new(); - // buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx (outer header) - // buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // counter (outer header) - // buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // version + reserved (inner prefix) - // buf_too_short.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // msg type - // // No payload, no trailer. Length = 12+4+2=18. Min size = 34. - // let result_too_short = parse_lp_packet(&buf_too_short, None); - // assert!(result_too_short.is_err()); - // assert!(matches!( - // result_too_short.unwrap_err(), - // LpError::InsufficientBufferSize - // )); - // - // // Test a buffer missing PART of the trailer - // let mut buf_partial_trailer = BytesMut::new(); - // buf_partial_trailer.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf_partial_trailer.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf_partial_trailer.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf_partial_trailer.extend_from_slice(&MessageType::Handshake.to_u32().to_le_bytes()); // Handshake type - // let payload = vec![42u8; 20]; // Assume 20 byte payload - // buf_partial_trailer.extend_from_slice(&payload); - // buf_partial_trailer.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer - // - // // Total length = 16 + 2 + 20 + 15 = 53. Min size = 34. This passes. - // // Payload size = 53 - 34 = 19. <--- THIS IS WRONG. The parser assumes the length dictates payload. - // // Let's fix the parser logic slightly. - // - // // The point is, parse_lp_packet expects a COMPLETE datagram. Providing less bytes - // // than LpHeader + Type + Trailer should fail. Providing *more* is also an issue unless - // // the length calculation works out perfectly. The most direct test is just < min_size. - // // Renaming test to reflect this. - // } - // - // #[test] - // fn test_parse_buffer_smaller_than_minimum() { - // // Test a buffer that's smaller than the smallest possible packet (LpHeader+Type+Trailer) - // let mut buf_too_short = BytesMut::new(); - // buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf_too_short.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Type - // buf_too_short.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer - // // Length = 16 + 2 + 15 = 33. Min Size = 34. - // let result_too_short = parse_lp_packet(&buf_too_short, None); - // assert!( - // result_too_short.is_err(), - // "Expected error for buffer size 33, min 34" - // ); - // assert!(matches!( - // result_too_short.unwrap_err(), - // LpError::InsufficientBufferSize - // )); - // } - // - // #[test] - // fn test_parse_invalid_message_type() { - // // Create a buffer with invalid message type - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf.extend_from_slice(&231u16.to_le_bytes()); // Invalid message type - // // Need payload and trailer to meet min_size requirement - // let payload_size = 10; // Arbitrary - // buf.extend_from_slice(&vec![0u8; payload_size]); // Some data - // buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - // - // // Attempt to parse - // let result = parse_lp_packet(&buf, None); - // assert!(result.is_err()); - // match result { - // Err(LpError::InvalidMessageType(231)) => {} // Expected error - // Err(e) => panic!("Expected InvalidMessageType error, got {:?}", e), - // Ok(_) => panic!("Expected error, but got Ok"), - // } - // } - // - // #[test] - // fn test_parse_incorrect_payload_size_for_busy() { - // // Create a Busy packet but *with* a payload (which is invalid) - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index - // buf.extend_from_slice(&123u64.to_le_bytes()); // Counter - // buf.extend_from_slice(&MessageType::Busy.to_u32().to_le_bytes()); // Busy type - // buf.extend_from_slice(&[42; 1]); // <<< Invalid 1-byte payload for Busy - // buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer - // - // // Total size = 16 + 2 + 1 + 16 = 35. Min size = 34. - // // Calculated payload size = 35 - 34 = 1. - // let result = parse_lp_packet(&buf, None); - // assert!(result.is_err()); - // assert!(matches!( - // result.unwrap_err(), - // LpError::InvalidPayloadSize { - // expected: 0, - // actual: 1 - // } - // )); - // } - // - // // Test multiple packets simulation isn't relevant for datagram parsing - // // #[test] - // // fn test_multiple_packets_in_buffer() { ... } - // - // - // - // #[test] - // fn test_forward_packet_encode_decode_roundtrip_v4() { - // let mut dst = BytesMut::new(); - // - // let forward_data = crate::message::ForwardPacketData { - // target_gateway_identity: [77u8; 32], - // target_lp_address: "1.2.3.4:41264".parse().unwrap(), - // inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 999, - // counter: 555, - // }, - // message: LpMessage::ForwardPacket(forward_data), - // trailer: [0xff; TRAILER_LEN], - // }; - // - // // Serialize - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // - // // Parse back - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // // Verify LP protocol handling works correctly - // assert_eq!(decoded.header.receiver_idx, 999); - // assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); - // - // if let LpMessage::ForwardPacket(data) = decoded.message { - // assert_eq!(data.target_gateway_identity, [77u8; 32]); - // assert_eq!(data.target_lp_address, "1.2.3.4:41264".parse().unwrap()); - // assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); - // } else { - // panic!("Expected ForwardPacket message"); - // } - // } - // - // #[test] - // fn test_forward_packet_encode_decode_roundtrip_v6() { - // let mut dst = BytesMut::new(); - // - // let forward_data = crate::message::ForwardPacketData { - // target_gateway_identity: [77u8; 32], - // target_lp_address: "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap(), - // inner_packet_bytes: vec![0xa, 0xb, 0xc, 0xd], - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 999, - // counter: 555, - // }, - // message: LpMessage::ForwardPacket(forward_data), - // trailer: [0xff; TRAILER_LEN], - // }; - // - // // Serialize - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // - // // Parse back - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // // Verify LP protocol handling works correctly - // assert_eq!(decoded.header.receiver_idx, 999); - // assert!(matches!(decoded.message.typ(), MessageType::ForwardPacket)); - // - // if let LpMessage::ForwardPacket(data) = decoded.message { - // assert_eq!(data.target_gateway_identity, [77u8; 32]); - // assert_eq!( - // data.target_lp_address, - // "[dead:beef:4242:c0ff:ee00::1111]:41264".parse().unwrap() - // ); - // assert_eq!(data.inner_packet_bytes, vec![0xa, 0xb, 0xc, 0xd]); - // } else { - // panic!("Expected ForwardPacket message"); - // } - // } - // - // // === Outer AEAD Tests === - // - // #[test] - // fn test_aead_roundtrip_with_key() { - // // Test that encrypt/decrypt roundtrip works with an AEAD key - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 999, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // // Parse back with the same key - // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - // - // assert_eq!(decoded.header.protocol_version, 1); - // assert_eq!(decoded.header.receiver_idx, 12345); - // assert_eq!(decoded.header.counter, 999); - // assert!(matches!(decoded.message, LpMessage::Busy)); - // } - // - // #[test] - // fn test_aead_ciphertext_differs_from_plaintext() { - // // Verify that encrypted payload differs from plaintext - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 999, - // }, - // message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload(vec![ - // 0xAA, 0xBB, 0xCC, 0xDD, - // ])), - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut cleartext = BytesMut::new(); - // serialize_lp_packet(&packet, &mut cleartext, None).unwrap(); - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // // Outer header (receiver_idx + counter) should be the same - always cleartext - // assert_eq!(&cleartext[..OUTER_HDR], &encrypted[..OUTER_HDR]); - // - // // Inner payload (proto + reserved + msg_type + content) should differ (encrypted) - // let payload_start = OUTER_HDR; - // let payload_end_cleartext = cleartext.len() - TRAILER_LEN; - // let payload_end_encrypted = encrypted.len() - TRAILER_LEN; - // - // assert_ne!( - // &cleartext[payload_start..payload_end_cleartext], - // &encrypted[payload_start..payload_end_encrypted], - // "Encrypted payload should differ from plaintext" - // ); - // - // // Trailer should differ (zeros vs AEAD tag) - // assert_ne!( - // &cleartext[payload_end_cleartext..], - // &encrypted[payload_end_encrypted..], - // "Encrypted trailer should be a tag, not zeros" - // ); - // } - // - // #[test] - // fn test_aead_tampered_tag_fails() { - // // Verify that tampering with the tag causes decryption failure - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 999, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // // Tamper with the tag (last byte) - // let last_idx = encrypted.len() - 1; - // encrypted[last_idx] ^= 0xFF; - // - // // Parsing should fail with AeadTagMismatch - // let result = parse_lp_packet(&encrypted, Some(&outer_key)); - // assert!(matches!(result, Err(LpError::AeadTagMismatch))); - // } - // - // #[test] - // fn test_aead_tampered_header_fails() { - // // Verify that tampering with the header (AAD) causes decryption failure - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 999, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // // Tamper with the outer header AAD (flip a bit in counter at byte 4) - // // New format: [receiver_idx(0-3), counter(4-11)], so byte 4 is counter's LSB - // encrypted[4] ^= 0x01; - // - // // Parsing should fail with AeadTagMismatch - // let result = parse_lp_packet(&encrypted, Some(&outer_key)); - // assert!(matches!(result, Err(LpError::AeadTagMismatch))); - // } - // - // #[test] - // fn test_aead_different_counters_produce_different_ciphertext() { - // // Verify that different counters (nonces) produce different ciphertexts - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let packet1 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 1, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let packet2 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 2, // Different counter - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted1 = BytesMut::new(); - // serialize_lp_packet(&packet1, &mut encrypted1, Some(&outer_key)).unwrap(); - // - // let mut encrypted2 = BytesMut::new(); - // serialize_lp_packet(&packet2, &mut encrypted2, Some(&outer_key)).unwrap(); - // - // // The encrypted inner payloads should differ even though the message is the same - // // (because nonce is different). Inner payload starts after outer header. - // let payload_start = OUTER_HDR; - // assert_ne!( - // &encrypted1[payload_start..], - // &encrypted2[payload_start..], - // "Different counters should produce different ciphertexts" - // ); - // } - // - // #[test] - // fn test_aead_wrong_key_fails() { - // // Verify that decryption with wrong key fails - // let psk1 = [42u8; 32]; - // let psk2 = [43u8; 32]; // Different PSK - // let outer_key1 = OuterAeadKey::from_psk(&psk1); - // let outer_key2 = OuterAeadKey::from_psk(&psk2); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 12345, - // counter: 999, - // }, - // message: LpMessage::Busy, - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key1)).unwrap(); - // - // // Parsing with wrong key should fail - // let result = parse_lp_packet(&encrypted, Some(&outer_key2)); - // assert!(matches!(result, Err(LpError::AeadTagMismatch))); - // } - // - // #[test] - // fn test_aead_encrypted_data_message_roundtrip() { - // // Test AEAD with EncryptedData message type (larger payload) - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let payload_data = vec![0xDE; 100]; - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 54321, - // counter: 12345678, - // }, - // message: LpMessage::EncryptedData(crate::message::EncryptedDataPayload( - // payload_data.clone(), - // )), - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - // - // match decoded.message { - // LpMessage::EncryptedData(data) => { - // assert_eq!(data.0, payload_data); - // } - // _ => panic!("Expected EncryptedData message"), - // } - // } - // - // #[test] - // fn test_aead_handshake_message_roundtrip() { - // // Test AEAD with Handshake message type - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let handshake_data = vec![0x01, 0x02, 0x03, 0x04, 0x05]; - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 99999, - // counter: 2, - // }, - // message: LpMessage::PSQRequest(PSQRequestData(handshake_data.clone())), - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - // - // match decoded.message { - // LpMessage::PSQResponse(data) => { - // assert_eq!(data.0, handshake_data); - // } - // _ => panic!("Expected Handshake message"), - // } - // } - // - // // === Subsession Message Tests === - // - // #[test] - // fn test_serialize_parse_subsession_request() { - // let mut dst = BytesMut::new(); - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 100, - // }, - // message: LpMessage::SubsessionRequest, - // trailer: [0; TRAILER_LEN], - // }; - // - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // assert_eq!(decoded.header.receiver_idx, 42); - // assert_eq!(decoded.header.counter, 100); - // assert!(matches!(decoded.message, LpMessage::SubsessionRequest)); - // } - // - // #[test] - // fn test_serialize_parse_subsession_kk1() { - // use crate::message::SubsessionKK1Data; - // - // let mut dst = BytesMut::new(); - // - // let kk1_data = SubsessionKK1Data { - // payload: vec![0xAA; 50], // 50 bytes KK payload - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 123, - // counter: 456, - // }, - // message: LpMessage::SubsessionKK1(kk1_data.clone()), - // trailer: [0; TRAILER_LEN], - // }; - // - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // assert_eq!(decoded.header.receiver_idx, 123); - // match decoded.message { - // LpMessage::SubsessionKK1(data) => { - // assert_eq!(data.payload, kk1_data.payload); - // } - // _ => panic!("Expected SubsessionKK1 message"), - // } - // } - // - // #[test] - // fn test_serialize_parse_subsession_kk2() { - // use crate::message::SubsessionKK2Data; - // - // let mut dst = BytesMut::new(); - // - // let kk2_data = SubsessionKK2Data { - // payload: vec![0x11; 60], // 60 bytes KK response payload - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 789, - // counter: 1000, - // }, - // message: LpMessage::SubsessionKK2(kk2_data.clone()), - // trailer: [0; TRAILER_LEN], - // }; - // - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // assert_eq!(decoded.header.receiver_idx, 789); - // match decoded.message { - // LpMessage::SubsessionKK2(data) => { - // assert_eq!(data.payload, kk2_data.payload); - // } - // _ => panic!("Expected SubsessionKK2 message"), - // } - // } - // - // #[test] - // fn test_serialize_parse_subsession_ready() { - // use crate::message::SubsessionReadyData; - // - // let mut dst = BytesMut::new(); - // - // let ready_data = SubsessionReadyData { - // receiver_index: 99999, - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 200, - // }, - // message: LpMessage::SubsessionReady(ready_data.clone()), - // trailer: [0; TRAILER_LEN], - // }; - // - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // assert_eq!(decoded.header.receiver_idx, 42); - // match decoded.message { - // LpMessage::SubsessionReady(data) => { - // assert_eq!(data.receiver_index, 99999); - // } - // _ => panic!("Expected SubsessionReady message"), - // } - // } - // - // #[test] - // fn test_subsession_request_with_payload_fails() { - // // SubsessionRequest should have no payload - // let mut buf = BytesMut::new(); - // buf.extend_from_slice(&42u32.to_le_bytes()); // receiver_idx - // buf.extend_from_slice(&123u64.to_le_bytes()); // counter - // buf.extend_from_slice(&[1, 0, 0, 0]); // version + reserved - // buf.extend_from_slice(&MessageType::SubsessionRequest.to_u32().to_le_bytes()); - // buf.extend_from_slice(&[0xFF]); // Invalid payload for SubsessionRequest - // buf.extend_from_slice(&[0; TRAILER_LEN]); - // - // let result = parse_lp_packet(&buf, None); - // assert!(matches!( - // result, - // Err(LpError::InvalidPayloadSize { - // expected: 0, - // actual: 1 - // }) - // )); - // } - // - // #[test] - // fn test_aead_subsession_roundtrip() { - // use crate::message::SubsessionKK1Data; - // - // let psk = [42u8; 32]; - // let outer_key = OuterAeadKey::from_psk(&psk); - // - // let kk1_data = SubsessionKK1Data { - // payload: vec![0xDE; 48], // 48 bytes KK payload - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 54321, - // counter: 999, - // }, - // message: LpMessage::SubsessionKK1(kk1_data.clone()), - // trailer: [0; TRAILER_LEN], - // }; - // - // let mut encrypted = BytesMut::new(); - // serialize_lp_packet(&packet, &mut encrypted, Some(&outer_key)).unwrap(); - // - // let decoded = parse_lp_packet(&encrypted, Some(&outer_key)).unwrap(); - // - // match decoded.message { - // LpMessage::SubsessionKK1(data) => { - // assert_eq!(data.payload, kk1_data.payload); - // } - // _ => panic!("Expected SubsessionKK1 message"), - // } - // } - // - // #[test] - // fn test_serialize_parse_error() { - // use crate::message::ErrorPacketData; - // - // let mut dst = BytesMut::new(); - // - // let error_data = ErrorPacketData { - // message: "this is an error".to_string(), - // }; - // - // let packet = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 200, - // }, - // message: LpMessage::Error(error_data.clone()), - // trailer: [0; TRAILER_LEN], - // }; - // - // serialize_lp_packet(&packet, &mut dst, None).unwrap(); - // let decoded = parse_lp_packet(&dst, None).unwrap(); - // - // assert_eq!(decoded.header.receiver_idx, 42); - // match decoded.message { - // LpMessage::Error(data) => { - // assert_eq!(data.message, "this is an error"); - // } - // _ => panic!("Expected Error message"), - // } - // } } diff --git a/common/nym-lp/src/error.rs b/common/nym-lp/src/error.rs index 0cfd56c18d..635db970f5 100644 --- a/common/nym-lp/src/error.rs +++ b/common/nym-lp/src/error.rs @@ -27,6 +27,9 @@ pub enum LpError { #[error("Attempted operation on closed session")] SessionClosed, + #[error("There already exists an LP session with id {0:?}")] + DuplicateSessionId(SessionId), + #[error("Internal error: {0}")] Internal(String), @@ -112,6 +115,9 @@ pub enum LpError { #[error("transport failure: {0}")] TransportFailure(#[from] LpTransportError), + + #[error("the current session is not in transport state")] + NotInTransport, } impl LpError { diff --git a/common/nym-lp/src/lib.rs b/common/nym-lp/src/lib.rs index b672ff8828..c4475239e7 100644 --- a/common/nym-lp/src/lib.rs +++ b/common/nym-lp/src/lib.rs @@ -135,240 +135,3 @@ pub fn sessions_for_tests() -> (LpSession, LpSession) { pub fn mock_session_for_test() -> LpSession { SessionsMock::mock_initiator() } - -#[cfg(test)] -mod tests { - use crate::session_manager::SessionManager; - use crate::{LpError, SessionsMock, mock_session_for_test}; - use bytes::BytesMut; - use nym_kkt_ciphersuite::{Ciphersuite, HashFunction, IntoEnumIterator, KEM, SignatureScheme}; - - // Import the new standalone functions - use crate::codec::serialize_lp_packet; - - #[test] - fn test_replay_protection_integration() { - todo!() - // for kem in kem_list() { - // // Create session - // let mut session = mock_session_for_test(); - // - // // === Packet 1 (Counter 0 - Should succeed) === - // let packet1 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, // Matches session's sending_index assumption for this test - // counter: 0, - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize packet - // let mut buf1 = BytesMut::new(); - // serialize_lp_packet(&packet1, &mut buf1, None).unwrap(); - // - // // Parse packet - // let parsed_packet1 = parse_lp_packet(&buf1, None).unwrap(); - // - // // Perform replay check (should pass) - // session - // .receiving_counter_quick_check(parsed_packet1.header.counter) - // .expect("Initial packet failed replay check"); - // - // // Mark received (simulating successful processing) - // session - // .receiving_counter_mark(parsed_packet1.header.counter) - // .expect("Failed to mark initial packet received"); - // - // // === Packet 2 (Counter 0 - Replay, should fail check) === - // let packet2 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 0, // Same counter as before (replay) - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize packet - // let mut buf2 = BytesMut::new(); - // serialize_lp_packet(&packet2, &mut buf2, None).unwrap(); - // - // // Parse packet - // let parsed_packet2 = parse_lp_packet(&buf2, None).unwrap(); - // - // // Perform replay check (should fail) - // let replay_result = - // session.receiving_counter_quick_check(parsed_packet2.header.counter); - // assert!(replay_result.is_err()); - // match replay_result.unwrap_err() { - // LpError::Replay(e) => { - // assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter)); - // } - // e => panic!("Expected replay error, got {:?}", e), - // } - // // Do not mark received as it failed validation - // - // // === Packet 3 (Counter 1 - Should succeed) === - // let packet3 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: 42, - // counter: 1, // Incremented counter - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize packet - // let mut buf3 = BytesMut::new(); - // serialize_lp_packet(&packet3, &mut buf3, None).unwrap(); - // - // // Parse packet - // let parsed_packet3 = parse_lp_packet(&buf3, None).unwrap(); - // - // // Perform replay check (should pass) - // session - // .receiving_counter_quick_check(parsed_packet3.header.counter) - // .expect("Packet 3 failed replay check"); - // - // // Mark received - // session - // .receiving_counter_mark(parsed_packet3.header.counter) - // .expect("Failed to mark packet 3 received"); - // - // // Verify validator state directly on the session - // let state = session.current_packet_cnt(); - // assert_eq!(state.0, 2); // Next expected counter (correct - was 1, now expects 2) - // assert_eq!(state.1, 2); // Total marked received (correct - packets 1 and 3) - // } - } - - #[test] - fn test_session_manager_integration() { - // Create session manager - let mut local_manager = SessionManager::new(); - let mut remote_manager = SessionManager::new(); - - for kem in KEM::iter() { - todo!() - // // Generate Ed25519 keypairs for PSQ authentication - // let (init, resp) = mock_peers(kem); - // - // let mut ciphersuite = Ciphersuite::resolve_ciphersuite( - // kem, - // HashFunction::Blake3, - // SignatureScheme::Ed25519, - // None, - // ); - // - // // Use fixed receiver_index for deterministic test - // let receiver_index: u32 = 54321; - // - // let sessions = SessionsMock::mock_post_handshake(receiver_index); - // let local_session = sessions.initiator; - // let remote_session = sessions.responder; - // - // // Create a session via manager - // let _ = local_manager.create_session_state_machine(local_session); - // let _ = remote_manager.create_session_state_machine(remote_session); - // - // // === Packet 1 (Counter 0 - Should succeed) === - // let packet1 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: receiver_index, - // counter: 0, - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize - // let mut buf1 = BytesMut::new(); - // serialize_lp_packet(&packet1, &mut buf1, None).unwrap(); - // - // // Parse - // let parsed_packet1 = parse_lp_packet(&buf1, None).unwrap(); - // - // // Process via SessionManager method (which should handle checks + marking) - // // NOTE: We might need a method on SessionManager/LpSession like `process_incoming_packet` - // // that encapsulates parse -> check -> process_noise -> mark. - // // For now, we simulate the steps using the retrieved session. - // - // // Perform replay check - // local_manager - // .receiving_counter_quick_check(receiver_index, parsed_packet1.header.counter) - // .expect("Packet 1 check failed"); - // // Mark received - // local_manager - // .receiving_counter_mark(receiver_index, parsed_packet1.header.counter) - // .expect("Packet 1 mark failed"); - // - // // === Packet 2 (Counter 1 - Should succeed on same session) === - // let packet2 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: receiver_index, - // counter: 1, - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize - // let mut buf2 = BytesMut::new(); - // serialize_lp_packet(&packet2, &mut buf2, None).unwrap(); - // - // // Parse - // let parsed_packet2 = parse_lp_packet(&buf2, None).unwrap(); - // - // // Perform replay check - // local_manager - // .receiving_counter_quick_check(receiver_index, parsed_packet2.header.counter) - // .expect("Packet 2 check failed"); - // // Mark received - // local_manager - // .receiving_counter_mark(receiver_index, parsed_packet2.header.counter) - // .expect("Packet 2 mark failed"); - // - // // === Packet 3 (Counter 0 - Replay, should fail check) === - // let packet3 = LpPacket { - // header: LpHeader { - // protocol_version: 1, - // reserved: [0u8; 3], - // receiver_idx: receiver_index, - // counter: 0, // Replay of first packet - // }, - // message: LpMessage::Busy, - // trailer: [0u8; TRAILER_LEN], - // }; - // - // // Serialize - // let mut buf3 = BytesMut::new(); - // serialize_lp_packet(&packet3, &mut buf3, None).unwrap(); - // - // // Parse - // let parsed_packet3 = parse_lp_packet(&buf3, None).unwrap(); - // - // // Perform replay check (should fail) - // let replay_result = local_manager - // .receiving_counter_quick_check(receiver_index, parsed_packet3.header.counter); - // assert!(replay_result.is_err()); - // match replay_result.unwrap_err() { - // LpError::Replay(e) => { - // assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter)); - // } - // e => panic!("Expected replay error for packet 3, got {:?}", e), - // } - // // Do not mark received - } - } -} diff --git a/common/nym-lp/src/packet.rs b/common/nym-lp/src/packet.rs index 230a6b2795..5cc04ba867 100644 --- a/common/nym-lp/src/packet.rs +++ b/common/nym-lp/src/packet.rs @@ -17,7 +17,7 @@ pub(crate) const UDP_PAYLOAD_SIZE: usize = MTU - UDP_OVERHEAD; pub use nym_lp_packet::version; -pub(crate) trait LpPacketReplayExt { +pub trait LpPacketReplayExt { /// Validate packet counter against a replay protection validator /// /// This performs a quick check to see if the packet counter is valid before diff --git a/common/nym-lp/src/peer.rs b/common/nym-lp/src/peer.rs index 5bfc78a6cb..18e66a7a07 100644 --- a/common/nym-lp/src/peer.rs +++ b/common/nym-lp/src/peer.rs @@ -152,7 +152,7 @@ pub fn mock_peer() -> LpLocalPeer { } #[cfg(any(feature = "mock", test))] -pub fn random_peer<'a, R: rand09::CryptoRng + rand09::RngCore>(rng: &mut R) -> LpLocalPeer { +pub fn random_peer(rng: &mut R) -> LpLocalPeer { let x25519 = Arc::new(nym_kkt::key_utils::generate_keypair_x25519(rng)); LpLocalPeer { diff --git a/common/nym-lp/src/psq/initiator.rs b/common/nym-lp/src/psq/initiator.rs index 05b81a854d..2004ee72a6 100644 --- a/common/nym-lp/src/psq/initiator.rs +++ b/common/nym-lp/src/psq/initiator.rs @@ -214,7 +214,6 @@ mod tests { let conn_resp = conn_resp.leak(); let (mut init, mut resp) = mock_peers(); - let init_remote = init.as_remote(); let resp_remote = resp.as_remote(); let ciphersuite = Ciphersuite::default().with_kem(kem); @@ -246,8 +245,8 @@ mod tests { let responder_x25519_keypair = resp.x25519(); let kkt_responder = KKTResponder::new( - &responder_x25519_keypair, - &resp_keys, + responder_x25519_keypair, + resp_keys, &supported_hash, &supported_sigs, &[1], @@ -282,11 +281,6 @@ mod tests { .await??; responder.read_message(&msg, &mut []).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") - }; - // 4 send PSQ response let mut buf = vec![0u8; PSQ_MSG2_SIZE]; let n = responder.write_message(&[], &mut buf).unwrap(); @@ -304,7 +298,7 @@ mod tests { let mut r_transport = responder.into_session().unwrap(); // test serialization, deserialization - let mut channel_i = session_init.active_transport(); + let channel_i = session_init.active_transport(); let mut channel_r = r_transport.transport_channel().unwrap(); assert_eq!(channel_i.identifier(), channel_r.identifier()); @@ -312,13 +306,13 @@ mod tests { let app_data_i = b"Derived session hey".as_slice(); let app_data_r = b"Derived session ho".as_slice(); - let ct_i = encrypt_data(app_data_i, &mut channel_i)?; + let ct_i = encrypt_data(app_data_i, channel_i)?; let pt_r = decrypt_data(&ct_i, &mut channel_r)?; assert_eq!(app_data_i, pt_r); let ct_r = encrypt_data(app_data_r, &mut channel_r)?; - let pt_i = decrypt_data(&ct_r, &mut channel_i)?; + let pt_i = decrypt_data(&ct_r, channel_i)?; assert_eq!(app_data_r, pt_i); } diff --git a/common/nym-lp/src/psq/mod.rs b/common/nym-lp/src/psq/mod.rs index 2dac8c1566..0aa9b89ed3 100644 --- a/common/nym-lp/src/psq/mod.rs +++ b/common/nym-lp/src/psq/mod.rs @@ -172,21 +172,21 @@ mod tests { ); // test serialization, deserialization - let mut channel_i = session_init.active_transport(); - let mut channel_r = session_resp.active_transport(); + let channel_i = session_init.active_transport(); + let channel_r = session_resp.active_transport(); assert_eq!(channel_i.identifier(), channel_r.identifier()); let app_data_i = b"Derived session hey".as_slice(); let app_data_r = b"Derived session ho".as_slice(); - let ct_i = encrypt_data(app_data_i, &mut channel_i)?; - let pt_r = decrypt_data(&ct_i, &mut channel_r)?; + let ct_i = encrypt_data(app_data_i, channel_i)?; + let pt_r = decrypt_data(&ct_i, channel_r)?; assert_eq!(app_data_i, pt_r); - let ct_r = encrypt_data(app_data_r, &mut channel_r)?; - let pt_i = decrypt_data(&ct_r, &mut channel_i)?; + let ct_r = encrypt_data(app_data_r, channel_r)?; + let pt_i = decrypt_data(&ct_r, channel_i)?; assert_eq!(app_data_r, pt_i); } @@ -218,8 +218,8 @@ mod tests { HashFunction::SHA256, ]; let kkt_responder = KKTResponder::new( - &responder_x25519_keypair, - &resp_keys, + responder_x25519_keypair, + resp_keys, &supported_hash, &supported_sigs, &[protocol_version], diff --git a/common/nym-lp/src/psq/responder.rs b/common/nym-lp/src/psq/responder.rs index 3b44f555d1..6cb7cb9048 100644 --- a/common/nym-lp/src/psq/responder.rs +++ b/common/nym-lp/src/psq/responder.rs @@ -233,7 +233,6 @@ mod tests { let conn_resp = conn_resp.leak(); let (mut init, mut resp) = mock_peers(); - let init_remote = init.as_remote(); let resp_remote = resp.as_remote(); let ciphersuite = Ciphersuite::default().with_kem(kem); @@ -314,7 +313,7 @@ mod tests { // test serialization, deserialization let mut channel_i = i_transport.transport_channel().unwrap(); - let mut channel_r = session_resp.active_transport(); + let channel_r = session_resp.active_transport(); assert_eq!(channel_i.identifier(), channel_r.identifier()); @@ -322,11 +321,11 @@ mod tests { let app_data_r = b"Derived session ho".as_slice(); let ct_i = encrypt_data(app_data_i, &mut channel_i)?; - let pt_r = decrypt_data(&ct_i, &mut channel_r)?; + let pt_r = decrypt_data(&ct_i, channel_r)?; assert_eq!(app_data_i, pt_r); - let ct_r = encrypt_data(app_data_r, &mut channel_r)?; + let ct_r = encrypt_data(app_data_r, channel_r)?; let pt_i = decrypt_data(&ct_r, &mut channel_i)?; assert_eq!(app_data_r, pt_i); diff --git a/common/nym-lp/src/replay/validator.rs b/common/nym-lp/src/replay/validator.rs index 3e7a0cfa5d..d3edb7b84b 100644 --- a/common/nym-lp/src/replay/validator.rs +++ b/common/nym-lp/src/replay/validator.rs @@ -38,6 +38,16 @@ const N_WORDS: usize = 16; /// Total number of bits in the bitmap const N_BITS: usize = WORD_SIZE * N_WORDS; +/// Current packet count statistics +#[derive(Copy, Clone, Debug, PartialEq)] +pub struct PacketCount { + /// the next expected counter value + pub next: u64, + + /// the total number of received packets + pub received: u64, +} + /// Validator for receiving key counters to prevent replay attacks. /// /// This structure maintains a bitmap of received packets and validates @@ -205,11 +215,14 @@ impl ReceivingKeyCounterValidator { /// Returns the current packet count statistics. /// - /// Returns a tuple of `(next, receive_cnt)` where: + /// Returns a struct consisting of `(next, receive_cnt)` where: /// - `next` is the next expected counter value /// - `receive_cnt` is the total number of received packets - pub fn current_packet_cnt(&self) -> (u64, u64) { - (self.next, self.receive_cnt) + pub fn current_packet_cnt(&self) -> PacketCount { + PacketCount { + next: self.next, + received: self.receive_cnt, + } } #[inline(always)] @@ -481,7 +494,10 @@ mod tests { let mut validator = ReceivingKeyCounterValidator::default(); // Initial state - let (next, count) = validator.current_packet_cnt(); + let PacketCount { + next, + received: count, + } = validator.current_packet_cnt(); assert_eq!(next, 0); assert_eq!(count, 0); @@ -490,21 +506,30 @@ mod tests { assert!(validator.mark_did_receive_branchless(1).is_ok()); assert!(validator.mark_did_receive_branchless(2).is_ok()); - let (next, count) = validator.current_packet_cnt(); + let PacketCount { + next, + received: count, + } = validator.current_packet_cnt(); assert_eq!(next, 3); assert_eq!(count, 3); // After an out of order packet assert!(validator.mark_did_receive_branchless(10).is_ok()); - let (next, count) = validator.current_packet_cnt(); + let PacketCount { + next, + received: count, + } = validator.current_packet_cnt(); assert_eq!(next, 11); assert_eq!(count, 4); // After a packet from the past (within window) assert!(validator.mark_did_receive_branchless(5).is_ok()); - let (next, count) = validator.current_packet_cnt(); + let PacketCount { + next, + received: count, + } = validator.current_packet_cnt(); assert_eq!(next, 11); // Next doesn't change assert_eq!(count, 5); // Count increases } @@ -553,7 +578,7 @@ mod tests { assert!(validator.mark_did_receive_branchless(first_jump).is_ok()); // Verify next counter is updated - let (next, _) = validator.current_packet_cnt(); + let PacketCount { next, .. } = validator.current_packet_cnt(); assert_eq!(next, first_jump + 1); // Second large jump, even further ahead @@ -561,7 +586,7 @@ mod tests { assert!(validator.mark_did_receive_branchless(second_jump).is_ok()); // Verify next counter is updated again - let (next, _) = validator.current_packet_cnt(); + let PacketCount { next, .. } = validator.current_packet_cnt(); assert_eq!(next, second_jump + 1); // Test packets within the new window @@ -726,10 +751,10 @@ mod tests { // Check final state of the validator let final_state = validator.lock().unwrap(); - let (_next, receive_cnt) = final_state.current_packet_cnt(); + let count = final_state.current_packet_cnt(); // Verify that the received count matches our successful operations - assert_eq!(receive_cnt, total_successes as u64); + assert_eq!(count.received, total_successes as u64); } #[test] diff --git a/common/nym-lp/src/session.rs b/common/nym-lp/src/session.rs index 5b3f6e79f4..336d28ed45 100644 --- a/common/nym-lp/src/session.rs +++ b/common/nym-lp/src/session.rs @@ -11,6 +11,7 @@ use crate::psq::{ InitiatorData, PSQHandshakeState, PSQHandshakeStateInitiator, PSQHandshakeStateResponder, ResponderData, }; +use crate::replay::validator::PacketCount; use crate::{LpError, LpMessage, LpPacket, ReceivingKeyCounterValidator}; use libcrux_psq::handshake::types::{Authenticator, DHPublicKey}; use libcrux_psq::session::{Session, SessionBinding}; @@ -162,7 +163,7 @@ impl LpSession { self.psq_session.identifier() } - pub fn id(&self) -> u32 { + pub fn receiver_index(&self) -> u32 { let TODO = "this is temporary..."; let id = self.psq_session.identifier(); u32::from_le_bytes([id[0], id[1], id[2], id[3]]) @@ -177,7 +178,12 @@ impl LpSession { pub fn next_packet(&mut self, message: LpMessage) -> Result { let counter = self.next_counter(); - let header = LpHeader::new(self.id(), counter, self.protocol_version, message.typ()); + let header = LpHeader::new( + self.receiver_index(), + counter, + self.protocol_version, + message.typ(), + ); let packet = LpPacket::new(header, message); Ok(packet) } @@ -236,7 +242,7 @@ impl LpSession { /// A tuple containing: /// * The next expected counter value for incoming packets /// * The total number of received packets - pub fn current_packet_cnt(&self) -> (u64, u64) { + pub fn current_packet_cnt(&self) -> PacketCount { self.receiving_counter.current_packet_cnt() } @@ -280,7 +286,7 @@ impl LpSession { #[cfg(test)] mod tests { use super::*; - use crate::{ReplayError, SessionsMock, sessions_for_tests}; + use crate::{ReplayError, SessionsMock}; use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; #[test] @@ -351,17 +357,17 @@ mod tests { 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); + let packet_count = session.current_packet_cnt(); + assert_eq!(packet_count.next, 0); + assert_eq!(packet_count.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); + let packet_count = session.current_packet_cnt(); + assert_eq!(packet_count.next, 2); + assert_eq!(packet_count.received, 2); } } } diff --git a/common/nym-lp/src/session_integration/mod.rs b/common/nym-lp/src/session_integration/mod.rs index 8934ec25a9..33fdf91b51 100644 --- a/common/nym-lp/src/session_integration/mod.rs +++ b/common/nym-lp/src/session_integration/mod.rs @@ -1,721 +1,393 @@ #[cfg(test)] mod tests { - use nym_lp_packet::{LpHeader, LpMessage, LpPacket}; + use crate::state_machine::{LpAction, LpData, LpInput, LpStateBare}; + use crate::{LpError, SessionManager, SessionsMock}; + use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; + use nym_lp_packet::EncryptedLpPacket; - // Function to create a test packet - similar to how it's done in codec.rs tests - fn create_test_packet( - protocol_version: u8, - receiver_idx: u32, - counter: u64, - message: LpMessage, - ) -> LpPacket { - // Create the header - let header = LpHeader::new(receiver_idx, counter, protocol_version, message.typ()); + // helpers to make tests smaller + trait ActionExtract { + fn ciphertext(self) -> EncryptedLpPacket; - // Create and return the packet directly - LpPacket::new(header, message) + fn data(self) -> LpData; } - /// Tests the complete session flow including: + impl ActionExtract for LpAction { + fn ciphertext(self) -> EncryptedLpPacket { + if let LpAction::SendPacket(packet) = self { + packet + } else { + panic!("invalid action"); + } + } + + fn data(self) -> LpData { + if let LpAction::DeliverData(data) = self { + data + } else { + panic!("invalid action"); + } + } + } + + /// Tests simultaneous bidirectional communication between sessions + #[test] + fn test_bidirectional_communication() { + // 1. Initialize session manager + let mut session_manager_1 = SessionManager::new(); + let mut session_manager_2 = SessionManager::new(); + + for kem in KEM::iter() { + let sessions = SessionsMock::mock_post_handshake(kem); + + // 2. Create sessions using the pre-built Noise states + let peer_a_sm = session_manager_1 + .create_session_state_machine(sessions.initiator) + .unwrap(); + let peer_b_sm = session_manager_2 + .create_session_state_machine(sessions.responder) + .unwrap(); + + // 3. Send multiple encrypted messages both ways + const NUM_MESSAGES: u64 = 5; + for i in 0..NUM_MESSAGES { + println!("Bidirectional test: Round {i}"); + // --- A sends to B --- + let plaintext_a = format!("A->B Message {i}").into_bytes(); + let ciphertext_a = session_manager_1 + .send_data(peer_a_sm, LpData::new_opaque(plaintext_a.clone())) + .unwrap() + .ciphertext(); + + // B parses and checks replay + let decrypted_payload = session_manager_2 + .receive_packet(peer_b_sm, ciphertext_a) + .unwrap() + .unwrap() + .data(); + assert_eq!(decrypted_payload.content, plaintext_a); + + // --- B sends to A --- + let plaintext_b = format!("B->A Message {i}").into_bytes(); + let ciphertext_b = session_manager_2 + .send_data(peer_b_sm, LpData::new_opaque(plaintext_b.clone())) + .unwrap() + .ciphertext(); + + // B parses and checks replay + let decrypted_payload = session_manager_1 + .receive_packet(peer_a_sm, ciphertext_b) + .unwrap() + .unwrap() + .data(); + assert_eq!(decrypted_payload.content, plaintext_b); + } + + // 5. Verify counter stats + // Note: current_packet_cnt() returns (next_expected_receive_counter, total_received) + let count_a = session_manager_1.current_packet_cnt(peer_a_sm).unwrap(); + let count_b = session_manager_2.current_packet_cnt(peer_b_sm).unwrap(); + + // Peer A sent handshake(0), handshake(1) + 5 data packets = 7 total. Next send counter = 7. + // Peer A received handshake(0) + 5 data packets = 6 total. Next expected recv counter = 6. + assert_eq!( + count_a.received, NUM_MESSAGES, + "Peer A total received count mismatch" + ); // Received 5 data + assert_eq!( + count_a.next, NUM_MESSAGES, + "Peer A next expected receive counter mismatch" + ); // Expected counter for msg from B + + // Peer B sent handshake(0) + 5 data packets = 6 total. Next send counter = 6. + // Peer B received handshake(0), handshake(1) + 5 data packets = 7 total. Next expected recv counter = 7. + assert_eq!( + count_b.received, NUM_MESSAGES, + "Peer B total received count mismatch" + ); // Received 5 data + assert_eq!( + count_b.next, NUM_MESSAGES, + "Peer B next expected receive counter mismatch" + ); // Expected counter for msg from A + + println!("Bidirectional test completed."); + } + } + + /// Tests error handling in session flow + #[test] + fn test_session_error_handling() { + for kem in KEM::iter() { + // 1. Initialize session manager + let mut session_manager = SessionManager::new(); + + let sessions = SessionsMock::mock_post_handshake(kem); + let session_id = *sessions.initiator.session_identifier(); + + let non_existent = [42u8; 32]; + // sanity check in case of the 1 in 2^256 + assert_ne!(session_id, non_existent); + + let session1 = sessions.initiator; + let session2 = sessions.responder; + + // 2. Create a session (using real noise state) + let _session = session_manager.create_session_state_machine(session1); + + // 3. Try to get a non-existent session + let result = session_manager.state_machine_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); + 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); + + assert!( + session_manager.remove_state_machine(session_id), + "Should remove the session" + ); + } + } + + /// Tests the complete session flow using ONLY the process_input interface: /// - Creation of sessions through session manager - /// - Packet encoding/decoding with the session - /// - Replay protection across the session - /// - Multiple sessions with unique indices - /// - Session removal and cleanup + /// - Data transfer driven by SendData, ReceivePacket inputs + /// - Actions like SendPacket, DeliverData handled from output + /// - Implicit replay protection via state machine logic + /// - Closing driven by Close input #[test] fn test_full_session_flow() { - todo!() + // 1. Initialize session managers + let mut session_manager_1 = SessionManager::new(); + let mut session_manager_2 = SessionManager::new(); - // // 1. Initialize session manager - // let mut session_manager_1 = SessionManager::new(); - // let mut session_manager_2 = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let receiver_index = 12345; - // let sessions = SessionsMock::mock_post_handshake(receiver_index); - // - // // 2. Create sessions using the pre-built Noise states - // let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); - // let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); - // - // // Verify session count - // assert_eq!(session_manager_1.session_count(), 1); - // assert_eq!(session_manager_2.session_count(), 1); - // - // // 3. Simulate Data Transfer (Post-Handshake) - // println!("Starting data transfer simulation..."); - // let plaintext_a_to_b = b"Hello from A!"; - // - // // A encrypts data - // let ciphertext_a_to_b = session_manager_1 - // .encrypt_data(peer_a_sm, plaintext_a_to_b) - // .expect("A encrypt failed"); - // - // // A prepares packet - // let counter_a = session_manager_1.next_counter(peer_a_sm).unwrap(); - // let message_a_to_b = create_test_packet(1, receiver_index, counter_a, ciphertext_a_to_b); - // let mut encoded_data_a_to_b = BytesMut::new(); - // serialize_lp_packet(&message_a_to_b, &mut encoded_data_a_to_b, None) - // .expect("A serialize data failed"); - // - // // B parses packet and checks replay - // let decoded_packet_b = - // parse_lp_packet(&encoded_data_a_to_b, None).expect("B parse data failed"); - // assert_eq!(decoded_packet_b.header.counter, counter_a); - // - // // Check replay before decrypting - // session_manager_2 - // .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) - // .expect("B data replay check failed (A->B)"); - // - // // B decrypts data - // let decrypted_payload = session_manager_2 - // .decrypt_data(peer_b_sm, &decoded_packet_b.message) - // .expect("B decrypt failed"); - // assert_eq!(decrypted_payload, plaintext_a_to_b); - // // Mark counter only after successful decryption - // session_manager_2 - // .receiving_counter_mark(peer_b_sm, decoded_packet_b.header.counter) - // .expect("B mark data counter failed"); - // println!( - // " A->B: Decrypted successfully: {:?}", - // String::from_utf8_lossy(&decrypted_payload) - // ); - // - // // B sends data to A - // let plaintext_b_to_a = b"Hello from B!"; - // let ciphertext_b_to_a = session_manager_2 - // .encrypt_data(peer_b_sm, plaintext_b_to_a) - // .expect("B encrypt failed"); - // let counter_b = session_manager_2.next_counter(peer_b_sm).unwrap(); - // let message_b_to_a = create_test_packet(1, receiver_index, counter_b, ciphertext_b_to_a); - // let mut encoded_data_b_to_a = BytesMut::new(); - // serialize_lp_packet(&message_b_to_a, &mut encoded_data_b_to_a, None) - // .expect("B serialize data failed"); - // - // // A parses packet and checks replay - // let decoded_packet_a = - // parse_lp_packet(&encoded_data_b_to_a, None).expect("A parse data failed"); - // assert_eq!(decoded_packet_a.header.counter, counter_b); - // - // // Check replay before decrypting - // session_manager_1 - // .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) - // .expect("A data replay check failed (B->A)"); - // - // // A decrypts data - // let decrypted_payload = session_manager_1 - // .decrypt_data(peer_a_sm, &decoded_packet_a.message) - // .expect("A decrypt failed"); - // assert_eq!(decrypted_payload, plaintext_b_to_a); - // // Mark counter only after successful decryption - // session_manager_1 - // .receiving_counter_mark(peer_a_sm, decoded_packet_a.header.counter) - // .expect("A mark data counter failed"); - // println!( - // " B->A: Decrypted successfully: {:?}", - // String::from_utf8_lossy(&decrypted_payload) - // ); - // - // println!("Data transfer simulation completed."); - // - // // 4. Replay Protection Test (Data Packet) - // println!("Testing data packet replay protection..."); - // // Try to replay the last message from B to A - // // Need to re-encode because decode consumes the buffer - // let message_b_to_a_replay = create_test_packet( - // 1, - // receiver_index, - // counter_b, - // LpMessage::EncryptedData(crate::message::EncryptedDataPayload( - // plaintext_b_to_a.to_vec(), - // )), // Using plaintext here, but content doesn't matter for replay check - // ); - // let mut encoded_data_b_to_a_replay = BytesMut::new(); - // serialize_lp_packet( - // &message_b_to_a_replay, - // &mut encoded_data_b_to_a_replay, - // None, - // ) - // .expect("B serialize replay failed"); - // - // let parsed_replay_packet = - // parse_lp_packet(&encoded_data_b_to_a_replay, None).expect("A parse replay failed"); - // let replay_result = session_manager_1 - // .receiving_counter_quick_check(peer_a_sm, parsed_replay_packet.header.counter); - // assert!(replay_result.is_err(), "Data replay should be prevented"); - // assert!( - // matches!(replay_result.unwrap_err(), LpError::Replay(_)), - // "Should be a replay protection error for data packet" - // ); - // println!("Data packet replay protection test passed."); - // - // // 5. Test out-of-order packet reception (send counter N+1 before counter N) - // println!("Testing out-of-order data packet reception..."); - // let counter_a_next = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 1 - // let counter_a_skip = session_manager_1.next_counter(peer_a_sm).unwrap(); // Should be counter_a + 2 - // - // // Prepare data for counter_a_skip (N+1) - // let plaintext_skip = b"Out of order message"; - // let ciphertext_skip = session_manager_1 - // .encrypt_data(peer_a_sm, plaintext_skip) - // .expect("A encrypt skip failed"); - // - // let message_a_to_b_skip = create_test_packet( - // 1, // protocol version - // receiver_index, - // counter_a_skip, // Send N+1 first - // ciphertext_skip, - // ); - // - // // Encode the skip message - // let mut encoded_skip = BytesMut::new(); - // serialize_lp_packet(&message_a_to_b_skip, &mut encoded_skip, None) - // .expect("Failed to serialize skip message"); - // - // // B parses skip message and checks replay - // let decoded_packet_skip = - // parse_lp_packet(&encoded_skip, None).expect("B parse skip failed"); - // session_manager_2 - // .receiving_counter_quick_check(peer_b_sm, decoded_packet_skip.header.counter) - // .expect("B replay check skip failed"); - // assert_eq!(decoded_packet_skip.header.counter, counter_a_skip); - // - // // B decrypts skip message - // let decrypted_payload = session_manager_2 - // .decrypt_data(peer_b_sm, &decoded_packet_skip.message) - // .expect("B decrypt skip failed"); - // assert_eq!(decrypted_payload, plaintext_skip); - // // Mark counter N+1 - // session_manager_2 - // .receiving_counter_mark(peer_b_sm, decoded_packet_skip.header.counter) - // .expect("B mark skip counter failed"); - // println!( - // " A->B (Counter {}): Decrypted successfully: {:?}", - // counter_a_skip, - // String::from_utf8_lossy(&decrypted_payload) - // ); - // - // // 6. Now send the skipped counter N message (should still work) - // println!("Testing delayed data packet reception..."); - // // Prepare data for counter_a_next (N) - // let plaintext_delayed = b"Delayed message"; - // let ciphertext_delayed = session_manager_1 - // .encrypt_data(peer_a_sm, plaintext_delayed) - // .expect("A encrypt delayed failed"); - // - // let message_a_to_b_delayed = create_test_packet( - // 1, // protocol version - // receiver_index, - // counter_a_next, // counter N (delayed packet) - // ciphertext_delayed, - // ); - // - // // Encode the delayed message - // let mut encoded_delayed = BytesMut::new(); - // serialize_lp_packet(&message_a_to_b_delayed, &mut encoded_delayed, None) - // .expect("Failed to serialize delayed message"); - // - // // Make a copy for replay test later - // let encoded_delayed_copy = encoded_delayed.clone(); - // - // // B parses delayed message and checks replay - // let decoded_packet_delayed = - // parse_lp_packet(&encoded_delayed, None).expect("B parse delayed failed"); - // session_manager_2 - // .receiving_counter_quick_check(peer_b_sm, decoded_packet_delayed.header.counter) - // .expect("B replay check delayed failed"); - // assert_eq!(decoded_packet_delayed.header.counter, counter_a_next); - // - // // B decrypts delayed message - // let decrypted_payload = session_manager_2 - // .decrypt_data(peer_b_sm, &decoded_packet_delayed.message) - // .expect("B decrypt delayed failed"); - // assert_eq!(decrypted_payload, plaintext_delayed); - // // Mark counter N - // session_manager_2 - // .receiving_counter_mark(peer_b_sm, decoded_packet_delayed.header.counter) - // .expect("B mark delayed counter failed"); - // println!( - // " A->B (Counter {}): Decrypted successfully: {:?}", - // counter_a_next, - // String::from_utf8_lossy(&decrypted_payload) - // ); - // - // println!("Delayed data packet reception test passed."); - // - // // 7. Try to replay message with counter N (should fail) - // println!("Testing replay of delayed packet..."); - // let parsed_delayed_replay = - // parse_lp_packet(&encoded_delayed_copy, None).expect("Parse delayed replay failed"); - // let result = session_manager_2 - // .receiving_counter_quick_check(peer_b_sm, parsed_delayed_replay.header.counter); - // assert!(result.is_err(), "Replay attack should be prevented"); - // assert!( - // matches!(result, Err(LpError::Replay(_))), - // "Should be a replay protection error" - // ); - // - // // 8. Session removal - // assert!(session_manager_1.remove_state_machine(receiver_index)); - // assert_eq!(session_manager_1.session_count(), 0); - // - // // Verify the session is gone - // let session = session_manager_1.state_machine_exists(receiver_index); - // assert!(!session, "Session should be removed"); - // - // // But the other session still exists - // let session = session_manager_2.state_machine_exists(receiver_index); - // assert!(session, "Session still exists in the other manager"); + for kem in KEM::iter() { + let sessions = SessionsMock::mock_post_handshake(kem); + let session_id = *sessions.responder.session_identifier(); + + // 2. Create sessions state machines + session_manager_1 + .create_session_state_machine(sessions.initiator) + .unwrap(); + session_manager_2 + .create_session_state_machine(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 + ); + + // --- 3. Simulate Data Transfer via process_input --- + println!("Starting data transfer simulation via process_input..."); + let plaintext_a_to_b = LpData::new_opaque(b"Hello from A via process_input!".to_vec()); + let plaintext_b_to_a = LpData::new_opaque(b"Hello from B via process_input!".to_vec()); + + // --- A sends to B --- + 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(); + + // B receives + 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 { + assert_eq!(data, plaintext_a_to_b, "Decrypted data mismatch A->B"); + println!( + " B successfully decrypted: {:?}", + String::from_utf8_lossy(&data.content) + ); + } else { + panic!("B ReceivePacket did not produce DeliverData"); + } + + // --- B sends to A --- + 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(); + + // Keep a copy for replay test + let data_packet_b_replay = data_packet_b.clone(); + + // A receives + 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 { + assert_eq!(data, plaintext_b_to_a, "Decrypted data mismatch B->A"); + println!( + " A successfully decrypted: {:?}", + String::from_utf8_lossy(&data.content) + ); + } else { + panic!("A ReceivePacket did not produce DeliverData"); + } + println!("Data transfer simulation completed."); + + // --- 4. Replay Protection Test --- + println!("Testing data packet replay protection via process_input..."); + let replay_result = session_manager_1 + .process_input(session_id, LpInput::ReceivePacket(data_packet_b_replay)); // Use cloned packet + + assert!(replay_result.is_err(), "Replay should produce Err(...)"); + let error = replay_result.err().unwrap(); + assert!( + matches!(error, LpError::Replay(_)), + "Expected Replay error, got {:?}", + error + ); + println!("Data packet replay protection test passed."); + + // --- 5. Out-of-Order Test --- + println!("Testing out-of-order reception via process_input..."); + + // A prepares N+1 then N + let data_n_plus_1 = LpData::new_opaque(b"Message N+1".to_vec()); + let data_n = LpData::new_opaque(b"Message N".to_vec()); + + 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, + _ => panic!("Expected SendPacket"), + }; + + 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, + _ => panic!("Expected SendPacket"), + }; + let packet_n_replay = packet_n.clone(); // For replay test + + // B receives N+1 first + 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"), + _ => panic!("Expected DeliverData for N+1"), + } + + // B receives N second (should work) + 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"), + _ => panic!("Expected DeliverData for N"), + } + + // B tries to replay N (should fail) + println!(" B tries to replay N"); + let replay_n_result = session_manager_2 + .process_input(session_id, LpInput::ReceivePacket(packet_n_replay)); + assert!(replay_n_result.is_err(), "Replay N should produce Err"); + assert!( + matches!(replay_n_result.err().unwrap(), LpError::Replay(_)), + "Expected Replay error for N" + ); + 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(LpData::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(LpData::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)); + assert_eq!(session_manager_1.session_count(), 0); + assert!(!session_manager_1.state_machine_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_eq!(session_manager_2.session_count(), 0); + assert!(!session_manager_2.state_machine_exists(&session_id)); + println!("Session removal test passed."); + } } - // - // /// Tests simultaneous bidirectional communication between sessions - // #[test] - // fn test_bidirectional_communication() { - // // 1. Initialize session manager - // let mut session_manager_1 = SessionManager::new(); - // let mut session_manager_2 = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let receiver_index = 12345; - // let sessions = SessionsMock::mock_post_handshake(receiver_index); - // - // // 2. Create sessions using the pre-built Noise states - // let peer_a_sm = session_manager_1.create_session_state_machine(sessions.initiator); - // let peer_b_sm = session_manager_2.create_session_state_machine(sessions.responder); - // - // // Counters after handshake - // let mut counter_a = 0; // Next counter for A to send - // let mut counter_b = 0; // Next counter for B to send - // - // // 3. Send multiple encrypted messages both ways - // const NUM_MESSAGES: u64 = 5; - // for i in 0..NUM_MESSAGES { - // println!("Bidirectional test: Round {}", i); - // // --- A sends to B --- - // let plaintext_a = format!("A->B Message {}", i).into_bytes(); - // let ciphertext_a = session_manager_1 - // .encrypt_data(peer_a_sm, &plaintext_a) - // .expect("A encrypt failed"); - // let current_counter_a = counter_a; - // counter_a += 1; - // - // let message_a = create_test_packet(1, receiver_index, current_counter_a, ciphertext_a); - // let mut encoded_a = BytesMut::new(); - // serialize_lp_packet(&message_a, &mut encoded_a, None).expect("A serialize failed"); - // - // // B parses and checks replay - // let decoded_packet_b = parse_lp_packet(&encoded_a, None).expect("B parse failed"); - // session_manager_2 - // .receiving_counter_quick_check(peer_b_sm, decoded_packet_b.header.counter) - // .expect("B replay check failed (A->B)"); - // assert_eq!(decoded_packet_b.header.counter, current_counter_a); - // let decrypted_payload = session_manager_2 - // .decrypt_data(peer_b_sm, &decoded_packet_b.message) - // .expect("B decrypt failed"); - // assert_eq!(decrypted_payload, plaintext_a); - // session_manager_2 - // .receiving_counter_mark(peer_b_sm, current_counter_a) - // .expect("B mark counter failed"); - // - // // --- B sends to A --- - // let plaintext_b = format!("B->A Message {}", i).into_bytes(); - // let ciphertext_b = session_manager_2 - // .encrypt_data(peer_b_sm, &plaintext_b) - // .expect("B encrypt failed"); - // let current_counter_b = counter_b; - // counter_b += 1; - // - // let message_b = create_test_packet(1, receiver_index, current_counter_b, ciphertext_b); - // let mut encoded_b = BytesMut::new(); - // serialize_lp_packet(&message_b, &mut encoded_b, None).expect("B serialize failed"); - // - // // A parses and checks replay - // let decoded_packet_a = parse_lp_packet(&encoded_b, None).expect("A parse failed"); - // session_manager_1 - // .receiving_counter_quick_check(peer_a_sm, decoded_packet_a.header.counter) - // .expect("A replay check failed (B->A)"); - // assert_eq!(decoded_packet_a.header.counter, current_counter_b); - // let decrypted_payload = session_manager_1 - // .decrypt_data(peer_a_sm, &decoded_packet_a.message) - // .expect("A decrypt failed"); - // assert_eq!(decrypted_payload, plaintext_b); - // session_manager_1 - // .receiving_counter_mark(peer_a_sm, current_counter_b) - // .expect("A mark counter failed"); - // } - // - // // 5. Verify counter stats - // // Note: current_packet_cnt() returns (next_expected_receive_counter, total_received) - // let (next_recv_a, total_recv_a) = session_manager_1.current_packet_cnt(peer_a_sm).unwrap(); - // let (next_recv_b, total_recv_b) = session_manager_2.current_packet_cnt(peer_b_sm).unwrap(); - // - // // Peer A sent handshake(0), handshake(1) + 5 data packets = 7 total. Next send counter = 7. - // // Peer A received handshake(0) + 5 data packets = 6 total. Next expected recv counter = 6. - // assert_eq!( - // counter_a, NUM_MESSAGES, - // "Peer A final send counter mismatch" - // ); - // assert_eq!( - // total_recv_a, NUM_MESSAGES, - // "Peer A total received count mismatch" - // ); // Received 5 data - // assert_eq!( - // next_recv_a, NUM_MESSAGES, - // "Peer A next expected receive counter mismatch" - // ); // Expected counter for msg from B - // - // // Peer B sent handshake(0) + 5 data packets = 6 total. Next send counter = 6. - // // Peer B received handshake(0), handshake(1) + 5 data packets = 7 total. Next expected recv counter = 7. - // assert_eq!( - // counter_b, NUM_MESSAGES, - // "Peer B final send counter mismatch" - // ); - // assert_eq!( - // total_recv_b, NUM_MESSAGES, - // "Peer B total received count mismatch" - // ); // Received 5 data - // assert_eq!( - // next_recv_b, NUM_MESSAGES, - // "Peer B next expected receive counter mismatch" - // ); // Expected counter for msg from A - // - // println!("Bidirectional test completed."); - // } - // - // /// Tests error handling in session flow - // #[test] - // fn test_session_error_handling() { - // // 1. Initialize session manager - // let mut session_manager = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // let receiver_index = 123; - // let session1 = SessionsMock::mock_post_handshake(receiver_index).initiator; - // let session2 = SessionsMock::mock_post_handshake(124).initiator; - // - // // 2. Create a session (using real noise state) - // let _session = session_manager.create_session_state_machine(session1); - // - // // 3. Try to get a non-existent session - // let result = session_manager.state_machine_exists(999); - // assert!(!result, "Non-existent session should return None"); - // - // // 4. Try to remove a non-existent session - // let result = session_manager.remove_state_machine(999); - // 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); - // - // assert!( - // session_manager.remove_state_machine(124), - // "Should remove the session" - // ); - // - // // 6. Create a codec and test error cases - // // let mut codec = LPCodec::new(session); - // - // // 7. Create an invalid message type packet - // let mut buf = BytesMut::new(); - // - // // Add header - // buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved - // buf.extend_from_slice(&receiver_index.to_le_bytes()); // Sender index - // buf.extend_from_slice(&0u64.to_le_bytes()); // Counter - // - // // Add invalid message type - // buf.extend_from_slice(&0xFFFFu16.to_le_bytes()); - // - // // Add some dummy data - // buf.extend_from_slice(&[0u8; 80]); - // - // // Add trailer - // buf.extend_from_slice(&[0u8; TRAILER_LEN]); - // - // // Try to parse the invalid message type - // let result = parse_lp_packet(&buf, None); - // assert!(result.is_err(), "Decoding invalid message type should fail"); - // - // // Add assertion for the specific error type - // assert!(matches!( - // result.unwrap_err(), - // LpError::InvalidMessageType(0xFFFF) - // )); - // - // // 8. Test partial packet decoding - // let partial_packet = &buf[0..10]; // Too short to be a valid packet - // let partial_bytes = BytesMut::from(partial_packet); - // - // let result = parse_lp_packet(&partial_bytes, None); - // assert!(result.is_err(), "Parsing partial packet should fail"); - // assert!(matches!( - // result.unwrap_err(), - // LpError::InsufficientBufferSize - // )); - // } - // // Remove unused imports if SessionManager methods are no longer direct dependencies - // // use crate::noise_protocol::{create_noise_state, create_noise_state_responder}; - // use crate::state_machine::LpData; - // use crate::state_machine::{LpAction, LpInput, LpStateBare}; - // // Use Bytes for SendData input - // - // // Keep helper function for creating test packets if needed, - // // but LpAction::SendPacket should provide the packets now. - // // fn create_test_packet(...) -> LpPacket { ... } - // - // /// Tests the complete session flow using ONLY the process_input interface: - // /// - Creation of sessions through session manager - // /// - Handshake driven by StartHandshake, ReceivePacket inputs - // /// - Data transfer driven by SendData, ReceivePacket inputs - // /// - Actions like SendPacket, DeliverData handled from output - // /// - Implicit replay protection via state machine logic - // /// - Closing driven by Close input - // #[test] - // fn test_full_session_flow_with_process_input() { - // // 1. Initialize session managers - // let mut session_manager_1 = SessionManager::new(); - // let mut session_manager_2 = SessionManager::new(); - // - // let TODO = " for kem in kem_list() {"; - // - // let receiver_index = 12345; - // let sessions = SessionsMock::mock_post_handshake(receiver_index); - // - // // 2. Create sessions state machines - // session_manager_1.create_session_state_machine(sessions.initiator); - // session_manager_2.create_session_state_machine(sessions.responder); - // - // assert_eq!(session_manager_1.session_count(), 1); - // assert_eq!(session_manager_2.session_count(), 1); - // assert!(session_manager_1.state_machine_exists(receiver_index)); - // assert!(session_manager_2.state_machine_exists(receiver_index)); - // - // // Verify initial states are Transport - // assert_eq!( - // session_manager_1.get_state(receiver_index).unwrap(), - // LpStateBare::Transport - // ); - // assert_eq!( - // session_manager_2.get_state(receiver_index).unwrap(), - // LpStateBare::Transport - // ); - // - // // --- 3. Simulate Data Transfer via process_input --- - // println!("Starting data transfer simulation via process_input..."); - // let plaintext_a_to_b = LpData::new_opaque(b"Hello from A via process_input!".to_vec()); - // let plaintext_b_to_a = LpData::new_opaque(b"Hello from B via process_input!".to_vec()); - // - // // --- A sends to B --- - // println!(" A sends to B"); - // let action_a_send = session_manager_1 - // .process_input(receiver_index, LpInput::SendData(plaintext_a_to_b.clone())) - // .expect("A SendData should produce action") - // .expect("A SendData failed"); - // - // let data_packet_a = if let LpAction::SendPacket(packet) = action_a_send { - // packet - // } else { - // panic!("A SendData did not produce SendPacket"); - // }; - // - // // Simulate network - // let mut buf_data_a = BytesMut::new(); - // serialize_lp_packet(&data_packet_a, &mut buf_data_a, None).unwrap(); - // let parsed_data_a = parse_lp_packet(&buf_data_a, None).unwrap(); - // - // // B receives - // println!(" B receives from A"); - // let action_b_recv = session_manager_2 - // .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_a)) - // .expect("B ReceivePacket (data) should produce action") - // .expect("B ReceivePacket (data) failed"); - // - // if let LpAction::DeliverData(data) = action_b_recv { - // assert_eq!(data, plaintext_a_to_b, "Decrypted data mismatch A->B"); - // println!( - // " B successfully decrypted: {:?}", - // String::from_utf8_lossy(&data.content) - // ); - // } else { - // panic!("B ReceivePacket did not produce DeliverData"); - // } - // - // // --- B sends to A --- - // println!(" B sends to A"); - // let action_b_send = session_manager_2 - // .process_input(receiver_index, LpInput::SendData(plaintext_b_to_a.clone())) - // .expect("B SendData should produce action") - // .expect("B SendData failed"); - // - // let data_packet_b = if let LpAction::SendPacket(packet) = action_b_send { - // packet - // } else { - // panic!("B SendData did not produce SendPacket"); - // }; - // // Keep a copy for replay test - // let data_packet_b_replay = data_packet_b.clone(); - // - // // Simulate network - // let mut buf_data_b = BytesMut::new(); - // serialize_lp_packet(&data_packet_b, &mut buf_data_b, None).unwrap(); - // let parsed_data_b = parse_lp_packet(&buf_data_b, None).unwrap(); - // - // // A receives - // println!(" A receives from B"); - // let action_a_recv = session_manager_1 - // .process_input(receiver_index, LpInput::ReceivePacket(parsed_data_b)) - // .expect("A ReceivePacket (data) should produce action") - // .expect("A ReceivePacket (data) failed"); - // - // if let LpAction::DeliverData(data) = action_a_recv { - // assert_eq!(data, plaintext_b_to_a, "Decrypted data mismatch B->A"); - // println!( - // " A successfully decrypted: {:?}", - // String::from_utf8_lossy(&data.content) - // ); - // } else { - // panic!("A ReceivePacket did not produce DeliverData"); - // } - // println!("Data transfer simulation completed."); - // - // // --- 4. Replay Protection Test --- - // println!("Testing data packet replay protection via process_input..."); - // let replay_result = session_manager_1 - // .process_input(receiver_index, LpInput::ReceivePacket(data_packet_b_replay)); // Use cloned packet - // - // assert!(replay_result.is_err(), "Replay should produce Err(...)"); - // let error = replay_result.err().unwrap(); - // assert!( - // matches!(error, LpError::Replay(_)), - // "Expected Replay error, got {:?}", - // error - // ); - // println!("Data packet replay protection test passed."); - // - // // --- 5. Out-of-Order Test --- - // println!("Testing out-of-order reception via process_input..."); - // - // // A prepares N+1 then N - // let data_n_plus_1 = LpData::new_opaque(b"Message N+1".to_vec()); - // let data_n = LpData::new_opaque(b"Message N".to_vec()); - // - // let action_send_n1 = session_manager_1 - // .process_input(receiver_index, LpInput::SendData(data_n_plus_1.clone())) - // .unwrap() - // .unwrap(); - // let packet_n1 = match action_send_n1 { - // LpAction::SendPacket(p) => p, - // _ => panic!("Expected SendPacket"), - // }; - // - // let action_send_n = session_manager_1 - // .process_input(receiver_index, LpInput::SendData(data_n.clone())) - // .unwrap() - // .unwrap(); - // let packet_n = match action_send_n { - // LpAction::SendPacket(p) => p, - // _ => panic!("Expected SendPacket"), - // }; - // let packet_n_replay = packet_n.clone(); // For replay test - // - // // B receives N+1 first - // println!(" B receives N+1"); - // let action_recv_n1 = session_manager_2 - // .process_input(receiver_index, LpInput::ReceivePacket(packet_n1)) - // .unwrap() - // .unwrap(); - // match action_recv_n1 { - // LpAction::DeliverData(d) => assert_eq!(d, data_n_plus_1, "Data N+1 mismatch"), - // _ => panic!("Expected DeliverData for N+1"), - // } - // - // // B receives N second (should work) - // println!(" B receives N"); - // let action_recv_n = session_manager_2 - // .process_input(receiver_index, LpInput::ReceivePacket(packet_n)) - // .unwrap() - // .unwrap(); - // match action_recv_n { - // LpAction::DeliverData(d) => assert_eq!(d, data_n, "Data N mismatch"), - // _ => panic!("Expected DeliverData for N"), - // } - // - // // B tries to replay N (should fail) - // println!(" B tries to replay N"); - // let replay_n_result = session_manager_2 - // .process_input(receiver_index, LpInput::ReceivePacket(packet_n_replay)); - // assert!(replay_n_result.is_err(), "Replay N should produce Err"); - // assert!( - // matches!(replay_n_result.err().unwrap(), LpError::Replay(_)), - // "Expected Replay error for N" - // ); - // 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(receiver_index, 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(receiver_index).unwrap(), - // LpStateBare::Closed - // ); - // - // // Further actions on A fail - // let send_after_close_a = session_manager_1.process_input( - // receiver_index, - // LpInput::SendData(LpData::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(receiver_index, 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(receiver_index).unwrap(), - // LpStateBare::Closed - // ); - // - // // Further actions on B fail - // let send_after_close_b = session_manager_2.process_input( - // receiver_index, - // LpInput::SendData(LpData::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(receiver_index)); - // assert_eq!(session_manager_1.session_count(), 0); - // assert!(!session_manager_1.state_machine_exists(receiver_index)); - // - // // B's session manager still has it until removed - // assert!(session_manager_2.state_machine_exists(receiver_index)); - // assert!(session_manager_2.remove_state_machine(receiver_index)); - // assert_eq!(session_manager_2.session_count(), 0); - // assert!(!session_manager_2.state_machine_exists(receiver_index)); - // println!("Session removal test passed."); - // } // ... other tests ... } diff --git a/common/nym-lp/src/session_manager.rs b/common/nym-lp/src/session_manager.rs index b3d60b2efc..a23164bd65 100644 --- a/common/nym-lp/src/session_manager.rs +++ b/common/nym-lp/src/session_manager.rs @@ -7,10 +7,13 @@ //! creation, retrieval, and storage of sessions. use crate::session::SessionId; -use crate::state_machine::{LpAction, LpInput, LpStateBare}; +use crate::state_machine::{LpAction, LpData, LpInput, LpStateBare}; use crate::{LpError, LpSession, LpStateMachine}; +use nym_lp_packet::EncryptedLpPacket; use std::collections::HashMap; +pub use crate::replay::validator::PacketCount; + /// Manages the lifecycle of Lewes Protocol sessions. /// /// The SessionManager is responsible for creating, storing, and retrieving sessions @@ -41,6 +44,19 @@ impl SessionManager { self.with_state_machine_mut(lp_id, |sm| sm.process_input(input).transpose())? } + pub fn send_data(&mut self, lp_id: SessionId, data: LpData) -> Result { + self.process_input(lp_id, LpInput::SendData(data))? + .ok_or(LpError::NotInTransport) + } + + pub fn receive_packet( + &mut self, + lp_id: SessionId, + packet: EncryptedLpPacket, + ) -> Result, LpError> { + self.process_input(lp_id, LpInput::ReceivePacket(packet)) + } + pub fn closed(&self, lp_id: SessionId) -> Result { Ok(self.get_state(lp_id)? == LpStateBare::Closed) } @@ -58,31 +74,7 @@ impl SessionManager { self.with_state_machine(lp_id, |sm| Ok(sm.bare_state()))? } - pub fn receiving_counter_quick_check( - &self, - lp_id: SessionId, - counter: u64, - ) -> Result<(), LpError> { - self.with_state_machine(lp_id, |sm| { - sm.session()?.receiving_counter_quick_check(counter) - })? - } - - pub fn receiving_counter_mark( - &mut self, - lp_id: SessionId, - counter: u64, - ) -> Result<(), LpError> { - self.with_state_machine_mut(lp_id, |sm| { - sm.session_mut()?.receiving_counter_mark(counter) - })? - } - - pub fn next_counter(&mut self, lp_id: SessionId) -> Result { - self.with_state_machine_mut(lp_id, |sm| Ok(sm.session_mut()?.next_counter()))? - } - - pub fn current_packet_cnt(&self, lp_id: SessionId) -> Result<(u64, u64), LpError> { + pub fn current_packet_cnt(&self, lp_id: SessionId) -> Result { self.with_state_machine(lp_id, |sm| Ok(sm.session()?.current_packet_cnt()))? } @@ -117,11 +109,19 @@ impl SessionManager { } } - pub fn create_session_state_machine(&mut self, lp_session: LpSession) -> SessionId { + pub fn create_session_state_machine( + &mut self, + lp_session: LpSession, + ) -> Result { let session_id = *lp_session.session_identifier(); + + if self.state_machines.contains_key(&session_id) { + return Err(LpError::DuplicateSessionId(session_id)); + } + let sm = LpStateMachine::new(lp_session); self.state_machines.insert(session_id, sm); - session_id + Ok(session_id) } /// Method to remove a state machine @@ -144,7 +144,7 @@ mod tests { let local_session = mock_session_for_test(); let id = *local_session.session_identifier(); - let sm_1_id = manager.create_session_state_machine(local_session); + let sm_1_id = manager.create_session_state_machine(local_session).unwrap(); assert_eq!(sm_1_id, id); let retrieved = manager.state_machine_exists(&id); @@ -158,7 +158,7 @@ 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); + let sm_1_id = manager.create_session_state_machine(local_session).unwrap(); let removed = manager.remove_state_machine(sm_1_id); assert!(removed); @@ -176,9 +176,9 @@ 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); - let sm_2 = manager.create_session_state_machine(session2); - let sm_3 = manager.create_session_state_machine(session3); + 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(); assert_eq!(manager.session_count(), 3); @@ -198,7 +198,7 @@ mod tests { let sesion = mock_session_for_test(); - let sm = manager.create_session_state_machine(sesion); + let sm = manager.create_session_state_machine(sesion).unwrap(); assert_eq!(manager.session_count(), 1); let retrieved = manager.get_state_machine_id(sm); diff --git a/common/nym-lp/src/state_machine.rs b/common/nym-lp/src/state_machine.rs index 971ec5ed23..8d0fbbfa4f 100644 --- a/common/nym-lp/src/state_machine.rs +++ b/common/nym-lp/src/state_machine.rs @@ -2,14 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 //! Lewes Protocol State Machine for managing connection lifecycle. -//! -//! LP protocol flow (KKT → PSQ → Noise): -//! 1. KKTExchange: Client requests gateway's KEM public key (signed for MITM protection) -//! 2. Handshaking: Noise XKpsk3 with PSQ-derived PSK embedded in handshake messages -//! - PSQ ciphertext piggybacked on ClientHello (no extra round-trip) -//! - PSK = Blake3(ECDH || PSQ_secret || salt) provides hybrid classical+PQ security -//! 3. Transport: ChaCha20-Poly1305 authenticated encryption with derived keys -//! //! State machine ensures protocol steps execute in correct order. Invalid transitions //! return LpError, preventing protocol violations. @@ -190,11 +182,6 @@ impl LpStateMachine { } } - pub fn id(&self) -> Result { - todo!() - // Ok(self.session()?.id()) - } - pub fn session_identifier(&self) -> Result { Ok(*self.session()?.session_identifier()) } @@ -217,7 +204,7 @@ impl LpStateMachine { match input { LpInput::ReceivePacket(packet) => { // Check if packet lp_id matches our session - if packet.outer_header().receiver_idx != session.id() { + if packet.outer_header().receiver_idx != session.receiver_index() { let result_action = Some(Err(LpError::UnknownSessionId( packet.outer_header().receiver_idx, ))); @@ -231,17 +218,17 @@ impl LpStateMachine { return (LpState::Transport(state), Some(Err(e))); } - // 2. Mark counter as received - if let Err(e) = session.receiving_counter_mark(ctr) { - return (LpState::Transport(state), Some(Err(e))); - } - - // 3. decrypt the packet and attempt to deliver data + // 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))); + } + // Check message type match packet.into_message() { // Normal encrypted data @@ -373,77 +360,79 @@ mod tests { let resp_session = responder_sm.session().unwrap(); // Check both state machines use the same receiver_index - assert_eq!(init_session.id(), resp_session.id()); + assert_eq!(init_session.receiver_index(), resp_session.receiver_index()); } } #[test] fn test_state_machine_simplified_flow() { - todo!() - // for kem in KEM::iter() { - // let receiver_index: u32 = 123; - // let mock_sessions = SessionsMock::mock_post_handshake(kem); - // - // // 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 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 = 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.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 = 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.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/common/registration/src/lib.rs b/common/registration/src/lib.rs index baab036160..5ad090488e 100644 --- a/common/registration/src/lib.rs +++ b/common/registration/src/lib.rs @@ -5,7 +5,7 @@ use nym_authenticator_requests::AuthenticatorVersion; use nym_crypto::asymmetric::x25519::serde_helpers::bs58_x25519_pubkey; use nym_crypto::asymmetric::{ed25519, x25519}; use nym_ip_packet_requests::IpPair; -use nym_kkt_ciphersuite::{Ciphersuite, KEM, KEMKeyDigests, SignatureScheme}; +use nym_kkt_ciphersuite::{Ciphersuite, KEM, KEMKeyDigests}; use nym_sphinx::addressing::Recipient; use serde::{Deserialize, Serialize}; use std::collections::HashMap; diff --git a/gateway/src/node/lp_listener/handler.rs b/gateway/src/node/lp_listener/handler.rs index 4d9d7046b2..6d31a6ac41 100644 --- a/gateway/src/node/lp_listener/handler.rs +++ b/gateway/src/node/lp_listener/handler.rs @@ -162,7 +162,7 @@ where } Ok(Ok(session)) => session, }; - let receiver_idx = session.id(); + let receiver_idx = session.receiver_index(); // 2. insert the state machine into the shared state let state_machine = LpStateMachine::new(session); @@ -742,7 +742,7 @@ mod tests { } fn add_dummy_lp_state(handler: &mut LpConnectionHandler, session: LpSession) { - let id = session.id(); + let id = session.receiver_index(); let state_machine = LpStateMachine::new(session); handler.bound_receiver_idx = Some(id); diff --git a/integration-tests/src/lp_registration.rs b/integration-tests/src/lp_registration.rs index 08d2e52719..2d11be5783 100644 --- a/integration-tests/src/lp_registration.rs +++ b/integration-tests/src/lp_registration.rs @@ -398,112 +398,115 @@ mod tests { #[cfg(test)] mod using_lp_registration_client { use super::*; + use nym_kkt_ciphersuite::{IntoEnumIterator, KEM}; use nym_registration_client::NestedLpSession; use nym_test_utils::helpers::u64_seeded_rng_09; use nym_wireguard::DefguardPeer; #[tokio::test] async fn test_basic_lp_entry_registration() -> anyhow::Result<()> { - let TODO = "test with different kems"; - let ciphersuite = Ciphersuite::default(); + for kem in KEM::iter() { + let ciphersuite = Ciphersuite::default().with_kem(kem); - // nym_test_utils::helpers::setup_test_logger(); - // initialise random, but deterministic, keys, addresses, etc. for the parties - let mut client_rng = u64_seeded_rng_09(0); - let mut gateway_rng = u64_seeded_rng_09(1); + // nym_test_utils::helpers::setup_test_logger(); + // initialise random, but deterministic, keys, addresses, etc. for the parties + let mut client_rng = u64_seeded_rng_09(0); + let mut gateway_rng = u64_seeded_rng_09(1); - let client_data = Client::mock(&mut client_rng); - let client_key = *client_data.base.x25519_wg_keys.public_key(); - let mut entry = Gateway::mock(&mut gateway_rng).await?; + let client_data = Client::mock(&mut client_rng); + let client_key = *client_data.base.x25519_wg_keys.public_key(); + let mut entry = Gateway::mock(&mut gateway_rng).await?; - let mut client = LpRegistrationClient::::new_with_default_config( - client_data.base.peer.x25519().clone(), - entry.base.peer.as_remote(), - entry.base.socket_addr, - ciphersuite, - entry.base.lp_version, - ); + let mut client = LpRegistrationClient::::new_with_default_config( + client_data.base.peer.x25519().clone(), + entry.base.peer.as_remote(), + entry.base.socket_addr, + ciphersuite, + entry.base.lp_version, + ); - // 1. establish mock connection between client and gateway and retrieve gateway's handle - client.ensure_connected().await?; - let gateway_conn = client - .connection() - .as_ref() - .context("mock connection has failed!")? - .try_get_remote_handle(); + // 1. establish mock connection between client and gateway and retrieve gateway's handle + client.ensure_connected().await?; + let gateway_conn = client + .connection() + .as_ref() + .context("mock connection has failed!")? + .try_get_remote_handle(); - // 2. create and spawn gateway handler for the client connection - entry.create_lp_handler(gateway_conn, client_data.base.socket_addr); - entry.spawn_lp_handler(); + // 2. create and spawn gateway handler for the client connection + entry.create_lp_handler(gateway_conn, client_data.base.socket_addr); + entry.spawn_lp_handler(); - // 3. register all needed responses for the dvpn registration that will reach the peer controller - // 1) peer registration - ip pair allocation - let ip_pair = entry.pre_allocate_ip_pair(); - let reg_res = Ok::<_, nym_wireguard::Error>(ip_pair); + // 3. register all needed responses for the dvpn registration that will reach the peer controller + // 1) peer registration - ip pair allocation + let ip_pair = entry.pre_allocate_ip_pair(); + let reg_res = Ok::<_, nym_wireguard::Error>(ip_pair); - entry - .register_peer_controller_response( - PeerControlRequestType::AllocatePeerIpPair {}, - reg_res, - ) - .await; + entry + .register_peer_controller_response( + PeerControlRequestType::AllocatePeerIpPair {}, + reg_res, + ) + .await; - // 2) new peer inclusion - in non-mock system it would spawn handlers, - // here we'll just set a flag and say it's all fine - let public_key = client_key.to_wg_key(); - let add_res = Ok::<_, nym_wireguard::Error>(()); - entry - .register_peer_controller_response( - PeerControlRequestType::AddPeer { public_key }, - add_res, - ) - .await; + // 2) new peer inclusion - in non-mock system it would spawn handlers, + // here we'll just set a flag and say it's all fine + let public_key = client_key.to_wg_key(); + let add_res = Ok::<_, nym_wireguard::Error>(()); + entry + .register_peer_controller_response( + PeerControlRequestType::AddPeer { public_key }, + add_res, + ) + .await; - // 3) peer query - check for prior registrations - let query_res = Ok::<_, nym_wireguard::Error>(Option::::None); - let key = client_key.to_wg_key(); - entry - .register_peer_controller_response( - PeerControlRequestType::QueryPeer { key }, - query_res, - ) - .await; + // 3) peer query - check for prior registrations + let query_res = Ok::<_, nym_wireguard::Error>(Option::::None); + let key = client_key.to_wg_key(); + entry + .register_peer_controller_response( + PeerControlRequestType::QueryPeer { key }, + query_res, + ) + .await; - // 4. spawn peer controller to be able to handle dvpn registration requests - entry.spawn_peer_controller(); + // 4. spawn peer controller to be able to handle dvpn registration requests + entry.spawn_peer_controller(); - // 5. perform client handshake - client.perform_handshake().timeboxed().await??; + // 5. perform client handshake + client.perform_handshake().timeboxed().await??; - // 6. perform registration with entry only - let wg_keypair = client_data.base.x25519_wg_keys; - let gateway_identity = entry.base.identity.public_key(); - let registration_result = client - .register_dvpn( - &mut client_rng, - &wg_keypair, - gateway_identity, - &client_data.ticket_provider, - TicketType::V1WireguardEntry, - ) - .timeboxed() - .await??; + // 6. perform registration with entry only + let wg_keypair = client_data.base.x25519_wg_keys; + let gateway_identity = entry.base.identity.public_key(); + let registration_result = client + .register_dvpn( + &mut client_rng, + &wg_keypair, + gateway_identity, + &client_data.ticket_provider, + TicketType::V1WireguardEntry, + ) + .timeboxed() + .await??; - // 7. verify registration result - let peers_guard = entry.mock_peer_controller_state.peers.read().await; - let peer = peers_guard.get_by_x25519_key(&client_key).unwrap().clone(); - drop(peers_guard); - assert!(peer.add_success); + // 7. verify registration result + let peers_guard = entry.mock_peer_controller_state.peers.read().await; + let peer = peers_guard.get_by_x25519_key(&client_key).unwrap().clone(); + drop(peers_guard); + assert!(peer.add_success); - assert_eq!(registration_result.private_ipv4, ip_pair.ipv4); - assert_eq!(registration_result.private_ipv6, ip_pair.ipv6); - assert_eq!( - registration_result.public_key, - *entry.base.x25519_wg_keys.public_key() - ); + assert_eq!(registration_result.private_ipv4, ip_pair.ipv4); + assert_eq!(registration_result.private_ipv6, ip_pair.ipv6); + assert_eq!( + registration_result.public_key, + *entry.base.x25519_wg_keys.public_key() + ); + + // 8. stop the gateway task and finish the test + entry.stop_tasks().await?; + } - // 8. stop the gateway task and finish the test - entry.stop_tasks().await?; Ok(()) } @@ -570,7 +573,7 @@ mod tests { #[tokio::test] async fn test_basic_lp_exit_registration() -> anyhow::Result<()> { - nym_test_utils::helpers::setup_test_logger(); + // nym_test_utils::helpers::setup_test_logger(); let TODO = "test with different kems"; let ciphersuite = Ciphersuite::default(); @@ -705,7 +708,6 @@ mod tests { let mut nested_session = NestedLpSession::new( exit.base.socket_addr, client_data.base.peer.x25519().clone(), - *exit.base.identity.public_key(), exit.base.peer.as_remote(), ciphersuite, exit.base.lp_version, diff --git a/nym-registration-client/src/clients/lp.rs b/nym-registration-client/src/clients/lp.rs index 1411ebfe45..3b6131db76 100644 --- a/nym-registration-client/src/clients/lp.rs +++ b/nym-registration-client/src/clients/lp.rs @@ -60,8 +60,6 @@ impl LpBasedRegistrationClient { let entry_address = entry_lp_data.address; let exit_address = exit_lp_data.address; - let exit_identity = self.config.exit.node.identity; - tracing::debug!("Entry gateway LP address: {entry_address}"); tracing::debug!("Exit gateway LP address: {exit_address}"); @@ -103,7 +101,6 @@ impl LpBasedRegistrationClient { let mut nested_session = NestedLpSession::new( exit_address, exit_lp_keypair, - exit_identity, exit_peer, exit_ciphersuite, exit_lp_protocol, diff --git a/nym-registration-client/src/lp_client/client.rs b/nym-registration-client/src/lp_client/client.rs index 5cfde7057c..3dee0bf25a 100644 --- a/nym-registration-client/src/lp_client/client.rs +++ b/nym-registration-client/src/lp_client/client.rs @@ -5,20 +5,16 @@ use super::config::LpRegistrationConfig; use super::error::{LpClientError, Result}; -use crate::lp_client::helpers::{ - LpDataDeliverExt, LpDataSendExt, convert_forward_data, try_convert_forward_response, -}; +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 bytes::BytesMut; -use futures::SinkExt; 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::peer::{DHKeyPair, LpLocalPeer, LpRemotePeer}; -use nym_lp::state_machine::{LpAction, LpData, LpInput, LpStateMachine}; -use nym_lp::{Ciphersuite, EncryptedLpPacket, ForwardPacketData, packet::version}; -use nym_lp::{LpPacket, LpSession}; +use nym_lp::state_machine::LpStateMachine; +use nym_lp::{Ciphersuite, EncryptedLpPacket, packet::version}; use nym_lp_transport::traits::LpTransportChannel; use nym_lp_transport::{LpHandshakeChannel, LpTransportError}; use nym_registration_common::dvpn::LpDvpnRegistrationResponseMessageContent; @@ -120,13 +116,8 @@ where } /// Attempt to use this `LpRegistrationClient` as transport for `NestedSession` - pub fn as_nested_connection( - &mut self, - exit_identity: ed25519::PublicKey, - exit_address: SocketAddr, - ) -> NestedConnection<'_, S> { + pub fn as_nested_connection(&mut self, exit_address: SocketAddr) -> NestedConnection<'_, S> { NestedConnection { - exit_identity, exit_address, outer_client: self, } @@ -723,7 +714,7 @@ where pub fn session_id(&self) -> Result { self.state_machine()? .session() - .map(|s| s.id()) + .map(|s| s.receiver_index()) .map_err(Into::into) } } diff --git a/nym-registration-client/src/lp_client/error.rs b/nym-registration-client/src/lp_client/error.rs index ca9a4bd330..e4ca5f63e4 100644 --- a/nym-registration-client/src/lp_client/error.rs +++ b/nym-registration-client/src/lp_client/error.rs @@ -6,8 +6,6 @@ use nym_lp::state_machine::{LpAction, LpDataKind}; use nym_lp::{LpError, MalformedLpPacketError}; use nym_lp_transport::LpTransportError; -use nym_registration_common::BincodeError; -use std::time::Duration; use thiserror::Error; /// Errors that can occur during LP client operations. @@ -74,49 +72,6 @@ pub enum LpClientError { #[error("{0}")] Other(String), - // /// Failed to send registration request - // #[error("Failed to send registration request: {0}")] - // SendRegistrationRequest(String), - // - // /// Failed to receive registration response - // #[error("Failed to receive registration response: {0}")] - // ReceiveRegistrationResponse(String), - // - // /// Registration was rejected by gateway - // #[error("Gateway rejected registration: {reason}")] - // RegistrationRejected { reason: String }, - // - // /// Failed to receive response within specified deadline - // #[error("Failed to receive response within the set timeout: {timeout:?}")] - // ResponseReceiveTimeout { timeout: Duration }, - // - // /// LP transport error - // #[error("LP transport error: {0}")] - // Transport(String), - // - // #[error(transparent)] - // LpTransportError(#[from] LpTransportError), - // - // /// Invalid LP address format - // #[error("Invalid LP address '{address}': {reason}")] - // InvalidAddress { address: String, reason: String }, - // - // /// Serialization/deserialization error - // #[error("Serialization error: {0}")] - // Serialization(#[from] BincodeError), - // - // /// Connection closed unexpectedly - // #[error("Connection closed unexpectedly")] - // ConnectionClosed, - // - // /// Timeout waiting for response - // #[error("Timeout waiting for {operation}")] - // Timeout { operation: String }, - - // - // /// Another uncategorized error - // #[error("{0}")] - // Other(String), } impl LpClientError { diff --git a/nym-registration-client/src/lp_client/helpers.rs b/nym-registration-client/src/lp_client/helpers.rs index c84d766034..4479389b4f 100644 --- a/nym-registration-client/src/lp_client/helpers.rs +++ b/nym-registration-client/src/lp_client/helpers.rs @@ -4,7 +4,6 @@ #![allow(dead_code)] use crate::LpClientError; -use nym_crypto::asymmetric::ed25519; use nym_lp::ForwardPacketData; use nym_lp::peer::LpRemotePeer; use nym_lp::state_machine::{LpAction, LpData, LpDataKind, LpInput}; 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 e0f9ecd2d2..b09e0372e8 100644 --- a/nym-registration-client/src/lp_client/nested_session/connection.rs +++ b/nym-registration-client/src/lp_client/nested_session/connection.rs @@ -4,7 +4,6 @@ use crate::lp_client::helpers::{convert_forward_data, try_convert_forward_response}; use crate::{LpClientError, LpRegistrationClient}; use bytes::{BufMut, BytesMut}; -use nym_crypto::asymmetric::ed25519; use nym_lp::state_machine::{LpAction, LpInput}; use nym_lp::{EncryptedLpPacket, ExpectedResponseSize, ForwardPacketData, KEM}; use nym_lp_transport::traits::{HandshakeMessage, LpTransportChannel}; @@ -14,9 +13,6 @@ use std::net::SocketAddr; /// Attempt to treat the inner client as a LP connection pub struct NestedConnection<'a, S> { - /// Remote Ed25519 public key - pub(crate) exit_identity: ed25519::PublicKey, - /// Exit gateway's LP address (e.g., "2.2.2.2:41264") pub(crate) exit_address: SocketAddr, 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 c04713e477..d62d5e8a60 100644 --- a/nym-registration-client/src/lp_client/nested_session/mod.rs +++ b/nym-registration-client/src/lp_client/nested_session/mod.rs @@ -28,7 +28,7 @@ use nym_crypto::asymmetric::{ed25519, x25519}; use nym_lp::packet::version; use nym_lp::peer::{DHKeyPair, LpLocalPeer, LpRemotePeer}; use nym_lp::state_machine::{LpData, LpStateMachine}; -use nym_lp::{Ciphersuite, EncryptedLpPacket, ForwardPacketData, LpPacket, LpSession}; +use nym_lp::{Ciphersuite, EncryptedLpPacket, LpSession}; use nym_lp_transport::LpHandshakeChannel; use nym_lp_transport::traits::LpTransportChannel; use nym_registration_common::dvpn::LpDvpnRegistrationResponseMessageContent; @@ -71,9 +71,6 @@ pub struct NestedLpSession { /// Encapsulates all the client keys needed for the Lewes Protocol. lp_local_peer: LpLocalPeer, - /// Identity of the exit gateway - gateway_identity: ed25519::PublicKey, - /// Encapsulates all the exit gateway keys needed for the Lewes Protocol. gateway_lp_peer: LpRemotePeer, @@ -97,7 +94,6 @@ impl NestedLpSession { pub fn new( exit_address: SocketAddr, client_keypair: Arc, - gateway_identity: ed25519::PublicKey, gateway_lp_peer: LpRemotePeer, ciphersuite: Ciphersuite, gateway_supported_lp_protocol_version: u8, @@ -117,19 +113,12 @@ impl NestedLpSession { Self { exit_address, lp_local_peer, - gateway_identity, gateway_lp_peer, gateway_supported_lp_protocol_version: lp_protocol, state_machine: None, } } - fn state_machine(&self) -> Result<&LpStateMachine> { - self.state_machine - .as_ref() - .ok_or(LpClientError::IncompleteHandshake) - } - fn state_machine_mut(&mut self) -> Result<&mut LpStateMachine> { self.state_machine .as_mut() @@ -178,8 +167,7 @@ impl NestedLpSession { self.exit_address ); - let mut nested_connection = - outer_client.as_nested_connection(self.gateway_identity, self.exit_address); + let mut nested_connection = outer_client.as_nested_connection(self.exit_address); let local_peer = self.lp_local_peer.clone(); let remote_peer = self.gateway_lp_peer.clone(); @@ -233,8 +221,7 @@ impl NestedLpSession { S: LpTransportChannel + LpHandshakeChannel + Unpin, { tracing::debug!("Acquiring bandwidth credential for registration"); - let mut nested_connection = - outer_client.as_nested_connection(self.gateway_identity, self.exit_address); + let mut nested_connection = outer_client.as_nested_connection(self.exit_address); // Step 1: Get bandwidth credential from controller let credential_spending = bandwidth_controller @@ -341,8 +328,7 @@ impl NestedLpSession { S: LpTransportChannel + LpHandshakeChannel + Unpin, R: RngCore + CryptoRng, { - let mut nested_connection = - outer_client.as_nested_connection(self.gateway_identity, self.exit_address); + let mut nested_connection = outer_client.as_nested_connection(self.exit_address); tracing::debug!("Building registration request for exit gateway"); @@ -441,8 +427,7 @@ impl NestedLpSession { /// - Forwarding through entry gateway fails /// - Response decryption/deserialization fails /// - Gateway rejects the registration - #[deprecated(note = "use explicit handshake and register methods")] - pub async fn handshake_and_register_dvpn( + pub(crate) async fn handshake_and_register_dvpn( &mut self, outer_client: &mut LpRegistrationClient, rng: &mut R, 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 d17773756a..5410a67a8d 100644 --- a/nym-registration-client/src/lp_client/state_machine_helpers.rs +++ b/nym-registration-client/src/lp_client/state_machine_helpers.rs @@ -3,7 +3,7 @@ use crate::LpClientError; use nym_lp::state_machine::{LpAction, LpData, LpInput}; -use nym_lp::{EncryptedLpPacket, LpPacket, LpStateMachine}; +use nym_lp::{EncryptedLpPacket, LpStateMachine}; /// Attempt to prepare the provided data for sending by wrapping it in appropriate `LpAction`, /// and attempting to extract `EncryptedLpPacket` from the provided state machine.