diff --git a/src/node/dataplane/peer_actions.rs b/src/node/dataplane/peer_actions.rs index de46ca6..1e2bc62 100644 --- a/src/node/dataplane/peer_actions.rs +++ b/src/node/dataplane/peer_actions.rs @@ -191,7 +191,6 @@ impl Node { // (`handle_msg1` L665): the send error text is surfaced // at the executor point where the failure is now handled. warn!(link_id = %link, error = %e, "Failed to send msg2"); - self.connections.remove(&link); self.links.remove(&link); self.addr_to_link .remove(&(ambient.transport_id, ambient.remote_addr.clone())); diff --git a/src/node/handlers/handshake.rs b/src/node/handlers/handshake.rs index f87c2d9..0707ddc 100644 --- a/src/node/handlers/handshake.rs +++ b/src/node/handlers/handshake.rs @@ -39,7 +39,7 @@ impl EstablishView for Node { rekey_in_progress: existing.map(|p| p.rekey_in_progress()).unwrap_or(false), existing_msg2: existing.and_then(|p| p.handshake_msg2().map(|m| m.to_vec())), at_max_peers: max_peers > 0 && self.peers.len() >= max_peers, - has_pending_outbound_to_peer: self.connections.values().any(|conn| { + has_pending_outbound_to_peer: self.connections().any(|conn| { conn.expected_identity() .map(|id| id.node_addr() == peer_addr) .unwrap_or(false) @@ -626,13 +626,12 @@ impl Node { ); self.links.insert(link_id, link); self.addr_to_link.insert(addr_key, link_id); - self.connections.insert(link_id, conn); let wire_msg2 = build_msg2(our_index, their_index, &msg2_payload); - if let Some(conn) = self.connections.get_mut(&link_id) { - conn.set_handshake_msg2(wire_msg2.clone()); - } + conn.set_handshake_msg2(wire_msg2.clone()); - // Register the machine (Promote/Restart tail only). + // Register the machine, carrying 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 @@ -781,14 +780,13 @@ impl Node { ); self.links.insert(link_id, link); self.addr_to_link.insert(addr_key, link_id); - self.connections.insert(link_id, conn); let wire_msg2 = build_msg2(our_index, their_index, &msg2_payload); - if let Some(conn) = self.connections.get_mut(&link_id) { - conn.set_handshake_msg2(wire_msg2.clone()); - } + conn.set_handshake_msg2(wire_msg2.clone()); - // Register the machine (Promote tail only — discarded on every - // reject/resend/rekey arm per the insertion discipline). + // Register the machine, carrying the connection (Promote tail + // only — discarded on every reject/resend/rekey arm per the + // insertion discipline). + machine.set_leg(conn); self.peer_machines.insert(link_id, machine); // Execute [SendHandshake, PromoteToActive]. The executor frames + @@ -842,7 +840,7 @@ impl Node { /// (if already promoted). fn find_stored_msg2(&self, link_id: LinkId) -> Option> { // Check pending connection first - if let Some(conn) = self.connections.get(&link_id) + if let Some(conn) = self.leg(&link_id) && let Some(msg2) = conn.handshake_msg2() { return Some(msg2.to_vec()); @@ -889,9 +887,17 @@ impl Node { }; // Check if this is a rekey msg2: the handshake state is on the - // ActivePeer (not a PeerConnection), so self.connections won't have it. - // Look for a peer with matching rekey_our_index. - if !self.connections.contains_key(&link_id) { + // ActivePeer (not a PeerConnection), so the link's machine — if one + // survives at all — carries no pending connection. A bare machine + // lookup would NOT discriminate here: an established peer's machine + // stays keyed by this link, so the pending connection's presence is + // what marks a fresh establish. Look for a peer with matching + // rekey_our_index. + if self + .peer_machines + .get(&link_id) + .is_none_or(|machine| machine.leg().is_none()) + { let noise_msg2 = &packet.data[header.noise_msg2_offset..]; // Find peer with rekey in progress for this index @@ -989,7 +995,7 @@ impl Node { } let (peer_identity, our_index) = { - let conn = self.connections.get_mut(&link_id).unwrap(); + let conn = self.leg_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) { @@ -1037,8 +1043,8 @@ impl Node { transport.close_connection(&addr).await; } } - self.connections.remove(&link_id); - // Drop the machine persisted at dial — this leg never promotes. + // Drop the machine persisted at dial — this leg never promotes, + // and its pending connection is dropped with it. self.remove_peer_machine(link_id); self.remove_link(&link_id); if let Some(idx) = our_index { @@ -1074,14 +1080,21 @@ impl Node { let out_snap = self.outbound_snapshot(&peer_node_addr); let out_decision = self.fmp.establish_outbound(&out_snap); if out_decision != OutboundDecision::Promote { - // The dial-persisted outbound machine is not consumed by the inline - // Swap/Keep resolution below (which mutates the existing promoted peer - // directly, with no machine). Drop it on entry so none of this block's - // exits leave a dangling machine — matching the pre-persistence path, + // Extract the outbound connection from its machine FIRST — the + // machine owns it, so disposing the machine before the take would + // destroy the connection. The dial-persisted outbound machine is + // not consumed by the inline Swap/Keep resolution below (which + // mutates the existing promoted peer directly, with no machine), + // so drop it right after the take — unconditionally, whether or + // not a connection was carried — so none of this block's exits + // leave a dangling machine. Matches the pre-persistence path, // which created no machine for a cross-connection. + let taken_conn = self + .peer_machines + .get_mut(&link_id) + .and_then(|machine| machine.take_leg()); self.remove_peer_machine(link_id); - // Extract the outbound connection - let mut conn = match self.connections.remove(&link_id) { + let mut conn = match taken_conn { Some(c) => c, None => { self.pending_outbound.remove(&key); @@ -1272,7 +1285,7 @@ impl Node { // The outbound Msg2 Promote step cancels the two dial-armed handshake // timers (the machine survives promotion, so they would otherwise linger // in `peer_timers` until `drive_peer_timers` lazily discards them — the - // promoted leg's `connections` entry is gone and the machine has left + // promoted leg's pending connection is consumed and the machine has left // `SentMsg1`, so they can no longer fire) and then promotes. // `PromoteToActive` is still what performs the promotion. debug_assert_eq!( @@ -1347,10 +1360,13 @@ impl Node { verified_identity: PeerIdentity, current_time_ms: u64, ) -> Result { - // Remove the connection from pending + // Take the pending connection off its control machine. The machine + // survives the promotion (it becomes the active peer's control + // machine), left with no pending connection. let mut connection = self - .connections - .remove(&link_id) + .peer_machines + .get_mut(&link_id) + .and_then(|machine| machine.take_leg()) .ok_or(NodeError::ConnectionNotFound(link_id))?; // Verify handshake is complete and extract session @@ -1522,14 +1538,13 @@ impl Node { // peer. The outbound will be cleaned up in handle_msg2 or by // the 30s handshake timeout. let pending_to_same_peer: Vec = self - .connections - .iter() - .filter(|(_, conn)| { + .connections() + .filter(|conn| { conn.expected_identity() .map(|id| *id.node_addr() == peer_node_addr) .unwrap_or(false) }) - .map(|(lid, _)| *lid) + .map(|conn| conn.link_id()) .collect(); for pending_link_id in &pending_to_same_peer { diff --git a/src/node/handlers/timeout.rs b/src/node/handlers/timeout.rs index 42e6e7c..e7e85a1 100644 --- a/src/node/handlers/timeout.rs +++ b/src/node/handlers/timeout.rs @@ -14,11 +14,12 @@ impl LifecycleView for Node { // `is_failed()` legs are always reaped here (~1s), as before. The // idle-timeout is reaped here ONLY for legs whose timeout is not already // driven by a machine `HandshakeTimeout` timer — i.e. inbound legs (IK - // inbound arms none) and machine-less test connections. Outbound legs with - // an armed timer are reaped by `drive_handshake_timeouts`, so excluding - // them here avoids a double reap. - self.connections + // inbound arms none). Outbound legs with an armed timer are reaped by + // `drive_handshake_timeouts`, so excluding them here avoids a double + // reap. + self.peer_machines .iter() + .filter_map(|(link_id, machine)| machine.leg().map(|conn| (link_id, conn))) .filter(|(link_id, conn)| { conn.is_failed() || (conn.is_timed_out(now_ms, timeout_ms) @@ -57,7 +58,7 @@ impl Node { /// the retry-then-teardown choreography is the pure /// [`Fmp::poll_timeouts`](crate::proto::fmp::Fmp::poll_timeouts) decision. pub(in crate::node) fn check_timeouts(&mut self) { - if self.connections.is_empty() { + if self.connection_count() == 0 { return; } @@ -70,7 +71,7 @@ impl Node { ConnAction::ScheduleRetry { peer } => self.note_handshake_timeout(peer, now_ms), ConnAction::Teardown { link } => { // Log before cleanup (needs live connection state). - if let Some(conn) = self.connections.get(&link) { + if let Some(conn) = self.leg(&link) { let direction = conn.direction(); if conn.is_failed() { debug!( @@ -101,14 +102,21 @@ impl Node { /// the link and address mapping. Does not log — callers provide context-appropriate /// log messages. pub(in crate::node) fn cleanup_stale_connection(&mut self, link_id: LinkId, _now_ms: u64) { - let conn = match self.connections.remove(&link_id) { + // Take the connection off its machine BEFORE disposing the machine + // (the machine owns it), keeping it readable for the index/link + // cleanup below. The machine shares the connection's `link_id` and + // lifetime; dropping it here means a reaped handshake leg leaves no + // dangling machine. A no-op for promoted peers — `promote_connection` + // already consumed their connection, so this reaper never runs for + // them. + let conn = match self + .peer_machines + .get_mut(&link_id) + .and_then(|machine| machine.take_leg()) + { Some(c) => c, None => return, }; - // A dial-persisted outbound control machine shares the connection's - // `link_id` and lifetime; drop it here so a reaped handshake leg leaves - // no dangling machine. A no-op for promoted peers — `promote_connection` - // already removed their connection, so this reaper never runs for them. self.remove_peer_machine(link_id); let transport_id = conn.transport_id(); @@ -158,7 +166,7 @@ impl Node { /// reflex, then `cleanup_stale_connection` (which drops the machine + timers). /// /// `check_timeouts` keeps reaping everything else — `is_failed()` legs and the - /// idle-timeout of legs without a machine timer (inbound / machine-less). + /// idle-timeout of legs without a machine timer (inbound legs). fn drive_handshake_timeouts(&mut self, now_ms: u64) { let timeout_ms = self.config().node.rate_limit.handshake_timeout_secs * 1000; let timer_links: Vec = self @@ -168,7 +176,7 @@ impl Node { .map(|(link, _)| *link) .collect(); for link in timer_links { - let (reap, retry_peer) = match self.connections.get(&link) { + let (reap, retry_peer) = match self.leg(&link) { Some(conn) if conn.is_timed_out(now_ms, timeout_ms) => { let retry_peer = if conn.is_outbound() { conn.expected_identity().map(|id| *id.node_addr()) @@ -251,7 +259,7 @@ impl Node { continue; } }; - match self.connections.get(&link).and_then(|c| c.handshake_msg1()) { + match self.leg(&link).and_then(|c| c.handshake_msg1()) { // Armed but the stored wire isn't there yet — leave the timer and // retry next tick (matches the old candidate filter skipping it). None => continue, @@ -283,7 +291,7 @@ impl Node { continue; }; - let (transport_id, remote_addr) = match self.connections.get(&link) { + let (transport_id, remote_addr) = match self.leg(&link) { Some(conn) => match (conn.transport_id(), conn.source_addr()) { (Some(tid), Some(addr)) => (tid, addr.clone()), _ => continue, diff --git a/src/node/lifecycle/mod.rs b/src/node/lifecycle/mod.rs index b3f35fb..fb9c2e0 100644 --- a/src/node/lifecycle/mod.rs +++ b/src/node/lifecycle/mod.rs @@ -372,7 +372,7 @@ impl Node { } fn is_connecting_to_peer(&self, peer_node_addr: &NodeAddr) -> bool { - self.connections.values().any(|conn| { + self.connections().any(|conn| { conn.expected_identity() .map(|id| id.node_addr() == peer_node_addr) .unwrap_or(false) @@ -385,7 +385,7 @@ impl Node { transport_id: TransportId, remote_addr: &TransportAddr, ) -> bool { - self.connections.values().any(|conn| { + self.connections().any(|conn| { conn.expected_identity() .map(|id| id.node_addr() == peer_node_addr) .unwrap_or(false) @@ -559,14 +559,14 @@ impl Node { /// Prepare an outbound Noise msg1 at dial: allocate the session index, run /// the Noise leaf, frame the wire, arm the shell-side resend, track - /// `pending_outbound`, and persist the connection. Returns `Err` on - /// index-allocation or Noise failure (cleaning the partial leg, including - /// the dial-time control machine), leaving the armed wire on the connection - /// for `send_stored_msg1` to transmit. Does NOT send — so the fallible setup - /// can propagate its error synchronously before any machine drive. The - /// control machine itself is persisted at dial in `initiate_connection`; - /// this function no longer touches `peer_machines` except to clean it up on - /// the failure paths. + /// `pending_outbound`, and embed the prepared connection on the dial-time + /// control machine. Returns `Err` on index-allocation or Noise failure + /// (cleaning the partial leg, including the dial-time control machine), + /// leaving the armed wire on the connection for `send_stored_msg1` to + /// transmit. Does NOT send — so the fallible setup can propagate its error + /// synchronously before any machine drive. The control machine itself is + /// persisted at dial in `initiate_connection`; this function hands it the + /// connection on success and cleans it up on the failure paths. pub(in crate::node) fn prepare_outbound_msg1( &mut self, link_id: LinkId, @@ -633,7 +633,19 @@ impl Node { // Track in pending_outbound for msg2 dispatch self.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); - self.connections.insert(link_id, connection); + + // 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" + ); + self.peer_machines + .entry(link_id) + .or_insert_with(|| PeerMachine::new_outbound(link_id, peer_identity, current_time_ms)) + .set_leg(connection); Ok(()) } @@ -650,15 +662,11 @@ impl Node { transport_id: TransportId, remote_addr: &TransportAddr, ) { - let wire_msg1 = match self - .connections - .get(&link_id) - .and_then(|c| c.handshake_msg1()) - { + let wire_msg1 = match self.leg(&link_id).and_then(|c| c.handshake_msg1()) { Some(w) => w.to_vec(), None => return, }; - let our_index = self.connections.get(&link_id).and_then(|c| c.our_index()); + let our_index = self.leg(&link_id).and_then(|c| c.our_index()); // Send the wire format handshake message if let Some(transport) = self.transports.get(&transport_id) { @@ -681,7 +689,7 @@ impl Node { ); // Mark connection as failed but don't remove it yet // The event loop can handle retry logic - if let Some(conn) = self.connections.get_mut(&link_id) { + if let Some(conn) = self.leg_mut(&link_id) { conn.mark_failed(); } } @@ -881,14 +889,13 @@ impl Node { ); let now_ms = Self::now_ms(); let stale: Vec = self - .connections - .iter() - .filter(|(_, conn)| { + .connections() + .filter(|conn| { conn.expected_identity() .map(|id| id.node_addr() == &peer_addr) .unwrap_or(false) }) - .map(|(link_id, _)| *link_id) + .map(|conn| conn.link_id()) .collect(); for link_id in stale { self.cleanup_stale_connection(link_id, now_ms); @@ -1821,7 +1828,7 @@ impl Node { info!("Node started:"); info!(" state: {}", self.supervisor.state); info!(" transports: {}", self.transports.len()); - info!(" connections: {}", self.connections.len()); + info!(" connections: {}", self.connection_count()); Ok(()) } @@ -2665,18 +2672,17 @@ impl Node { /// The `connected` / `connecting` sets gate the floor, retry-dial, overlay, /// and LAN layers; `in_flight_by_peer` feeds the opportunistic layer's /// per-peer parallel cap, computed exactly as the deleted - /// `path_candidate_attempt_budget` did: `connections(expected == addr) + + /// `path_candidate_attempt_budget` did: `pending legs(expected == addr) + /// pending_connects(addr)`. The scalar counts stay unpopulated at the /// ceiling-only posture (no layer reads them; see [`Observed`]). pub(in crate::node) fn observe_peering(&self) -> Observed { let connected: HashSet = self.peers.keys().copied().collect(); let connecting: HashSet = self - .connections - .values() + .connections() .filter_map(|conn| conn.expected_identity().map(|id| *id.node_addr())) .collect(); let mut in_flight_by_peer: HashMap = HashMap::new(); - for conn in self.connections.values() { + for conn in self.connections() { if let Some(id) = conn.expected_identity() { *in_flight_by_peer.entry(*id.node_addr()).or_default() += 1; } @@ -2715,9 +2721,11 @@ impl Node { } fn outbound_handshake_slots(&self) -> usize { + // Count pending connections (legs), not machines: a dial-born machine + // in the connect window has no connection yet but DOES hold a + // `pending_connects` slot, so counting machines would tally it twice. let used = self - .connections - .len() + .connection_count() .saturating_add(self.peering.pending_connects.len()); if self.max_connections() == 0 { usize::MAX @@ -2743,8 +2751,7 @@ impl Node { } let in_flight_for_peer = self - .connections - .values() + .connections() .filter(|conn| { conn.expected_identity() .map(|identity| identity.node_addr() == peer_node_addr) diff --git a/src/node/mod.rs b/src/node/mod.rs index d220085..ad1bf6d 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -281,7 +281,8 @@ struct PendingConnect { /// ## Peer Lifecycle /// /// Peers go through two phases: -/// 1. **Connection phase** (`connections`): Handshake in progress, indexed by LinkId +/// 1. **Connection phase**: Handshake in progress; the pending connection is +/// carried by its per-peer control machine (`peer_machines`), indexed by LinkId /// 2. **Active phase** (`peers`): Authenticated, indexed by NodeAddr /// /// The `addr_to_link` map enables dispatching incoming packets to the right @@ -348,16 +349,12 @@ pub struct Node { /// rx_loop select arm that feeds `Event::ChildExited` to the supervisor FSM. child_exit_rx: Option>, - // === Connections (Handshake Phase) === - /// Pending connections (handshake in progress). - /// Indexed by LinkId since we don't know the peer's identity yet. - connections: HashMap, - // === Per-Peer Control Machines === /// Per-peer lifecycle control FSMs, keyed by the stable `LinkId` that spans - /// the handshake→active lifetime. A parallel structure introduced by the - /// node-runtime decomposition: `connections`/`peers` stay byte-unchanged (hot - /// path pristine) and are cut over to this machine home path-by-path. + /// the handshake→active lifetime. Each machine owns its pending handshake + /// connection (the `PeerConnection` leg) while the handshake is in + /// progress — the single LinkId-keyed per-peer map on `Node`; `peers` + /// stays byte-unchanged (hot path pristine). /// Machines are inserted at dial and inbound msg1, and stepped in production /// by the handshake handlers, the rekey-cadence and liveness-reap routers, /// and the lifecycle paths, with the executor (`dataplane/peer_actions.rs`) @@ -643,7 +640,6 @@ impl Node { packet_rx: None, child_exit_tx: None, child_exit_rx: None, - connections: HashMap::new(), peer_machines: HashMap::new(), peer_timers: HashMap::new(), peers: HashMap::new(), @@ -792,7 +788,6 @@ impl Node { packet_rx: None, child_exit_tx: None, child_exit_rx: None, - connections: HashMap::new(), peer_machines: HashMap::new(), peer_timers: HashMap::new(), peers: HashMap::new(), @@ -1520,7 +1515,7 @@ impl Node { tun_state: self.tun_state, tun_name: self.tun_name.clone(), effective_ipv6_mtu: self.effective_ipv6_mtu(), - connection_count: self.connections.len(), + connection_count: self.connection_count(), peer_count: self.peers.len(), link_count: self.links.len(), transport_count: self.transports.len(), @@ -2152,7 +2147,10 @@ impl Node { /// Number of pending connections (handshake in progress). pub fn connection_count(&self) -> usize { - self.connections.len() + self.peer_machines + .values() + .filter(|machine| machine.leg().is_some()) + .count() } /// Number of authenticated peers. @@ -2273,30 +2271,19 @@ impl Node { self.peer_timers.remove(&link); } - /// Gates the leg→machine direction of `debug_assert_peer_maps_coherent`. - /// Asserting it requires the convention that every `connections` insert is - /// paired with a `peer_machines` insert before the next await point — - /// which holds here: inbound msg1 births a machine alongside the window - /// leg, and dials birth one before the leg exists. Where handshake-window - /// legs legitimately run machine-less, flip this to `false` rather than - /// weakening the machine→carrier direction. - #[cfg(debug_assertions)] - const CHECK_LEGS_HAVE_MACHINES: bool = true; - /// Debug-build coherence sweep over the peer-lifecycle maps, run once per /// rx-loop tick and invoked directly by unit tests. /// /// Machine→carrier: every `peer_machines` entry must have a live carrier — - /// a pending connection on the same link, an active peer on it, or a + /// its own embedded pending connection, an active peer on its link, or a /// pending connect still resolving toward it. A machine with none of these - /// is unreachable by every teardown path and has leaked. - /// - /// Leg→machine: every pending connection carries a control machine (gated - /// by `CHECK_LEGS_HAVE_MACHINES` above). + /// is unreachable by every teardown path and has leaked. (The pending + /// connection lives inside the machine, so no separate connection→machine + /// direction exists to check.) #[cfg(debug_assertions)] pub(in crate::node) fn debug_assert_peer_maps_coherent(&self) { - for link in self.peer_machines.keys() { - let has_carrier = self.connections.contains_key(link) + for (link, machine) in &self.peer_machines { + let has_carrier = machine.leg().is_some() || self.peers.values().any(|peer| peer.link_id() == *link) || self .peering @@ -2309,22 +2296,13 @@ impl Node { (no pending connection, active peer, or pending connect)" ); } - - if Self::CHECK_LEGS_HAVE_MACHINES { - for link in self.connections.keys() { - assert!( - self.peer_machines.contains_key(link), - "pending connection on link {link} has no control machine" - ); - } - } } /// Operator-visible msg1 resend count for a pending handshake `link`, read /// from the per-peer machine (the counter's home once the resend drive moved - /// off the shell connection). Machine-less connections (inbound legs that - /// never resend, and test-created connections) report 0, matching what the - /// shell connection reported before the counter moved. + /// off the shell connection). A link with no machine reports 0, and inbound + /// machines never resend, matching what the shell connection reported + /// before the counter moved. pub(crate) fn connection_resend_count(&self, link: LinkId) -> u32 { self.peer_machines .get(&link) @@ -2345,8 +2323,7 @@ impl Node { .values() .any(|link| link.transport_id() == transport_id) || self - .connections - .values() + .connections() .any(|conn| conn.transport_id() == Some(transport_id)) || self .peers @@ -2381,59 +2358,87 @@ impl Node { // === Connection Management (Handshake Phase) === + /// The pending connection for `link_id`, read through the control machine + /// that carries it. + fn leg(&self, link_id: &LinkId) -> Option<&PeerConnection> { + self.peer_machines + .get(link_id) + .and_then(|machine| machine.leg()) + } + + /// Mutable access to the pending connection for `link_id`. + fn leg_mut(&mut self, link_id: &LinkId) -> Option<&mut PeerConnection> { + self.peer_machines + .get_mut(link_id) + .and_then(|machine| machine.leg_mut()) + } + /// Add a pending connection. /// - /// Also seeds a control machine for the leg when none exists yet, keeping - /// the leg→machine invariant (`debug_assert_peer_maps_coherent`) intact - /// for callers that insert a connection directly rather than through the - /// dial or inbound-msg1 paths, which pair the two inserts themselves. + /// Seeds a control machine for the leg when none exists yet and embeds the + /// connection on it, for callers that insert a connection directly rather + /// than through the dial or inbound-msg1 paths, which build the machine + /// themselves. pub fn add_connection(&mut self, connection: PeerConnection) -> Result<(), NodeError> { let link_id = connection.link_id(); - if self.connections.contains_key(&link_id) { + if self + .peer_machines + .get(&link_id) + .is_some_and(|machine| machine.leg().is_some()) + { return Err(NodeError::ConnectionAlreadyExists(link_id)); } - if self.max_connections() > 0 && self.connections.len() >= self.max_connections() { + if self.max_connections() > 0 && self.connection_count() >= self.max_connections() { return Err(NodeError::MaxConnectionsExceeded { max: self.max_connections(), }); } - self.peer_machines.entry(link_id).or_insert_with(|| { - let now = connection.started_at(); - match connection.expected_identity() { - Some(identity) if connection.is_outbound() => { - PeerMachine::new_outbound(link_id, *identity, now) + self.peer_machines + .entry(link_id) + .or_insert_with(|| { + let now = connection.started_at(); + match connection.expected_identity() { + Some(identity) if connection.is_outbound() => { + PeerMachine::new_outbound(link_id, *identity, now) + } + _ => PeerMachine::new_inbound(link_id, now), } - _ => PeerMachine::new_inbound(link_id, now), - } - }); - - self.connections.insert(link_id, connection); + }) + .set_leg(connection); Ok(()) } /// Get a connection by LinkId. pub fn get_connection(&self, link_id: &LinkId) -> Option<&PeerConnection> { - self.connections.get(link_id) + self.leg(link_id) } /// Get a mutable connection by LinkId. pub fn get_connection_mut(&mut self, link_id: &LinkId) -> Option<&mut PeerConnection> { - self.connections.get_mut(link_id) + self.leg_mut(link_id) } /// Remove a connection, disposing its control machine alongside - /// (the disposal complement of `add_connection`'s machine seeding). + /// (the disposal complement of `add_connection`'s machine seeding). The + /// connection is taken off the machine BEFORE the machine is dropped, so + /// the caller still receives it. pub fn remove_connection(&mut self, link_id: &LinkId) -> Option { + let connection = self + .peer_machines + .get_mut(link_id) + .and_then(|machine| machine.take_leg()); self.remove_peer_machine(*link_id); - self.connections.remove(link_id) + connection } /// Iterate over all connections. pub fn connections(&self) -> impl Iterator { - self.connections.values() + self.peer_machines + .values() + .filter_map(|machine| machine.leg()) } // === Peer Management (Active Phase) === diff --git a/src/node/peering/reconcile.rs b/src/node/peering/reconcile.rs index 2e850d7..82db277 100644 --- a/src/node/peering/reconcile.rs +++ b/src/node/peering/reconcile.rs @@ -147,7 +147,7 @@ pub(crate) struct Observed { #[allow(dead_code)] // ceiling-only posture: no layer reads the scalar counts (future set-point) pub peers: usize, - /// `self.connections.len()`. + /// `self.connection_count()`. #[allow(dead_code)] // ceiling-only posture: no layer reads the scalar counts (future set-point) pub connections: usize, diff --git a/src/node/tests/acl.rs b/src/node/tests/acl.rs index 5670d0f..b28c77a 100644 --- a/src/node/tests/acl.rs +++ b/src/node/tests/acl.rs @@ -104,7 +104,7 @@ async fn test_outbound_msg2_denied_after_acl_reload() { node_a .addr_to_link .insert((transport_id, remote_addr.clone()), link_id_a); - node_a.connections.insert(link_id_a, conn_a); + node_a.add_connection(conn_a).unwrap(); node_a .pending_outbound .insert((transport_id, our_index_a.as_u32()), link_id_a); diff --git a/src/node/tests/establish_chartests.rs b/src/node/tests/establish_chartests.rs index 989a8cc..02cae39 100644 --- a/src/node/tests/establish_chartests.rs +++ b/src/node/tests/establish_chartests.rs @@ -206,7 +206,7 @@ async fn chartest_msg1_duplicate_pending_resends_stored_msg2() { node.links.insert(link_id, link); node.addr_to_link .insert((transport_id, peer_addr.clone()), link_id); - node.connections.insert(link_id, conn); + node.add_connection(conn).unwrap(); assert_eq!(node.peer_count(), 0); let before_pending = node.msg1_rate_limiter.pending_count(); @@ -235,7 +235,7 @@ async fn chartest_msg1_duplicate_pending_resends_stored_msg2() { ); assert_eq!(node.peer_count(), 0, "duplicate msg1 promotes nothing"); assert!( - node.connections.contains_key(&link_id), + node.get_connection(&link_id).is_some(), "pending connection is left intact" ); assert_eq!( @@ -341,7 +341,7 @@ async fn chartest_msg1_inbound_promote_defers_pending_outbound_to_same_identity( node.links.insert(out_link, out_l); node.addr_to_link .insert((transport_id, out_addr.clone()), out_link); - node.connections.insert(out_link, out_conn); + node.add_connection(out_conn).unwrap(); node.pending_outbound .insert((transport_id, out_index.as_u32()), out_link); assert_eq!(node.peer_count(), 0); @@ -362,7 +362,7 @@ async fn chartest_msg1_inbound_promote_defers_pending_outbound_to_same_identity( assert!(peer.has_session()); assert_eq!(node.peer_count(), 1); assert!( - node.connections.contains_key(&out_link), + node.get_connection(&out_link).is_some(), "pending outbound to the same identity must be preserved (deferred cleanup)" ); assert!( @@ -406,7 +406,7 @@ async fn chartest_msg1_at_cap_with_pending_outbound_bypasses_early_gate() { out_conn.set_our_index(out_index); out_conn.set_transport_id(transport_id); out_conn.set_source_addr(out_addr.clone()); - node.connections.insert(out_link, out_conn); + node.add_connection(out_conn).unwrap(); node.pending_outbound .insert((transport_id, out_index.as_u32()), out_link); @@ -516,7 +516,7 @@ async fn chartest_cross_connection_tiebreak_winner_and_loser() { node_a .addr_to_link .insert((transport_id_a, remote_addr_b.clone()), link_a_out); - node_a.connections.insert(link_a_out, conn_a); + node_a.add_connection(conn_a).unwrap(); node_a .pending_outbound .insert((transport_id_a, out_index_a.as_u32()), link_a_out); @@ -545,7 +545,7 @@ async fn chartest_cross_connection_tiebreak_winner_and_loser() { node_b .addr_to_link .insert((transport_id_b, remote_addr_a.clone()), link_b_out); - node_b.connections.insert(link_b_out, conn_b); + node_b.add_connection(conn_b).unwrap(); node_b .pending_outbound .insert((transport_id_b, out_index_b.as_u32()), link_b_out); diff --git a/src/node/tests/handshake.rs b/src/node/tests/handshake.rs index dd241bb..18db30f 100644 --- a/src/node/tests/handshake.rs +++ b/src/node/tests/handshake.rs @@ -78,7 +78,7 @@ async fn test_two_node_handshake_udp() { Duration::from_millis(100), ); node_a.links.insert(link_id_a, link_a); - node_a.connections.insert(link_id_a, conn_a); + node_a.add_connection(conn_a).unwrap(); node_a .pending_outbound .insert((transport_id_a, our_index_a.as_u32()), link_id_a); @@ -315,7 +315,7 @@ async fn test_run_rx_loop_handshake() { Duration::from_millis(100), ); node_a.links.insert(link_id_a, link_a); - node_a.connections.insert(link_id_a, conn_a); + node_a.add_connection(conn_a).unwrap(); node_a .pending_outbound .insert((transport_id_a, our_index_a.as_u32()), link_id_a); @@ -506,7 +506,7 @@ async fn test_cross_connection_both_initiate() { node_a .addr_to_link .insert((transport_id_a, remote_addr_b.clone()), link_id_a_out); - node_a.connections.insert(link_id_a_out, conn_a); + node_a.add_connection(conn_a).unwrap(); node_a .pending_outbound .insert((transport_id_a, our_index_a.as_u32()), link_id_a_out); @@ -536,7 +536,7 @@ async fn test_cross_connection_both_initiate() { node_b .addr_to_link .insert((transport_id_b, remote_addr_a.clone()), link_id_b_out); - node_b.connections.insert(link_id_b_out, conn_b); + node_b.add_connection(conn_b).unwrap(); node_b .pending_outbound .insert((transport_id_b, our_index_b.as_u32()), link_id_b_out); @@ -685,7 +685,7 @@ async fn test_stale_connection_cleanup() { node.links.insert(link_id, link); node.addr_to_link .insert((transport_id, remote_addr.clone()), link_id); - node.connections.insert(link_id, conn); + node.add_connection(conn).unwrap(); node.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); @@ -763,7 +763,7 @@ async fn test_failed_connection_cleanup() { node.links.insert(link_id, link); node.addr_to_link .insert((transport_id, remote_addr.clone()), link_id); - node.connections.insert(link_id, conn); + node.add_connection(conn).unwrap(); node.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); @@ -859,11 +859,11 @@ async fn test_resend_scheduling() { .insert((transport_id, remote_addr.clone()), link_id); node.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); - node.connections.insert(link_id, conn); - // The msg1-resend counter and its due timer live on the per-peer machine. - // Dial it to `SentMsg1` (connectionless: no connect step) and arm its - // retransmit timer at now + 1000ms, mirroring what a real dial arms. + // The msg1-resend counter and its due timer live on the per-peer machine, + // which also carries the pending connection. Dial it to `SentMsg1` + // (connectionless: no connect step) and arm its retransmit timer at + // now + 1000ms, mirroring what a real dial arms. let mut machine = crate::peer::machine::PeerMachine::new_outbound(link_id, peer_identity, now_ms); let _ = machine.step( @@ -876,6 +876,7 @@ async fn test_resend_scheduling() { now_ms, &mut node.index_allocator, ); + machine.set_leg(conn); node.peer_machines.insert(link_id, machine); node.peer_timers.entry(link_id).or_default().insert( crate::peer::machine::TimerKind::HandshakeRetransmit, @@ -934,11 +935,11 @@ async fn test_handshake_timeout_drive() { node.links.insert(link_id, link); node.addr_to_link .insert((transport_id, remote_addr.clone()), link_id); - node.connections.insert(link_id, conn); node.pending_outbound .insert((transport_id, our_index.as_u32()), link_id); - // Machine in SentMsg1 with a HandshakeTimeout timer armed at dial + 30s. + // Machine in SentMsg1, carrying the pending connection, with a + // HandshakeTimeout timer armed at dial + 30s. let mut machine = crate::peer::machine::PeerMachine::new_outbound(link_id, peer_identity, dial_ms); let _ = machine.step( @@ -951,6 +952,7 @@ async fn test_handshake_timeout_drive() { dial_ms, &mut node.index_allocator, ); + machine.set_leg(conn); 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/spanning_tree.rs b/src/node/tests/spanning_tree.rs index 85de6ee..12b323e 100644 --- a/src/node/tests/spanning_tree.rs +++ b/src/node/tests/spanning_tree.rs @@ -141,7 +141,7 @@ pub(super) async fn initiate_handshake(nodes: &mut [TestNode], i: usize, j: usiz .node .addr_to_link .insert((transport_id, responder_addr.clone()), link_id); - initiator.node.connections.insert(link_id, conn); + initiator.node.add_connection(conn).unwrap(); initiator .node .pending_outbound diff --git a/src/node/tests/unit.rs b/src/node/tests/unit.rs index 1d46941..ad4cbc8 100644 --- a/src/node/tests/unit.rs +++ b/src/node/tests/unit.rs @@ -137,8 +137,7 @@ async fn test_try_peer_addresses_races_all_concrete_udp_candidates() { .unwrap(); let mut addrs = node - .connections - .values() + .connections() .filter_map(|conn| conn.source_addr().and_then(|addr| addr.as_str())) .collect::>(); addrs.sort(); @@ -733,7 +732,7 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() { node.links.insert(pending_link_id, pending_link); node.addr_to_link .insert((transport_id, pending_addr.clone()), pending_link_id); - node.connections.insert(pending_link_id, pending_conn); + node.add_connection(pending_conn).unwrap(); node.pending_outbound .insert((transport_id, pending_index.as_u32()), pending_link_id); @@ -1252,8 +1251,7 @@ async fn update_peers_races_new_alternative_without_dropping_active_peer() { assert_eq!(node.peer_count(), 1, "existing link must stay live"); assert_eq!(node.connection_count(), 1); assert_eq!( - node.connections - .values() + node.connections() .next() .and_then(|conn| conn.source_addr()), Some(&new_addr) diff --git a/src/peer/machine.rs b/src/peer/machine.rs index 2b546b8..36d375f 100644 --- a/src/peer/machine.rs +++ b/src/peer/machine.rs @@ -56,6 +56,7 @@ #![allow(dead_code)] +use crate::peer::PeerConnection; use crate::proto::fmp::{ ConnAction, ConnSnapshot, ConnectionState, EstablishSnapshot, Fmp, InboundDecision, OutboundDecision, OutboundSnapshot, PeerSnapshot, PromotionResult, RekeyCfg, @@ -370,6 +371,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 + /// 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. + leg: Option, /// Pure handshake-phase bookkeeping (link/direction/indices/transport/ /// stored handshake bytes/epoch). Reused verbatim from the FMP state core. conn: ConnectionState, @@ -415,6 +422,7 @@ impl PeerMachine { state: PeerState::Discovered, link, identity: Some(identity), + leg: None, conn: ConnectionState::outbound(link, identity, now), remote_epoch: None, pending_msg2_payload: None, @@ -441,6 +449,7 @@ impl PeerMachine { }, link, identity: None, + leg: None, conn: ConnectionState::inbound(link, now), remote_epoch: None, pending_msg2_payload: None, @@ -463,6 +472,28 @@ impl PeerMachine { self.state } + /// The pending handshake connection, if this leg is still in the + /// handshake window. + pub(crate) fn leg(&self) -> Option<&PeerConnection> { + self.leg.as_ref() + } + + /// Mutable access to the pending handshake connection. + pub(crate) fn leg_mut(&mut self) -> Option<&mut PeerConnection> { + self.leg.as_mut() + } + + /// Take the pending handshake connection off the machine (promotion and + /// teardown consume it by value). + pub(crate) fn take_leg(&mut self) -> Option { + self.leg.take() + } + + /// Embed a pending handshake connection on the machine. + pub(crate) fn set_leg(&mut self, leg: PeerConnection) { + self.leg = Some(leg); + } + /// The index we allocated for this peer's inbound session, once Phase 2 /// (`on_authorized`) has run. `None` before allocation (and after a /// rejected/unauthorized msg1). The inbound cutover reads this to perform