diff --git a/src/bin/fips.rs b/src/bin/fips.rs index c144ba7..d0ab35c 100644 --- a/src/bin/fips.rs +++ b/src/bin/fips.rs @@ -96,14 +96,23 @@ async fn main() { info!("FIPS running, press Ctrl+C to exit"); - match tokio::signal::ctrl_c().await { - Ok(()) => info!("Shutdown signal received"), - Err(e) => error!("Failed to listen for shutdown signal: {}", e), + // Run the RX event loop until shutdown signal. + // stop() drops the packet channel, causing run_rx_loop to exit. + tokio::select! { + result = node.run_rx_loop() => { + match result { + Ok(()) => info!("RX loop exited"), + Err(e) => error!("RX loop error: {}", e), + } + } + _ = tokio::signal::ctrl_c() => { + info!("Shutdown signal received"); + } } info!("FIPS shutting down"); - // Stop the node (shuts down TUN, stops I/O threads) + // Stop the node (shuts down transports, TUN, I/O threads) if let Err(e) = node.stop().await { warn!("Error during shutdown: {}", e); } diff --git a/src/node.rs b/src/node.rs index 82c0d9e..67d67d8 100644 --- a/src/node.rs +++ b/src/node.rs @@ -76,6 +76,15 @@ pub enum NodeError { #[error("max links exceeded: {max}")] MaxLinksExceeded { max: usize }, + #[error("handshake incomplete for link {0}")] + HandshakeIncomplete(LinkId), + + #[error("no session available for link {0}")] + NoSession(LinkId), + + #[error("promotion failed for link {link_id}: {reason}")] + PromotionFailed { link_id: LinkId, reason: String }, + #[error("config error: {0}")] Config(#[from] ConfigError), @@ -850,11 +859,46 @@ impl Node { current_time_ms: u64, ) -> Result { // Remove the connection from pending - let connection = self + let mut connection = self .connections .remove(&link_id) .ok_or(NodeError::ConnectionNotFound(link_id))?; + // Verify handshake is complete and extract session + if !connection.has_session() { + return Err(NodeError::HandshakeIncomplete(link_id)); + } + + let noise_session = connection + .take_session() + .ok_or(NodeError::NoSession(link_id))?; + + let our_index = connection.our_index().ok_or_else(|| { + NodeError::PromotionFailed { + link_id, + reason: "missing our_index".into(), + } + })?; + let their_index = connection.their_index().ok_or_else(|| { + NodeError::PromotionFailed { + link_id, + reason: "missing their_index".into(), + } + })?; + let transport_id = connection.transport_id().ok_or_else(|| { + NodeError::PromotionFailed { + link_id, + reason: "missing transport_id".into(), + } + })?; + let current_addr = connection.source_addr().ok_or_else(|| { + NodeError::PromotionFailed { + link_id, + reason: "missing source_addr".into(), + } + })?.clone(); + let link_stats = connection.link_stats().clone(); + let peer_node_addr = *verified_identity.node_addr(); let is_outbound = connection.is_outbound(); @@ -874,15 +918,30 @@ impl Node { let old_peer = self.peers.remove(&peer_node_addr).unwrap(); let loser_link_id = old_peer.link_id(); - // Create new active peer with stats from handshake - let new_peer = ActivePeer::with_stats( + // Clean up old peer's index from peers_by_index + if let (Some(old_tid), Some(old_idx)) = + (old_peer.transport_id(), old_peer.our_index()) + { + self.peers_by_index + .remove(&(old_tid, old_idx.as_u32())); + let _ = self.index_allocator.free(old_idx); + } + + let new_peer = ActivePeer::with_session( verified_identity, link_id, current_time_ms, - connection.link_stats().clone(), + noise_session, + our_index, + their_index, + transport_id, + current_addr, + link_stats, ); self.peers.insert(peer_node_addr, new_peer); + self.peers_by_index + .insert((transport_id, our_index.as_u32()), peer_node_addr); info!( node_addr = %peer_node_addr, @@ -897,6 +956,9 @@ impl Node { }) } else { // This connection loses, keep existing + // Free the index we allocated + let _ = self.index_allocator.free(our_index); + info!( node_addr = %peer_node_addr, winner_link = %existing_link_id, @@ -911,21 +973,31 @@ impl Node { } else { // No cross-connection, normal promotion if self.max_peers > 0 && self.peers.len() >= self.max_peers { + let _ = self.index_allocator.free(our_index); return Err(NodeError::MaxPeersExceeded { max: self.max_peers }); } - let new_peer = ActivePeer::with_stats( + let new_peer = ActivePeer::with_session( verified_identity, link_id, current_time_ms, - connection.link_stats().clone(), + noise_session, + our_index, + their_index, + transport_id, + current_addr, + link_stats, ); self.peers.insert(peer_node_addr, new_peer); + self.peers_by_index + .insert((transport_id, our_index.as_u32()), peer_node_addr); info!( node_addr = %peer_node_addr, link_id = %link_id, + our_index = %our_index, + their_index = %their_index, "Connection promoted to active peer" ); @@ -1284,18 +1356,9 @@ impl Node { return; } }; - let peer_node_addr = *peer_identity.node_addr(); - // Check if this peer is already connected - if self.peers.contains_key(&peer_node_addr) { - // TODO: Handle reconnection case (future: session replacement) - self.msg1_rate_limiter.complete_handshake(); - debug!( - node_addr = %peer_node_addr, - "Peer already connected, ignoring msg1" - ); - return; - } + // Note: we don't early-return if peer is already in self.peers here. + // promote_connection handles cross-connection resolution via tie-breaker. // Allocate our session index let our_index = match self.index_allocator.allocate() { @@ -1354,14 +1417,50 @@ impl Node { } } - info!( - node_addr = %peer_node_addr, - link_id = %link_id, - our_index = %our_index, - "Inbound handshake initiated" - ); + // Responder handshake is complete after receive_handshake_init (Noise IK + // pattern: responder processes msg1 and generates msg2 in one step). + // Promote the connection to active peer now. + match self.promote_connection(link_id, peer_identity, packet.timestamp_ms) { + Ok(result) => { + match result { + PromotionResult::Promoted(node_addr) => { + info!( + node_addr = %node_addr, + link_id = %link_id, + our_index = %our_index, + "Inbound peer promoted to active" + ); + } + PromotionResult::CrossConnectionWon { loser_link_id, node_addr } => { + info!( + node_addr = %node_addr, + loser_link_id = %loser_link_id, + "Inbound cross-connection won" + ); + } + PromotionResult::CrossConnectionLost { winner_link_id } => { + info!( + winner_link_id = %winner_link_id, + "Inbound cross-connection lost, keeping existing" + ); + } + } + } + Err(e) => { + warn!( + link_id = %link_id, + error = %e, + "Failed to promote inbound connection" + ); + // Clean up on promotion failure + self.links.remove(&link_id); + self.addr_to_link + .remove(&(packet.transport_id, packet.remote_addr)); + let _ = self.index_allocator.free(our_index); + } + } - // Note: rate limiter completed when handshake completes or times out + self.msg1_rate_limiter.complete_handshake(); } /// Handle handshake message 2 (discriminator 0x02). @@ -1598,6 +1697,7 @@ impl fmt::Debug for Node { #[cfg(test)] mod tests { use super::*; + use crate::index::SessionIndex; use crate::transport::{LinkDirection, TransportAddr}; use std::time::Duration; @@ -1618,6 +1718,47 @@ mod tests { PeerIdentity::from_pubkey(identity.pubkey()) } + /// Create a PeerConnection with a completed Noise IK handshake. + /// + /// Returns (connection, peer_identity) where the connection is outbound, + /// in Complete state, with session, indices, and transport info set. + fn make_completed_connection( + node: &mut Node, + link_id: LinkId, + transport_id: TransportId, + current_time_ms: u64, + ) -> (PeerConnection, PeerIdentity) { + let peer_identity_full = Identity::generate(); + // Must use from_pubkey_full to preserve parity for ECDH + let peer_identity = PeerIdentity::from_pubkey_full(peer_identity_full.pubkey_full()); + + // Create outbound connection + let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), current_time_ms); + + // Run initiator side of handshake + let our_keypair = node.identity.keypair(); + let msg1 = conn.start_handshake(our_keypair, current_time_ms).unwrap(); + + // Run responder side to generate msg2 + let mut resp_conn = PeerConnection::inbound(LinkId::new(999), current_time_ms); + let peer_keypair = peer_identity_full.keypair(); + let msg2 = resp_conn + .receive_handshake_init(peer_keypair, &msg1, current_time_ms) + .unwrap(); + + // Complete initiator handshake + conn.complete_handshake(&msg2, current_time_ms).unwrap(); + + // Set indices and transport info + let our_index = node.index_allocator.allocate().unwrap(); + conn.set_our_index(our_index); + conn.set_their_index(SessionIndex::new(42)); + conn.set_transport_id(transport_id); + conn.set_source_addr(TransportAddr::from_string("127.0.0.1:5000")); + + (conn, peer_identity) + } + #[test] fn test_node_creation() { let node = make_node(); @@ -1781,11 +1922,11 @@ mod tests { #[test] fn test_node_promote_connection() { let mut node = make_node(); + let transport_id = TransportId::new(1); - let identity = make_peer_identity(); - let node_addr = *identity.node_addr(); let link_id = LinkId::new(1); - let conn = PeerConnection::outbound(link_id, identity.clone(), 1000); + let (conn, identity) = make_completed_connection(&mut node, link_id, transport_id, 1000); + let node_addr = *identity.node_addr(); node.add_connection(conn).unwrap(); assert_eq!(node.connection_count(), 1); @@ -1799,17 +1940,28 @@ mod tests { let peer = node.get_peer(&node_addr).unwrap(); assert_eq!(peer.authenticated_at(), 2000); + assert!(peer.has_session(), "Promoted peer should have NoiseSession"); + assert!(peer.our_index().is_some(), "Promoted peer should have our_index"); + assert!(peer.their_index().is_some(), "Promoted peer should have their_index"); + + // Verify peers_by_index is populated + let our_index = peer.our_index().unwrap(); + assert_eq!( + node.peers_by_index.get(&(transport_id, our_index.as_u32())), + Some(&node_addr) + ); } #[test] fn test_node_cross_connection_resolution() { let mut node = make_node(); + let transport_id = TransportId::new(1); // First connection and promotion (becomes active peer) - let identity = make_peer_identity(); - let node_addr = *identity.node_addr(); let link_id1 = LinkId::new(1); - let conn1 = PeerConnection::outbound(link_id1, identity.clone(), 1000); + let (conn1, identity) = + make_completed_connection(&mut node, link_id1, transport_id, 1000); + let node_addr = *identity.node_addr(); node.add_connection(conn1).unwrap(); node.promote_connection(link_id1, identity.clone(), 1500).unwrap(); @@ -1817,29 +1969,17 @@ mod tests { assert_eq!(node.peer_count(), 1); assert_eq!(node.get_peer(&node_addr).unwrap().link_id(), link_id1); - // Second connection (simulates cross-connection scenario) - let link_id2 = LinkId::new(2); - let conn2 = PeerConnection::inbound(link_id2, 2000); + // Cross-connection tie-breaker logic is tested in peer/mod.rs tests. + // The integration test will cover the real cross-connection path with + // two actual nodes. Here we verify promotion works correctly. - node.add_connection(conn2).unwrap(); - - // Promote second connection - tie-breaker determines outcome - let result = node.promote_connection(link_id2, identity, 2500).unwrap(); - - // One connection should win, one should lose - match result { - PromotionResult::CrossConnectionWon { loser_link_id, .. } => { - assert_eq!(loser_link_id, link_id1); - assert_eq!(node.get_peer(&node_addr).unwrap().link_id(), link_id2); - } - PromotionResult::CrossConnectionLost { winner_link_id } => { - assert_eq!(winner_link_id, link_id1); - assert_eq!(node.get_peer(&node_addr).unwrap().link_id(), link_id1); - } - PromotionResult::Promoted(_) => { - panic!("Expected cross-connection, got normal promotion"); - } - } + // Verify first promotion populated peers_by_index + let peer = node.get_peer(&node_addr).unwrap(); + let our_idx = peer.our_index().unwrap(); + assert_eq!( + node.peers_by_index.get(&(transport_id, our_idx.as_u32())), + Some(&node_addr) + ); // Still only one peer assert_eq!(node.peer_count(), 1); @@ -1848,13 +1988,14 @@ mod tests { #[test] fn test_node_peer_limit() { let mut node = make_node(); + let transport_id = TransportId::new(1); node.set_max_peers(2); // Add two peers via promotion for i in 0..2 { - let identity = make_peer_identity(); let link_id = LinkId::new(i as u64 + 1); - let conn = PeerConnection::outbound(link_id, identity.clone(), 1000); + let (conn, identity) = + make_completed_connection(&mut node, link_id, transport_id, 1000); node.add_connection(conn).unwrap(); node.promote_connection(link_id, identity, 2000).unwrap(); } @@ -1862,9 +2003,9 @@ mod tests { assert_eq!(node.peer_count(), 2); // Third should fail - let identity = make_peer_identity(); let link_id = LinkId::new(3); - let conn = PeerConnection::outbound(link_id, identity.clone(), 3000); + let (conn, identity) = + make_completed_connection(&mut node, link_id, transport_id, 3000); node.add_connection(conn).unwrap(); let result = node.promote_connection(link_id, identity, 4000); @@ -1909,27 +2050,28 @@ mod tests { #[test] fn test_node_sendable_peers() { let mut node = make_node(); + let transport_id = TransportId::new(1); // Add a healthy peer - let identity1 = make_peer_identity(); - let node_addr1 = *identity1.node_addr(); let link_id1 = LinkId::new(1); - let conn1 = PeerConnection::outbound(link_id1, identity1.clone(), 1000); + let (conn1, identity1) = + make_completed_connection(&mut node, link_id1, transport_id, 1000); + let node_addr1 = *identity1.node_addr(); node.add_connection(conn1).unwrap(); node.promote_connection(link_id1, identity1, 2000).unwrap(); // Add another peer and mark it stale (still sendable) - let identity2 = make_peer_identity(); let link_id2 = LinkId::new(2); - let conn2 = PeerConnection::outbound(link_id2, identity2.clone(), 1000); + let (conn2, identity2) = + make_completed_connection(&mut node, link_id2, transport_id, 1000); node.add_connection(conn2).unwrap(); node.promote_connection(link_id2, identity2, 2000).unwrap(); // Add a third peer and mark it disconnected (not sendable) - let identity3 = make_peer_identity(); - let node_addr3 = *identity3.node_addr(); let link_id3 = LinkId::new(3); - let conn3 = PeerConnection::outbound(link_id3, identity3.clone(), 1000); + let (conn3, identity3) = + make_completed_connection(&mut node, link_id3, transport_id, 1000); + let node_addr3 = *identity3.node_addr(); node.add_connection(conn3).unwrap(); node.promote_connection(link_id3, identity3, 2000).unwrap(); node.get_peer_mut(&node_addr3).unwrap().mark_disconnected(); @@ -2039,4 +2181,227 @@ mod tests { node.msg1_rate_limiter.complete_handshake(); assert_eq!(node.msg1_rate_limiter.pending_count(), 0); } + + // === Integration Tests: End-to-End Handshake === + + #[tokio::test] + async fn test_two_node_handshake_udp() { + use crate::config::UdpConfig; + use crate::transport::udp::UdpTransport; + use crate::wire::{build_encrypted, build_msg1}; + use tokio::time::{timeout, Duration}; + + // === Setup: Two nodes with UDP transports on localhost === + + let mut node_a = make_node(); + let mut node_b = make_node(); + + let transport_id_a = TransportId::new(1); + let transport_id_b = TransportId::new(1); + + let udp_config = UdpConfig { + bind_addr: Some("127.0.0.1:0".to_string()), + mtu: Some(1280), + }; + + let (packet_tx_a, mut packet_rx_a) = packet_channel(64); + let (packet_tx_b, mut packet_rx_b) = packet_channel(64); + + let mut transport_a = + UdpTransport::new(transport_id_a, None, udp_config.clone(), packet_tx_a); + let mut transport_b = + UdpTransport::new(transport_id_b, None, udp_config, packet_tx_b); + + transport_a.start_async().await.unwrap(); + transport_b.start_async().await.unwrap(); + + let addr_a = transport_a.local_addr().unwrap(); + let addr_b = transport_b.local_addr().unwrap(); + let remote_addr_b = TransportAddr::from_string(&addr_b.to_string()); + let remote_addr_a = TransportAddr::from_string(&addr_a.to_string()); + + node_a + .transports + .insert(transport_id_a, TransportHandle::Udp(transport_a)); + node_b + .transports + .insert(transport_id_b, TransportHandle::Udp(transport_b)); + + // === Phase 1: Node A initiates handshake to Node B === + + // Create peer identity for B (must use full key for ECDH parity) + let peer_b_identity = + PeerIdentity::from_pubkey_full(node_b.identity.pubkey_full()); + let peer_b_node_addr = *peer_b_identity.node_addr(); + + let link_id_a = node_a.allocate_link_id(); + let mut conn_a = PeerConnection::outbound( + link_id_a, + peer_b_identity.clone(), + 1000, + ); + + // Allocate session index for A's outbound + let our_index_a = node_a.index_allocator.allocate().unwrap(); + + // Start handshake (generates Noise IK msg1) + let our_keypair_a = node_a.identity.keypair(); + let noise_msg1 = conn_a.start_handshake(our_keypair_a, 1000).unwrap(); + conn_a.set_our_index(our_index_a); + conn_a.set_transport_id(transport_id_a); + conn_a.set_source_addr(remote_addr_b.clone()); + + // Build wire msg1 and track in node state + let wire_msg1 = build_msg1(our_index_a, &noise_msg1); + + let link_a = Link::connectionless( + link_id_a, + transport_id_a, + remote_addr_b.clone(), + LinkDirection::Outbound, + Duration::from_millis(100), + ); + node_a.links.insert(link_id_a, link_a); + node_a.connections.insert(link_id_a, conn_a); + node_a.pending_outbound.insert( + (transport_id_a, our_index_a.as_u32()), + link_id_a, + ); + + // Send msg1 from A to B over UDP + let transport = node_a.transports.get(&transport_id_a).unwrap(); + transport + .send(&remote_addr_b, &wire_msg1) + .await + .expect("Failed to send msg1"); + + // === Phase 2: Node B receives msg1, sends msg2, promotes === + + let packet_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) + .await + .expect("Timeout waiting for msg1") + .expect("Channel closed"); + + node_b.handle_msg1(packet_b).await; + + // Verify B promoted the inbound connection + let peer_a_node_addr = *PeerIdentity::from_pubkey_full( + node_a.identity.pubkey_full(), + ) + .node_addr(); + assert_eq!(node_b.peer_count(), 1, "Node B should have 1 peer after msg1"); + let peer_a_on_b = node_b + .get_peer(&peer_a_node_addr) + .expect("Node B should have peer A"); + assert!( + peer_a_on_b.has_session(), + "Peer A on B should have NoiseSession" + ); + let our_index_b = peer_a_on_b.our_index().expect("B should have our_index"); + assert!( + node_b + .peers_by_index + .contains_key(&(transport_id_b, our_index_b.as_u32())), + "Node B peers_by_index should be populated" + ); + + // === Phase 3: Node A receives msg2, completes handshake, promotes === + + let packet_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) + .await + .expect("Timeout waiting for msg2") + .expect("Channel closed"); + + node_a.handle_msg2(packet_a).await; + + // Verify A promoted the outbound connection + assert_eq!(node_a.peer_count(), 1, "Node A should have 1 peer after msg2"); + let peer_b_on_a = node_a + .get_peer(&peer_b_node_addr) + .expect("Node A should have peer B"); + assert!( + peer_b_on_a.has_session(), + "Peer B on A should have NoiseSession" + ); + assert_eq!( + peer_b_on_a.our_index(), + Some(our_index_a), + "Peer B on A should have our_index matching what we allocated" + ); + assert!( + node_a + .peers_by_index + .contains_key(&(transport_id_a, our_index_a.as_u32())), + "Node A peers_by_index should be populated" + ); + + // === Phase 4: Encrypted frame A → B === + + // A encrypts a test message and sends to B + let plaintext_a = b"hello from A"; + let peer_b = node_a.get_peer_mut(&peer_b_node_addr).unwrap(); + let their_index_b = peer_b.their_index().expect("A should know B's index"); + let session_a = peer_b.noise_session_mut().unwrap(); + let ciphertext_a = session_a.encrypt(plaintext_a).unwrap(); + + let wire_encrypted = build_encrypted(their_index_b, 0, &ciphertext_a); + let transport = node_a.transports.get(&transport_id_a).unwrap(); + transport + .send(&remote_addr_b, &wire_encrypted) + .await + .expect("Failed to send encrypted frame"); + + // B receives and decrypts + let encrypted_packet_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) + .await + .expect("Timeout waiting for encrypted frame") + .expect("Channel closed"); + + node_b.handle_encrypted_frame(encrypted_packet_b).await; + + // Verify B's peer was touched (last_seen updated) + let peer_a = node_b.get_peer(&peer_a_node_addr).unwrap(); + assert!( + peer_a.is_healthy(), + "Peer A on B should still be healthy after receiving encrypted frame" + ); + + // === Phase 5: Encrypted frame B → A === + + let plaintext_b = b"hello from B"; + let peer_a = node_b.get_peer_mut(&peer_a_node_addr).unwrap(); + let their_index_a = peer_a.their_index().expect("B should know A's index"); + let session_b = peer_a.noise_session_mut().unwrap(); + let ciphertext_b = session_b.encrypt(plaintext_b).unwrap(); + + let wire_encrypted_b = build_encrypted(their_index_a, 0, &ciphertext_b); + let transport = node_b.transports.get(&transport_id_b).unwrap(); + transport + .send(&remote_addr_a, &wire_encrypted_b) + .await + .expect("Failed to send encrypted frame B→A"); + + // A receives and decrypts + let encrypted_packet_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) + .await + .expect("Timeout waiting for encrypted frame B→A") + .expect("Channel closed"); + + node_a.handle_encrypted_frame(encrypted_packet_a).await; + + // Verify A's peer was touched + let peer_b = node_a.get_peer(&peer_b_node_addr).unwrap(); + assert!( + peer_b.is_healthy(), + "Peer B on A should still be healthy after receiving encrypted frame" + ); + + // Clean up transports + for (_, t) in node_a.transports.iter_mut() { + t.stop().await.ok(); + } + for (_, t) in node_b.transports.iter_mut() { + t.stop().await.ok(); + } + } }