diff --git a/src/node/handlers/discovery.rs b/src/node/handlers/discovery.rs index 138e803..74742de 100644 --- a/src/node/handlers/discovery.rs +++ b/src/node/handlers/discovery.rs @@ -289,11 +289,4 @@ impl Node { .retain(|_, entry| !entry.is_expired(current_time_ms)); } - /// Get current time in milliseconds since Unix epoch. - fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) - } } diff --git a/src/node/handlers/forwarding.rs b/src/node/handlers/forwarding.rs index 8ac13cd..466d884 100644 --- a/src/node/handlers/forwarding.rs +++ b/src/node/handlers/forwarding.rs @@ -40,14 +40,10 @@ impl Node { // Coordinate cache warming from plaintext session-layer headers self.try_warm_coord_cache(&datagram); - // Local delivery check + // Local delivery: dispatch to session layer handlers if datagram.dest_addr == *self.node_addr() { - debug!( - src = %datagram.src_addr, - payload_len = datagram.payload.len(), - "SessionDatagram delivered locally" - ); - // TODO: deliver to session layer + self.handle_session_payload(&datagram.src_addr, &datagram.payload) + .await; return; } diff --git a/src/node/handlers/mod.rs b/src/node/handlers/mod.rs index de6338f..87baa82 100644 --- a/src/node/handlers/mod.rs +++ b/src/node/handlers/mod.rs @@ -6,4 +6,5 @@ mod encrypted; mod forwarding; mod handshake; mod rx_loop; +mod session; mod timeout; diff --git a/src/node/handlers/session.rs b/src/node/handlers/session.rs new file mode 100644 index 0000000..75af5df --- /dev/null +++ b/src/node/handlers/session.rs @@ -0,0 +1,499 @@ +//! End-to-end session message handlers. +//! +//! Handles locally-delivered session payloads from SessionDatagram envelopes. +//! Dispatches based on session message type to specific handlers for +//! SessionSetup (Noise IK msg1), SessionAck (msg2), DataPacket, and +//! error signals (CoordsRequired, PathBroken). + +use crate::node::session::{EndToEndState, SessionEntry}; +use crate::node::{Node, NodeError}; +use crate::noise::{HandshakeState, HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE}; +use crate::protocol::{ + CoordsRequired, DataPacket, PathBroken, SessionAck, SessionDatagram, SessionMessageType, + SessionSetup, +}; +use crate::NodeAddr; +use secp256k1::PublicKey; +use tracing::debug; + +impl Node { + /// Handle a locally-delivered session datagram payload. + /// + /// Called from `handle_session_datagram()` when `dest_addr == self.node_addr()`. + /// Dispatches to the appropriate handler based on the session message type byte. + pub(in crate::node) async fn handle_session_payload( + &mut self, + src_addr: &NodeAddr, + payload: &[u8], + ) { + if payload.is_empty() { + debug!("Empty session payload"); + return; + } + + let msg_type = payload[0]; + let inner = &payload[1..]; + + match SessionMessageType::from_byte(msg_type) { + Some(SessionMessageType::SessionSetup) => { + self.handle_session_setup(src_addr, inner).await; + } + Some(SessionMessageType::SessionAck) => { + self.handle_session_ack(src_addr, inner).await; + } + Some(SessionMessageType::DataPacket) => { + self.handle_data_packet(src_addr, inner).await; + } + Some(SessionMessageType::CoordsRequired) => { + self.handle_coords_required(inner); + } + Some(SessionMessageType::PathBroken) => { + self.handle_path_broken(inner); + } + None => { + debug!(msg_type, "Unknown session message type"); + } + } + } + + /// Handle an incoming SessionSetup (Noise IK msg1). + /// + /// The remote node wants to establish an end-to-end session with us. + /// We create a responder handshake, process msg1, send SessionAck with msg2. + async fn handle_session_setup(&mut self, src_addr: &NodeAddr, inner: &[u8]) { + let setup = match SessionSetup::decode(inner) { + Ok(s) => s, + Err(e) => { + debug!(error = %e, "Malformed SessionSetup"); + return; + } + }; + + if setup.handshake_payload.len() != HANDSHAKE_MSG1_SIZE { + debug!( + len = setup.handshake_payload.len(), + expected = HANDSHAKE_MSG1_SIZE, + "Invalid handshake payload size in SessionSetup" + ); + return; + } + + // Check for existing session with this remote + if let Some(existing) = self.sessions.get(src_addr) { + match existing.state() { + EndToEndState::Initiating(_) => { + // Simultaneous initiation: smaller NodeAddr wins as initiator + if self.identity.node_addr() < src_addr { + // We win — drop their setup, they'll process ours + debug!( + src = %src_addr, + "Simultaneous session initiation: we win (smaller addr), dropping their setup" + ); + return; + } + // We lose — discard our pending handshake, become responder below + debug!( + src = %src_addr, + "Simultaneous session initiation: we lose, becoming responder" + ); + } + EndToEndState::Responding(_) => { + // Duplicate setup while we already responded — drop + debug!(src = %src_addr, "Duplicate SessionSetup, already responding"); + return; + } + EndToEndState::Established(_) => { + // Re-establishment: replace existing session below + debug!(src = %src_addr, "Session re-establishment from peer"); + } + } + } + + // Create responder handshake and process msg1 + let our_keypair = self.identity.keypair(); + let mut handshake = HandshakeState::new_responder(our_keypair); + + if let Err(e) = handshake.read_message_1(&setup.handshake_payload) { + debug!(error = %e, "Failed to process Noise IK msg1 in SessionSetup"); + return; + } + + // Extract the initiator's static public key (learned from msg1) + let remote_pubkey = match handshake.remote_static() { + Some(pk) => *pk, + None => { + debug!("No remote static key after processing msg1"); + return; + } + }; + + // Generate msg2 + let msg2 = match handshake.write_message_2() { + Ok(m) => m, + Err(e) => { + debug!(error = %e, "Failed to generate Noise IK msg2 for SessionAck"); + return; + } + }; + + // Build and send SessionAck + let our_coords = self.tree_state.my_coords().clone(); + let ack = SessionAck::new(our_coords).with_handshake(msg2); + let my_addr = *self.node_addr(); + let datagram = SessionDatagram::new(my_addr, *src_addr, ack.encode()); + + // Route the ack back to the initiator + if let Err(e) = self.send_session_datagram(&datagram).await { + debug!(error = %e, dest = %src_addr, "Failed to send SessionAck"); + return; + } + + // Store session entry in Responding state + let now_ms = Self::now_ms(); + let entry = SessionEntry::new(*src_addr, remote_pubkey, EndToEndState::Responding(handshake), now_ms); + self.sessions.insert(*src_addr, entry); + + debug!(src = %src_addr, "SessionSetup processed, SessionAck sent"); + } + + /// Handle an incoming SessionAck (Noise IK msg2). + /// + /// Completes our initiated handshake, transitions to Established. + async fn handle_session_ack(&mut self, src_addr: &NodeAddr, inner: &[u8]) { + let ack = match SessionAck::decode(inner) { + Ok(a) => a, + Err(e) => { + debug!(error = %e, "Malformed SessionAck"); + return; + } + }; + + if ack.handshake_payload.len() != HANDSHAKE_MSG2_SIZE { + debug!( + len = ack.handshake_payload.len(), + expected = HANDSHAKE_MSG2_SIZE, + "Invalid handshake payload size in SessionAck" + ); + return; + } + + // Remove the entry to take ownership of the handshake state + let mut entry = match self.sessions.remove(src_addr) { + Some(e) => e, + None => { + debug!(src = %src_addr, "SessionAck for unknown session"); + return; + } + }; + + // Must be in Initiating state + let handshake = match entry.take_state() { + Some(EndToEndState::Initiating(hs)) => hs, + _ => { + debug!(src = %src_addr, "SessionAck but session not in Initiating state"); + // Put it back + self.sessions.insert(*src_addr, entry); + return; + } + }; + + // Complete the handshake + let session = match Self::complete_initiator_handshake(handshake, &ack.handshake_payload) { + Ok(s) => s, + Err(e) => { + debug!(error = %e, "Failed to complete session handshake"); + return; // Entry was already removed, don't put back a broken session + } + }; + + entry.set_state(EndToEndState::Established(session)); + entry.touch(Self::now_ms()); + self.sessions.insert(*src_addr, entry); + + // Cache the responder's coordinates + let now_ms = Self::now_ms(); + self.coord_cache.insert(*src_addr, ack.src_coords, now_ms); + + debug!(src = %src_addr, "Session established (initiator)"); + } + + /// Complete an initiator-side Noise IK handshake given msg2. + fn complete_initiator_handshake( + mut handshake: HandshakeState, + msg2: &[u8], + ) -> Result { + handshake + .read_message_2(msg2) + .map_err(|e| format!("read_message_2 failed: {}", e))?; + handshake + .into_session() + .map_err(|e| format!("into_session failed: {}", e)) + } + + /// Handle an incoming DataPacket. + /// + /// Decrypts the payload using the established session key and delivers + /// to the TUN interface. + async fn handle_data_packet(&mut self, src_addr: &NodeAddr, inner: &[u8]) { + let packet = match DataPacket::decode(inner) { + Ok(p) => p, + Err(e) => { + debug!(error = %e, "Malformed DataPacket"); + return; + } + }; + + // Remove entry to take ownership for potential state transition + let mut entry = match self.sessions.remove(src_addr) { + Some(e) => e, + None => { + debug!(src = %src_addr, "DataPacket for unknown session"); + return; + } + }; + + // If in Responding state, transition to Established first + // (responder wrote msg2, handshake is complete from our side) + if entry.state().is_responding() { + let old_state = entry.take_state(); + let handshake = match old_state { + Some(EndToEndState::Responding(hs)) => hs, + _ => { + debug!(src = %src_addr, "Unexpected state in DataPacket handler"); + return; + } + }; + let noise_session = match handshake.into_session() { + Ok(s) => s, + Err(e) => { + debug!(error = %e, "Failed to create session from responding handshake"); + return; + } + }; + entry.set_state(EndToEndState::Established(noise_session)); + debug!(src = %src_addr, "Session established (responder, on first data)"); + } + + // Decrypt + let session = match entry.state_mut() { + EndToEndState::Established(s) => s, + _ => { + debug!(src = %src_addr, "DataPacket but session not established"); + self.sessions.insert(*src_addr, entry); + return; + } + }; + + let plaintext = match session.decrypt(&packet.payload) { + Ok(pt) => pt, + Err(e) => { + debug!(error = %e, src = %src_addr, "Session decryption failed"); + self.sessions.insert(*src_addr, entry); + return; + } + }; + + entry.touch(Self::now_ms()); + self.sessions.insert(*src_addr, entry); + + // Deliver to TUN + if let Some(tun_tx) = &self.tun_tx { + if let Err(e) = tun_tx.send(plaintext) { + debug!(error = %e, "Failed to deliver decrypted packet to TUN"); + } + } else { + debug!( + src = %src_addr, + "DataPacket decrypted (no TUN interface, plaintext dropped)" + ); + } + } + + /// Handle a CoordsRequired error signal from a transit router. + /// + /// The router couldn't route our packet because it lacks cached + /// coordinates for the destination. Future packets should include + /// coordinates (set COORDS_PRESENT flag). + fn handle_coords_required(&mut self, inner: &[u8]) { + let msg = match CoordsRequired::decode(inner) { + Ok(m) => m, + Err(e) => { + debug!(error = %e, "Malformed CoordsRequired"); + return; + } + }; + + debug!( + dest = %msg.dest_addr, + reporter = %msg.reporter, + "CoordsRequired: transit router needs coordinates" + ); + } + + /// Handle a PathBroken error signal from a transit router. + /// + /// The router has coordinates but still can't route to the destination. + /// Invalidate cached coordinates and consider re-discovery. + fn handle_path_broken(&mut self, inner: &[u8]) { + let msg = match PathBroken::decode(inner) { + Ok(m) => m, + Err(e) => { + debug!(error = %e, "Malformed PathBroken"); + return; + } + }; + + debug!( + dest = %msg.dest_addr, + reporter = %msg.reporter, + "PathBroken: transit router reports routing failure" + ); + + // Invalidate stale cached coordinates + self.coord_cache.remove(&msg.dest_addr); + } + + // === Session Initiation (Send Path) === + + /// Initiate an end-to-end session with a remote node. + /// + /// Creates a Noise IK handshake as initiator, wraps msg1 in a + /// SessionSetup, encapsulates in a SessionDatagram, and routes + /// toward the destination. + pub(in crate::node) async fn initiate_session( + &mut self, + dest_addr: NodeAddr, + dest_pubkey: PublicKey, + ) -> Result<(), NodeError> { + // Check for existing session + if let Some(existing) = self.sessions.get(&dest_addr) { + if existing.state().is_established() || existing.state().is_initiating() { + return Ok(()); + } + } + + // Create Noise IK initiator handshake + let our_keypair = self.identity.keypair(); + let mut handshake = HandshakeState::new_initiator(our_keypair, dest_pubkey); + let msg1 = handshake.write_message_1().map_err(|e| NodeError::SendFailed { + node_addr: dest_addr, + reason: format!("Noise msg1 generation failed: {}", e), + })?; + + // Build SessionSetup with coordinates + let our_coords = self.tree_state.my_coords().clone(); + let dest_coords = self.get_dest_coords(&dest_addr); + let setup = SessionSetup::new(our_coords, dest_coords) + .with_handshake(msg1); + + // Wrap in SessionDatagram + let my_addr = *self.node_addr(); + let datagram = SessionDatagram::new(my_addr, dest_addr, setup.encode()); + + // Route toward destination + self.send_session_datagram(&datagram).await?; + + // Store session entry + let now_ms = Self::now_ms(); + let entry = SessionEntry::new(dest_addr, dest_pubkey, EndToEndState::Initiating(handshake), now_ms); + self.sessions.insert(dest_addr, entry); + + debug!(dest = %dest_addr, "Session initiation started"); + Ok(()) + } + + /// Send application data over an established session. + /// + /// Encrypts the payload with the session key, wraps in DataPacket + /// and SessionDatagram, routes toward destination. + pub(in crate::node) async fn send_session_data( + &mut self, + dest_addr: &NodeAddr, + plaintext: &[u8], + ) -> Result<(), NodeError> { + let entry = self.sessions.get_mut(dest_addr).ok_or_else(|| NodeError::SendFailed { + node_addr: *dest_addr, + reason: "no session".into(), + })?; + + let session = match entry.state_mut() { + EndToEndState::Established(s) => s, + _ => { + return Err(NodeError::SendFailed { + node_addr: *dest_addr, + reason: "session not established".into(), + }); + } + }; + + // Encrypt with session key + let ciphertext = session.encrypt(plaintext).map_err(|e| NodeError::SendFailed { + node_addr: *dest_addr, + reason: format!("session encrypt failed: {}", e), + })?; + + // Build DataPacket and wrap in SessionDatagram + let data_packet = DataPacket::new(ciphertext); + let my_addr = *self.node_addr(); + let datagram = SessionDatagram::new(my_addr, *dest_addr, data_packet.encode()); + + self.send_session_datagram(&datagram).await?; + + // Re-borrow after send (which borrowed &mut self) + if let Some(entry) = self.sessions.get_mut(dest_addr) { + entry.touch(Self::now_ms()); + } + + Ok(()) + } + + /// Route and send a SessionDatagram through the mesh. + /// + /// Finds the next hop for the destination and sends the datagram + /// as an encrypted link message. + async fn send_session_datagram( + &mut self, + datagram: &SessionDatagram, + ) -> Result<(), NodeError> { + let next_hop_addr = match self.find_next_hop(&datagram.dest_addr) { + Some(peer) => *peer.node_addr(), + None => { + return Err(NodeError::SendFailed { + node_addr: datagram.dest_addr, + reason: "no route to destination".into(), + }); + } + }; + + let encoded = datagram.encode(); + self.send_encrypted_link_message(&next_hop_addr, &encoded).await + } + + /// Look up destination coordinates from available caches. + /// + /// Returns our own coordinates as a fallback (the SessionSetup will + /// carry src_coords for return path routing; empty dest_coords + /// would fail wire encoding since TreeCoordinate requires ≥1 entry). + fn get_dest_coords(&self, dest: &NodeAddr) -> crate::tree::TreeCoordinate { + let now_ms = Self::now_ms(); + if let Some(coords) = self.coord_cache.get(dest, now_ms) { + return coords.clone(); + } + if let Some(cached) = self.route_cache.get(dest) { + return cached.coords().clone(); + } + // Fallback: use our own coordinates. The SessionSetup dest_coords + // field cannot be empty (wire format requires ≥1 entry). Using our + // own coords is safe — transit routers will still cache them, and + // the destination will return its actual coords in the SessionAck. + self.tree_state.my_coords().clone() + } + + /// Current Unix time in milliseconds. + pub(in crate::node) fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) + } +} diff --git a/src/node/mod.rs b/src/node/mod.rs index ab8db98..83f1034 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -8,6 +8,7 @@ mod bloom; mod handlers; mod lifecycle; mod retry; +pub(crate) mod session; mod tree; #[cfg(test)] mod tests; @@ -15,6 +16,7 @@ mod tests; use crate::bloom::BloomState; use crate::cache::{CoordCache, RouteCache}; use crate::index::IndexAllocator; +use crate::node::session::SessionEntry; use crate::peer::{ActivePeer, PeerConnection}; use crate::rate_limit::HandshakeRateLimiter; use crate::transport::{ @@ -241,6 +243,11 @@ pub struct Node { /// Indexed by NodeAddr (verified identity). peers: HashMap, + // === End-to-End Sessions === + /// Session table for end-to-end encrypted sessions. + /// Keyed by remote NodeAddr. + sessions: HashMap, + // === Resource Limits === /// Maximum connections (0 = unlimited). max_connections: usize, @@ -335,6 +342,7 @@ impl Node { packet_rx: None, connections: HashMap::new(), peers: HashMap::new(), + sessions: HashMap::new(), max_connections: 256, max_peers: 128, max_links: 256, @@ -386,6 +394,7 @@ impl Node { packet_rx: None, connections: HashMap::new(), peers: HashMap::new(), + sessions: HashMap::new(), max_connections: 256, max_peers: 128, max_links: 256, @@ -754,6 +763,28 @@ impl Node { self.peers.values().filter(|p| p.can_send()).count() } + // === End-to-End Sessions === + + /// Get a session by remote NodeAddr. + pub(crate) fn get_session(&self, remote: &NodeAddr) -> Option<&SessionEntry> { + self.sessions.get(remote) + } + + /// Get a mutable session by remote NodeAddr. + pub(crate) fn get_session_mut(&mut self, remote: &NodeAddr) -> Option<&mut SessionEntry> { + self.sessions.get_mut(remote) + } + + /// Remove a session. + pub(crate) fn remove_session(&mut self, remote: &NodeAddr) -> Option { + self.sessions.remove(remote) + } + + /// Number of end-to-end sessions. + pub fn session_count(&self) -> usize { + self.sessions.len() + } + // === Routing === /// Find next hop for a destination node address. diff --git a/src/node/session.rs b/src/node/session.rs new file mode 100644 index 0000000..bf93638 --- /dev/null +++ b/src/node/session.rs @@ -0,0 +1,125 @@ +//! End-to-end session state. +//! +//! Tracks Noise IK sessions between this node and remote endpoints. +//! Sessions are established via SessionSetup/SessionAck handshake +//! messages carried inside SessionDatagram envelopes through the mesh. + +use crate::noise::{HandshakeState, NoiseSession}; +use crate::NodeAddr; +use secp256k1::PublicKey; + +/// State machine for an end-to-end session. +pub(crate) enum EndToEndState { + /// We initiated: sent SessionSetup with Noise IK msg1, awaiting SessionAck. + Initiating(HandshakeState), + /// We are responding: received msg1, sent SessionAck with msg2. + Responding(HandshakeState), + /// Handshake complete, NoiseSession available for encrypt/decrypt. + Established(NoiseSession), +} + +impl EndToEndState { + /// Check if the session is established and ready for data. + pub(crate) fn is_established(&self) -> bool { + matches!(self, EndToEndState::Established(_)) + } + + /// Check if we are the initiator (waiting for ack). + pub(crate) fn is_initiating(&self) -> bool { + matches!(self, EndToEndState::Initiating(_)) + } + + /// Check if we are the responder (sent ack, waiting for data). + pub(crate) fn is_responding(&self) -> bool { + matches!(self, EndToEndState::Responding(_)) + } +} + +/// A single end-to-end session with a remote node. +/// +/// The state is wrapped in `Option` to allow taking ownership of the +/// handshake state during transitions without placeholder values. +/// The state is `None` only transiently during handler processing. +pub(crate) struct SessionEntry { + /// Remote node's address (session table key). + remote_addr: NodeAddr, + /// Remote node's static public key (for Noise IK). + remote_pubkey: PublicKey, + /// Current session state. `None` only during state transitions. + state: Option, + /// When the session was created (Unix milliseconds). + created_at: u64, + /// Last activity timestamp (Unix milliseconds). + last_activity: u64, +} + +impl SessionEntry { + /// Create a new session entry. + pub(crate) fn new( + remote_addr: NodeAddr, + remote_pubkey: PublicKey, + state: EndToEndState, + now_ms: u64, + ) -> Self { + Self { + remote_addr, + remote_pubkey, + state: Some(state), + created_at: now_ms, + last_activity: now_ms, + } + } + + /// Get the remote node address. + pub(crate) fn remote_addr(&self) -> &NodeAddr { + &self.remote_addr + } + + /// Get the remote node's public key. + pub(crate) fn remote_pubkey(&self) -> &PublicKey { + &self.remote_pubkey + } + + /// Get the current session state. + pub(crate) fn state(&self) -> &EndToEndState { + self.state.as_ref().expect("session state taken but not restored") + } + + /// Get mutable access to the session state. + pub(crate) fn state_mut(&mut self) -> &mut EndToEndState { + self.state.as_mut().expect("session state taken but not restored") + } + + /// Replace the session state. + pub(crate) fn set_state(&mut self, state: EndToEndState) { + self.state = Some(state); + } + + /// Take the state out, leaving `None`. + /// + /// The caller must call `set_state()` to restore a valid state, + /// or discard the entry entirely. + pub(crate) fn take_state(&mut self) -> Option { + self.state.take() + } + + /// Update the last activity timestamp. + pub(crate) fn touch(&mut self, now_ms: u64) { + self.last_activity = now_ms; + } + + /// Check if the session is established. + pub(crate) fn is_established(&self) -> bool { + self.state.as_ref().map_or(false, |s| s.is_established()) + } + + /// Get creation time. + pub(crate) fn created_at(&self) -> u64 { + self.created_at + } + + /// Get last activity time. + pub(crate) fn last_activity(&self) -> u64 { + self.last_activity + } +} diff --git a/src/node/tests/mod.rs b/src/node/tests/mod.rs index fd2cf48..1bbba46 100644 --- a/src/node/tests/mod.rs +++ b/src/node/tests/mod.rs @@ -10,6 +10,7 @@ mod discovery; mod forwarding; mod handshake; mod routing; +mod session; mod spanning_tree; mod unit; diff --git a/src/node/tests/session.rs b/src/node/tests/session.rs new file mode 100644 index 0000000..990eb16 --- /dev/null +++ b/src/node/tests/session.rs @@ -0,0 +1,872 @@ +//! End-to-end session establishment tests. + +use super::*; +use crate::node::session::EndToEndState; +use crate::node::tests::spanning_tree::{ + cleanup_nodes, generate_random_edges, process_available_packets, run_tree_test, + verify_tree_convergence, TestNode, +}; +use crate::protocol::{SessionAck, SessionDatagram}; + +/// Populate all nodes' coordinate caches with each other's coords. +/// +/// This enables routing between non-adjacent nodes (bloom filter + tree +/// routing both require cached destination coordinates). +fn populate_all_coord_caches(nodes: &mut [TestNode]) { + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + + let all_coords: Vec<(NodeAddr, crate::tree::TreeCoordinate)> = nodes + .iter() + .map(|tn| { + ( + *tn.node.node_addr(), + tn.node.tree_state().my_coords().clone(), + ) + }) + .collect(); + + for tn in nodes.iter_mut() { + for (addr, coords) in &all_coords { + if addr != tn.node.node_addr() { + tn.node + .coord_cache_mut() + .insert(*addr, coords.clone(), now_ms); + } + } + } +} + +// ============================================================================ +// Unit tests: SessionEntry data structure +// ============================================================================ + +#[test] +fn test_session_entry_new_initiating() { + use crate::noise::HandshakeState; + + let identity_a = Identity::generate(); + let identity_b = Identity::generate(); + + let handshake = HandshakeState::new_initiator( + identity_a.keypair(), + identity_b.pubkey_full(), + ); + + let entry = crate::node::session::SessionEntry::new( + *identity_b.node_addr(), + identity_b.pubkey_full(), + EndToEndState::Initiating(handshake), + 1000, + ); + + assert!(entry.state().is_initiating()); + assert!(!entry.state().is_established()); + assert!(!entry.state().is_responding()); + assert_eq!(entry.created_at(), 1000); + assert_eq!(entry.last_activity(), 1000); +} + +#[test] +fn test_session_entry_touch() { + use crate::noise::HandshakeState; + + let identity_a = Identity::generate(); + let identity_b = Identity::generate(); + + let handshake = HandshakeState::new_initiator( + identity_a.keypair(), + identity_b.pubkey_full(), + ); + + let mut entry = crate::node::session::SessionEntry::new( + *identity_b.node_addr(), + identity_b.pubkey_full(), + EndToEndState::Initiating(handshake), + 1000, + ); + + entry.touch(2000); + assert_eq!(entry.last_activity(), 2000); + assert_eq!(entry.created_at(), 1000); +} + +#[test] +fn test_session_table_operations() { + use crate::noise::HandshakeState; + + let mut node = make_node(); + let identity_b = Identity::generate(); + + let handshake = HandshakeState::new_initiator( + node.identity().keypair(), + identity_b.pubkey_full(), + ); + + let dest_addr = *identity_b.node_addr(); + let entry = crate::node::session::SessionEntry::new( + dest_addr, + identity_b.pubkey_full(), + EndToEndState::Initiating(handshake), + 1000, + ); + + node.sessions.insert(dest_addr, entry); + assert_eq!(node.session_count(), 1); + assert!(node.get_session(&dest_addr).is_some()); + assert!(node.get_session(&make_node_addr(0xFF)).is_none()); + + let removed = node.remove_session(&dest_addr); + assert!(removed.is_some()); + assert_eq!(node.session_count(), 0); +} + +// ============================================================================ +// Integration tests: 2-node direct session establishment +// ============================================================================ + +#[tokio::test] +async fn test_session_direct_peer_handshake() { + // Two directly connected nodes: A initiates a session with B + let edges = vec![(0, 1)]; + let mut nodes = run_tree_test(2, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let node0_addr = *nodes[0].node.node_addr(); + let node1_addr = *nodes[1].node.node_addr(); + let node1_pubkey = nodes[1].node.identity().pubkey_full(); + + // Node 0 initiates session with Node 1 + nodes[0] + .node + .initiate_session(node1_addr, node1_pubkey) + .await + .expect("initiate_session failed"); + + // Node 0 should have a session in Initiating state + assert_eq!(nodes[0].node.session_count(), 1); + assert!(nodes[0] + .node + .get_session(&node1_addr) + .unwrap() + .state() + .is_initiating()); + + // Process packets: SessionSetup arrives at Node 1 + tokio::time::sleep(Duration::from_millis(20)).await; + let count = process_available_packets(&mut nodes).await; + assert!(count > 0, "Expected SessionSetup packet to arrive"); + + // Node 1 should now have a session in Responding state + assert_eq!(nodes[1].node.session_count(), 1); + assert!(nodes[1] + .node + .get_session(&node0_addr) + .unwrap() + .state() + .is_responding()); + + // Process packets: SessionAck arrives at Node 0 + tokio::time::sleep(Duration::from_millis(20)).await; + let count = process_available_packets(&mut nodes).await; + assert!(count > 0, "Expected SessionAck packet to arrive"); + + // Node 0 should now be Established + assert!(nodes[0] + .node + .get_session(&node1_addr) + .unwrap() + .state() + .is_established()); + + cleanup_nodes(&mut nodes).await; +} + +#[tokio::test] +async fn test_session_direct_peer_data_transfer() { + // Two nodes: establish session, then send data + let edges = vec![(0, 1)]; + let mut nodes = run_tree_test(2, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let node0_addr = *nodes[0].node.node_addr(); + let node1_addr = *nodes[1].node.node_addr(); + let node1_pubkey = nodes[1].node.identity().pubkey_full(); + + // Establish session + nodes[0] + .node + .initiate_session(node1_addr, node1_pubkey) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; // Setup → Node 1 + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; // Ack → Node 0 + + assert!(nodes[0] + .node + .get_session(&node1_addr) + .unwrap() + .state() + .is_established()); + + // Send data from Node 0 to Node 1 + let test_data = b"Hello, FIPS session!"; + nodes[0] + .node + .send_session_data(&node1_addr, test_data) + .await + .expect("send_session_data failed"); + + // Process packets: DataPacket arrives at Node 1 + tokio::time::sleep(Duration::from_millis(20)).await; + let count = process_available_packets(&mut nodes).await; + assert!(count > 0, "Expected DataPacket to arrive"); + + // Node 1's session should now be Established (was Responding, transitions on first data) + assert!(nodes[1] + .node + .get_session(&node0_addr) + .unwrap() + .state() + .is_established()); + + cleanup_nodes(&mut nodes).await; +} + +// ============================================================================ +// Integration tests: 3-node forwarded session +// ============================================================================ + +#[tokio::test] +async fn test_session_3node_forwarded_handshake() { + // A—B—C: Node A initiates session with Node C through transit node B + let edges = vec![(0, 1), (1, 2)]; + let mut nodes = run_tree_test(3, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let node0_addr = *nodes[0].node.node_addr(); + let node2_addr = *nodes[2].node.node_addr(); + let node2_pubkey = nodes[2].node.identity().pubkey_full(); + + // Node 0 initiates session with Node 2 + nodes[0] + .node + .initiate_session(node2_addr, node2_pubkey) + .await + .expect("initiate_session failed"); + + // Process: SessionSetup: 0→1 (forwarded by transit B) + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + + // Process: SessionSetup: 1→2 (arrives at destination C) + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + + // Node 2 should have a Responding session + assert!( + nodes[2].node.get_session(&node0_addr).is_some(), + "Node 2 should have a session entry for Node 0" + ); + assert!(nodes[2] + .node + .get_session(&node0_addr) + .unwrap() + .state() + .is_responding()); + + // Process: SessionAck: 2→1 (forwarded by transit B) + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + + // Process: SessionAck: 1→0 (arrives at initiator A) + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + + // Node 0 should now be Established + assert!(nodes[0] + .node + .get_session(&node2_addr) + .unwrap() + .state() + .is_established()); + + // Transit node B should NOT have a session + assert_eq!( + nodes[1].node.session_count(), + 0, + "Transit node should have no sessions" + ); + + cleanup_nodes(&mut nodes).await; +} + +#[tokio::test] +async fn test_session_3node_forwarded_data() { + // A—B—C: Establish session, send data end-to-end + let edges = vec![(0, 1), (1, 2)]; + let mut nodes = run_tree_test(3, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let node0_addr = *nodes[0].node.node_addr(); + let node2_addr = *nodes[2].node.node_addr(); + let node2_pubkey = nodes[2].node.identity().pubkey_full(); + + // Establish session (needs more hops) + nodes[0] + .node + .initiate_session(node2_addr, node2_pubkey) + .await + .unwrap(); + + // Drain packets until handshake completes (multi-hop needs several rounds) + for _ in 0..10 { + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + } + + assert!( + nodes[0] + .node + .get_session(&node2_addr) + .map(|s| s.state().is_established()) + .unwrap_or(false), + "Session should be established after handshake rounds" + ); + + // Send data + let test_data = b"End-to-end through transit node B"; + nodes[0] + .node + .send_session_data(&node2_addr, test_data) + .await + .expect("send_session_data failed"); + + // Drain data packet through transit node + for _ in 0..5 { + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + } + + // Node 2 should have transitioned to Established on first data + assert!(nodes[2] + .node + .get_session(&node0_addr) + .unwrap() + .state() + .is_established()); + + cleanup_nodes(&mut nodes).await; +} + +// ============================================================================ +// Edge cases +// ============================================================================ + +#[tokio::test] +async fn test_session_initiate_idempotent() { + // Calling initiate_session twice should be idempotent + let edges = vec![(0, 1)]; + let mut nodes = run_tree_test(2, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let node1_addr = *nodes[1].node.node_addr(); + let node1_pubkey = nodes[1].node.identity().pubkey_full(); + + // First call + nodes[0] + .node + .initiate_session(node1_addr, node1_pubkey) + .await + .unwrap(); + assert_eq!(nodes[0].node.session_count(), 1); + + // Second call should be a no-op + nodes[0] + .node + .initiate_session(node1_addr, node1_pubkey) + .await + .unwrap(); + assert_eq!(nodes[0].node.session_count(), 1); + + cleanup_nodes(&mut nodes).await; +} + +#[tokio::test] +async fn test_session_send_data_no_session_fails() { + let mut node = make_node(); + let fake_addr = make_node_addr(0xAA); + + let result = node.send_session_data(&fake_addr, b"test").await; + assert!(result.is_err(), "Should fail with no session"); +} + +#[tokio::test] +async fn test_session_ack_for_unknown_session() { + // Receiving a SessionAck when we have no Initiating session should be dropped + let edges = vec![(0, 1)]; + let mut nodes = run_tree_test(2, &edges, false).await; + verify_tree_convergence(&nodes); + + let node0_addr = *nodes[0].node.node_addr(); + let node1_addr = *nodes[1].node.node_addr(); + + // Fabricate a SessionAck and deliver directly + let coords = nodes[1].node.tree_state().my_coords().clone(); + let ack = SessionAck::new(coords).with_handshake(vec![0u8; 33]); + let datagram = SessionDatagram::new(node1_addr, node0_addr, ack.encode()); + + // Send through link layer + let encoded = datagram.encode(); + nodes[1] + .node + .send_encrypted_link_message(&node0_addr, &encoded) + .await + .unwrap(); + + tokio::time::sleep(Duration::from_millis(20)).await; + process_available_packets(&mut nodes).await; + + // Node 0 should have no sessions (ack was for unknown session) + assert_eq!(nodes[0].node.session_count(), 0); + + cleanup_nodes(&mut nodes).await; +} + +// ============================================================================ +// Large-scale test: 100-node session establishment + bidirectional data +// ============================================================================ + +/// Drain packets until quiescent (2 consecutive idle rounds). +async fn drain_to_quiescence(nodes: &mut [TestNode]) { + let mut idle_rounds = 0; + for _ in 0..40 { + tokio::time::sleep(Duration::from_millis(10)).await; + let count = process_available_packets(nodes).await; + if count == 0 { + idle_rounds += 1; + if idle_rounds >= 2 { + break; + } + } else { + idle_rounds = 0; + } + } +} + +#[tokio::test] +async fn test_session_100_nodes() { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + use std::sync::mpsc; + use std::time::Instant; + + // Same random topology as other 100-node tests + const NUM_NODES: usize = 100; + const TARGET_EDGES: usize = 250; + const SEED: u64 = 42; + + let start = Instant::now(); + + let edges = generate_random_edges(NUM_NODES, TARGET_EDGES, SEED); + let mut nodes = run_tree_test(NUM_NODES, &edges, false).await; + verify_tree_convergence(&nodes); + populate_all_coord_caches(&mut nodes); + + let setup_time = start.elapsed(); + + // Collect identities: (node_addr, pubkey) for all nodes + let all_info: Vec<(NodeAddr, secp256k1::PublicKey)> = nodes + .iter() + .map(|tn| { + ( + *tn.node.node_addr(), + tn.node.identity().pubkey_full(), + ) + }) + .collect(); + + // Each node picks one random target for its outbound session. + // Use deterministic RNG so failures are reproducible. + let mut rng = StdRng::seed_from_u64(SEED + 1); + let mut session_pairs: Vec<(usize, usize)> = Vec::with_capacity(NUM_NODES); + for src in 0..NUM_NODES { + let mut dst = rng.gen_range(0..NUM_NODES); + while dst == src { + dst = rng.gen_range(0..NUM_NODES); + } + session_pairs.push((src, dst)); + } + + // === Phase 1: Establish all sessions === + + let session_start = Instant::now(); + + for &(src, dst) in &session_pairs { + let (dest_addr, dest_pubkey) = all_info[dst]; + + nodes[src] + .node + .initiate_session(dest_addr, dest_pubkey) + .await + .expect("initiate_session failed"); + + drain_to_quiescence(&mut nodes).await; + } + + drain_to_quiescence(&mut nodes).await; + let session_time = session_start.elapsed(); + + // Verify all initiator sessions reached Established before data phase + let mut handshake_failures: Vec<(usize, usize)> = Vec::new(); + for &(src, dst) in &session_pairs { + let dest_addr = all_info[dst].0; + let ok = nodes[src] + .node + .get_session(&dest_addr) + .map(|e| e.state().is_established()) + .unwrap_or(false); + if !ok { + handshake_failures.push((src, dst)); + } + } + assert!( + handshake_failures.is_empty(), + "Handshake failed for {} pairs (first: {:?})", + handshake_failures.len(), + handshake_failures.first() + ); + + // === Phase 2: Inject TUN receivers and snapshot link stats === + + // Install a tun_tx on every node so delivered datagrams can be counted. + let mut tun_receivers: Vec>> = Vec::with_capacity(NUM_NODES); + for tn in nodes.iter_mut() { + let (tx, rx) = mpsc::channel(); + tn.node.tun_tx = Some(tx); + tun_receivers.push(rx); + } + + // Snapshot per-peer link stats before data phase + let link_pkts_sent_before: Vec> = nodes + .iter() + .map(|tn| { + tn.node + .peers() + .map(|p| (*p.node_addr(), p.link_stats().packets_sent)) + .collect() + }) + .collect(); + + // === Phase 3: Bidirectional data transfer === + // + // For each session pair: + // 1. Initiator sends one datagram to responder + // (this also transitions responder from Responding → Established) + // 2. Responder sends one datagram back to initiator + // + // Batched per pair with draining between each. + + let data_start = Instant::now(); + let mut send_forward_ok = 0usize; + let mut send_forward_err = 0usize; + let mut send_reverse_ok = 0usize; + let mut send_reverse_err = 0usize; + + for (pair_idx, &(src, dst)) in session_pairs.iter().enumerate() { + let dest_addr = all_info[dst].0; + let src_addr = all_info[src].0; + + // Forward: initiator → responder + let fwd_payload = format!("fwd-{}", pair_idx).into_bytes(); + match nodes[src] + .node + .send_session_data(&dest_addr, &fwd_payload) + .await + { + Ok(()) => send_forward_ok += 1, + Err(_) => send_forward_err += 1, + } + + drain_to_quiescence(&mut nodes).await; + + // Reverse: responder → initiator + // (Responder should now be Established after receiving the forward datagram) + let rev_payload = format!("rev-{}", pair_idx).into_bytes(); + match nodes[dst] + .node + .send_session_data(&src_addr, &rev_payload) + .await + { + Ok(()) => send_reverse_ok += 1, + Err(_) => send_reverse_err += 1, + } + + drain_to_quiescence(&mut nodes).await; + } + + let data_time = data_start.elapsed(); + + // === Phase 4: Collect delivered datagrams from TUN receivers === + + let mut delivered_per_node: Vec>> = Vec::with_capacity(NUM_NODES); + for rx in tun_receivers.iter_mut() { + let mut packets = Vec::new(); + while let Ok(pkt) = rx.try_recv() { + packets.push(pkt); + } + delivered_per_node.push(packets); + } + + let total_delivered: usize = delivered_per_node.iter().map(|v| v.len()).sum(); + + // Verify each pair's forward and reverse datagrams arrived + let mut fwd_delivered = 0usize; + let mut rev_delivered = 0usize; + let mut fwd_missing: Vec<(usize, usize)> = Vec::new(); + let mut rev_missing: Vec<(usize, usize)> = Vec::new(); + + for (pair_idx, &(src, dst)) in session_pairs.iter().enumerate() { + let fwd_payload = format!("fwd-{}", pair_idx).into_bytes(); + let rev_payload = format!("rev-{}", pair_idx).into_bytes(); + + if delivered_per_node[dst].iter().any(|p| *p == fwd_payload) { + fwd_delivered += 1; + } else if fwd_missing.len() < 20 { + fwd_missing.push((src, dst)); + } + + if delivered_per_node[src].iter().any(|p| *p == rev_payload) { + rev_delivered += 1; + } else if rev_missing.len() < 20 { + rev_missing.push((src, dst)); + } + } + + // === Phase 5: Final session state === + + let mut total_established = 0usize; + let mut total_responding = 0usize; + let mut total_initiating = 0usize; + let mut fully_established_nodes = 0usize; + + for tn in &nodes { + let mut all_est = true; + for (_, entry) in tn.node.sessions.iter() { + if entry.state().is_established() { + total_established += 1; + } else if entry.state().is_responding() { + total_responding += 1; + all_est = false; + } else { + total_initiating += 1; + all_est = false; + } + } + if tn.node.session_count() > 0 && all_est { + fully_established_nodes += 1; + } + } + + let session_counts: Vec = nodes + .iter() + .map(|tn| tn.node.session_count()) + .collect(); + let total_sessions: usize = session_counts.iter().sum(); + let min_sessions = *session_counts.iter().min().unwrap(); + let max_sessions = *session_counts.iter().max().unwrap(); + + // === Phase 6: Link and routing statistics === + + // Link stats delta: packets sent during data phase + let mut data_link_pkts_sent: u64 = 0; + let mut total_link_pkts_sent: u64 = 0; + let mut total_link_pkts_recv: u64 = 0; + let mut total_link_bytes_sent: u64 = 0; + let mut total_link_bytes_recv: u64 = 0; + + for (i, tn) in nodes.iter().enumerate() { + for peer in tn.node.peers() { + let stats = peer.link_stats(); + // Delta for this peer since before data phase + let before = link_pkts_sent_before[i] + .iter() + .find(|(addr, _)| addr == peer.node_addr()) + .map(|(_, pkts)| *pkts) + .unwrap_or(0); + data_link_pkts_sent += stats.packets_sent.saturating_sub(before); + + // Totals (cumulative since node creation) + total_link_pkts_sent += stats.packets_sent; + total_link_pkts_recv += stats.packets_recv; + total_link_bytes_sent += stats.bytes_sent; + total_link_bytes_recv += stats.bytes_recv; + } + } + + // Estimate average hop count from link packet overhead. + // Each data datagram traverses N link hops, each producing 1 link send. + // We sent 200 datagrams total (100 forward + 100 reverse). + let total_data_datagrams = (send_forward_ok + send_reverse_ok) as u64; + let avg_hops = if total_data_datagrams > 0 { + data_link_pkts_sent as f64 / total_data_datagrams as f64 + } else { + 0.0 + }; + + // Coord cache stats + let coord_cache_sizes: Vec = nodes + .iter() + .map(|tn| tn.node.coord_cache().len()) + .collect(); + let total_coord_entries: usize = coord_cache_sizes.iter().sum(); + let min_coord = *coord_cache_sizes.iter().min().unwrap(); + let max_coord = *coord_cache_sizes.iter().max().unwrap(); + + let route_cache_sizes: Vec = nodes + .iter() + .map(|tn| tn.node.route_cache().len()) + .collect(); + let total_route_entries: usize = route_cache_sizes.iter().sum(); + + // === Report === + + eprintln!("\n === Session 100-Node Test ==="); + eprintln!( + " Topology: {} nodes, {} edges (seed {})", + NUM_NODES, + edges.len(), + SEED + ); + eprintln!( + " Session pairs: {} (1 outbound per node, random target)", + session_pairs.len() + ); + + eprintln!("\n --- Handshake ---"); + eprintln!( + " Initiator established: {}/{}", + session_pairs.len(), + session_pairs.len() + ); + + eprintln!("\n --- Data Transfer ---"); + eprintln!( + " Forward (initiator->responder): {} sent, {} errors", + send_forward_ok, send_forward_err + ); + eprintln!( + " Reverse (responder->initiator): {} sent, {} errors", + send_reverse_ok, send_reverse_err + ); + eprintln!( + " TUN delivery: {} total ({} expected)", + total_delivered, + send_forward_ok + send_reverse_ok + ); + eprintln!( + " Forward delivered: {}/{} | Reverse delivered: {}/{}", + fwd_delivered, send_forward_ok, rev_delivered, send_reverse_ok + ); + + eprintln!("\n --- Final Session State ---"); + eprintln!( + " Entries: {} total ({} established, {} responding, {} initiating)", + total_sessions, total_established, total_responding, total_initiating + ); + eprintln!( + " Per node: min={} max={} avg={:.1}", + min_sessions, + max_sessions, + total_sessions as f64 / NUM_NODES as f64 + ); + eprintln!( + " All-established nodes: {}/{}", + fully_established_nodes, NUM_NODES + ); + + eprintln!("\n --- Routing ---"); + eprintln!( + " Data-phase link hops: {} ({:.1} avg hops/datagram over {} datagrams)", + data_link_pkts_sent, avg_hops, total_data_datagrams + ); + eprintln!( + " Lifetime link totals: {} pkts sent, {} pkts recv, {:.1} KB sent, {:.1} KB recv", + total_link_pkts_sent, + total_link_pkts_recv, + total_link_bytes_sent as f64 / 1024.0, + total_link_bytes_recv as f64 / 1024.0 + ); + eprintln!( + " Coord cache: total={} min={} max={} avg={:.1}", + total_coord_entries, + min_coord, + max_coord, + total_coord_entries as f64 / NUM_NODES as f64 + ); + eprintln!(" Route cache: total={}", total_route_entries); + + eprintln!("\n --- Timing ---"); + eprintln!( + " Setup: {:.1}s | Handshake: {:.1}s | Data: {:.1}s | Total: {:.1}s", + setup_time.as_secs_f64(), + session_time.as_secs_f64(), + data_time.as_secs_f64(), + start.elapsed().as_secs_f64() + ); + + if !fwd_missing.is_empty() { + eprintln!( + "\n First {} undelivered forward datagrams:", + fwd_missing.len() + ); + for &(src, dst) in &fwd_missing { + eprintln!(" node {} -> node {}", src, dst); + } + } + if !rev_missing.is_empty() { + eprintln!( + "\n First {} undelivered reverse datagrams:", + rev_missing.len() + ); + for &(src, dst) in &rev_missing { + eprintln!(" node {} <- node {}", src, dst); + } + } + + // === Assertions === + + assert_eq!( + send_forward_err, 0, + "All forward sends should succeed" + ); + assert_eq!( + send_reverse_err, 0, + "All reverse sends should succeed (responder Established after forward data)" + ); + assert_eq!( + fwd_delivered, send_forward_ok, + "All forward datagrams should be delivered to responder TUN" + ); + assert_eq!( + rev_delivered, send_reverse_ok, + "All reverse datagrams should be delivered to initiator TUN" + ); + assert_eq!( + total_established, total_sessions, + "All {} session entries should be Established, \ + but {} responding, {} initiating", + total_sessions, total_responding, total_initiating + ); + + cleanup_nodes(&mut nodes).await; +}