From 347cbe60bd692fa3d54283948e96b7459874381d Mon Sep 17 00:00:00 2001 From: Johnathan Corgan Date: Sat, 18 Jul 2026 22:55:38 +0000 Subject: [PATCH] peer: move the Noise handshake operations onto the control machine The pending connection drove its own Noise handshake while the control machine held it, so the crypto and the state it produces lived on the carrier that is going away. Move the six operations onto the machine: starting and completing an initiation, processing an inbound initiation, taking the session, testing for one, and dropping the handle on failure. Bodies are unchanged apart from reaching the handles through the attached connection. Each operation now records its results on both carriers. The learned identity, the remote epoch, and the activity stamp are written to the machine's own bookkeeping at the same point and with the same value as they are written to the connection's. The connection's copies still have readers until those reads are repointed, so both have to be written; the machine's copy of the learned identity was previously never populated for an inbound connection, which is why an inbound pending row showed no expected peer. Driving the handshake from the machine means the machine has to exist before the crypto runs. On the inbound path it is built above the message-1 processing and still kept local, so a rejected message leaves no registry entry and allocates no index. On the outbound path it already existed from the dial, so it simply takes the connection before the index allocation, and both failure arms unwind it as they did. Completing message 2 no longer mirrors the activity stamp separately, since the completion itself now writes both carriers at the point the mirror was approximating. Three tests cover what the compiler cannot. A connection whose handshake failed holds neither Noise handle yet must stay visible to the stale-connection sweep, or every failed connection would leak and hold a peering-budget slot forever. A message 1 rejected by the crypto or by the ACL must leave no machine, no index, and the same rejection count as before. A dial whose message-1 preparation fails must unwind the machine registered at dial time; that test drives the index-allocation failure rather than a crypto failure, which is the arm this change actually widens, since the allocation now happens with the connection already attached. --- src/node/handlers/handshake.rs | 97 ++++--- src/node/lifecycle/mod.rs | 85 +++--- src/node/tests/acl.rs | 53 +++- src/node/tests/establish_chartests.rs | 14 +- src/node/tests/handshake.rs | 56 ++-- src/node/tests/mod.rs | 32 ++- src/node/tests/spanning_tree.rs | 3 +- src/node/tests/unit.rs | 140 +++++++++- src/peer/connection.rs | 303 ++------------------- src/peer/machine.rs | 375 +++++++++++++++++++++++++- 10 files changed, 747 insertions(+), 411 deletions(-) diff --git a/src/node/handlers/handshake.rs b/src/node/handlers/handshake.rs index 3b78bf6..c5ebd3c 100644 --- a/src/node/handlers/handshake.rs +++ b/src/node/handlers/handshake.rs @@ -266,16 +266,28 @@ impl Node { // === CRYPTO COST PAID HERE === let link_id = self.allocate_link_id(); - let mut conn = PeerConnection::inbound_with_transport( + let conn = PeerConnection::inbound_with_transport( link_id, packet.transport_id, packet.remote_addr.clone(), packet.timestamp_ms, ); + // The control machine drives the handshake, so it is built here, above + // the crypto, carrying the pending connection. It stays a local: it + // enters `peer_machines` only at the promote tails, so a rejected msg1 + // still leaves no registry trace and allocates no index. + let mut machine = PeerMachine::new_inbound(link_id, packet.timestamp_ms); + // The inbound connection carries the transport ID from msg1, but the + // machine's carrier is only written on the outbound dial. Seed it here + // so the promotion hand-off reads it from the surviving carrier, + // matching the connection's own inbound seed. + machine.set_conn_transport_id(packet.transport_id); + machine.set_leg(conn); + let our_keypair = self.identity().keypair(); let noise_msg1 = &packet.data[header.noise_msg1_offset..]; - let msg2_response = match conn.receive_handshake_init( + let msg2_response = match machine.receive_handshake_init( our_keypair, self.startup_epoch(), noise_msg1, @@ -295,7 +307,11 @@ impl Node { }; // Learn peer identity from msg1 - let peer_identity = match conn.expected_identity() { + let peer_identity = match machine + .leg() + .expect("pending connection attached above") + .expected_identity() + { Some(id) => *id, None => { self.msg1_rate_limiter.complete_handshake(); @@ -314,7 +330,10 @@ impl Node { // state; from here the decision reads only `wire` and the snapshot. let wire = WireOutcome { peer_identity, - remote_epoch: conn.remote_epoch(), + remote_epoch: machine + .leg() + .expect("pending connection attached above") + .remote_epoch(), their_index: header.sender_idx, msg2_payload: msg2_response, }; @@ -336,12 +355,6 @@ impl Node { // shared authorize → allocate → send-msg2 → promote tail; the other // variants complete the rate-limiter and return here. The local machine // enters `peer_machines` only at the promote tails. - let mut machine = PeerMachine::new_inbound(link_id, packet.timestamp_ms); - // The inbound leg carries the transport ID from msg1, but the machine's - // carrier is only written on the outbound dial. Seed it here so the - // promotion hand-off reads it from the surviving carrier, matching the - // leg's inbound seed. - machine.set_conn_transport_id(packet.transport_id); let (decision, actions) = machine.inbound_msg1(link_id, &wire, est, packet.timestamp_ms); match decision { InboundDecision::Reject { @@ -443,7 +456,7 @@ impl Node { } // Rekey: process as responder, store new session as pending. - let noise_session = conn.take_session(); + let noise_session = machine.take_session(); let our_new_index = match self.index_allocator.allocate() { Ok(idx) => idx, Err(e) => { @@ -623,7 +636,10 @@ impl Node { // connection, then build + store the framed msg2. The old index was // already freed by `remove_active_peer` above, BEFORE this fresh // allocation — matching the pre-refactor allocation sequence. - conn.set_our_index(our_index); + machine + .leg_mut() + .expect("pending connection attached above") + .set_our_index(our_index); let link = Link::connectionless( link_id, packet.transport_id, @@ -638,9 +654,8 @@ impl Node { // msg1 resend while the connection is still pending. machine.set_conn_handshake_msg2(wire_msg2.clone()); - // Register the machine, carrying the connection + // Register the machine, which already carries the connection // (Promote/Restart tail only). - machine.set_leg(conn); self.peer_machines.insert(link_id, machine); // Execute [SendHandshake, PromoteToActive]. Because the old peer was @@ -771,7 +786,10 @@ impl Node { // Shell registry surgery, in the pre-refactor order: // set indices on the shell connection, insert link / reverse map / // connection, then build + store the framed msg2. - conn.set_our_index(our_index); + machine + .leg_mut() + .expect("pending connection attached above") + .set_our_index(our_index); let link = Link::connectionless( link_id, packet.transport_id, @@ -786,10 +804,9 @@ impl Node { // msg1 resend while the connection is still pending. machine.set_conn_handshake_msg2(wire_msg2.clone()); - // Register the machine, carrying the connection (Promote tail - // only — discarded on every reject/resend/rekey arm per the - // insertion discipline). - machine.set_leg(conn); + // Register the machine, which already carries the connection + // (Promote tail only — discarded on every reject/resend/rekey + // arm per the insertion discipline). self.peer_machines.insert(link_id, machine); // Execute [SendHandshake, PromoteToActive]. The executor frames + @@ -1001,32 +1018,32 @@ impl Node { } let (peer_identity, our_index) = { - let conn = self.leg_mut(&link_id).unwrap(); + let machine = self.peer_machines.get_mut(&link_id).unwrap(); let noise_msg2 = &packet.data[header.noise_msg2_offset..]; - if let Err(e) = conn.complete_handshake(noise_msg2, packet.timestamp_ms) { + if let Err(e) = machine.complete_handshake(noise_msg2, packet.timestamp_ms) { warn!( link_id = %link_id, error = %e, "Handshake completion failed" ); - // Drop the leg's Noise handle (byte-identical point) and record - // the failure on the control machine as `send_failed` — the - // failure state's new home. The machine PHASE stays exactly - // where the old leg-carried failure left it (`SentMsg1`): the - // stale-connection sweep reclaims the leg unconditionally via - // the machine `is_failed()` at the next tick, before any - // projection or resend, so the phase in that window is - // byte-identical to the pre-collapse machine. - conn.mark_failed(); - if let Some(machine) = self.peer_machines.get_mut(&link_id) { - machine.mark_send_failed(); - } + // Drop the Noise handle (byte-identical point) and record the + // failure on the control machine as `send_failed` — the + // failure state's home. The machine PHASE stays exactly where + // the old failure left it (`SentMsg1`): the stale-connection + // sweep reclaims the connection unconditionally via the + // machine `is_failed()` at the next tick, before any + // projection or resend. + machine.mark_failed(); + machine.mark_send_failed(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; } + let conn = machine + .leg_mut() + .expect("pending connection present for msg2 completion"); conn.set_source_addr(packet.remote_addr.clone()); let peer_identity = match conn.expected_identity() { @@ -1042,13 +1059,6 @@ impl Node { (peer_identity, conn.our_index()) }; - // Mirror the leg's completion `touch` on the surviving carrier so the - // connection's last-activity advances at msg2 completion, matching the - // leg's clock. - if let Some(machine) = self.peer_machines.get_mut(&link_id) { - machine.touch_conn(packet.timestamp_ms); - } - if self .authorize_peer( &peer_identity, @@ -1173,7 +1183,7 @@ impl Node { // We're the smaller node. Swap to outbound session + indices. // The peer will keep their inbound session (complement of ours). let outbound_our_index = conn.our_index(); - let outbound_session = conn.take_session(); + let outbound_session = conn.noise_session.take(); let (outbound_session, outbound_our_index) = match ( outbound_session, @@ -1391,12 +1401,13 @@ impl Node { let link_stats = machine.conn_link_stats().clone(); // Verify handshake is complete and extract session - if !connection.has_session() { + if connection.noise_session.is_none() { return Err(NodeError::HandshakeIncomplete(link_id)); } let noise_session = connection - .take_session() + .noise_session + .take() .ok_or(NodeError::NoSession(link_id))?; let our_index = connection diff --git a/src/node/lifecycle/mod.rs b/src/node/lifecycle/mod.rs index 22fa219..f8652d6 100644 --- a/src/node/lifecycle/mod.rs +++ b/src/node/lifecycle/mod.rs @@ -585,7 +585,20 @@ impl Node { // Create connection in handshake phase (outbound knows expected identity) let current_time_ms = Self::now_ms(); - let mut connection = PeerConnection::outbound(link_id, peer_identity, current_time_ms); + let connection = PeerConnection::outbound(link_id, peer_identity, current_time_ms); + + // The control machine drives the handshake, so it takes the connection + // before the crypto runs. The machine was born at dial and persisted in + // `initiate_connection`, so every live caller already has one; recover + // with a fresh one if a direct caller ever skips the dial. + debug_assert!( + self.peer_machines.contains_key(&link_id), + "outbound msg1 prepared for link {link_id} with no dial-time machine" + ); + self.peer_machines + .entry(link_id) + .or_insert_with(|| PeerMachine::new_outbound(link_id, peer_identity, current_time_ms)) + .set_leg(connection); // Allocate a session index for this handshake let our_index = match self.index_allocator.allocate() { @@ -602,23 +615,35 @@ impl Node { // Start the Noise handshake and get message 1 let our_keypair = self.identity().keypair(); - let noise_msg1 = - match connection.start_handshake(our_keypair, self.startup_epoch(), current_time_ms) { - Ok(msg) => msg, - Err(e) => { - // Clean up the index, link, and dial-time machine - let _ = self.index_allocator.free(our_index); - self.links.remove(&link_id); - self.addr_to_link - .remove(&(transport_id, remote_addr.clone())); - self.remove_peer_machine(link_id); - return Err(NodeError::HandshakeFailed(e.to_string())); - } - }; + let startup_epoch = self.startup_epoch(); + let noise_msg1 = match self + .peer_machines + .get_mut(&link_id) + .expect("dial-time machine carries the connection") + .start_handshake(our_keypair, startup_epoch, current_time_ms) + { + Ok(msg) => msg, + Err(e) => { + // Clean up the index, link, and dial-time machine + let _ = self.index_allocator.free(our_index); + self.links.remove(&link_id); + self.addr_to_link + .remove(&(transport_id, remote_addr.clone())); + self.remove_peer_machine(link_id); + return Err(NodeError::HandshakeFailed(e.to_string())); + } + }; // Set index and transport info on the connection - connection.set_our_index(our_index); - connection.set_source_addr(remote_addr.clone()); + { + let conn = self + .peer_machines + .get_mut(&link_id) + .and_then(|machine| machine.leg_mut()) + .expect("dial-time machine carries the connection"); + conn.set_our_index(our_index); + conn.set_source_addr(remote_addr.clone()); + } // Build wire format msg1: [0x01][sender_idx:4 LE][noise_msg1:82] let wire_msg1 = build_msg1(our_index, &noise_msg1); @@ -641,21 +666,13 @@ impl Node { self.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); - // The dial-born machine (persisted in `initiate_connection`) carries - // the prepared connection from here. Every live caller dialed first, - // so the machine exists; recover with a fresh one if a direct caller - // ever skips the dial. - debug_assert!( - self.peer_machines.contains_key(&link_id), - "outbound msg1 prepared for link {link_id} with no dial-time machine" - ); let machine = self .peer_machines - .entry(link_id) - .or_insert_with(|| PeerMachine::new_outbound(link_id, peer_identity, current_time_ms)); + .get_mut(&link_id) + .expect("dial-time machine carries the connection"); // The dial-born machine carrier was stamped at dial; re-stamp it with the - // leg's msg1-prep clock so the surviving `started_at`/`last_activity` - // carry the leg's provenance. The two clocks differ when a connect + // msg1-prep clock so the surviving `started_at`/`last_activity` carry + // the preparation's provenance. The two clocks differ when a connect // round-trip separates dial from msg1 preparation. machine.set_conn_started_at(current_time_ms); machine.touch_conn(current_time_ms); @@ -663,14 +680,14 @@ impl Node { // projects it to the promotion hand-off); holds even if a direct caller // reached here without the dial-time `on_dial` write. machine.set_conn_transport_id(transport_id); - // Record our session index on the surviving carrier — the same index just - // written on the leg above — so the carrier is the single index home on - // the outbound path (the inbound path writes it at authorize). + // Record our session index on the surviving carrier — the same index + // just written on the connection above — so the carrier is the single + // index home on the outbound path (the inbound path writes it at + // authorize). machine.set_conn_our_index(our_index); - // Store the msg1 wire on the surviving carrier (the leg no longer holds - // the resend source); the retransmit driver reads it from here. + // Store the msg1 wire on the surviving carrier (the connection does not + // hold the resend source); the retransmit driver reads it from here. machine.set_conn_handshake_msg1(wire_msg1, first_resend_at_ms); - machine.set_leg(connection); Ok(()) } diff --git a/src/node/tests/acl.rs b/src/node/tests/acl.rs index 1ce5dbc..5a87260 100644 --- a/src/node/tests/acl.rs +++ b/src/node/tests/acl.rs @@ -56,7 +56,7 @@ async fn test_inbound_msg1_denied_by_acl() { node_b.reload_peer_acl().await; let peer_b_identity = PeerIdentity::from_pubkey_full(node_b.identity().pubkey_full()); - let mut conn_a = PeerConnection::outbound(LinkId::new(1), peer_b_identity, 1000); + let mut conn_a = outbound_leg(LinkId::new(1), peer_b_identity, 1000); let noise_msg1 = conn_a .start_handshake(node_a.identity().keypair(), node_a.startup_epoch(), 1000) .unwrap(); @@ -96,7 +96,8 @@ async fn test_outbound_msg2_denied_after_acl_reload() { let keypair_a = node_a.identity().keypair(); let epoch_a = node_a.startup_epoch(); let noise_msg1 = node_a - .get_connection_mut(&link_id_a) + .peer_machines + .get_mut(&link_id_a) .unwrap() .start_handshake(keypair_a, epoch_a, 1000) .unwrap(); @@ -116,7 +117,7 @@ async fn test_outbound_msg2_denied_after_acl_reload() { .pending_outbound .insert((transport_id, our_index_a.as_u32()), link_id_a); - let mut conn_b = PeerConnection::inbound(LinkId::new(2), 1000); + let mut conn_b = inbound_leg(LinkId::new(2), 1000); let responder_epoch = [0x11; 8]; let noise_msg2 = conn_b .receive_handshake_init( @@ -185,3 +186,49 @@ async fn test_outbound_connect_not_denied_by_allowlist_miss() { assert!(!matches!(result, Err(NodeError::AccessDenied(_)))); } + +/// The ACL-rejected arm of the same property the Noise-failure arm pins in +/// `unit.rs`: a msg1 that is admitted by the crypto but denied by the ACL +/// still leaves nothing behind. The control machine is built above the crypto +/// so it can drive the handshake, but it stays a local until a promote tail +/// inserts it, so a denial drops it. +#[tokio::test] +async fn test_acl_rejected_msg1_leaves_no_registry_trace() { + let (dir, mut node_b) = make_acl_node(); + let node_a = make_node(); + + std::fs::write(deny_path(&dir), format!("{}\n", node_a.npub())).unwrap(); + node_b.reload_peer_acl().await; + + let peer_b_identity = PeerIdentity::from_pubkey_full(node_b.identity().pubkey_full()); + let mut conn_a = outbound_leg(LinkId::new(1), peer_b_identity, 1000); + let noise_msg1 = conn_a + .start_handshake(node_a.identity().keypair(), node_a.startup_epoch(), 1000) + .unwrap(); + let wire_msg1 = build_msg1(SessionIndex::new(7), &noise_msg1); + let packet = ReceivedPacket::with_timestamp( + TransportId::new(1), + TransportAddr::from_string("127.0.0.1:5000"), + wire_msg1, + 1000, + ); + + node_b.handle_msg1(packet).await; + + assert!( + node_b.peer_machines.is_empty(), + "an ACL-denied msg1 must leave no control machine behind" + ); + assert_eq!(node_b.connection_count(), 0); + assert_eq!(node_b.peer_count(), 0); + assert_eq!(node_b.link_count(), 0); + assert!( + node_b.peers_by_index.is_empty(), + "an ACL-denied msg1 must allocate no session index" + ); + assert_eq!( + node_b.stats().handshake.bad_state, + 1, + "the denial is attributed to the handshake state-machine counter" + ); +} diff --git a/src/node/tests/establish_chartests.rs b/src/node/tests/establish_chartests.rs index 995e887..42fd01f 100644 --- a/src/node/tests/establish_chartests.rs +++ b/src/node/tests/establish_chartests.rs @@ -49,7 +49,7 @@ fn craft_msg1_wire( use crate::proto::fmp::wire::build_msg1; let peer_b_identity = PeerIdentity::from_pubkey_full(node.identity().pubkey_full()); let link_id = LinkId::new(0x0BAD_C0DE); - let mut conn = PeerConnection::outbound(link_id, peer_b_identity, ts); + let mut conn = outbound_leg(link_id, peer_b_identity, ts); let noise_msg1 = conn .start_handshake(sender.keypair(), epoch, ts) .expect("start_handshake produces noise msg1"); @@ -341,7 +341,8 @@ async fn chartest_msg1_inbound_promote_defers_pending_outbound_to_same_identity( let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let _ = node - .get_connection_mut(&out_link) + .peer_machines + .get_mut(&out_link) .unwrap() .start_handshake(our_keypair, startup_epoch, 1000) .unwrap(); @@ -421,7 +422,8 @@ async fn chartest_msg1_at_cap_with_pending_outbound_bypasses_early_gate() { let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let _ = node - .get_connection_mut(&out_link) + .peer_machines + .get_mut(&out_link) .unwrap() .start_handshake(our_keypair, startup_epoch, 1000) .unwrap(); @@ -524,7 +526,8 @@ async fn chartest_cross_connection_tiebreak_winner_and_loser() { let keypair_a = node_a.identity().keypair(); let epoch_a = node_a.startup_epoch(); let noise_msg1_a = node_a - .get_connection_mut(&link_a_out) + .peer_machines + .get_mut(&link_a_out) .unwrap() .start_handshake(keypair_a, epoch_a, 1000) .unwrap(); @@ -560,7 +563,8 @@ async fn chartest_cross_connection_tiebreak_winner_and_loser() { let keypair_b = node_b.identity().keypair(); let epoch_b = node_b.startup_epoch(); let noise_msg1_b = node_b - .get_connection_mut(&link_b_out) + .peer_machines + .get_mut(&link_b_out) .unwrap() .start_handshake(keypair_b, epoch_b, 1000) .unwrap(); diff --git a/src/node/tests/handshake.rs b/src/node/tests/handshake.rs index da5e087..dae87c6 100644 --- a/src/node/tests/handshake.rs +++ b/src/node/tests/handshake.rs @@ -70,7 +70,8 @@ async fn test_two_node_handshake_udp() { let our_keypair_a = node_a.identity().keypair(); let startup_epoch_a = node_a.startup_epoch(); let noise_msg1 = node_a - .get_connection_mut(&link_id_a) + .peer_machines + .get_mut(&link_id_a) .unwrap() .start_handshake(our_keypair_a, startup_epoch_a, 1000) .unwrap(); @@ -314,7 +315,8 @@ async fn test_run_rx_loop_handshake() { let our_keypair_a = node_a.identity().keypair(); let startup_epoch_a = node_a.startup_epoch(); let noise_msg1 = node_a - .get_connection_mut(&link_id_a) + .peer_machines + .get_mut(&link_id_a) .unwrap() .start_handshake(our_keypair_a, startup_epoch_a, 1000) .unwrap(); @@ -508,7 +510,8 @@ async fn test_cross_connection_both_initiate() { let our_keypair_a = node_a.identity().keypair(); let startup_epoch_a = node_a.startup_epoch(); let noise_msg1_a = node_a - .get_connection_mut(&link_id_a_out) + .peer_machines + .get_mut(&link_id_a_out) .unwrap() .start_handshake(our_keypair_a, startup_epoch_a, 1000) .unwrap(); @@ -544,7 +547,8 @@ async fn test_cross_connection_both_initiate() { let our_keypair_b = node_b.identity().keypair(); let startup_epoch_b = node_b.startup_epoch(); let noise_msg1_b = node_b - .get_connection_mut(&link_id_b_out) + .peer_machines + .get_mut(&link_id_b_out) .unwrap() .start_handshake(our_keypair_b, startup_epoch_b, 1000) .unwrap(); @@ -700,7 +704,8 @@ async fn test_stale_connection_cleanup() { let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let _noise_msg1 = node - .get_connection_mut(&link_id) + .peer_machines + .get_mut(&link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, past_time_ms) .unwrap(); @@ -783,7 +788,8 @@ async fn test_failed_connection_cleanup() { let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let _noise_msg1 = node - .get_connection_mut(&link_id) + .peer_machines + .get_mut(&link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, now_ms) .unwrap(); @@ -853,24 +859,26 @@ async fn test_msg1_stored_for_resend() { .map(|d| d.as_millis() as u64) .unwrap_or(0); let link_id = node.allocate_link_id(); - let mut conn = PeerConnection::outbound(link_id, peer_identity, now_ms); + let mut conn = outbound_leg(link_id, peer_identity, now_ms); let our_index = node.index_allocator.allocate().unwrap(); let our_keypair = node.identity().keypair(); let noise_msg1 = conn .start_handshake(our_keypair, node.startup_epoch(), now_ms) .unwrap(); - conn.set_our_index(our_index); - conn.set_transport_id(transport_id); - conn.set_source_addr(remote_addr.clone()); + conn.leg_mut().unwrap().set_our_index(our_index); + conn.leg_mut().unwrap().set_transport_id(transport_id); + conn.leg_mut().unwrap().set_source_addr(remote_addr.clone()); // Build wire msg1 and store it (as initiate_peer_connection does) let wire_msg1 = build_msg1(our_index, &noise_msg1); let resend_interval = node.config().node.rate_limit.handshake_resend_interval_ms; - conn.set_handshake_msg1(wire_msg1.clone(), now_ms + resend_interval); + conn.leg_mut() + .unwrap() + .set_handshake_msg1(wire_msg1.clone(), now_ms + resend_interval); // Verify stored msg1 matches what was built - assert_eq!(conn.handshake_msg1().unwrap(), &wire_msg1); + assert_eq!(conn.leg().unwrap().handshake_msg1().unwrap(), &wire_msg1); } /// Test that resend scheduling respects max_resends and backoff. @@ -884,20 +892,22 @@ async fn test_resend_scheduling() { let now_ms = 100_000u64; // Use a fixed time for predictable testing let link_id = node.allocate_link_id(); - let mut conn = PeerConnection::outbound(link_id, peer_identity, now_ms); + let mut conn = outbound_leg(link_id, peer_identity, now_ms); let our_index = node.index_allocator.allocate().unwrap(); let our_keypair = node.identity().keypair(); let noise_msg1 = conn .start_handshake(our_keypair, node.startup_epoch(), now_ms) .unwrap(); - conn.set_our_index(our_index); - conn.set_transport_id(transport_id); - conn.set_source_addr(remote_addr.clone()); + conn.leg_mut().unwrap().set_our_index(our_index); + conn.leg_mut().unwrap().set_transport_id(transport_id); + conn.leg_mut().unwrap().set_source_addr(remote_addr.clone()); // Store msg1 with first resend at now + 1000ms let wire_msg1 = crate::proto::fmp::wire::build_msg1(our_index, &noise_msg1); - conn.set_handshake_msg1(wire_msg1.clone(), now_ms + 1000); + conn.leg_mut() + .unwrap() + .set_handshake_msg1(wire_msg1.clone(), now_ms + 1000); let link = Link::connectionless( link_id, @@ -931,7 +941,7 @@ async fn test_resend_scheduling() { // The msg1 wire lives on the machine's carrier (the retransmit driver's // resend source), mirroring `prepare_outbound_msg1`. machine.set_conn_handshake_msg1(wire_msg1, now_ms + 1000); - machine.set_leg(conn); + machine.set_leg(conn.take_leg().unwrap()); node.peer_machines.insert(link_id, machine); node.peer_timers.entry(link_id).or_default().insert( crate::peer::machine::TimerKind::HandshakeRetransmit, @@ -970,15 +980,15 @@ async fn test_handshake_timeout_drive() { let dial_ms = 1000u64; let link_id = node.allocate_link_id(); - let mut conn = PeerConnection::outbound(link_id, peer_identity, dial_ms); + let mut conn = outbound_leg(link_id, peer_identity, dial_ms); let our_index = node.index_allocator.allocate().unwrap(); let our_keypair = node.identity().keypair(); let _ = conn .start_handshake(our_keypair, node.startup_epoch(), dial_ms) .unwrap(); - conn.set_our_index(our_index); - conn.set_transport_id(transport_id); - conn.set_source_addr(remote_addr.clone()); + conn.leg_mut().unwrap().set_our_index(our_index); + conn.leg_mut().unwrap().set_transport_id(transport_id); + conn.leg_mut().unwrap().set_source_addr(remote_addr.clone()); let link = Link::connectionless( link_id, @@ -1007,7 +1017,7 @@ async fn test_handshake_timeout_drive() { dial_ms, &mut node.index_allocator, ); - machine.set_leg(conn); + machine.set_leg(conn.take_leg().unwrap()); node.peer_machines.insert(link_id, machine); node.peer_timers.entry(link_id).or_default().insert( crate::peer::machine::TimerKind::HandshakeTimeout, diff --git a/src/node/tests/mod.rs b/src/node/tests/mod.rs index dc75785..a302e26 100644 --- a/src/node/tests/mod.rs +++ b/src/node/tests/mod.rs @@ -84,6 +84,30 @@ pub(super) fn make_peer_identity() -> PeerIdentity { PeerIdentity::from_pubkey(identity.pubkey()) } +/// A control machine carrying a fresh outbound connection, for tests that +/// drive one end of a handshake without a whole node behind it. The machine +/// owns the handshake operations, so it is what the crypto runs on. +pub(super) fn outbound_leg( + link_id: LinkId, + expected_identity: PeerIdentity, + current_time_ms: u64, +) -> PeerMachine { + let mut machine = PeerMachine::new_outbound(link_id, expected_identity, current_time_ms); + machine.set_leg(PeerConnection::outbound( + link_id, + expected_identity, + current_time_ms, + )); + machine +} + +/// The responder twin of [`outbound_leg`]. +pub(super) fn inbound_leg(link_id: LinkId, current_time_ms: u64) -> PeerMachine { + let mut machine = PeerMachine::new_inbound(link_id, current_time_ms); + machine.set_leg(PeerConnection::inbound(link_id, current_time_ms)); + machine +} + /// Seed a control machine whose leg carries a completed Noise IK handshake. /// /// Returns the peer identity. The leg is outbound, in Complete state, with @@ -138,13 +162,14 @@ pub(super) fn seed_completed_connection_with( let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let msg1 = node - .get_connection_mut(&link_id) + .peer_machines + .get_mut(&link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, current_time_ms) .unwrap(); // Run responder side to generate msg2 - let mut resp_conn = PeerConnection::inbound(LinkId::new(999), current_time_ms); + let mut resp_conn = inbound_leg(LinkId::new(999), current_time_ms); let peer_keypair = peer_identity_full.keypair(); let mut resp_epoch = [0u8; 8]; rand::Rng::fill_bytes(&mut rand::rng(), &mut resp_epoch); @@ -153,7 +178,8 @@ pub(super) fn seed_completed_connection_with( .unwrap(); // Complete initiator handshake - node.get_connection_mut(&link_id) + node.peer_machines + .get_mut(&link_id) .unwrap() .complete_handshake(&msg2, current_time_ms) .unwrap(); diff --git a/src/node/tests/spanning_tree.rs b/src/node/tests/spanning_tree.rs index cb6e8eb..84af809 100644 --- a/src/node/tests/spanning_tree.rs +++ b/src/node/tests/spanning_tree.rs @@ -131,7 +131,8 @@ pub(super) async fn initiate_handshake(nodes: &mut [TestNode], i: usize, j: usiz let startup_epoch = initiator.node.startup_epoch(); let noise_msg1 = initiator .node - .get_connection_mut(&link_id) + .peer_machines + .get_mut(&link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, 1000) .unwrap(); diff --git a/src/node/tests/unit.rs b/src/node/tests/unit.rs index da4b34a..adccce9 100644 --- a/src/node/tests/unit.rs +++ b/src/node/tests/unit.rs @@ -808,7 +808,8 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() { let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let _msg1 = node - .get_connection_mut(&pending_link_id) + .peer_machines + .get_mut(&pending_link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, pending_time_ms) .unwrap(); @@ -850,13 +851,14 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() { let our_keypair = node.identity().keypair(); let startup_epoch = node.startup_epoch(); let msg1 = node - .get_connection_mut(&completing_link_id) + .peer_machines + .get_mut(&completing_link_id) .unwrap() .start_handshake(our_keypair, startup_epoch, completing_time_ms) .unwrap(); // B responds - let mut resp_conn = PeerConnection::inbound(LinkId::new(999), completing_time_ms); + let mut resp_conn = inbound_leg(LinkId::new(999), completing_time_ms); let peer_keypair = peer_b_full.keypair(); let mut resp_epoch = [0u8; 8]; rand::Rng::fill_bytes(&mut rand::rng(), &mut resp_epoch); @@ -864,7 +866,8 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() { .receive_handshake_init(peer_keypair, resp_epoch, &msg1, completing_time_ms) .unwrap(); - node.get_connection_mut(&completing_link_id) + node.peer_machines + .get_mut(&completing_link_id) .unwrap() .complete_handshake(&msg2, completing_time_ms) .unwrap(); @@ -2083,7 +2086,7 @@ async fn craft_and_send_msg1( let sender_node_addr = *sender_pubkey_id.node_addr(); let link_id = LinkId::new(0xDEAD_BEEF); - let mut conn = PeerConnection::outbound(link_id, peer_b_identity, timestamp_ms); + let mut conn = outbound_leg(link_id, peer_b_identity, timestamp_ms); let sender_keypair = sender_identity.keypair(); let mut startup_epoch = [0u8; 8]; @@ -2345,3 +2348,130 @@ async fn start_skips_system_tun_when_app_owned() { node.stop().await.unwrap(); } + +/// A connection whose handshake failed is retained with BOTH Noise handles +/// empty, and the stale-connection sweep depends on that: presence of the +/// pending connection — not presence of a handle — is what marks a machine as +/// handshake-phase. If presence were ever derived from the handles, every +/// failed connection would become invisible to the sweep and leak forever, +/// holding a peering-budget slot and a wrong `connection_count` permanently. +#[test] +fn test_failed_connection_is_retained_and_reaped() { + use crate::proto::fmp::LifecycleView; + + let mut node = make_node(); + let link_id = LinkId::new(1); + let peer_identity = make_peer_identity(); + + node.seed_handshake_machine(HandshakeSeed::outbound(link_id, peer_identity, 1000)) + .unwrap(); + let our_keypair = node.identity().keypair(); + let startup_epoch = node.startup_epoch(); + node.peer_machines + .get_mut(&link_id) + .unwrap() + .start_handshake(our_keypair, startup_epoch, 1000) + .unwrap(); + + // The send of that stored initiation fails. + let machine = node.peer_machines.get_mut(&link_id).unwrap(); + machine.mark_failed(); + machine.mark_send_failed(); + + // Both handles are now empty — the initiation handle was dropped and no + // session was ever reached — yet the connection is deliberately retained. + let leg = node.peer_machines.get(&link_id).unwrap().leg().unwrap(); + assert!( + leg.noise_handshake.is_none() && leg.noise_session.is_none(), + "a failed connection holds neither handle" + ); + + // (a) it still counts while it waits for the sweep + assert_eq!( + node.connection_count(), + 1, + "a failed connection stays counted until it is reaped" + ); + + // (b) the sweep yields it + let stale = node.stale_connections(2000, 30_000); + assert_eq!( + stale.len(), + 1, + "the sweep must see a failed connection despite its empty handles" + ); + assert_eq!(stale[0].link, link_id); + + // (c) reaping it clears the carrier + node.remove_peer_machine(link_id); + assert_eq!(node.connection_count(), 0); +} + +/// A msg1 that fails Noise processing must leave no trace in the registry. +/// The control machine is built above the crypto so it can drive the +/// handshake, but it stays a local until a promote tail inserts it — a +/// rejected msg1 drops it. +#[tokio::test] +async fn test_rejected_msg1_leaves_no_registry_trace() { + let mut node = make_node(); + + // Well-formed framing, garbage Noise payload: processing fails. + let wire_msg1 = crate::proto::fmp::wire::build_msg1( + SessionIndex::new(7), + &[0u8; crate::noise::HANDSHAKE_MSG1_SIZE], + ); + let packet = ReceivedPacket::with_timestamp( + TransportId::new(1), + TransportAddr::from_string("127.0.0.1:5000"), + wire_msg1, + 1000, + ); + + node.handle_msg1(packet).await; + + assert!( + node.peer_machines.is_empty(), + "a rejected msg1 must leave no control machine behind" + ); + assert_eq!(node.connection_count(), 0); + assert_eq!(node.peer_count(), 0); + assert_eq!(node.link_count(), 0); + assert!( + node.peers_by_index.is_empty(), + "a rejected msg1 must allocate no session index" + ); + assert_eq!( + node.stats().handshake.bad_state, + 1, + "the rejection is attributed to the handshake state-machine counter" + ); +} + +/// The outbound path registers its control machine at dial, before msg1 is +/// prepared, so a preparation failure has to unwind that registration rather +/// than drop a local. +#[tokio::test] +async fn test_failed_msg1_preparation_unwinds_the_dial_machine() { + let mut node = make_node(); + let link_id = LinkId::new(1); + let transport_id = TransportId::new(1); + let remote_addr = TransportAddr::from_string("127.0.0.1:5000"); + let peer_identity = make_peer_identity(); + + // Stand in for the dial: the machine exists before msg1 is prepared. + node.peer_machines.insert( + link_id, + PeerMachine::new_outbound(link_id, peer_identity, 1000), + ); + // Force the index allocation inside msg1 preparation to fail. + node.index_allocator = crate::utils::index::IndexAllocator::with_max_attempts(0); + + let result = node.prepare_outbound_msg1(link_id, transport_id, &remote_addr, peer_identity); + + assert!(matches!(result, Err(NodeError::IndexAllocationFailed(_)))); + assert!( + !node.peer_machines.contains_key(&link_id), + "a failed msg1 preparation must unwind the dial-time machine" + ); + assert_eq!(node.connection_count(), 0); +} diff --git a/src/peer/connection.rs b/src/peer/connection.rs index ea689bd..6eb00d0 100644 --- a/src/peer/connection.rs +++ b/src/peer/connection.rs @@ -2,17 +2,16 @@ //! //! Represents an in-progress connection before authentication completes. //! PeerConnection tracks the Noise IK handshake and transitions to -//! ActivePeer upon successful authentication. The handshake *phase* (initial / -//! sent_msg1 / complete / failed) is no longer tracked here — it lives on the -//! per-peer control machine; the leg's crypto methods gate on the presence of -//! their Noise handles (`noise_handshake` / `noise_session`) directly. +//! ActivePeer upon successful authentication. Neither the handshake *phase* +//! (initial / sent_msg1 / complete / failed) nor the handshake operations are +//! tracked here — both live on the per-peer control machine, which drives the +//! Noise handles held below. use crate::PeerIdentity; -use crate::noise::{self, NoiseError, NoiseSession}; +use crate::noise::{self, NoiseSession}; use crate::proto::fmp::ConnectionState; use crate::transport::{LinkDirection, LinkId, TransportAddr, TransportId}; use crate::utils::index::SessionIndex; -use secp256k1::Keypair; use std::fmt; /// A connection in the handshake phase, before authentication completes. @@ -23,17 +22,21 @@ use std::fmt; /// This is the shell holder for the FMP crypto/state split: the pure /// connection bookkeeping lives in [`ConnectionState`] (`proto::fmp::state`), /// and the two Noise crypto handles stay here beside it. Pure public methods -/// delegate to `self.state`; the XX transition methods drive the crypto and -/// write results back through `self.state`'s setters. +/// delegate to `self.state`; the control machine drives the handles and +/// records each result here and on its own carrier. pub struct PeerConnection { /// Pure, runtime-agnostic connection bookkeeping. state: ConnectionState, /// Noise handshake state (consumes on completion). - noise_handshake: Option, + /// + /// Driven by the control machine, which owns the handshake operations. + pub(crate) noise_handshake: Option, /// Completed Noise session (available after handshake complete). - noise_session: Option, + /// + /// Driven by the control machine, which owns the handshake operations. + pub(crate) noise_session: Option, } impl PeerConnection { @@ -202,146 +205,13 @@ impl PeerConnection { self.state.handshake_msg2() } - // === Noise Handshake Operations (shell: drives crypto, updates pure state) === + // === Crypto handle plumbing (the control machine drives the handshake) === - /// Start the handshake as initiator and generate message 1. - /// - /// For outbound connections only. Returns the handshake message to send. - /// The epoch is our startup epoch, encrypted into msg1 for restart detection. - pub fn start_handshake( - &mut self, - our_keypair: Keypair, - epoch: [u8; 8], - current_time_ms: u64, - ) -> Result, NoiseError> { - if self.state.direction() != LinkDirection::Outbound { - return Err(NoiseError::WrongState { - expected: "outbound connection".to_string(), - got: "inbound connection".to_string(), - }); - } - - let remote_static = self - .state - .expected_identity() - .expect("outbound must have expected identity") - .pubkey_full(); - - let mut hs = noise::HandshakeState::new_initiator(our_keypair, remote_static); - hs.set_local_epoch(epoch); - let msg1 = hs.write_message_1()?; - - self.noise_handshake = Some(hs); - self.state.touch(current_time_ms); - - Ok(msg1) - } - - /// Initialize responder and process incoming message 1. - /// - /// For inbound connections only. Returns the handshake message 2 to send. - /// The epoch is our startup epoch, encrypted into msg2 for restart detection. - pub fn receive_handshake_init( - &mut self, - our_keypair: Keypair, - epoch: [u8; 8], - message: &[u8], - current_time_ms: u64, - ) -> Result, NoiseError> { - if self.state.direction() != LinkDirection::Inbound { - return Err(NoiseError::WrongState { - expected: "inbound connection".to_string(), - got: "outbound connection".to_string(), - }); - } - - let mut hs = noise::HandshakeState::new_responder(our_keypair); - hs.set_local_epoch(epoch); - - // Process message 1 (this reveals the initiator's identity and epoch) - hs.read_message_1(message)?; - - // Extract the discovered identity from the crypto and record it as - // pure data on the state. - let remote_static = *hs - .remote_static() - .expect("remote static available after msg1"); - self.state - .set_expected_identity(PeerIdentity::from_pubkey_full(remote_static)); - - // Capture remote epoch from msg1 - self.state.set_remote_epoch(hs.remote_epoch()); - - // Generate message 2 - let msg2 = hs.write_message_2()?; - - // Handshake is complete for responder - let session = hs.into_session()?; - self.noise_session = Some(session); - self.state.touch(current_time_ms); - - Ok(msg2) - } - - /// Complete the handshake by processing message 2. - /// - /// For outbound connections only (initiator completing handshake). - pub fn complete_handshake( - &mut self, - message: &[u8], - current_time_ms: u64, - ) -> Result<(), NoiseError> { - // The leg is at `SentMsg1` iff its Noise handshake handle is present - // (set by `start_handshake`, taken here on completion). Gate on the - // handle directly now that the phase enum is gone — byte-equivalent to - // the old `!= SentMsg1` guard for every reachable transition. - if self.noise_handshake.is_none() { - return Err(NoiseError::WrongState { - expected: "sent_msg1 state".to_string(), - got: "no active handshake".to_string(), - }); - } - - let mut hs = self - .noise_handshake - .take() - .expect("noise handshake must exist in SentMsg1 state"); - - hs.read_message_2(message)?; - - // Capture remote epoch from msg2 - self.state.set_remote_epoch(hs.remote_epoch()); - - let session = hs.into_session()?; - self.noise_session = Some(session); - self.state.touch(current_time_ms); - - Ok(()) - } - - /// Take the completed Noise session. - /// - /// Returns the NoiseSession for use in ActivePeer. Can only be called - /// once after handshake completes. - pub fn take_session(&mut self) -> Option { - // The session exists iff the handshake reached `Complete`, so taking it - // unconditionally is byte-equivalent to the old `== Complete` gate. - self.noise_session.take() - } - - /// Check if we have a completed session ready to take. - pub fn has_session(&self) -> bool { - self.noise_session.is_some() - } - - // === State Transitions (for manual control if needed) === - - /// Drop the shell-owned crypto handshake handle. The failure *state* now - /// lives on the control machine (`PeerMachine`); this only releases the - /// leg's Noise handle at the identical point it was released before, so a - /// subsequent `complete_handshake` on this leg still reports `WrongState`. - pub fn mark_failed(&mut self) { - self.noise_handshake = None; + /// Mutable access to the pure bookkeeping, so the control machine's + /// handshake operations can record their results here as well as on the + /// surviving carrier. + pub(crate) fn state_mut(&mut self) -> &mut ConnectionState { + &mut self.state } // === Validation === @@ -373,108 +243,12 @@ impl fmt::Debug for PeerConnection { mod tests { use super::*; use crate::Identity; - use rand::Rng; fn make_peer_identity() -> PeerIdentity { let identity = Identity::generate(); PeerIdentity::from_pubkey(identity.pubkey()) } - fn make_keypair() -> Keypair { - let identity = Identity::generate(); - identity.keypair() - } - - fn make_epoch() -> [u8; 8] { - let mut epoch = [0u8; 8]; - rand::rng().fill_bytes(&mut epoch); - epoch - } - - #[test] - fn test_outbound_connection() { - let identity = make_peer_identity(); - let conn = PeerConnection::outbound(LinkId::new(1), identity, 1000); - - assert!(conn.is_outbound()); - assert!(!conn.is_inbound()); - assert!(!conn.has_session()); - assert!(conn.expected_identity().is_some()); - assert_eq!(conn.started_at(), 1000); - } - - #[test] - fn test_inbound_connection() { - let conn = PeerConnection::inbound(LinkId::new(2), 2000); - - assert!(conn.is_inbound()); - assert!(!conn.is_outbound()); - assert!(!conn.has_session()); - assert!(conn.expected_identity().is_none()); - assert_eq!(conn.started_at(), 2000); - } - - #[test] - fn test_full_handshake_flow() { - // Create identities - let initiator_identity = Identity::generate(); - let responder_identity = Identity::generate(); - - let initiator_keypair = initiator_identity.keypair(); - let responder_keypair = responder_identity.keypair(); - let initiator_epoch = make_epoch(); - let responder_epoch = make_epoch(); - - // Use from_pubkey_full to preserve parity for ECDH - let responder_peer_id = PeerIdentity::from_pubkey_full(responder_identity.pubkey_full()); - - // Create connections - let mut initiator_conn = PeerConnection::outbound(LinkId::new(1), responder_peer_id, 1000); - let mut responder_conn = PeerConnection::inbound(LinkId::new(2), 1000); - - // Initiator starts handshake - let msg1 = initiator_conn - .start_handshake(initiator_keypair, initiator_epoch, 1100) - .unwrap(); - // Post-msg1 the initiator holds an in-flight handshake, not yet a session. - assert!(!initiator_conn.has_session()); - - // Responder processes msg1 and sends msg2 - let msg2 = responder_conn - .receive_handshake_init(responder_keypair, responder_epoch, &msg1, 1200) - .unwrap(); - // The IK responder completes in one step: it now holds a session. - assert!(responder_conn.has_session()); - - // Responder learned initiator's identity - let discovered = responder_conn.expected_identity().unwrap(); - assert_eq!(discovered.pubkey(), initiator_identity.pubkey()); - - // Responder learned initiator's epoch - assert_eq!(responder_conn.remote_epoch(), Some(initiator_epoch)); - - // Initiator completes handshake - initiator_conn.complete_handshake(&msg2, 1300).unwrap(); - assert!(initiator_conn.has_session()); - - // Initiator learned responder's epoch - assert_eq!(initiator_conn.remote_epoch(), Some(responder_epoch)); - - // Both have sessions - assert!(initiator_conn.has_session()); - assert!(responder_conn.has_session()); - - // Take and verify sessions work - let mut init_session = initiator_conn.take_session().unwrap(); - let mut resp_session = responder_conn.take_session().unwrap(); - - // Encrypt/decrypt test - let plaintext = b"test message"; - let ciphertext = init_session.encrypt(plaintext).unwrap(); - let decrypted = resp_session.decrypt(&ciphertext).unwrap(); - assert_eq!(decrypted, plaintext); - } - #[test] fn test_connection_timing() { let identity = make_peer_identity(); @@ -485,43 +259,4 @@ mod tests { assert!(!conn.is_timed_out(1500, 1000)); assert!(conn.is_timed_out(2500, 1000)); } - - #[test] - fn test_connection_failure() { - // `mark_failed` releases the leg's Noise handshake handle. The failure - // *state* now lives on the control machine, but the leg-local effect is - // still observable: a completion attempt afterward reports `WrongState` - // (the handle-presence gate) and no session is produced. - let identity = make_peer_identity(); - let keypair = make_keypair(); - let mut conn = PeerConnection::outbound(LinkId::new(1), identity, 1000); - conn.start_handshake(keypair, make_epoch(), 1100).unwrap(); - - conn.mark_failed(); - - assert!(!conn.has_session()); - assert!(conn.complete_handshake(&[0u8; 96], 1200).is_err()); - } - - #[test] - fn test_wrong_direction_errors() { - let identity = make_peer_identity(); - let keypair = make_keypair(); - - // Outbound can't receive_handshake_init - let mut outbound = PeerConnection::outbound(LinkId::new(1), identity, 1000); - assert!( - outbound - .receive_handshake_init(keypair, make_epoch(), &[0u8; 106], 1100) - .is_err() - ); - - // Inbound can't start_handshake - let mut inbound = PeerConnection::inbound(LinkId::new(2), 1000); - assert!( - inbound - .start_handshake(keypair, make_epoch(), 1100) - .is_err() - ); - } } diff --git a/src/peer/machine.rs b/src/peer/machine.rs index 91941ea..6c3157c 100644 --- a/src/peer/machine.rs +++ b/src/peer/machine.rs @@ -56,6 +56,7 @@ #![allow(dead_code)] +use crate::noise::{self, NoiseError, NoiseSession}; use crate::peer::PeerConnection; use crate::proto::fmp::{ ConnAction, ConnSnapshot, ConnectionState, EstablishSnapshot, Fmp, InboundDecision, @@ -63,9 +64,10 @@ use crate::proto::fmp::{ RekeyResendSnapshot, WireOutcome, }; use crate::proto::link::LinkMessageType; -use crate::transport::{LinkId, LinkStats, TransportAddr, TransportId}; +use crate::transport::{LinkDirection, LinkId, LinkStats, TransportAddr, TransportId}; use crate::utils::index::{IndexAllocator, SessionIndex}; use crate::{NodeAddr, PeerIdentity}; +use secp256k1::Keypair; // ============================================================================ // Timing placeholders @@ -406,6 +408,15 @@ pub(crate) enum PeerAction { // The machine (control tier) // ============================================================================ +/// The handshake operations are only reachable while a pending connection is +/// attached; every path that drives the crypto attaches it first. +fn no_pending_connection() -> NoiseError { + NoiseError::WrongState { + expected: "attached connection".to_string(), + got: "no connection".to_string(), + } +} + /// Per-peer control FSM. Holds control-tier lifecycle state only; the /// send-critical state is published as `PeerSendState` and mutated via the /// emitted [`PeerAction`]s. @@ -413,11 +424,12 @@ pub(crate) struct PeerMachine { state: PeerState, link: LinkId, identity: Option, - /// The pending handshake connection this machine owns while the leg is in - /// the handshake window. `None` before the connection is built (the dial + /// The pending handshake connection this machine owns while it is in the + /// handshake window. `None` before the connection is built (the dial /// window) and after promotion consumes it (the machine survives as the - /// active peer's control machine). Pure storage — the machine never reads - /// or drives it; the shell reaches it through the accessors below. + /// active peer's control machine). Its bookkeeping is storage the shell + /// reaches through the accessors below; its Noise handles are driven by + /// this machine's handshake operations. leg: Option, /// Pure handshake-phase bookkeeping (link/direction/indices/transport/ /// stored handshake bytes/epoch). Reused verbatim from the FMP state core. @@ -542,6 +554,181 @@ impl PeerMachine { self.leg = Some(leg); } + // === Noise handshake operations === + // + // Mechanism, not decision: these are called by the shell, are never + // reached from `step()`, and no event triggers them. Each drives the + // Noise crypto on the pending connection and records the results on both + // that connection and the surviving carrier, so a reader of either sees + // the same value at the same point. + + /// Start the handshake as initiator and generate message 1. + /// + /// For outbound connections only. Returns the handshake message to send. + /// The epoch is our startup epoch, encrypted into msg1 for restart detection. + pub(crate) fn start_handshake( + &mut self, + our_keypair: Keypair, + epoch: [u8; 8], + current_time_ms: u64, + ) -> Result, NoiseError> { + let msg1 = { + let leg = self.leg.as_mut().ok_or_else(no_pending_connection)?; + + if leg.direction() != LinkDirection::Outbound { + return Err(NoiseError::WrongState { + expected: "outbound connection".to_string(), + got: "inbound connection".to_string(), + }); + } + + let remote_static = leg + .expected_identity() + .expect("outbound must have expected identity") + .pubkey_full(); + + let mut hs = noise::HandshakeState::new_initiator(our_keypair, remote_static); + hs.set_local_epoch(epoch); + let msg1 = hs.write_message_1()?; + + leg.noise_handshake = Some(hs); + leg.state_mut().touch(current_time_ms); + + msg1 + }; + self.conn.touch(current_time_ms); + + Ok(msg1) + } + + /// Initialize responder and process incoming message 1. + /// + /// For inbound connections only. Returns the handshake message 2 to send. + /// The epoch is our startup epoch, encrypted into msg2 for restart detection. + pub(crate) fn receive_handshake_init( + &mut self, + our_keypair: Keypair, + epoch: [u8; 8], + message: &[u8], + current_time_ms: u64, + ) -> Result, NoiseError> { + let (msg2, learned_identity, remote_epoch) = { + let leg = self.leg.as_mut().ok_or_else(no_pending_connection)?; + + if leg.direction() != LinkDirection::Inbound { + return Err(NoiseError::WrongState { + expected: "inbound connection".to_string(), + got: "outbound connection".to_string(), + }); + } + + let mut hs = noise::HandshakeState::new_responder(our_keypair); + hs.set_local_epoch(epoch); + + // Process message 1 (this reveals the initiator's identity and epoch) + hs.read_message_1(message)?; + + // Extract the discovered identity from the crypto and record it as + // pure data on the state. + let remote_static = *hs + .remote_static() + .expect("remote static available after msg1"); + let learned_identity = PeerIdentity::from_pubkey_full(remote_static); + leg.state_mut().set_expected_identity(learned_identity); + + // Capture remote epoch from msg1 + let remote_epoch = hs.remote_epoch(); + leg.state_mut().set_remote_epoch(remote_epoch); + + // Generate message 2 + let msg2 = hs.write_message_2()?; + + // Handshake is complete for responder + let session = hs.into_session()?; + leg.noise_session = Some(session); + leg.state_mut().touch(current_time_ms); + + (msg2, learned_identity, remote_epoch) + }; + self.conn.set_expected_identity(learned_identity); + self.conn.set_remote_epoch(remote_epoch); + self.conn.touch(current_time_ms); + + Ok(msg2) + } + + /// Complete the handshake by processing message 2. + /// + /// For outbound connections only (initiator completing handshake). + pub(crate) fn complete_handshake( + &mut self, + message: &[u8], + current_time_ms: u64, + ) -> Result<(), NoiseError> { + let remote_epoch = { + let leg = self.leg.as_mut().ok_or_else(no_pending_connection)?; + + // The connection is at `SentMsg1` iff its Noise handshake handle is + // present (set by `start_handshake`, taken here on completion). + // Gating on the handle directly is byte-equivalent to the old + // `!= SentMsg1` guard for every reachable transition. + if leg.noise_handshake.is_none() { + return Err(NoiseError::WrongState { + expected: "sent_msg1 state".to_string(), + got: "no active handshake".to_string(), + }); + } + + let mut hs = leg + .noise_handshake + .take() + .expect("noise handshake must exist in SentMsg1 state"); + + hs.read_message_2(message)?; + + // Capture remote epoch from msg2 + let remote_epoch = hs.remote_epoch(); + leg.state_mut().set_remote_epoch(remote_epoch); + + let session = hs.into_session()?; + leg.noise_session = Some(session); + leg.state_mut().touch(current_time_ms); + + remote_epoch + }; + self.conn.set_remote_epoch(remote_epoch); + self.conn.touch(current_time_ms); + + Ok(()) + } + + /// Take the completed Noise session. + /// + /// Returns the NoiseSession for use in ActivePeer. Can only be called + /// once after the handshake completes. + pub(crate) fn take_session(&mut self) -> Option { + // The session exists iff the handshake reached `Complete`, so taking it + // unconditionally is byte-equivalent to the old `== Complete` gate. + self.leg.as_mut().and_then(|leg| leg.noise_session.take()) + } + + /// Check if we have a completed session ready to take. + pub(crate) fn has_session(&self) -> bool { + self.leg + .as_ref() + .is_some_and(|leg| leg.noise_session.is_some()) + } + + /// Drop the crypto handshake handle. The failure *state* lives on this + /// machine; this only releases the Noise handle at the identical point it + /// was released before, so a subsequent `complete_handshake` still reports + /// `WrongState`. + pub(crate) fn mark_failed(&mut self) { + if let Some(leg) = self.leg.as_mut() { + leg.noise_handshake = None; + } + } + /// The session index we allocated for this peer, read from the surviving /// carrier. Populated once the index is allocated on either establish path /// (inbound at `on_authorized`, outbound at msg1 preparation). `None` before @@ -861,11 +1048,9 @@ impl PeerMachine { /// (`is_handshaking_sent_msg1`) survives until the sweep, and no timer /// actions are emitted. fn on_handshake_send_failed(&mut self) -> Vec { - if let Some(leg) = self.leg.as_mut() { - // Drop the leg's Noise handshake handle at the identical point as - // before; the failure *state* is recorded on the machine. - leg.mark_failed(); - } + // Drop the Noise handshake handle at the identical point as before; + // the failure *state* is recorded on the machine. + self.mark_failed(); self.send_failed = true; Vec::new() } @@ -2970,6 +3155,176 @@ mod tests { assert_eq!(m.state(), PeerState::Active { addr }); assert_eq!(alloc.count(), 0); } + + // ---- Moved from `PeerConnection`'s own test module --------------------- + // These exercise the Noise handshake operations, which now live on the + // control machine. Constructed as a machine with a leg attached; the + // assertions are unchanged. + + fn make_peer_identity() -> PeerIdentity { + let identity = Identity::generate(); + PeerIdentity::from_pubkey(identity.pubkey()) + } + + fn make_keypair() -> Keypair { + let identity = Identity::generate(); + identity.keypair() + } + + fn make_epoch() -> [u8; 8] { + let mut epoch = [0u8; 8]; + rand::Rng::fill_bytes(&mut rand::rng(), &mut epoch); + epoch + } + + fn outbound_leg( + link_id: LinkId, + expected_identity: PeerIdentity, + current_time_ms: u64, + ) -> PeerMachine { + let mut machine = PeerMachine::new_outbound(link_id, expected_identity, current_time_ms); + machine.set_leg(PeerConnection::outbound( + link_id, + expected_identity, + current_time_ms, + )); + machine + } + + fn inbound_leg(link_id: LinkId, current_time_ms: u64) -> PeerMachine { + let mut machine = PeerMachine::new_inbound(link_id, current_time_ms); + machine.set_leg(PeerConnection::inbound(link_id, current_time_ms)); + machine + } + + #[test] + fn test_outbound_connection() { + let identity = make_peer_identity(); + let conn = outbound_leg(LinkId::new(1), identity, 1000); + + assert!(conn.leg().unwrap().is_outbound()); + assert!(!conn.leg().unwrap().is_inbound()); + assert!(!conn.has_session()); + assert!(conn.leg().unwrap().expected_identity().is_some()); + assert_eq!(conn.leg().unwrap().started_at(), 1000); + } + + #[test] + fn test_inbound_connection() { + let conn = inbound_leg(LinkId::new(2), 2000); + + assert!(conn.leg().unwrap().is_inbound()); + assert!(!conn.leg().unwrap().is_outbound()); + assert!(!conn.has_session()); + assert!(conn.leg().unwrap().expected_identity().is_none()); + assert_eq!(conn.leg().unwrap().started_at(), 2000); + } + + #[test] + fn test_full_handshake_flow() { + // Create identities + let initiator_identity = Identity::generate(); + let responder_identity = Identity::generate(); + + let initiator_keypair = initiator_identity.keypair(); + let responder_keypair = responder_identity.keypair(); + let initiator_epoch = make_epoch(); + let responder_epoch = make_epoch(); + + // Use from_pubkey_full to preserve parity for ECDH + let responder_peer_id = PeerIdentity::from_pubkey_full(responder_identity.pubkey_full()); + + // Create connections + let mut initiator_conn = outbound_leg(LinkId::new(1), responder_peer_id, 1000); + let mut responder_conn = inbound_leg(LinkId::new(2), 1000); + + // Initiator starts handshake + let msg1 = initiator_conn + .start_handshake(initiator_keypair, initiator_epoch, 1100) + .unwrap(); + // Post-msg1 the initiator holds an in-flight handshake, not yet a session. + assert!(!initiator_conn.has_session()); + + // Responder processes msg1 and sends msg2 + let msg2 = responder_conn + .receive_handshake_init(responder_keypair, responder_epoch, &msg1, 1200) + .unwrap(); + // The IK responder completes in one step: it now holds a session. + assert!(responder_conn.has_session()); + + // Responder learned initiator's identity + let discovered = responder_conn.leg().unwrap().expected_identity().unwrap(); + assert_eq!(discovered.pubkey(), initiator_identity.pubkey()); + + // Responder learned initiator's epoch + assert_eq!( + responder_conn.leg().unwrap().remote_epoch(), + Some(initiator_epoch) + ); + + // Initiator completes handshake + initiator_conn.complete_handshake(&msg2, 1300).unwrap(); + assert!(initiator_conn.has_session()); + + // Initiator learned responder's epoch + assert_eq!( + initiator_conn.leg().unwrap().remote_epoch(), + Some(responder_epoch) + ); + + // Both have sessions + assert!(initiator_conn.has_session()); + assert!(responder_conn.has_session()); + + // Take and verify sessions work + let mut init_session = initiator_conn.take_session().unwrap(); + let mut resp_session = responder_conn.take_session().unwrap(); + + // Encrypt/decrypt test + let plaintext = b"test message"; + let ciphertext = init_session.encrypt(plaintext).unwrap(); + let decrypted = resp_session.decrypt(&ciphertext).unwrap(); + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_connection_failure() { + // `mark_failed` releases the leg's Noise handshake handle. The failure + // *state* now lives on the control machine, but the leg-local effect is + // still observable: a completion attempt afterward reports `WrongState` + // (the handle-presence gate) and no session is produced. + let identity = make_peer_identity(); + let keypair = make_keypair(); + let mut conn = outbound_leg(LinkId::new(1), identity, 1000); + conn.start_handshake(keypair, make_epoch(), 1100).unwrap(); + + conn.mark_failed(); + + assert!(!conn.has_session()); + assert!(conn.complete_handshake(&[0u8; 96], 1200).is_err()); + } + + #[test] + fn test_wrong_direction_errors() { + let identity = make_peer_identity(); + let keypair = make_keypair(); + + // Outbound can't receive_handshake_init + let mut outbound = outbound_leg(LinkId::new(1), identity, 1000); + assert!( + outbound + .receive_handshake_init(keypair, make_epoch(), &[0u8; 106], 1100) + .is_err() + ); + + // Inbound can't start_handshake + let mut inbound = inbound_leg(LinkId::new(2), 1000); + assert!( + inbound + .start_handshake(keypair, make_epoch(), 1100) + .is_err() + ); + } } /// T-SANSIO: the action vocabulary must stay plain, comparable data.