diff --git a/src/node/handlers/mmp.rs b/src/node/handlers/mmp.rs index f20539e..0cff5e1 100644 --- a/src/node/handlers/mmp.rs +++ b/src/node/handlers/mmp.rs @@ -547,6 +547,7 @@ impl Node { let now = Instant::now(); let heartbeat_interval = Duration::from_secs(self.config().node.heartbeat_interval_secs); let dead_timeout = Duration::from_secs(self.config().node.link_dead_timeout_secs); + let max_resends = self.config().node.rate_limit.handshake_max_resends; let heartbeat_msg = [LinkMessageType::Heartbeat.to_byte()]; // Collect heartbeats to send and dead peers to remove @@ -556,7 +557,7 @@ impl Node { for (node_addr, peer) in self.peers.iter() { // Check liveness via MMP receiver last_recv_time. // Fall back to session_start for peers that never sent data. - let is_dead = if let Some(mmp) = peer.mmp() { + let time_dead = if let Some(mmp) = peer.mmp() { let reference_time = mmp .receiver .last_recv_time() @@ -565,6 +566,19 @@ impl Node { } else { false }; + + // Suppress teardown while an FMP rekey is genuinely in flight with + // budget left: a rekey-handshake link is not silent. The msg1 + // resend cap guarantees this terminates (abandon on exhaustion or + // cutover on completion clears `rekey_in_progress`), so a truly + // dead link is reaped on the next cycle. + let rekey_active = peer.rekey_in_progress() + && peer.rekey_msg1_resend_count() < max_resends + && peer.rekey_msg1().is_some() + || (peer.rekey_msg3_payload().is_some() + && peer.rekey_msg3_resend_count() < max_resends); + + let is_dead = time_dead && !rekey_active; if is_dead { dead_peers.push(*node_addr); continue; diff --git a/src/node/handlers/rekey.rs b/src/node/handlers/rekey.rs index 09d5ab4..2223deb 100644 --- a/src/node/handlers/rekey.rs +++ b/src/node/handlers/rekey.rs @@ -282,14 +282,21 @@ impl Node { } let interval_ms = self.config().node.rate_limit.handshake_resend_interval_ms; + let backoff = self.config().node.rate_limit.handshake_resend_backoff; + let max_resends = self.config().node.rate_limit.handshake_max_resends; // Collect peers needing action let mut to_resend: Vec<(NodeAddr, Vec)> = Vec::new(); + let mut to_abandon: Vec = Vec::new(); for (node_addr, peer) in &self.peers { if !peer.rekey_in_progress() || peer.rekey_msg1().is_none() { continue; } + if peer.rekey_msg1_resend_count() >= max_resends { + to_abandon.push(*node_addr); + continue; + } if peer.needs_msg1_resend(now_ms) && let Some(msg1) = peer.rekey_msg1() { @@ -297,6 +304,17 @@ impl Node { } } + // Abandon rekey cycles that exhausted their retransmission budget. + for node_addr in to_abandon { + if let Some(peer) = self.peers.get_mut(&node_addr) { + peer.abandon_rekey(); + } + warn!( + peer = %self.peer_display_name(&node_addr), + "FMP rekey aborted: msg1 unconfirmed after max retransmissions, abandoning cycle" + ); + } + for (node_addr, msg1_bytes) in to_resend { let (transport_id, remote_addr) = match self.peers.get(&node_addr) { Some(p) => match (p.transport_id(), p.current_addr()) { @@ -312,12 +330,13 @@ impl Node { false }; - if sent { - if let Some(peer) = self.peers.get_mut(&node_addr) { - peer.set_msg1_next_resend(now_ms + interval_ms); - } + if sent && let Some(peer) = self.peers.get_mut(&node_addr) { + let count = peer.rekey_msg1_resend_count() + 1; + let next = now_ms + (interval_ms as f64 * backoff.powi(count as i32)) as u64; + peer.record_rekey_msg1_resend(next); trace!( peer = %self.peer_display_name(&node_addr), + resend = count, "Resent rekey msg1" ); } diff --git a/src/node/tests/heartbeat.rs b/src/node/tests/heartbeat.rs new file mode 100644 index 0000000..3c52b24 --- /dev/null +++ b/src/node/tests/heartbeat.rs @@ -0,0 +1,162 @@ +//! Link-dead heartbeat rekey-awareness integration tests. +//! +//! `check_link_heartbeats()` reaps a peer after the link-dead timeout, +//! but suppresses teardown while an FMP rekey is genuinely in flight with +//! its msg1 retransmission budget unexhausted. These tests drive a real +//! two-node UDP peering, inject rekey state on the peer, and verify the +//! suppress / resume / regression behaviors. `link_dead_timeout_secs` is +//! set to 0 so the elapsed-time predicate is always satisfied and the only +//! variable is the rekey-active guard. + +use super::spanning_tree::*; +use super::*; +use crate::Identity; +use crate::noise::HandshakeState; +use crate::utils::index::SessionIndex; + +/// Arm a real (initiator) FMP rekey on the peer the given node holds for +/// `peer_addr`, so the msg1 resend budget can be exercised. +fn arm_rekey(node: &mut crate::node::Node, peer_addr: &NodeAddr) { + let local = Identity::generate(); + let hs = HandshakeState::new_initiator(local.keypair()); + let peer = node.get_peer_mut(peer_addr).expect("peer present"); + peer.set_rekey_state(hs, SessionIndex::new(7), vec![0xAB; 64], 0); +} + +/// Set `link_dead_timeout_secs` on an already-constructed node via the +/// sole-store copy-on-write context swap (immutable state is no longer a +/// directly-pokeable field; `config()` is a read-only accessor). +fn set_link_dead_timeout(node: &mut crate::node::Node, secs: u64) { + node.replace_context(|ctx| { + let mut cfg = (*ctx.config).clone(); + cfg.node.link_dead_timeout_secs = secs; + ctx.config = std::sync::Arc::new(cfg); + }); +} + +/// A peer past the link-dead timeout is NOT reaped while an FMP rekey is in +/// progress with its msg1 budget unexhausted. +#[tokio::test] +async fn heartbeat_suppressed_during_rekey() { + let mut nodes = run_tree_test(2, &[(0, 1)], false).await; + verify_tree_convergence(&nodes); + + let addr_1 = *nodes[1].node.node_addr(); + assert!(nodes[0].node.get_peer(&addr_1).is_some()); + + // Force every link to read as dead on elapsed time alone. + set_link_dead_timeout(&mut nodes[0].node, 0); + + // Arm a rekey with budget left (count 0 < max_resends default 5). + arm_rekey(&mut nodes[0].node, &addr_1); + assert!(nodes[0].node.get_peer(&addr_1).unwrap().rekey_in_progress()); + + nodes[0].node.check_link_heartbeats().await; + + assert!( + nodes[0].node.get_peer(&addr_1).is_some(), + "peer reaped despite an in-flight rekey with budget remaining" + ); + + cleanup_nodes(&mut nodes).await; +} + +/// Once the msg1 budget is exhausted the rekey-active guard no longer +/// holds, so a peer past the link-dead timeout IS reaped. +#[tokio::test] +async fn heartbeat_resumes_after_budget_exhausted() { + let mut nodes = run_tree_test(2, &[(0, 1)], false).await; + verify_tree_convergence(&nodes); + + let addr_1 = *nodes[1].node.node_addr(); + assert!(nodes[0].node.get_peer(&addr_1).is_some()); + + set_link_dead_timeout(&mut nodes[0].node, 0); + let max_resends = nodes[0].node.config().node.rate_limit.handshake_max_resends; + + arm_rekey(&mut nodes[0].node, &addr_1); + + // Exhaust the budget: count reaches max_resends, guard goes false. + let peer = nodes[0].node.get_peer_mut(&addr_1).unwrap(); + for i in 0..max_resends { + peer.record_rekey_msg1_resend(1000 + i as u64 * 100); + } + assert_eq!( + nodes[0] + .node + .get_peer(&addr_1) + .unwrap() + .rekey_msg1_resend_count(), + max_resends + ); + + nodes[0].node.check_link_heartbeats().await; + + assert!( + nodes[0].node.get_peer(&addr_1).is_none(), + "peer not reaped after its rekey budget was exhausted" + ); + + cleanup_nodes(&mut nodes).await; +} + +/// Next-only: a peer past the link-dead timeout is NOT reaped while a +/// retained FMP rekey msg3 is still being retransmitted with its budget +/// unexhausted. This is the XX/next msg3-liveness arm of the suppression +/// predicate, distinct from the msg1 arm above. +#[tokio::test] +async fn heartbeat_suppressed_during_msg3_retransmit() { + let mut nodes = run_tree_test(2, &[(0, 1)], false).await; + verify_tree_convergence(&nodes); + + let addr_1 = *nodes[1].node.node_addr(); + assert!(nodes[0].node.get_peer(&addr_1).is_some()); + + set_link_dead_timeout(&mut nodes[0].node, 0); + + // Retain a msg3 payload with budget left (count 0 < max_resends), with + // no msg1 rekey in flight: only the msg3 arm of the guard can hold. + let peer = nodes[0].node.get_peer_mut(&addr_1).unwrap(); + peer.set_rekey_msg3_payload(vec![0xCD; 64], 1000); + assert!(!peer.rekey_in_progress()); + assert_eq!( + nodes[0] + .node + .get_peer(&addr_1) + .unwrap() + .rekey_msg3_resend_count(), + 0 + ); + + nodes[0].node.check_link_heartbeats().await; + + assert!( + nodes[0].node.get_peer(&addr_1).is_some(), + "peer reaped despite a retained rekey msg3 with budget remaining" + ); + + cleanup_nodes(&mut nodes).await; +} + +/// Regression guard: with no rekey in flight, a peer past the link-dead +/// timeout is reaped exactly as before. +#[tokio::test] +async fn heartbeat_unaffected_without_rekey() { + let mut nodes = run_tree_test(2, &[(0, 1)], false).await; + verify_tree_convergence(&nodes); + + let addr_1 = *nodes[1].node.node_addr(); + assert!(nodes[0].node.get_peer(&addr_1).is_some()); + assert!(!nodes[0].node.get_peer(&addr_1).unwrap().rekey_in_progress()); + + set_link_dead_timeout(&mut nodes[0].node, 0); + + nodes[0].node.check_link_heartbeats().await; + + assert!( + nodes[0].node.get_peer(&addr_1).is_none(), + "dead peer with no rekey in flight should be reaped" + ); + + cleanup_nodes(&mut nodes).await; +} diff --git a/src/node/tests/mod.rs b/src/node/tests/mod.rs index 57394b1..2506958 100644 --- a/src/node/tests/mod.rs +++ b/src/node/tests/mod.rs @@ -18,6 +18,7 @@ mod discovery; mod ethernet; mod forwarding; mod handshake; +mod heartbeat; mod routing; mod session; mod spanning_tree; diff --git a/src/peer/active.rs b/src/peer/active.rs index e4f68a8..77e4bd3 100644 --- a/src/peer/active.rs +++ b/src/peer/active.rs @@ -209,6 +209,9 @@ pub struct ActivePeer { rekey_msg1: Option>, /// In-progress rekey: next resend timestamp (Unix ms). rekey_msg1_next_resend: u64, + /// In-progress rekey: number of msg1 retransmissions performed so far. + rekey_msg1_resend_count: u32, + // === Rekey Responder State (XX pattern) === /// In-progress rekey responder: Noise handshake state awaiting msg3. rekey_responder_handshake: Option, @@ -299,6 +302,7 @@ impl ActivePeer { rekey_our_index: None, rekey_msg1: None, rekey_msg1_next_resend: 0, + rekey_msg1_resend_count: 0, rekey_responder_handshake: None, rekey_responder_our_index: None, rekey_msg3_payload: None, @@ -400,6 +404,7 @@ impl ActivePeer { rekey_our_index: None, rekey_msg1: None, rekey_msg1_next_resend: 0, + rekey_msg1_resend_count: 0, rekey_responder_handshake: None, rekey_responder_our_index: None, rekey_msg3_payload: None, @@ -1042,6 +1047,7 @@ impl ActivePeer { self.rekey_handshake = None; self.rekey_msg1 = None; self.rekey_msg1_next_resend = 0; + self.rekey_msg1_resend_count = 0; } /// Cut over to the pending new session (initiator side). @@ -1069,6 +1075,7 @@ impl ActivePeer { self.session_established_at = Instant::now(); self.session_start = Instant::now(); self.rekey_in_progress = false; + self.rekey_msg1_resend_count = 0; self.rekey_jitter_secs = draw_rekey_jitter(); self.reset_replay_suppressed(); @@ -1105,6 +1112,7 @@ impl ActivePeer { self.session_established_at = Instant::now(); self.session_start = Instant::now(); self.rekey_in_progress = false; + self.rekey_msg1_resend_count = 0; self.rekey_jitter_secs = draw_rekey_jitter(); self.reset_replay_suppressed(); @@ -1149,6 +1157,7 @@ impl ActivePeer { self.rekey_handshake = None; self.rekey_msg1 = None; self.rekey_msg1_next_resend = 0; + self.rekey_msg1_resend_count = 0; self.rekey_in_progress = false; self.clear_rekey_msg3_payload(); // Return whichever index needs freeing @@ -1173,6 +1182,7 @@ impl ActivePeer { self.rekey_our_index = Some(our_index); self.rekey_msg1 = Some(wire_msg1); self.rekey_msg1_next_resend = next_resend_ms; + self.rekey_msg1_resend_count = 0; self.rekey_in_progress = true; } @@ -1226,6 +1236,7 @@ impl ActivePeer { // Clear msg1 resend state self.rekey_msg1 = None; self.rekey_msg1_next_resend = 0; + self.rekey_msg1_resend_count = 0; Ok((msg3, session, remote_epoch)) } @@ -1322,6 +1333,17 @@ impl ActivePeer { self.rekey_msg1_next_resend = next_ms; } + /// Number of rekey msg1 retransmissions performed so far. + pub fn rekey_msg1_resend_count(&self) -> u32 { + self.rekey_msg1_resend_count + } + + /// Record a rekey msg1 retransmission and schedule the next one. + pub fn record_rekey_msg1_resend(&mut self, next_ms: u64) { + self.rekey_msg1_resend_count += 1; + self.rekey_msg1_next_resend = next_ms; + } + // === Rekey Responder State (XX pattern) === /// Whether this peer has a rekey responder handshake awaiting msg3. @@ -1610,4 +1632,51 @@ mod tests { n ); } + + /// Put a peer into a rekey-in-progress state with a real (initiator) + /// handshake so the msg1 resend budget can be exercised. + fn arm_rekey(peer: &mut ActivePeer) { + let local = Identity::generate(); + let hs = NoiseHandshakeState::new_initiator(local.keypair()); + peer.set_rekey_state(hs, SessionIndex::new(7), vec![0xAB; 64], 0); + } + + #[test] + fn rekey_msg1_resend_count_increments_and_caps() { + let identity = make_peer_identity(); + let mut peer = ActivePeer::new(identity, LinkId::new(1), 1000); + arm_rekey(&mut peer); + + assert!(peer.rekey_in_progress()); + assert_eq!(peer.rekey_msg1_resend_count(), 0); + assert!(peer.rekey_msg1().is_some()); + + // The driver records one resend per call; the count tracks them. + let max_resends: u32 = 5; + for i in 0..max_resends { + peer.record_rekey_msg1_resend(1000 + i as u64 * 100); + assert_eq!(peer.rekey_msg1_resend_count(), i + 1); + } + assert_eq!(peer.rekey_msg1_resend_count(), max_resends); + } + + #[test] + fn rekey_msg1_budget_exhaustion_abandons_cleanly() { + let identity = make_peer_identity(); + let mut peer = ActivePeer::new(identity, LinkId::new(1), 1000); + arm_rekey(&mut peer); + + // Simulate the driver exhausting its budget. + let max_resends: u32 = 5; + for i in 0..max_resends { + peer.record_rekey_msg1_resend(1000 + i as u64 * 100); + } + assert_eq!(peer.rekey_msg1_resend_count(), max_resends); + + // Budget exhausted -> abandon: state clears and the counter resets. + peer.abandon_rekey(); + assert!(!peer.rekey_in_progress()); + assert!(peer.rekey_msg1().is_none()); + assert_eq!(peer.rekey_msg1_resend_count(), 0); + } }