diff --git a/src/node/handlers/rx_loop.rs b/src/node/handlers/rx_loop.rs index fcee5f5..a88de14 100644 --- a/src/node/handlers/rx_loop.rs +++ b/src/node/handlers/rx_loop.rs @@ -80,6 +80,7 @@ impl Node { .map(|d| d.as_millis() as u64) .unwrap_or(0); self.resend_pending_handshakes(now_ms).await; + self.resend_pending_session_handshakes(now_ms).await; self.purge_idle_sessions(now_ms); self.process_pending_retries(now_ms).await; self.check_tree_state().await; diff --git a/src/node/handlers/session.rs b/src/node/handlers/session.rs index 524ee13..b56869b 100644 --- a/src/node/handlers/session.rs +++ b/src/node/handlers/session.rs @@ -161,6 +161,7 @@ impl Node { entry.set_coords_warmup_remaining(self.config.node.session.coords_warmup_packets); entry.mark_established(Self::now_ms()); entry.init_mmp(&self.config.node.session_mmp); + entry.clear_handshake_payload(); info!(src = %self.peer_display_name(src_addr), "Session established (responder, on first encrypted message)"); } @@ -316,8 +317,18 @@ impl Node { ); } EndToEndState::Responding(_) => { - // Duplicate setup while we already responded — drop - debug!(src = %self.peer_display_name(src_addr), "Duplicate SessionSetup, already responding"); + // Duplicate setup while we already responded — resend stored ack + if let Some(payload) = existing.handshake_payload() { + debug!(src = %self.peer_display_name(src_addr), "Duplicate SessionSetup, resending SessionAck"); + let my_addr = *self.node_addr(); + let mut datagram = SessionDatagram::new(my_addr, *src_addr, payload.to_vec()) + .with_ttl(self.config.node.session.default_ttl); + if let Err(e) = self.send_session_datagram(&mut datagram).await { + debug!(error = %e, dest = %self.peer_display_name(src_addr), "Failed to resend SessionAck"); + } + } else { + debug!(src = %self.peer_display_name(src_addr), "Duplicate SessionSetup, no stored ack to resend"); + } return; } EndToEndState::Established(_) => { @@ -360,8 +371,9 @@ impl Node { // Build and send SessionAck let our_coords = self.tree_state.my_coords().clone(); let ack = SessionAck::new(our_coords).with_handshake(msg2); + let ack_payload = ack.encode(); let my_addr = *self.node_addr(); - let mut datagram = SessionDatagram::new(my_addr, *src_addr, ack.encode()) + let mut datagram = SessionDatagram::new(my_addr, *src_addr, ack_payload.clone()) .with_ttl(self.config.node.session.default_ttl); // Route the ack back to the initiator @@ -370,9 +382,11 @@ impl Node { return; } - // Store session entry in Responding state + // Store session entry in Responding state with ack payload for potential resend let now_ms = Self::now_ms(); - let entry = SessionEntry::new(*src_addr, remote_pubkey, EndToEndState::Responding(handshake), now_ms, false); + let resend_interval = self.config.node.rate_limit.handshake_resend_interval_ms; + let mut entry = SessionEntry::new(*src_addr, remote_pubkey, EndToEndState::Responding(handshake), now_ms, false); + entry.set_handshake_payload(ack_payload, now_ms + resend_interval); self.sessions.insert(*src_addr, entry); debug!(src = %self.peer_display_name(src_addr), "SessionSetup processed, SessionAck sent"); @@ -433,6 +447,7 @@ impl Node { entry.set_coords_warmup_remaining(self.config.node.session.coords_warmup_packets); entry.mark_established(now_ms); entry.init_mmp(&self.config.node.session_mmp); + entry.clear_handshake_payload(); entry.touch(now_ms); self.sessions.insert(*src_addr, entry); self.coord_cache.insert(*src_addr, ack.src_coords, now_ms); @@ -723,10 +738,11 @@ impl Node { let dest_coords = self.get_dest_coords(&dest_addr); let setup = SessionSetup::new(our_coords, dest_coords) .with_handshake(msg1); + let setup_payload = setup.encode(); // Wrap in SessionDatagram let my_addr = *self.node_addr(); - let mut datagram = SessionDatagram::new(my_addr, dest_addr, setup.encode()) + let mut datagram = SessionDatagram::new(my_addr, dest_addr, setup_payload.clone()) .with_ttl(self.config.node.session.default_ttl); // Route toward destination @@ -735,9 +751,11 @@ impl Node { // Register destination identity for TUN → session routing self.register_identity(dest_addr, dest_pubkey); - // Store session entry + // Store session entry with handshake payload for potential resend let now_ms = Self::now_ms(); - let entry = SessionEntry::new(dest_addr, dest_pubkey, EndToEndState::Initiating(handshake), now_ms, true); + let resend_interval = self.config.node.rate_limit.handshake_resend_interval_ms; + let mut entry = SessionEntry::new(dest_addr, dest_pubkey, EndToEndState::Initiating(handshake), now_ms, true); + entry.set_handshake_payload(setup_payload, now_ms + resend_interval); self.sessions.insert(dest_addr, entry); info!(dest = %self.peer_display_name(&dest_addr), "Session initiation started"); @@ -1027,7 +1045,7 @@ impl Node { /// /// Finds the next hop for the destination, seeds path_mtu from the /// first-hop transport MTU, and sends as an encrypted link message. - async fn send_session_datagram( + pub(in crate::node) async fn send_session_datagram( &mut self, datagram: &mut SessionDatagram, ) -> Result<(), NodeError> { diff --git a/src/node/handlers/timeout.rs b/src/node/handlers/timeout.rs index 4e2bc0e..55b6382 100644 --- a/src/node/handlers/timeout.rs +++ b/src/node/handlers/timeout.rs @@ -151,6 +151,84 @@ impl Node { } } + /// Resend session-layer handshake messages and timeout stale handshakes. + /// + /// For sessions in Initiating or Responding state: + /// - If the handshake has exceeded the timeout window, remove the session. + /// - If a resend is due and under max resends, resend the stored payload + /// wrapped in a fresh SessionDatagram (so routing can adapt). + pub(in crate::node) async fn resend_pending_session_handshakes(&mut self, now_ms: u64) { + if self.sessions.is_empty() { + return; + } + + let timeout_ms = self.config.node.rate_limit.handshake_timeout_secs * 1000; + let max_resends = self.config.node.rate_limit.handshake_max_resends; + let interval_ms = self.config.node.rate_limit.handshake_resend_interval_ms; + let backoff = self.config.node.rate_limit.handshake_resend_backoff; + let ttl = self.config.node.session.default_ttl; + + // First pass: find timed-out sessions to remove + let timed_out: Vec = self.sessions.iter() + .filter(|(_, entry)| { + !entry.is_established() + && now_ms.saturating_sub(entry.last_activity()) > timeout_ms + }) + .map(|(addr, _)| *addr) + .collect(); + + for addr in &timed_out { + let name = self.peer_display_name(addr); + info!(dest = %name, "Session handshake timed out, removing"); + self.sessions.remove(addr); + self.pending_tun_packets.remove(addr); + } + + // Second pass: collect resend candidates + let my_addr = *self.node_addr(); + let candidates: Vec<(crate::NodeAddr, Vec)> = self.sessions.iter() + .filter(|(_, entry)| { + !entry.is_established() + && entry.handshake_payload().is_some() + && entry.resend_count() < max_resends + && entry.next_resend_at_ms() > 0 + && now_ms >= entry.next_resend_at_ms() + }) + .map(|(addr, entry)| (*addr, entry.handshake_payload().unwrap().to_vec())) + .collect(); + + for (dest_addr, payload) in candidates { + use crate::protocol::SessionDatagram; + + let mut datagram = SessionDatagram::new(my_addr, dest_addr, payload) + .with_ttl(ttl); + let sent = match self.send_session_datagram(&mut datagram).await { + Ok(_) => true, + Err(e) => { + debug!( + dest = %self.peer_display_name(&dest_addr), + error = %e, + "Session handshake resend failed" + ); + false + } + }; + + if sent + && let Some(entry) = self.sessions.get_mut(&dest_addr) + { + let count = entry.resend_count() + 1; + let next = now_ms + (interval_ms as f64 * backoff.powi(count as i32)) as u64; + entry.record_resend(next); + debug!( + dest = %self.peer_display_name(&dest_addr), + resend = count, + "Resent session handshake" + ); + } + } + } + /// Remove established sessions that have been idle too long. /// /// Only targets sessions in the Established state. Initiating/Responding diff --git a/src/node/session.rs b/src/node/session.rs index 6bc91fd..d8ac764 100644 --- a/src/node/session.rs +++ b/src/node/session.rs @@ -70,6 +70,15 @@ pub(crate) struct SessionEntry { is_initiator: bool, /// Session-layer MMP state. Initialized on Established transition. mmp: Option, + + // === Handshake Resend === + /// Encoded session-layer payload for resend (SessionSetup or SessionAck). + /// Cleared on Established transition. + handshake_payload: Option>, + /// Number of resends performed. + resend_count: u32, + /// When the next resend should fire (Unix ms). 0 = no resend scheduled. + next_resend_at_ms: u64, } impl SessionEntry { @@ -91,6 +100,9 @@ impl SessionEntry { coords_warmup_remaining: 0, is_initiator, mmp: None, + handshake_payload: None, + resend_count: 0, + next_resend_at_ms: 0, } } @@ -193,4 +205,45 @@ impl SessionEntry { pub(crate) fn init_mmp(&mut self, config: &SessionMmpConfig) { self.mmp = Some(MmpSessionState::new(config, self.is_initiator)); } + + // === Handshake Resend === + + /// Store the encoded session-layer payload for potential resend. + /// + /// For initiators, this is the SessionSetup payload bytes. + /// For responders, this is the SessionAck payload bytes. + /// The payload is re-wrapped in a fresh SessionDatagram on each resend + /// so routing can adapt to topology changes. + pub(crate) fn set_handshake_payload(&mut self, payload: Vec, next_resend_at_ms: u64) { + self.handshake_payload = Some(payload); + self.resend_count = 0; + self.next_resend_at_ms = next_resend_at_ms; + } + + /// Get the stored handshake payload for resend. + pub(crate) fn handshake_payload(&self) -> Option<&[u8]> { + self.handshake_payload.as_deref() + } + + /// Clear the stored handshake payload (called on Established transition). + pub(crate) fn clear_handshake_payload(&mut self) { + self.handshake_payload = None; + self.next_resend_at_ms = 0; + } + + /// Number of resends performed so far. + pub(crate) fn resend_count(&self) -> u32 { + self.resend_count + } + + /// When the next resend should fire (Unix ms). 0 = no resend scheduled. + pub(crate) fn next_resend_at_ms(&self) -> u64 { + self.next_resend_at_ms + } + + /// Record a resend and schedule the next one. + pub(crate) fn record_resend(&mut self, next_resend_at_ms: u64) { + self.resend_count += 1; + self.next_resend_at_ms = next_resend_at_ms; + } } diff --git a/src/node/tests/session.rs b/src/node/tests/session.rs index fd29991..74b3f32 100644 --- a/src/node/tests/session.rs +++ b/src/node/tests/session.rs @@ -1480,3 +1480,184 @@ fn test_identity_cache_lookup() { assert_eq!(addr, remote_addr); assert_eq!(pk, remote.pubkey_full()); } + +// ============================================================================ +// Session-layer handshake resend tests +// ============================================================================ + +/// Test that SessionEntry handshake payload storage works correctly. +#[test] +fn test_session_entry_handshake_payload_storage() { + 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, + true, + ); + + // Initially no handshake payload + assert!(entry.handshake_payload().is_none()); + assert_eq!(entry.resend_count(), 0); + assert_eq!(entry.next_resend_at_ms(), 0); + + // Store a handshake payload + let payload = vec![0x01, 0x02, 0x03, 0x04]; + entry.set_handshake_payload(payload.clone(), 2000); + + assert_eq!(entry.handshake_payload().unwrap(), &payload); + assert_eq!(entry.resend_count(), 0); + assert_eq!(entry.next_resend_at_ms(), 2000); +} + +/// Test that resend_count and next_resend_at_ms track correctly on SessionEntry. +#[test] +fn test_session_entry_resend_tracking() { + 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, + true, + ); + + entry.set_handshake_payload(vec![0x01], 2000); + + // Record first resend + entry.record_resend(4000); + assert_eq!(entry.resend_count(), 1); + assert_eq!(entry.next_resend_at_ms(), 4000); + + // Record second resend + entry.record_resend(8000); + assert_eq!(entry.resend_count(), 2); + assert_eq!(entry.next_resend_at_ms(), 8000); +} + +/// Test that clear_handshake_payload clears payload and resets timer. +#[test] +fn test_session_entry_clear_handshake_payload() { + 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, + true, + ); + + entry.set_handshake_payload(vec![0x01, 0x02], 2000); + entry.record_resend(4000); + assert!(entry.handshake_payload().is_some()); + assert_eq!(entry.resend_count(), 1); + + // Clear on Established transition + entry.clear_handshake_payload(); + assert!(entry.handshake_payload().is_none()); + assert_eq!(entry.next_resend_at_ms(), 0); + // resend_count is NOT reset — it's a historical record + assert_eq!(entry.resend_count(), 1); +} + +/// Test that session handshake timeout removes stale Initiating sessions. +#[tokio::test] +async fn test_session_handshake_timeout() { + 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(); + + // Create a session at time 1000 + let entry = crate::node::session::SessionEntry::new( + dest_addr, + identity_b.pubkey_full(), + EndToEndState::Initiating(handshake), + 1000, + true, + ); + node.sessions.insert(dest_addr, entry); + + assert!(node.sessions.contains_key(&dest_addr)); + + // Before timeout: session should remain + let timeout_secs = node.config.node.rate_limit.handshake_timeout_secs; + let before_timeout = 1000 + timeout_secs * 1000 - 1; + node.resend_pending_session_handshakes(before_timeout).await; + assert!(node.sessions.contains_key(&dest_addr), "Session should survive before timeout"); + + // After timeout: session should be removed + let after_timeout = 1000 + timeout_secs * 1000 + 1; + node.resend_pending_session_handshakes(after_timeout).await; + assert!(!node.sessions.contains_key(&dest_addr), "Timed-out session should be removed"); +} + +/// Test that session handshake timeout removes stale Responding sessions. +#[tokio::test] +async fn test_session_responding_timeout() { + use crate::noise::HandshakeState; + + let mut node = make_node(); + + let identity_a = Identity::generate(); + let identity_b = Identity::generate(); + + let handshake = HandshakeState::new_responder( + identity_b.keypair(), + ); + + let src_addr = *identity_a.node_addr(); + + // Create a Responding session at time 1000 + let entry = crate::node::session::SessionEntry::new( + src_addr, + identity_a.pubkey_full(), + EndToEndState::Responding(handshake), + 1000, + false, + ); + node.sessions.insert(src_addr, entry); + + assert!(node.sessions.contains_key(&src_addr)); + + // After timeout: session should be removed + let timeout_secs = node.config.node.rate_limit.handshake_timeout_secs; + let after_timeout = 1000 + timeout_secs * 1000 + 1; + node.resend_pending_session_handshakes(after_timeout).await; + assert!(!node.sessions.contains_key(&src_addr), "Timed-out Responding session should be removed"); +}