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.