diff --git a/src/node/handlers/handshake.rs b/src/node/handlers/handshake.rs index 18bafe2..4a96beb 100644 --- a/src/node/handlers/handshake.rs +++ b/src/node/handlers/handshake.rs @@ -1,18 +1,28 @@ //! Handshake handlers and connection promotion. +//! +//! Implements the Noise XX 3-message handshake for FMP link establishment: +//! - msg1 (initiator → responder): ephemeral only, no identity +//! - msg2 (responder → initiator): responder identity + epoch + negotiation +//! - msg3 (initiator → responder): initiator identity + epoch + negotiation -use crate::PeerIdentity; -use crate::node::wire::{Msg1Header, Msg2Header, build_msg2}; use crate::node::{Node, NodeError}; -use crate::peer::{ActivePeer, PeerConnection, PromotionResult, cross_connection_winner}; +use crate::peer::{ + cross_connection_winner, ActivePeer, PeerConnection, PromotionResult, +}; +use crate::protocol::NegotiationPayload; use crate::transport::{Link, LinkDirection, LinkId, ReceivedPacket}; +use crate::node::wire::{build_msg2, build_msg3, Msg1Header, Msg2Header, Msg3Header}; +use crate::PeerIdentity; use std::time::Duration; use tracing::{debug, info, warn}; impl Node { /// Handle handshake message 1 (phase 0x1). /// - /// This creates a new inbound connection. Rate limiting is applied - /// before any expensive crypto operations. + /// With Noise XX, msg1 contains only the initiator's ephemeral key. + /// No identity is learned. The responder processes msg1, sends msg2 + /// (revealing its own identity), and stores the connection in + /// pending_inbound to await msg3. pub(in crate::node) async fn handle_msg1(&mut self, packet: ReceivedPacket) { // === RATE LIMITING (before any processing) === if !self.msg1_rate_limiter.start_handshake() { @@ -44,29 +54,22 @@ impl Node { // Check for existing connection from this address. // - // If we already have an *inbound* link from this address, this could be: - // 1. A duplicate msg1 (our msg2 was lost) — resend msg2 - // 2. A restarted peer (different epoch) — tear down and reprocess - // - // If we have an *outbound* link to this address (we initiated to them - // AND they initiated to us), this is a cross-connection — allow it. - // - // Epoch-based restart detection: if the sender already has an inbound - // link AND is an active peer in self.peers, fall through to decrypt - // the msg1 and check the epoch. Otherwise, treat as duplicate. + // With XX, we can't do identity-based checks in msg1 (no identity yet). + // We can only detect duplicates by address: if we already have an inbound + // link from this address with a pending connection, resend msg2. + // If we have an active peer on this address, it could be a restart or + // rekey — but we can't tell until msg3 reveals identity. For now, allow + // the new handshake to proceed. Identity-based checks happen in handle_msg3. let addr_key = (packet.transport_id, packet.remote_addr.clone()); - let mut possible_restart = false; if let Some(&existing_link_id) = self.addr_to_link.get(&addr_key) && let Some(link) = self.links.get(&existing_link_id) { if link.direction() == LinkDirection::Inbound { // Check if this link belongs to an already-promoted active peer - let is_active_peer = self.peers.values().any(|p| p.link_id() == existing_link_id); + let is_active_peer = self.peers.values() + .any(|p| p.link_id() == existing_link_id); - if is_active_peer { - // Possible restart — fall through to decrypt and check epoch - possible_restart = true; - } else { + if !is_active_peer { // Genuinely pending handshake — resend msg2 let msg2_bytes = self.find_stored_msg2(existing_link_id); if let Some(msg2) = msg2_bytes { @@ -92,22 +95,23 @@ impl Node { self.msg1_rate_limiter.complete_handshake(); return; } + // Active peer on this address — allow the new handshake. + // Identity checks (restart, rekey) deferred to handle_msg3. + debug!( + transport_id = %packet.transport_id, + remote_addr = %packet.remote_addr, + existing_link_id = %existing_link_id, + "XX msg1 from address with active peer — proceeding (identity check deferred to msg3)" + ); } else { - // Outbound link to this address. If it belongs to an active - // peer, this may be a rekey msg1 (same epoch) or a - // restart (different epoch). Set possible_restart to enable - // the epoch/rekey check below. - let is_active_peer = self.peers.values().any(|p| p.link_id() == existing_link_id); - if is_active_peer { - possible_restart = true; - } else { - debug!( - transport_id = %packet.transport_id, - remote_addr = %packet.remote_addr, - existing_link_id = %existing_link_id, - "Cross-connection detected: have outbound, received inbound msg1" - ); - } + // Outbound link to this address — cross-connection. + // Allow the inbound handshake to proceed. + debug!( + transport_id = %packet.transport_id, + remote_addr = %packet.remote_addr, + existing_link_id = %existing_link_id, + "Cross-connection detected: have outbound, received inbound msg1" + ); } } @@ -120,12 +124,16 @@ impl Node { packet.timestamp_ms, ); + // Create negotiation payload for msg2 + let neg_payload = NegotiationPayload::new(1, 1, 0).encode(); + let our_keypair = self.identity.keypair(); let noise_msg1 = &packet.data[header.noise_msg1_offset..]; let msg2_response = match conn.receive_handshake_init( our_keypair, self.startup_epoch, noise_msg1, + Some(&neg_payload), packet.timestamp_ms, ) { Ok(m) => m, @@ -139,213 +147,8 @@ impl Node { } }; - // Learn peer identity from msg1 - let peer_identity = match conn.expected_identity() { - Some(id) => *id, - None => { - self.msg1_rate_limiter.complete_handshake(); - warn!("Identity not learned from msg1"); - return; - } - }; - - let peer_node_addr = *peer_identity.node_addr(); - - // Identity-based restart/rekey detection: if the peer is already - // active but addr_to_link didn't match (different source address, e.g., - // TCP from a different port), we still need to check for restart/rekey. - if !possible_restart && self.peers.contains_key(&peer_node_addr) { - possible_restart = true; - } - - // Epoch-based restart detection and duplicate msg1 handling. - // - // If we fell through from the addr_to_link check above with - // possible_restart=true, we now have the decrypted epoch from msg1. - // Compare it against the stored epoch for this peer. - if possible_restart && let Some(existing_peer) = self.peers.get(&peer_node_addr) { - let new_epoch = conn.remote_epoch(); - let existing_epoch = existing_peer.remote_epoch(); - - match (existing_epoch, new_epoch) { - (Some(existing), Some(new)) if existing != new => { - // Epoch mismatch — peer restarted. Tear down stale session. - info!( - peer = %self.peer_display_name(&peer_node_addr), - "Peer restart detected (epoch mismatch), removing stale session" - ); - self.remove_active_peer(&peer_node_addr); - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - self.schedule_reconnect(peer_node_addr, now_ms); - // Fall through to process as new connection - } - _ => { - // Same epoch (or no epoch stored). - // If the peer has an active session and rekey is enabled, - // this is a rekey msg1 (not a duplicate initial msg1). - // Guard: the session must be at least 30s old to avoid - // misidentifying a cross-connection msg1 as a rekey. - // During simultaneous connection, both sides promote - // within the same tick and the peer's msg1 arrives - // immediately — a genuine rekey can't fire that fast. - let session_age_secs = - existing_peer.session_established_at().elapsed().as_secs(); - if self.config.node.rekey.enabled - && existing_peer.has_session() - && existing_peer.is_healthy() - && session_age_secs >= 30 - { - // Guard: already have a pending session from a completed - // rekey (waiting for K-bit cutover). Don't overwrite it - // with a new handshake — drop this msg1. - if existing_peer.pending_new_session().is_some() { - debug!( - peer = %self.peer_display_name(&peer_node_addr), - "Rekey msg1 received but already have pending session, dropping" - ); - self.connections.remove(&link_id); - self.links.remove(&link_id); - self.msg1_rate_limiter.complete_handshake(); - return; - } - - // Dual-initiation detection: both sides sent msg1 - // simultaneously. Apply tie-breaker — smaller NodeAddr - // wins as initiator (same as cross-connection resolution). - if existing_peer.rekey_in_progress() { - let our_addr = self.identity.node_addr(); - if our_addr < &peer_node_addr { - // We win as initiator — drop their msg1. - // Our msg2 will arrive at peer, who completes - // as our responder. - debug!( - peer = %self.peer_display_name(&peer_node_addr), - "Dual rekey initiation: we win (smaller addr), dropping their msg1" - ); - self.connections.remove(&link_id); - self.links.remove(&link_id); - self.msg1_rate_limiter.complete_handshake(); - return; - } - // We lose — abandon our rekey, become responder below. - debug!( - peer = %self.peer_display_name(&peer_node_addr), - "Dual rekey initiation: we lose (larger addr), abandoning ours" - ); - if let Some(peer) = self.peers.get_mut(&peer_node_addr) - && let Some(idx) = peer.abandon_rekey() - { - if let Some(tid) = peer.transport_id() { - self.peers_by_index.remove(&(tid, idx.as_u32())); - self.pending_outbound.remove(&(tid, idx.as_u32())); - } - let _ = self.index_allocator.free(idx); - } - // Fall through to respond as responder - } - - // Rekey: process as responder, store new session as pending - let noise_session = conn.take_session(); - let our_new_index = match self.index_allocator.allocate() { - Ok(idx) => idx, - Err(e) => { - warn!(error = %e, "Failed to allocate index for rekey"); - self.msg1_rate_limiter.complete_handshake(); - return; - } - }; - - let noise_session = match noise_session { - Some(s) => s, - None => { - warn!("Rekey msg1: no session from handshake"); - let _ = self.index_allocator.free(our_new_index); - self.msg1_rate_limiter.complete_handshake(); - return; - } - }; - - // Send msg2 response using the new handshake - let wire_msg2 = - build_msg2(our_new_index, header.sender_idx, &msg2_response); - if let Some(transport) = self.transports.get(&packet.transport_id) { - match transport.send(&packet.remote_addr, &wire_msg2).await { - Ok(_) => { - debug!( - peer = %self.peer_display_name(&peer_node_addr), - new_our_index = %our_new_index, - "Sent rekey msg2 response" - ); - } - Err(e) => { - warn!( - peer = %self.peer_display_name(&peer_node_addr), - error = %e, - "Failed to send rekey msg2" - ); - let _ = self.index_allocator.free(our_new_index); - self.msg1_rate_limiter.complete_handshake(); - return; - } - } - } - - // Store pending session on the existing peer - if let Some(peer) = self.peers.get_mut(&peer_node_addr) { - peer.set_pending_session( - noise_session, - our_new_index, - header.sender_idx, - ); - peer.record_peer_rekey(); - } - - // Register new index in peers_by_index - self.peers_by_index.insert( - (packet.transport_id, our_new_index.as_u32()), - peer_node_addr, - ); - - // Clean up: remove the temporary connection/link we created. - // Do NOT remove addr_to_link — the entry must remain pointing - // to the original link so future msg1s from this address are - // recognized as rekeys (not new connections). - self.connections.remove(&link_id); - self.links.remove(&link_id); - - self.msg1_rate_limiter.complete_handshake(); - return; - } - - // Not a rekey — duplicate msg1. Resend stored msg2. - if let Some(msg2) = existing_peer.handshake_msg2().map(|m| m.to_vec()) - && let Some(transport) = self.transports.get(&packet.transport_id) - { - match transport.send(&packet.remote_addr, &msg2).await { - Ok(_) => debug!( - peer = %self.peer_display_name(&peer_node_addr), - "Resent msg2 for duplicate msg1 (same epoch)" - ), - Err(e) => debug!( - peer = %self.peer_display_name(&peer_node_addr), - error = %e, - "Failed to resend msg2" - ), - } - } - self.msg1_rate_limiter.complete_handshake(); - return; - } - } - } - // If possible_restart was true but peer is no longer in self.peers - // (removed by another path), fall through to process as new connection. - - // Note: we don't early-return if peer is already in self.peers here. - // promote_connection handles cross-connection resolution via tie-breaker. + // XX: identity is NOT learned from msg1 (only ephemeral exchange). + // Identity will be learned from msg3 in handle_msg3. // Allocate our session index let our_index = match self.index_allocator.allocate() { @@ -399,8 +202,7 @@ impl Node { // Clean up on failure self.connections.remove(&link_id); self.links.remove(&link_id); - self.addr_to_link - .remove(&(packet.transport_id, packet.remote_addr)); + self.addr_to_link.remove(&(packet.transport_id, packet.remote_addr)); let _ = self.index_allocator.free(our_index); self.msg1_rate_limiter.complete_handshake(); return; @@ -408,92 +210,12 @@ impl Node { } } - // 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) => { - // Store msg2 on peer for resend on duplicate msg1 - if let Some(peer) = self.peers.get_mut(&node_addr) { - peer.set_handshake_msg2(wire_msg2.clone()); - } - debug!( - peer = %self.peer_display_name(&node_addr), - link_id = %link_id, - our_index = %our_index, - "Inbound peer promoted to active" - ); - // Send initial tree announce to new peer - if let Err(e) = self.send_tree_announce_to_peer(&node_addr).await { - debug!(peer = %self.peer_display_name(&node_addr), error = %e, "Failed to send initial TreeAnnounce"); - } - // Schedule filter announce (sent on next tick via debounce) - self.bloom_state.mark_update_needed(node_addr); - self.reset_discovery_backoff(); - } - PromotionResult::CrossConnectionWon { - loser_link_id, - node_addr, - } => { - // Store msg2 on peer for resend on duplicate msg1 - if let Some(peer) = self.peers.get_mut(&node_addr) { - peer.set_handshake_msg2(wire_msg2.clone()); - } - // Close the losing TCP connection (no-op for connectionless) - if let Some(loser_link) = self.links.get(&loser_link_id) { - let loser_tid = loser_link.transport_id(); - let loser_addr = loser_link.remote_addr().clone(); - if let Some(transport) = self.transports.get(&loser_tid) { - transport.close_connection(&loser_addr).await; - } - } - // Clean up the losing connection's link - self.remove_link(&loser_link_id); - debug!( - peer = %self.peer_display_name(&node_addr), - loser_link_id = %loser_link_id, - "Inbound cross-connection won, loser link cleaned up" - ); - // Send initial tree announce to peer (new or reconnected) - if let Err(e) = self.send_tree_announce_to_peer(&node_addr).await { - debug!(peer = %self.peer_display_name(&node_addr), error = %e, "Failed to send initial TreeAnnounce"); - } - // Schedule filter announce (sent on next tick via debounce) - self.bloom_state.mark_update_needed(node_addr); - self.reset_discovery_backoff(); - } - PromotionResult::CrossConnectionLost { winner_link_id } => { - // Close the losing TCP connection (no-op for connectionless) - if let Some(transport) = self.transports.get(&packet.transport_id) { - transport.close_connection(&packet.remote_addr).await; - } - // This connection lost — clean up its link - self.remove_link(&link_id); - // Restore addr_to_link for the winner's link - self.addr_to_link.insert( - (packet.transport_id, packet.remote_addr.clone()), - winner_link_id, - ); - debug!( - 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.remove_link(&link_id); - let _ = self.index_allocator.free(our_index); - } - } + // XX: handshake NOT complete yet — need msg3. + // Store in pending_inbound for msg3 dispatch. + self.pending_inbound.insert( + (packet.transport_id, our_index.as_u32()), + link_id, + ); self.msg1_rate_limiter.complete_handshake(); } @@ -522,7 +244,10 @@ impl Node { /// Handle handshake message 2 (phase 0x2). /// - /// This completes an outbound handshake we initiated. + /// With Noise XX, processing msg2 learns the responder's identity and + /// generates msg3 which must be sent before the handshake is complete. + /// After sending msg3, the initiator's handshake is complete and the + /// connection is promoted. pub(in crate::node) async fn handle_msg2(&mut self, packet: ReceivedPacket) { // Parse header let header = match Msg2Header::parse(&packet.data) { @@ -554,7 +279,9 @@ impl Node { // Find peer with rekey in progress for this index let peer_addr = self.peers.iter().find_map(|(addr, peer)| { - if peer.rekey_in_progress() && peer.rekey_our_index() == Some(header.receiver_idx) { + if peer.rekey_in_progress() + && peer.rekey_our_index() == Some(header.receiver_idx) + { Some(*addr) } else { None @@ -565,23 +292,69 @@ impl Node { let display_name = self.peer_display_name(&peer_node_addr); // Complete the rekey handshake on the ActivePeer + // XX: complete_rekey_msg2 processes msg2 and generates msg3 + let transport_id = self.peers.get(&peer_node_addr) + .and_then(|p| p.transport_id()); + let remote_addr = self.peers.get(&peer_node_addr) + .and_then(|p| p.current_addr().cloned()); + if let Some(peer) = self.peers.get_mut(&peer_node_addr) { match peer.complete_rekey_msg2(noise_msg2) { - Ok(session) => { - let our_index = peer.rekey_our_index().unwrap_or(header.receiver_idx); - peer.set_pending_session(session, our_index, header.sender_idx); + Ok((msg3_bytes, session)) => { + let our_index = peer.rekey_our_index() + .unwrap_or(header.receiver_idx); - if let Some(transport_id) = peer.transport_id() { - self.peers_by_index - .insert((transport_id, our_index.as_u32()), peer_node_addr); + // Send msg3 before setting pending session + let wire_msg3 = build_msg3(our_index, header.sender_idx, &msg3_bytes); + let msg3_sent = if let (Some(tid), Some(addr)) = (transport_id, &remote_addr) + && let Some(transport) = self.transports.get(&tid) + { + match transport.send(addr, &wire_msg3).await { + Ok(_) => { + debug!( + peer = %display_name, + "Sent rekey msg3" + ); + true + } + Err(e) => { + warn!( + peer = %display_name, + error = %e, + "Failed to send rekey msg3" + ); + false + } + } + } else { + false + }; + + if msg3_sent { + peer.set_pending_session(session, our_index, header.sender_idx); + + if let Some(tid) = transport_id { + self.peers_by_index.insert( + (tid, our_index.as_u32()), + peer_node_addr, + ); + } + + debug!( + peer = %display_name, + new_our_index = %our_index, + new_their_index = %header.sender_idx, + "Rekey completed (initiator), pending K-bit cutover" + ); + } else { + // msg3 send failed — abandon rekey + if let Some(idx) = peer.abandon_rekey() { + if let Some(tid) = peer.transport_id() { + self.peers_by_index.remove(&(tid, idx.as_u32())); + } + let _ = self.index_allocator.free(idx); + } } - - debug!( - peer = %display_name, - new_our_index = %our_index, - new_their_index = %header.sender_idx, - "Rekey completed (initiator), pending K-bit cutover" - ); } Err(e) => { warn!( @@ -610,23 +383,29 @@ impl Node { let conn = self.connections.get_mut(&link_id).unwrap(); - // Process Noise msg2 + // Create negotiation payload for msg3 + let neg_payload = NegotiationPayload::new(1, 1, 0).encode(); + + // Process Noise msg2 and generate msg3 let noise_msg2 = &packet.data[header.noise_msg2_offset..]; - if let Err(e) = conn.complete_handshake(noise_msg2, packet.timestamp_ms) { - warn!( - link_id = %link_id, - error = %e, - "Handshake completion failed" - ); - conn.mark_failed(); - return; - } + let (msg3_bytes, _received_negotiation) = match conn.complete_handshake(noise_msg2, Some(&neg_payload), packet.timestamp_ms) { + Ok(result) => result, + Err(e) => { + warn!( + link_id = %link_id, + error = %e, + "Handshake completion failed" + ); + conn.mark_failed(); + return; + } + }; // Store their index conn.set_their_index(header.sender_idx); conn.set_source_addr(packet.remote_addr.clone()); - // Get peer identity for promotion + // Get peer identity for promotion (learned from msg2 in XX) let peer_identity = match conn.expected_identity() { Some(id) => *id, None => { @@ -637,6 +416,35 @@ impl Node { let peer_node_addr = *peer_identity.node_addr(); + // Build and send msg3 + let our_index = conn.our_index().unwrap_or(header.receiver_idx); + let wire_msg3 = build_msg3(our_index, header.sender_idx, &msg3_bytes); + + if let Some(transport) = self.transports.get(&packet.transport_id) { + match transport.send(&packet.remote_addr, &wire_msg3).await { + Ok(bytes) => { + debug!( + peer = %self.peer_display_name(&peer_node_addr), + link_id = %link_id, + their_index = %header.sender_idx, + bytes, + "Sent msg3, outbound handshake completing" + ); + } + Err(e) => { + warn!( + link_id = %link_id, + error = %e, + "Failed to send msg3" + ); + if let Some(conn) = self.connections.get_mut(&link_id) { + conn.mark_failed(); + } + return; + } + } + } + debug!( peer = %self.peer_display_name(&peer_node_addr), link_id = %link_id, @@ -645,7 +453,7 @@ impl Node { ); // Cross-connection resolution: if the peer was already promoted via - // our inbound handshake (we processed their msg1), both nodes initially + // our inbound handshake (we processed their msg3), both nodes initially // use mismatched sessions. The tie-breaker determines which handshake // wins: smaller node_addr's outbound. // @@ -676,17 +484,15 @@ impl Node { let outbound_our_index = conn.our_index(); let outbound_session = conn.take_session(); - let (outbound_session, outbound_our_index) = match ( - outbound_session, - outbound_our_index, - ) { - (Some(s), Some(idx)) => (s, idx), - _ => { - warn!(peer = %self.peer_display_name(&peer_node_addr), "Incomplete outbound connection"); - self.pending_outbound.remove(&key); - return; - } - }; + let (outbound_session, outbound_our_index) = + match (outbound_session, outbound_our_index) { + (Some(s), Some(idx)) => (s, idx), + _ => { + warn!(peer = %self.peer_display_name(&peer_node_addr), "Incomplete outbound connection"); + self.pending_outbound.remove(&key); + return; + } + }; if let Some(peer) = self.peers.get_mut(&peer_node_addr) { let suppressed = peer.replay_suppressed_count(); @@ -699,12 +505,13 @@ impl Node { // Update peers_by_index: remove old inbound index, add outbound let transport_id = peer.transport_id().unwrap(); if let Some(old_idx) = old_our_index { - self.peers_by_index - .remove(&(transport_id, old_idx.as_u32())); + self.peers_by_index.remove(&(transport_id, old_idx.as_u32())); let _ = self.index_allocator.free(old_idx); } - self.peers_by_index - .insert((transport_id, outbound_our_index.as_u32()), peer_node_addr); + self.peers_by_index.insert( + (transport_id, outbound_our_index.as_u32()), + peer_node_addr, + ); if suppressed > 0 { debug!( @@ -770,6 +577,12 @@ impl Node { } // Normal path: promote to active peer + debug!( + peer = %self.peer_display_name(&peer_node_addr), + link_id = %link_id, + "handle_msg2: promoting outbound, peers_has_key={}", + self.peers.contains_key(&peer_node_addr), + ); match self.promote_connection(link_id, peer_identity, packet.timestamp_ms) { Ok(result) => { // Clean up pending_outbound @@ -789,10 +602,7 @@ impl Node { self.bloom_state.mark_update_needed(node_addr); self.reset_discovery_backoff(); } - PromotionResult::CrossConnectionWon { - loser_link_id, - node_addr, - } => { + PromotionResult::CrossConnectionWon { loser_link_id, node_addr } => { // Close the losing TCP connection (no-op for connectionless) if let Some(loser_link) = self.links.get(&loser_link_id) { let loser_tid = loser_link.transport_id(); @@ -804,8 +614,10 @@ impl Node { // Clean up the losing connection's link self.remove_link(&loser_link_id); // Ensure addr_to_link points to the winning link - self.addr_to_link - .insert((packet.transport_id, packet.remote_addr.clone()), link_id); + self.addr_to_link.insert( + (packet.transport_id, packet.remote_addr.clone()), + link_id, + ); debug!( peer = %self.peer_display_name(&node_addr), loser_link_id = %loser_link_id, @@ -848,6 +660,385 @@ impl Node { } } + /// Handle handshake message 3 (phase 0x3). + /// + /// Completes the XX handshake on the responder side. Processes msg3 to + /// learn the initiator's identity and epoch, then performs identity-based + /// checks (restart detection, rekey detection, cross-connection resolution) + /// and promotes the connection to active peer. + pub(in crate::node) async fn handle_msg3(&mut self, packet: ReceivedPacket) { + // Parse header + let header = match Msg3Header::parse(&packet.data) { + Some(h) => h, + None => { + debug!("Invalid msg3 header"); + return; + } + }; + + // Look up our pending inbound handshake by our index (receiver_idx in msg3) + let key = (packet.transport_id, header.receiver_idx.as_u32()); + let link_id = match self.pending_inbound.remove(&key) { + Some(id) => id, + None => { + // Check if this is a rekey msg3 for an active peer + self.handle_rekey_msg3(&packet, &header).await; + return; + } + }; + + // Get the pending connection + let conn = match self.connections.get_mut(&link_id) { + Some(c) => c, + None => { + debug!( + link_id = %link_id, + "No pending connection for msg3" + ); + return; + } + }; + + // Process msg3 — learns initiator's identity and epoch + let noise_msg3 = &packet.data[header.noise_msg3_offset..]; + let _received_negotiation = match conn.complete_handshake_msg3(noise_msg3, packet.timestamp_ms) { + Ok(neg) => neg, + Err(e) => { + warn!( + link_id = %link_id, + error = %e, + "Msg3 processing failed" + ); + // Clean up + self.connections.remove(&link_id); + self.remove_link(&link_id); + if let Some(idx) = self.connections.get(&link_id).and_then(|c| c.our_index()) { + let _ = self.index_allocator.free(idx); + } + return; + } + }; + + // Learn peer identity from msg3 + let peer_identity = match conn.expected_identity() { + Some(id) => *id, + None => { + warn!("Identity not learned from msg3"); + self.connections.remove(&link_id); + self.remove_link(&link_id); + return; + } + }; + + let peer_node_addr = *peer_identity.node_addr(); + let our_index = conn.our_index().unwrap_or(header.receiver_idx); + + // Identity-based restart/rekey detection. + // + // Now that we know the initiator's identity from msg3, perform the + // same checks that IK handle_msg1 used to do after decrypting msg1. + if let Some(existing_peer) = self.peers.get(&peer_node_addr) { + let new_epoch = conn.remote_epoch(); + let existing_epoch = existing_peer.remote_epoch(); + + match (existing_epoch, new_epoch) { + (Some(existing), Some(new)) if existing != new => { + // Epoch mismatch — peer restarted. Tear down stale session. + info!( + peer = %self.peer_display_name(&peer_node_addr), + "Peer restart detected (epoch mismatch), removing stale session" + ); + self.remove_active_peer(&peer_node_addr); + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + self.schedule_reconnect(peer_node_addr, now_ms); + // Fall through to process as new connection + } + _ => { + // Same epoch (or no epoch stored). + // Check for rekey: session must be at least 30s old. + let session_age_secs = existing_peer + .session_established_at() + .elapsed() + .as_secs(); + if self.config.node.rekey.enabled + && existing_peer.has_session() + && existing_peer.is_healthy() + && session_age_secs >= 30 + { + // Guard: already have a pending session from a completed + // rekey (waiting for K-bit cutover). Don't overwrite. + if existing_peer.pending_new_session().is_some() { + debug!( + peer = %self.peer_display_name(&peer_node_addr), + "Rekey msg3 received but already have pending session, dropping" + ); + self.connections.remove(&link_id); + self.links.remove(&link_id); + return; + } + + // Dual-initiation detection: both sides sent msg1 + // simultaneously. Apply tie-breaker. + if existing_peer.rekey_in_progress() { + let our_addr = self.identity.node_addr(); + if our_addr < &peer_node_addr { + // We win as initiator — drop their msg3/handshake. + debug!( + peer = %self.peer_display_name(&peer_node_addr), + "Dual rekey initiation: we win (smaller addr), dropping their msg3" + ); + self.connections.remove(&link_id); + self.links.remove(&link_id); + return; + } + // We lose — abandon our rekey, become responder below. + debug!( + peer = %self.peer_display_name(&peer_node_addr), + "Dual rekey initiation: we lose (larger addr), abandoning ours" + ); + if let Some(peer) = self.peers.get_mut(&peer_node_addr) + && let Some(idx) = peer.abandon_rekey() + { + if let Some(tid) = peer.transport_id() { + self.peers_by_index.remove(&(tid, idx.as_u32())); + self.pending_outbound.remove(&(tid, idx.as_u32())); + } + let _ = self.index_allocator.free(idx); + } + // Fall through to respond as responder + } + + // Rekey: process as responder, store new session as pending + let noise_session = { + let conn = self.connections.get_mut(&link_id).unwrap(); + conn.take_session() + }; + let our_new_index = our_index; + + let noise_session = match noise_session { + Some(s) => s, + None => { + warn!("Rekey msg3: no session from handshake"); + self.connections.remove(&link_id); + self.links.remove(&link_id); + return; + } + }; + + // Store pending session on the existing peer + if let Some(peer) = self.peers.get_mut(&peer_node_addr) { + peer.set_pending_session( + noise_session, + our_new_index, + header.sender_idx, + ); + peer.record_peer_rekey(); + } + + // Register new index in peers_by_index + self.peers_by_index.insert( + (packet.transport_id, our_new_index.as_u32()), + peer_node_addr, + ); + + // Clean up: remove the temporary connection/link. + // Do NOT remove addr_to_link — the entry must remain pointing + // to the original link. + self.connections.remove(&link_id); + self.links.remove(&link_id); + + debug!( + peer = %self.peer_display_name(&peer_node_addr), + new_our_index = %our_new_index, + "Rekey completed (responder), pending K-bit cutover" + ); + return; + } + + // Not a rekey — duplicate handshake from same epoch. + // Resend stored msg2. + if let Some(msg2) = existing_peer.handshake_msg2().map(|m| m.to_vec()) + && let Some(transport) = self.transports.get(&packet.transport_id) + { + match transport.send(&packet.remote_addr, &msg2).await { + Ok(_) => debug!( + peer = %self.peer_display_name(&peer_node_addr), + "Resent msg2 for duplicate handshake (same epoch)" + ), + Err(e) => debug!( + peer = %self.peer_display_name(&peer_node_addr), + error = %e, + "Failed to resend msg2" + ), + } + } + self.connections.remove(&link_id); + self.links.remove(&link_id); + return; + } + } + } + + // Promote the connection to active peer. + let wire_msg2 = self.connections.get(&link_id) + .and_then(|c| c.handshake_msg2().map(|m| m.to_vec())); + + debug!( + peer = %self.peer_display_name(&peer_node_addr), + link_id = %link_id, + our_index = %our_index, + "handle_msg3: promoting inbound, peers_has_key={}", + self.peers.contains_key(&peer_node_addr), + ); + match self.promote_connection(link_id, peer_identity, packet.timestamp_ms) { + Ok(result) => { + match result { + PromotionResult::Promoted(node_addr) => { + // Store msg2 on peer for resend on duplicate msg1 + if let (Some(peer), Some(msg2)) = (self.peers.get_mut(&node_addr), wire_msg2) { + peer.set_handshake_msg2(msg2); + } + debug!( + peer = %self.peer_display_name(&node_addr), + link_id = %link_id, + our_index = %our_index, + "Inbound peer promoted to active" + ); + // Send initial tree announce to new peer + if let Err(e) = self.send_tree_announce_to_peer(&node_addr).await { + debug!(peer = %self.peer_display_name(&node_addr), error = %e, "Failed to send initial TreeAnnounce"); + } + // Schedule filter announce (sent on next tick via debounce) + self.bloom_state.mark_update_needed(node_addr); + self.reset_discovery_backoff(); + } + PromotionResult::CrossConnectionWon { loser_link_id, node_addr } => { + // Store msg2 on peer for resend on duplicate msg1 + if let (Some(peer), Some(msg2)) = (self.peers.get_mut(&node_addr), wire_msg2) { + peer.set_handshake_msg2(msg2); + } + // Close the losing TCP connection (no-op for connectionless) + if let Some(loser_link) = self.links.get(&loser_link_id) { + let loser_tid = loser_link.transport_id(); + let loser_addr = loser_link.remote_addr().clone(); + if let Some(transport) = self.transports.get(&loser_tid) { + transport.close_connection(&loser_addr).await; + } + } + // Clean up the losing connection's link + self.remove_link(&loser_link_id); + debug!( + peer = %self.peer_display_name(&node_addr), + loser_link_id = %loser_link_id, + "Inbound cross-connection won, loser link cleaned up" + ); + if let Err(e) = self.send_tree_announce_to_peer(&node_addr).await { + debug!(peer = %self.peer_display_name(&node_addr), error = %e, "Failed to send initial TreeAnnounce"); + } + self.bloom_state.mark_update_needed(node_addr); + self.reset_discovery_backoff(); + } + PromotionResult::CrossConnectionLost { winner_link_id } => { + // Close the losing TCP connection (no-op for connectionless) + if let Some(transport) = self.transports.get(&packet.transport_id) { + transport.close_connection(&packet.remote_addr).await; + } + // This connection lost — clean up its link + self.remove_link(&link_id); + // Restore addr_to_link for the winner's link + self.addr_to_link.insert( + (packet.transport_id, packet.remote_addr.clone()), + winner_link_id, + ); + debug!( + 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.remove_link(&link_id); + let _ = self.index_allocator.free(our_index); + } + } + } + + /// Handle a rekey msg3 for an already-active peer. + /// + /// When a rekey is in progress (responder side), the ActivePeer holds the + /// handshake state. This processes msg3 to complete the rekey responder + /// handshake. + async fn handle_rekey_msg3(&mut self, packet: &ReceivedPacket, header: &Msg3Header) { + // Look for a peer expecting a rekey msg3 as responder. + // The responder's rekey handshake state is stored after processing + // the initiator's rekey msg1+msg2 exchange via the existing link. + let peer_addr = self.peers.iter().find_map(|(addr, peer)| { + if peer.has_rekey_responder_handshake() + && peer.rekey_responder_our_index() == Some(header.receiver_idx) + { + Some(*addr) + } else { + None + } + }); + + let peer_node_addr = match peer_addr { + Some(addr) => addr, + None => { + debug!( + receiver_idx = %header.receiver_idx, + "No pending inbound or rekey state for msg3" + ); + return; + } + }; + + let display_name = self.peer_display_name(&peer_node_addr); + let noise_msg3 = &packet.data[header.noise_msg3_offset..]; + + if let Some(peer) = self.peers.get_mut(&peer_node_addr) { + match peer.complete_rekey_msg3(noise_msg3) { + Ok(session) => { + let our_index = peer.rekey_responder_our_index() + .unwrap_or(header.receiver_idx); + peer.set_pending_session(session, our_index, header.sender_idx); + peer.record_peer_rekey(); + + if let Some(transport_id) = peer.transport_id() { + self.peers_by_index.insert( + (transport_id, our_index.as_u32()), + peer_node_addr, + ); + } + + debug!( + peer = %display_name, + new_our_index = %our_index, + "Rekey msg3 completed (responder), pending K-bit cutover" + ); + } + Err(e) => { + warn!( + peer = %display_name, + error = %e, + "Rekey msg3 processing failed" + ); + peer.clear_rekey_responder(); + } + } + } + } + /// Promote a connection to active peer after successful authentication. /// /// Handles cross-connection detection and resolution using tie-breaker rules. @@ -872,31 +1063,30 @@ impl Node { .take_session() .ok_or(NodeError::NoSession(link_id))?; - let our_index = connection - .our_index() - .ok_or_else(|| NodeError::PromotionFailed { + 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 { + } + })?; + 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 { + } + })?; + 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 { + } + })?; + let current_addr = connection.source_addr().ok_or_else(|| { + NodeError::PromotionFailed { link_id, reason: "missing source_addr".into(), - })? - .clone(); + } + })?.clone(); let link_stats = connection.link_stats().clone(); let remote_epoch = connection.remote_epoch(); @@ -908,8 +1098,11 @@ impl Node { let existing_link_id = existing_peer.link_id(); // Determine which connection wins - let this_wins = - cross_connection_winner(self.identity.node_addr(), &peer_node_addr, is_outbound); + let this_wins = cross_connection_winner( + self.identity.node_addr(), + &peer_node_addr, + is_outbound, + ); if this_wins { // This connection wins, replace the existing peer @@ -920,7 +1113,8 @@ impl Node { 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())); + self.peers_by_index + .remove(&(old_tid, old_idx.as_u32())); let _ = self.index_allocator.free(old_idx); } @@ -938,9 +1132,7 @@ impl Node { &self.config.node.mmp, remote_epoch, ); - new_peer.set_tree_announce_min_interval_ms( - self.config.node.tree.announce_min_interval_ms, - ); + new_peer.set_tree_announce_min_interval_ms(self.config.node.tree.announce_min_interval_ms); self.peers.insert(peer_node_addr, new_peer); self.peers_by_index @@ -1006,17 +1198,13 @@ impl Node { // 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, - }); + return Err(NodeError::MaxPeersExceeded { max: self.max_peers }); } // Preserve tree announce rate-limit state from old peer (if reconnecting). // Without this, reconnection resets the rate limit window to zero, // allowing an immediate announce that can feed an announce loop. - let old_announce_ts = self - .peers - .get(&peer_node_addr) + let old_announce_ts = self.peers.get(&peer_node_addr) .map(|p| p.last_tree_announce_sent_ms()); let mut new_peer = ActivePeer::with_session( @@ -1033,8 +1221,7 @@ impl Node { &self.config.node.mmp, remote_epoch, ); - new_peer - .set_tree_announce_min_interval_ms(self.config.node.tree.announce_min_interval_ms); + new_peer.set_tree_announce_min_interval_ms(self.config.node.tree.announce_min_interval_ms); if let Some(ts) = old_announce_ts { new_peer.set_last_tree_announce_sent_ms(ts); } diff --git a/src/node/handlers/rekey.rs b/src/node/handlers/rekey.rs index 44f4c8b..ac0fd7f 100644 --- a/src/node/handlers/rekey.rs +++ b/src/node/handlers/rekey.rs @@ -128,7 +128,7 @@ impl Node { /// Initiate an outbound rekey to a peer. /// - /// Creates a new IK handshake as initiator, sends msg1 over the existing + /// Creates a new XX handshake as initiator, sends msg1 over the existing /// link (same transport, same remote address), and stores the handshake /// state on the ActivePeer. No new Link or PeerConnection is created. async fn initiate_rekey(&mut self, node_addr: &NodeAddr) { @@ -146,7 +146,6 @@ impl Node { None => return, }; let link_id = peer.link_id(); - let peer_pubkey = peer.identity().pubkey_full(); // Allocate a new session index for the rekey let our_index = match self.index_allocator.allocate() { @@ -161,12 +160,12 @@ impl Node { } }; - // Create IK initiator handshake directly (no PeerConnection) + // Create XX initiator handshake directly (no PeerConnection) let our_keypair = self.identity.keypair(); - let mut hs = HandshakeState::new_initiator(our_keypair, peer_pubkey); + let mut hs = HandshakeState::new_xx_initiator(our_keypair); hs.set_local_epoch(self.startup_epoch); - let noise_msg1 = match hs.write_message_1() { + let noise_msg1 = match hs.write_xx_message_1() { Ok(msg) => msg, Err(e) => { warn!( diff --git a/src/node/handlers/rx_loop.rs b/src/node/handlers/rx_loop.rs index 225bb72..325f75d 100644 --- a/src/node/handlers/rx_loop.rs +++ b/src/node/handlers/rx_loop.rs @@ -1,12 +1,10 @@ //! RX event loop and packet dispatch. +use crate::control::{commands, ControlSocket}; use crate::control::queries; -use crate::control::{ControlSocket, commands}; -use crate::node::wire::{ - COMMON_PREFIX_SIZE, CommonPrefix, FMP_VERSION, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2, -}; use crate::node::{Node, NodeError}; use crate::transport::ReceivedPacket; +use crate::node::wire::{CommonPrefix, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2, PHASE_MSG3, FMP_VERSION, COMMON_PREFIX_SIZE}; use std::time::Duration; use tracing::{debug, info, warn}; @@ -18,6 +16,7 @@ impl Node { /// - Phase 0x0: Encrypted frame (session data) /// - Phase 0x1: Handshake message 1 (initiator -> responder) /// - Phase 0x2: Handshake message 2 (responder -> initiator) + /// - Phase 0x3: Handshake message 3 (initiator -> responder, XX completion) /// /// Also processes outbound IPv6 packets from the TUN reader for session /// encapsulation and routing through the mesh. @@ -31,7 +30,8 @@ impl Node { /// This method takes ownership of the packet_rx channel and runs /// until the channel is closed (typically when stop() is called). pub async fn run_rx_loop(&mut self) -> Result<(), NodeError> { - let mut packet_rx = self.packet_rx.take().ok_or(NodeError::NotStarted)?; + let mut packet_rx = self.packet_rx.take() + .ok_or(NodeError::NotStarted)?; // Take the TUN outbound receiver, or create a dummy channel that never // produces messages (when TUN is disabled). Holding the sender prevents @@ -54,12 +54,12 @@ impl Node { } }; - let mut tick = - tokio::time::interval(Duration::from_secs(self.config.node.tick_interval_secs)); + let mut tick = tokio::time::interval(Duration::from_secs(self.config.node.tick_interval_secs)); // Set up control socket channel - let (control_tx, mut control_rx) = - tokio::sync::mpsc::channel::(32); + let (control_tx, mut control_rx) = tokio::sync::mpsc::channel::< + crate::control::ControlMessage, + >(32); if self.config.node.control.enabled { let config = self.config.node.control.clone(); @@ -173,6 +173,9 @@ impl Node { PHASE_MSG2 => { self.handle_msg2(packet).await; } + PHASE_MSG3 => { + self.handle_msg3(packet).await; + } _ => { debug!( phase = prefix.phase, diff --git a/src/node/handlers/timeout.rs b/src/node/handlers/timeout.rs index e5c6dbc..547c494 100644 --- a/src/node/handlers/timeout.rs +++ b/src/node/handlers/timeout.rs @@ -71,10 +71,11 @@ impl Node { None => return, }; - // Free session index and pending_outbound if allocated + // Free session index and pending_outbound/pending_inbound if allocated if let Some(idx) = conn.our_index() { if let Some(tid) = conn.transport_id() { self.pending_outbound.remove(&(tid, idx.as_u32())); + self.pending_inbound.remove(&(tid, idx.as_u32())); } let _ = self.index_allocator.free(idx); } diff --git a/src/node/mod.rs b/src/node/mod.rs index b5f0c76..4114811 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -390,6 +390,10 @@ pub struct Node { /// Pending outbound handshakes by our sender_idx. /// Tracks which LinkId corresponds to which session index. pending_outbound: HashMap<(TransportId, u32), LinkId>, + /// Pending inbound connections awaiting msg3 (XX pattern), keyed by + /// (transport_id, our_index). The responder stores the connection here + /// after sending msg2 and awaits msg3 to learn the initiator's identity. + pending_inbound: HashMap<(TransportId, u32), LinkId>, // === Rate Limiting === /// Rate limiter for msg1 processing (DoS protection). @@ -547,6 +551,7 @@ impl Node { index_allocator: IndexAllocator::new(), peers_by_index: HashMap::new(), pending_outbound: HashMap::new(), + pending_inbound: HashMap::new(), msg1_rate_limiter, icmp_rate_limiter: IcmpRateLimiter::new(), routing_error_rate_limiter: RoutingErrorRateLimiter::new(), @@ -656,6 +661,7 @@ impl Node { index_allocator: IndexAllocator::new(), peers_by_index: HashMap::new(), pending_outbound: HashMap::new(), + pending_inbound: HashMap::new(), msg1_rate_limiter, icmp_rate_limiter: IcmpRateLimiter::new(), routing_error_rate_limiter: RoutingErrorRateLimiter::new(), diff --git a/src/node/tests/handshake.rs b/src/node/tests/handshake.rs index c42949e..7c6c31d 100644 --- a/src/node/tests/handshake.rs +++ b/src/node/tests/handshake.rs @@ -1,15 +1,13 @@ -//! Integration tests for end-to-end Noise IK handshake scenarios. +//! Integration tests for end-to-end Noise XX handshake scenarios. use super::*; #[tokio::test] async fn test_two_node_handshake_udp() { use crate::config::UdpConfig; - use crate::node::wire::{ - build_encrypted, build_established_header, build_msg1, prepend_inner_header, - }; use crate::transport::udp::UdpTransport; - use tokio::time::{Duration, timeout}; + use crate::node::wire::{build_encrypted, build_established_header, build_msg1, prepend_inner_header}; + use tokio::time::{timeout, Duration}; // === Setup: Two nodes with UDP transports on localhost === @@ -28,8 +26,10 @@ async fn test_two_node_handshake_udp() { 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); + 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(); @@ -49,20 +49,23 @@ async fn test_two_node_handshake_udp() { // === 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_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, 1000); + let mut conn_a = PeerConnection::outbound( + link_id_a, + peer_b_identity, + 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, node_a.startup_epoch, 1000) - .unwrap(); + let noise_msg1 = conn_a.start_handshake(our_keypair_a, node_a.startup_epoch, 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()); @@ -79,9 +82,10 @@ async fn test_two_node_handshake_udp() { ); 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); + 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(); @@ -90,7 +94,7 @@ async fn test_two_node_handshake_udp() { .await .expect("Failed to send msg1"); - // === Phase 2: Node B receives msg1, sends msg2, promotes === + // === Phase 2: Node B receives msg1, sends msg2 (XX: does NOT promote yet) === let packet_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) .await @@ -99,30 +103,16 @@ async fn test_two_node_handshake_udp() { 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" - ); + let peer_a_node_addr = *PeerIdentity::from_pubkey_full( + node_a.identity.pubkey_full(), + ) + .node_addr(); - // === Phase 3: Node A receives msg2, completes handshake, promotes === + // XX: B has NOT promoted yet (needs msg3) + assert_eq!(node_b.peer_count(), 0, "Node B should have 0 peers after msg1 (XX awaits msg3)"); + assert_eq!(node_b.connections.len(), 1, "Node B should have 1 pending connection awaiting msg3"); + + // === Phase 3: Node A receives msg2, sends msg3, promotes === let packet_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) .await @@ -132,11 +122,7 @@ async fn test_two_node_handshake_udp() { 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" - ); + 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"); @@ -156,6 +142,32 @@ async fn test_two_node_handshake_udp() { "Node A peers_by_index should be populated" ); + // === Phase 4: Node B receives msg3, promotes === + + let packet_b_msg3 = timeout(Duration::from_secs(1), packet_rx_b.recv()) + .await + .expect("Timeout waiting for msg3") + .expect("Channel closed"); + + node_b.handle_msg3(packet_b_msg3).await; + + // Verify B promoted after msg3 + assert_eq!(node_b.peer_count(), 1, "Node B should have 1 peer after msg3"); + 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 4: Encrypted frame A → B === // A encrypts a test message and sends to B @@ -243,8 +255,8 @@ async fn test_two_node_handshake_udp() { #[tokio::test] async fn test_run_rx_loop_handshake() { use crate::config::UdpConfig; - use crate::node::wire::build_msg1; use crate::transport::udp::UdpTransport; + use crate::node::wire::build_msg1; use tokio::time::Duration; // === Setup: Two nodes with UDP transports on localhost === @@ -264,8 +276,10 @@ async fn test_run_rx_loop_handshake() { let (packet_tx_a, packet_rx_a) = packet_channel(64); let (packet_tx_b, 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); + 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(); @@ -290,17 +304,20 @@ async fn test_run_rx_loop_handshake() { // === Phase 1: Node A initiates handshake to Node B === - let peer_b_identity = PeerIdentity::from_pubkey_full(node_b.identity.pubkey_full()); + 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, 1000); + let mut conn_a = PeerConnection::outbound( + link_id_a, + peer_b_identity, + 1000, + ); let our_index_a = node_a.index_allocator.allocate().unwrap(); let our_keypair_a = node_a.identity.keypair(); - let noise_msg1 = conn_a - .start_handshake(our_keypair_a, node_a.startup_epoch, 1000) - .unwrap(); + let noise_msg1 = conn_a.start_handshake(our_keypair_a, node_a.startup_epoch, 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()); @@ -316,9 +333,10 @@ async fn test_run_rx_loop_handshake() { ); 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); + node_a.pending_outbound.insert( + (transport_id_a, our_index_a.as_u32()), + link_id_a, + ); // Send msg1 from A to B over real UDP let transport = node_a.transports.get(&transport_id_a).unwrap(); @@ -330,11 +348,16 @@ async fn test_run_rx_loop_handshake() { // Small delay to ensure msg1 is received by B's transport tokio::time::sleep(Duration::from_millis(50)).await; - // === Phase 2: Run Node B's rx loop (processes msg1, sends msg2) === + // === Phase 2: Run Node B's rx loop (processes msg1 and later msg3) === // // This is the key difference from test_two_node_handshake_udp: // instead of calling handle_msg1() directly, we run the full rx loop // which dispatches based on the common prefix phase field. + // + // With XX, the rx loop will process msg1 (sending msg2) but NOT + // promote B yet (needs msg3). We run the rx loop once for msg1, + // then later use direct handler calls for msg3 (since run_rx_loop + // takes packet_rx and can't be called twice). tokio::select! { result = node_b.run_rx_loop() => { @@ -345,38 +368,11 @@ async fn test_run_rx_loop_handshake() { } } - // Verify Node B promoted the inbound connection via rx loop dispatch - let peer_a_node_addr = - *PeerIdentity::from_pubkey_full(node_a.identity.pubkey_full()).node_addr(); + // XX: Node B has NOT promoted yet (needs msg3) + assert_eq!(node_b.peer_count(), 0, "Node B should have 0 peers after rx loop processed msg1 (XX awaits msg3)"); + assert_eq!(node_b.connections.len(), 1, "Node B should have 1 pending connection"); - assert_eq!( - node_b.peer_count(), - 1, - "Node B should have 1 peer after rx loop processed 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!( - peer_a_on_b.their_index().is_some(), - "B should have their_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: Run Node A's rx loop (processes msg2) === - // - // msg2 was sent by Node B during its rx loop processing of msg1. - // It arrived at A's UDP transport, which forwarded it to A's packet channel. + // === Phase 3: Run Node A's rx loop (processes msg2, sends msg3) === tokio::select! { result = node_a.run_rx_loop() => { @@ -387,12 +383,8 @@ async fn test_run_rx_loop_handshake() { } } - // Verify Node A promoted the outbound connection via rx loop dispatch - assert_eq!( - node_a.peer_count(), - 1, - "Node A should have 1 peer after rx loop processed msg2" - ); + // Verify Node A promoted after processing msg2 + assert_eq!(node_a.peer_count(), 1, "Node A should have 1 peer after rx loop processed msg2"); let peer_b_on_a = node_a .get_peer(&peer_b_node_addr) .expect("Node A should have peer B"); @@ -416,6 +408,13 @@ async fn test_run_rx_loop_handshake() { "Node A peers_by_index should be populated" ); + // Note: Phase 4 (msg3 → B promotes) cannot be tested via run_rx_loop + // because it consumes packet_rx on first call. The msg3 dispatch is + // verified by test_two_node_handshake_udp which uses direct handler calls. + // This test verifies rx_loop correctly dispatches PHASE_MSG1 (Phase 2) + // and PHASE_MSG2 (Phase 3). B still has a pending connection awaiting msg3. + assert_eq!(node_b.connections.len(), 1, "Node B should still have pending connection awaiting msg3"); + // Clean up transports for (_, t) in node_a.transports.iter_mut() { t.stop().await.ok(); @@ -434,9 +433,9 @@ async fn test_run_rx_loop_handshake() { #[tokio::test] async fn test_cross_connection_both_initiate() { use crate::config::UdpConfig; - use crate::node::wire::build_msg1; use crate::transport::udp::UdpTransport; - use tokio::time::{Duration, timeout}; + use crate::node::wire::build_msg1; + use tokio::time::{timeout, Duration}; // === Setup: Two nodes with UDP transports on localhost === @@ -455,8 +454,10 @@ async fn test_cross_connection_both_initiate() { 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); + 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(); @@ -474,9 +475,11 @@ async fn test_cross_connection_both_initiate() { .insert(transport_id_b, TransportHandle::Udp(transport_b)); // Peer identities (must use full key for ECDH parity) - let peer_b_identity = PeerIdentity::from_pubkey_full(node_b.identity.pubkey_full()); + let peer_b_identity = + PeerIdentity::from_pubkey_full(node_b.identity.pubkey_full()); let peer_b_node_addr = *peer_b_identity.node_addr(); - let peer_a_identity = PeerIdentity::from_pubkey_full(node_a.identity.pubkey_full()); + let peer_a_identity = + PeerIdentity::from_pubkey_full(node_a.identity.pubkey_full()); let peer_a_node_addr = *peer_a_identity.node_addr(); // === Phase 1: Both nodes initiate handshakes (simulate auto_connect) === @@ -486,9 +489,7 @@ async fn test_cross_connection_both_initiate() { let mut conn_a = PeerConnection::outbound(link_id_a_out, peer_b_identity, 1000); let our_index_a = node_a.index_allocator.allocate().unwrap(); let our_keypair_a = node_a.identity.keypair(); - let noise_msg1_a = conn_a - .start_handshake(our_keypair_a, node_a.startup_epoch, 1000) - .unwrap(); + let noise_msg1_a = conn_a.start_handshake(our_keypair_a, node_a.startup_epoch, 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()); @@ -496,29 +497,20 @@ async fn test_cross_connection_both_initiate() { let wire_msg1_a = build_msg1(our_index_a, &noise_msg1_a); let link_a_out = Link::connectionless( - link_id_a_out, - transport_id_a, - remote_addr_b.clone(), - LinkDirection::Outbound, - Duration::from_millis(100), + link_id_a_out, transport_id_a, remote_addr_b.clone(), + LinkDirection::Outbound, Duration::from_millis(100), ); node_a.links.insert(link_id_a_out, link_a_out); - node_a - .addr_to_link - .insert((transport_id_a, remote_addr_b.clone()), link_id_a_out); + 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 - .pending_outbound - .insert((transport_id_a, our_index_a.as_u32()), link_id_a_out); + node_a.pending_outbound.insert((transport_id_a, our_index_a.as_u32()), link_id_a_out); // Node B initiates to Node A let link_id_b_out = node_b.allocate_link_id(); let mut conn_b = PeerConnection::outbound(link_id_b_out, peer_a_identity, 1000); let our_index_b = node_b.index_allocator.allocate().unwrap(); let our_keypair_b = node_b.identity.keypair(); - let noise_msg1_b = conn_b - .start_handshake(our_keypair_b, node_b.startup_epoch, 1000) - .unwrap(); + let noise_msg1_b = conn_b.start_handshake(our_keypair_b, node_b.startup_epoch, 1000).unwrap(); conn_b.set_our_index(our_index_b); conn_b.set_transport_id(transport_id_b); conn_b.set_source_addr(remote_addr_a.clone()); @@ -526,111 +518,77 @@ async fn test_cross_connection_both_initiate() { let wire_msg1_b = build_msg1(our_index_b, &noise_msg1_b); let link_b_out = Link::connectionless( - link_id_b_out, - transport_id_b, - remote_addr_a.clone(), - LinkDirection::Outbound, - Duration::from_millis(100), + link_id_b_out, transport_id_b, remote_addr_a.clone(), + LinkDirection::Outbound, Duration::from_millis(100), ); node_b.links.insert(link_id_b_out, link_b_out); - node_b - .addr_to_link - .insert((transport_id_b, remote_addr_a.clone()), link_id_b_out); + 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 - .pending_outbound - .insert((transport_id_b, our_index_b.as_u32()), link_id_b_out); + node_b.pending_outbound.insert((transport_id_b, our_index_b.as_u32()), link_id_b_out); // Both send msg1 over UDP let transport = node_a.transports.get(&transport_id_a).unwrap(); - transport - .send(&remote_addr_b, &wire_msg1_a) - .await - .expect("A send msg1"); + transport.send(&remote_addr_b, &wire_msg1_a).await.expect("A send msg1"); let transport = node_b.transports.get(&transport_id_b).unwrap(); - transport - .send(&remote_addr_a, &wire_msg1_b) - .await - .expect("B send msg1"); + transport.send(&remote_addr_a, &wire_msg1_b).await.expect("B send msg1"); - // === Phase 2: Both nodes receive the other's msg1 === - // Before the fix, addr_to_link would reject these because outbound links - // already exist for these addresses. + // === Phase 2: Both nodes receive the other's msg1 (XX: no promotion yet) === // B receives A's msg1 let packet_at_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) - .await - .expect("Timeout") - .expect("Channel closed"); + .await.expect("Timeout").expect("Channel closed"); node_b.handle_msg1(packet_at_b).await; - // B should have promoted the inbound connection - assert_eq!( - node_b.peer_count(), - 1, - "Node B should have 1 peer after processing A's msg1" - ); - assert!( - node_b.get_peer(&peer_a_node_addr).is_some(), - "Node B should have peer A" - ); + // XX: B has NOT promoted yet (needs msg3 from A) + assert_eq!(node_b.peer_count(), 0, "Node B should have 0 peers after processing A's msg1 (XX)"); // A receives B's msg1 let packet_at_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) - .await - .expect("Timeout") - .expect("Channel closed"); + .await.expect("Timeout").expect("Channel closed"); node_a.handle_msg1(packet_at_a).await; - // A should have promoted the inbound connection - assert_eq!( - node_a.peer_count(), - 1, - "Node A should have 1 peer after processing B's msg1" - ); - assert!( - node_a.get_peer(&peer_b_node_addr).is_some(), - "Node A should have peer B" - ); + // XX: A has NOT promoted yet (needs msg3 from B) + assert_eq!(node_a.peer_count(), 0, "Node A should have 0 peers after processing B's msg1 (XX)"); - // === Phase 3: Both nodes receive msg2 responses === - // The msg2 was sent during handle_msg1 processing. When handle_msg2 - // processes it, it will detect the cross-connection and resolve. + // === Phase 3: Both nodes receive msg2 + send msg3, initiator side promotes === - // A receives B's msg2 (response to A's original msg1) + // A receives B's msg2 (response to A's original msg1) → A sends msg3, A promotes let msg2_at_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) - .await - .expect("Timeout waiting for msg2 at A") - .expect("Channel closed"); + .await.expect("Timeout waiting for msg2 at A").expect("Channel closed"); node_a.handle_msg2(msg2_at_a).await; - // B receives A's msg2 (response to B's original msg1) + // A promoted as initiator + assert_eq!(node_a.peer_count(), 1, "Node A should have 1 peer after processing msg2"); + + // B receives A's msg2 (response to B's original msg1) → B sends msg3, B promotes let msg2_at_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) - .await - .expect("Timeout waiting for msg2 at B") - .expect("Channel closed"); + .await.expect("Timeout waiting for msg2 at B").expect("Channel closed"); node_b.handle_msg2(msg2_at_b).await; + // B promoted as initiator + assert_eq!(node_b.peer_count(), 1, "Node B should have 1 peer after processing msg2"); + + // === Phase 4: Both nodes receive msg3, responder side completes === + // Cross-connection resolution happens here (or in Phase 3 promotion). + + // A receives B's msg3 (B completing A's inbound handshake) + let msg3_at_a = timeout(Duration::from_secs(1), packet_rx_a.recv()) + .await.expect("Timeout waiting for msg3 at A").expect("Channel closed"); + node_a.handle_msg3(msg3_at_a).await; + + // B receives A's msg3 (A completing B's inbound handshake) + let msg3_at_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) + .await.expect("Timeout waiting for msg3 at B").expect("Channel closed"); + node_b.handle_msg3(msg3_at_b).await; + // === Verification === // Both nodes should have exactly 1 peer each after cross-connection resolution - assert_eq!( - node_a.peer_count(), - 1, - "Node A should have exactly 1 peer after cross-connection" - ); - assert_eq!( - node_b.peer_count(), - 1, - "Node B should have exactly 1 peer after cross-connection" - ); + assert_eq!(node_a.peer_count(), 1, "Node A should have exactly 1 peer after cross-connection"); + assert_eq!(node_b.peer_count(), 1, "Node B should have exactly 1 peer after cross-connection"); - let peer_b_on_a = node_a - .get_peer(&peer_b_node_addr) - .expect("A should have peer B"); - let peer_a_on_b = node_b - .get_peer(&peer_a_node_addr) - .expect("B should have peer A"); + let peer_b_on_a = node_a.get_peer(&peer_b_node_addr).expect("A should have peer B"); + let peer_a_on_b = node_b.get_peer(&peer_a_node_addr).expect("B should have peer A"); assert!(peer_b_on_a.has_session(), "Peer B on A should have session"); assert!(peer_a_on_b.has_session(), "Peer A on B should have session"); @@ -667,35 +625,25 @@ async fn test_stale_connection_cleanup() { // Allocate session index and set transport info 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, past_time_ms) - .unwrap(); + let _noise_msg1 = conn.start_handshake(our_keypair, node.startup_epoch, past_time_ms).unwrap(); conn.set_our_index(our_index); conn.set_transport_id(transport_id); conn.set_source_addr(remote_addr.clone()); // Set up all the state that initiate_peer_connection would create let link = Link::connectionless( - link_id, - transport_id, - remote_addr.clone(), - LinkDirection::Outbound, - Duration::from_millis(100), + link_id, transport_id, remote_addr.clone(), + LinkDirection::Outbound, Duration::from_millis(100), ); node.links.insert(link_id, link); - node.addr_to_link - .insert((transport_id, remote_addr.clone()), link_id); + 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); + node.pending_outbound.insert((transport_id, our_index.as_u32()), link_id); // Verify state before timeout check assert_eq!(node.connection_count(), 1); assert_eq!(node.link_count(), 1); - assert!( - node.pending_outbound - .contains_key(&(transport_id, our_index.as_u32())) - ); + assert!(node.pending_outbound.contains_key(&(transport_id, our_index.as_u32()))); assert_eq!(node.index_allocator.count(), 1); // Connection was created at time 1000ms. check_timeouts uses SystemTime::now(), @@ -703,27 +651,13 @@ async fn test_stale_connection_cleanup() { node.check_timeouts(); // Verify everything was cleaned up - assert_eq!( - node.connection_count(), - 0, - "Stale connection should be removed" - ); + assert_eq!(node.connection_count(), 0, "Stale connection should be removed"); assert_eq!(node.link_count(), 0, "Stale link should be removed"); - assert!( - !node - .pending_outbound - .contains_key(&(transport_id, our_index.as_u32())), - "pending_outbound should be cleaned up" - ); - assert_eq!( - node.index_allocator.count(), - 0, - "Session index should be freed" - ); - assert!( - !node.addr_to_link.contains_key(&(transport_id, remote_addr)), - "addr_to_link should be cleaned up" - ); + assert!(!node.pending_outbound.contains_key(&(transport_id, our_index.as_u32())), + "pending_outbound should be cleaned up"); + assert_eq!(node.index_allocator.count(), 0, "Session index should be freed"); + assert!(!node.addr_to_link.contains_key(&(transport_id, remote_addr)), + "addr_to_link should be cleaned up"); } /// Test that failed connections are cleaned up by check_timeouts(). @@ -745,44 +679,29 @@ async fn test_failed_connection_cleanup() { 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(); + 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.mark_failed(); // Simulate send failure let link = Link::connectionless( - link_id, - transport_id, - remote_addr.clone(), - LinkDirection::Outbound, - Duration::from_millis(100), + link_id, transport_id, remote_addr.clone(), + LinkDirection::Outbound, Duration::from_millis(100), ); node.links.insert(link_id, link); - node.addr_to_link - .insert((transport_id, remote_addr.clone()), link_id); + 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); + node.pending_outbound.insert((transport_id, our_index.as_u32()), link_id); assert_eq!(node.connection_count(), 1); // Failed connections should be cleaned up immediately regardless of age node.check_timeouts(); - assert_eq!( - node.connection_count(), - 0, - "Failed connection should be removed" - ); + assert_eq!(node.connection_count(), 0, "Failed connection should be removed"); assert_eq!(node.link_count(), 0, "Failed link should be removed"); - assert_eq!( - node.index_allocator.count(), - 0, - "Session index should be freed" - ); + assert_eq!(node.index_allocator.count(), 0, "Session index should be freed"); } /// Test that msg1 bytes are stored on connection for resend. @@ -805,9 +724,7 @@ async fn test_msg1_stored_for_resend() { 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(); + 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()); @@ -838,9 +755,7 @@ async fn test_resend_scheduling() { 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(); + 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()); @@ -850,17 +765,12 @@ async fn test_resend_scheduling() { conn.set_handshake_msg1(wire_msg1, now_ms + 1000); let link = Link::connectionless( - link_id, - transport_id, - remote_addr.clone(), - LinkDirection::Outbound, - Duration::from_millis(100), + link_id, transport_id, remote_addr.clone(), + LinkDirection::Outbound, Duration::from_millis(100), ); node.links.insert(link_id, link); - node.addr_to_link - .insert((transport_id, remote_addr), link_id); - node.pending_outbound - .insert((transport_id, our_index.as_u32()), link_id); + node.addr_to_link.insert((transport_id, remote_addr), link_id); + node.pending_outbound.insert((transport_id, our_index.as_u32()), link_id); node.connections.insert(link_id, conn); // Before resend time: nothing should happen (no transport = can't send, @@ -876,11 +786,7 @@ async fn test_resend_scheduling() { // No transport registered, so send fails — count stays 0. // That's the expected behavior (transport absence is a transient condition). let conn = node.connections.get(&link_id).unwrap(); - assert_eq!( - conn.resend_count(), - 0, - "No transport means no resend recorded" - ); + assert_eq!(conn.resend_count(), 0, "No transport means no resend recorded"); } /// Test that msg2 is stored on PeerConnection for responder resend. @@ -934,8 +840,8 @@ async fn test_duplicate_msg2_dropped() { let receiver_idx = SessionIndex::new(42); let sender_idx = SessionIndex::new(99); - // Build a fake msg2 packet - let fake_noise_msg2 = vec![0u8; 57]; // Noise IK msg2 is 57 bytes (33 ephem + 24 encrypted epoch) + // Build a fake msg2 packet (XX msg2 is at least 106 bytes) + let fake_noise_msg2 = vec![0u8; 106]; let wire_msg2 = build_msg2(sender_idx, receiver_idx, &fake_noise_msg2); let packet = ReceivedPacket { diff --git a/src/node/tests/mod.rs b/src/node/tests/mod.rs index 9898e42..57f5c9e 100644 --- a/src/node/tests/mod.rs +++ b/src/node/tests/mod.rs @@ -64,11 +64,14 @@ pub(super) fn make_completed_connection( let mut resp_epoch = [0u8; 8]; rand::Rng::fill_bytes(&mut rand::rng(), &mut resp_epoch); let msg2 = resp_conn - .receive_handshake_init(peer_keypair, resp_epoch, &msg1, current_time_ms) + .receive_handshake_init(peer_keypair, resp_epoch, &msg1, None, current_time_ms) .unwrap(); - // Complete initiator handshake - conn.complete_handshake(&msg2, current_time_ms).unwrap(); + // Complete initiator handshake (XX: generates msg3) + let (msg3, _neg) = conn.complete_handshake(&msg2, None, current_time_ms).unwrap(); + + // Complete responder handshake (XX: processes msg3) + resp_conn.complete_handshake_msg3(&msg3, current_time_ms).unwrap(); // Set indices and transport info let our_index = node.index_allocator.allocate().unwrap(); diff --git a/src/node/tests/spanning_tree.rs b/src/node/tests/spanning_tree.rs index bcc85e8..f14951e 100644 --- a/src/node/tests/spanning_tree.rs +++ b/src/node/tests/spanning_tree.rs @@ -69,9 +69,7 @@ pub(super) async fn initiate_handshake(nodes: &mut [TestNode], i: usize, j: usiz let our_index = initiator.node.index_allocator.allocate().unwrap(); let our_keypair = initiator.node.identity().keypair(); - let noise_msg1 = conn - .start_handshake(our_keypair, initiator.node.startup_epoch, 1000) - .unwrap(); + let noise_msg1 = conn.start_handshake(our_keypair, initiator.node.startup_epoch, 1000).unwrap(); conn.set_our_index(our_index); conn.set_transport_id(transport_id); conn.set_source_addr(responder_addr.clone()); @@ -186,12 +184,7 @@ pub(super) fn print_tree_snapshot(label: &str, nodes: &[TestNode]) { .count(); eprintln!( " node[{}] root=node[{}] depth={} parent=node[{}] peers={} pending={}", - i, - root_idx, - ts.my_coords().depth(), - parent_idx, - tn.node.peer_count(), - pending, + i, root_idx, ts.my_coords().depth(), parent_idx, tn.node.peer_count(), pending, ); } } else if correct_root_count < nodes.len() { @@ -216,9 +209,7 @@ pub(super) fn print_tree_snapshot(label: &str, nodes: &[TestNode]) { /// /// Returns the number of packets processed. pub(super) async fn process_available_packets(nodes: &mut [TestNode]) -> usize { - use crate::node::wire::{ - COMMON_PREFIX_SIZE, CommonPrefix, FMP_VERSION, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2, - }; + use crate::node::wire::{CommonPrefix, FMP_VERSION, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2, PHASE_MSG3, COMMON_PREFIX_SIZE}; let mut count = 0; for node in nodes.iter_mut() { @@ -233,7 +224,10 @@ pub(super) async fn process_available_packets(nodes: &mut [TestNode]) -> usize { match prefix.phase { PHASE_MSG1 => node.node.handle_msg1(packet).await, PHASE_MSG2 => node.node.handle_msg2(packet).await, - PHASE_ESTABLISHED => node.node.handle_encrypted_frame(packet).await, + PHASE_MSG3 => node.node.handle_msg3(packet).await, + PHASE_ESTABLISHED => { + node.node.handle_encrypted_frame(packet).await + } _ => {} } count += 1; @@ -326,11 +320,7 @@ pub(super) async fn drain_all_packets(nodes: &mut [TestNode], verbose: bool) -> /// /// First builds a random spanning tree to ensure connectivity, /// then adds extra edges up to the target count. -pub(super) fn generate_random_edges( - n: usize, - target_edges: usize, - seed: u64, -) -> Vec<(usize, usize)> { +pub(super) fn generate_random_edges(n: usize, target_edges: usize, seed: u64) -> Vec<(usize, usize)> { use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; @@ -384,7 +374,11 @@ pub(super) fn verify_tree_convergence(nodes: &[TestNode]) { assert!(n > 0); // Find the expected root (smallest NodeAddr across all nodes) - let expected_root = nodes.iter().map(|tn| *tn.node.node_addr()).min().unwrap(); + let expected_root = nodes + .iter() + .map(|tn| *tn.node.node_addr()) + .min() + .unwrap(); // All nodes should agree on the root for (i, tn) in nodes.iter().enumerate() { @@ -634,16 +628,12 @@ pub(super) async fn run_tree_test_with_mtus( assert!( nodes[i].node.get_peer(&j_addr).is_some(), "Node {} should have peer {} (node {})", - i, - j_addr, - j + i, j_addr, j ); assert!( nodes[j].node.get_peer(&i_addr).is_some(), "Node {} should have peer {} (node {})", - j, - i_addr, - i + j, i_addr, i ); } diff --git a/src/node/tests/unit.rs b/src/node/tests/unit.rs index 015762e..bb33448 100644 --- a/src/node/tests/unit.rs +++ b/src/node/tests/unit.rs @@ -505,12 +505,13 @@ fn test_promote_cleans_up_pending_outbound_to_same_peer() { let mut resp_epoch = [0u8; 8]; rand::Rng::fill_bytes(&mut rand::rng(), &mut resp_epoch); let msg2 = resp_conn - .receive_handshake_init(peer_keypair, resp_epoch, &msg1, completing_time_ms) + .receive_handshake_init(peer_keypair, resp_epoch, &msg1, None, completing_time_ms) .unwrap(); - completing_conn - .complete_handshake(&msg2, completing_time_ms) + let (msg3, _neg) = completing_conn + .complete_handshake(&msg2, None, completing_time_ms) .unwrap(); + resp_conn.complete_handshake_msg3(&msg3, completing_time_ms).unwrap(); let completing_index = node.index_allocator.allocate().unwrap(); completing_conn.set_our_index(completing_index); diff --git a/src/node/wire.rs b/src/node/wire.rs index 21bc149..428368c 100644 --- a/src/node/wire.rs +++ b/src/node/wire.rs @@ -14,39 +14,48 @@ //! | Phase | Type | Size | Description | //! |-------|-----------------|------------|--------------------------------| //! | 0x0 | Encrypted frame | 32+ bytes | Post-handshake encrypted data | -//! | 0x1 | Noise IK msg1 | 114 bytes | Handshake initiation | -//! | 0x2 | Noise IK msg2 | 69 bytes | Handshake response | +//! | 0x1 | Noise XX msg1 | 41 bytes | Handshake initiation | +//! | 0x2 | Noise XX msg2 | 118+ bytes | Handshake response | +//! | 0x3 | Noise XX msg3 | 85+ bytes | Handshake completion | -use crate::noise::{HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE, TAG_SIZE}; use crate::utils::index::SessionIndex; +use crate::noise::{XX_HANDSHAKE_MSG1_SIZE, XX_HANDSHAKE_MSG2_SIZE, XX_HANDSHAKE_MSG3_SIZE, TAG_SIZE}; // ============================================================================ // Constants // ============================================================================ /// FMP protocol version (4 high bits of byte 0). -pub const FMP_VERSION: u8 = 0; +pub const FMP_VERSION: u8 = 1; /// Phase value for established (encrypted) frames. pub const PHASE_ESTABLISHED: u8 = 0x0; -/// Phase value for Noise IK message 1 (handshake initiation). +/// Phase value for handshake message 1 (initiation). pub const PHASE_MSG1: u8 = 0x1; -/// Phase value for Noise IK message 2 (handshake response). +/// Phase value for handshake message 2 (response). pub const PHASE_MSG2: u8 = 0x2; +/// Phase value for handshake message 3 (completion, XX only). +pub const PHASE_MSG3: u8 = 0x3; + /// Size of the common packet prefix (all packet types). pub const COMMON_PREFIX_SIZE: usize = 4; /// Size of the full established frame header (prefix + receiver_idx + counter). pub const ESTABLISHED_HEADER_SIZE: usize = 16; -/// Size of Noise IK message 1 wire packet: prefix + sender_idx + noise_msg1. -pub const MSG1_WIRE_SIZE: usize = COMMON_PREFIX_SIZE + 4 + HANDSHAKE_MSG1_SIZE; // 114 bytes +/// Size of handshake msg1 wire packet: prefix + sender_idx + noise_msg1. +pub const MSG1_WIRE_SIZE: usize = COMMON_PREFIX_SIZE + 4 + XX_HANDSHAKE_MSG1_SIZE; // 41 bytes -/// Size of Noise IK message 2 wire packet: prefix + sender_idx + receiver_idx + noise_msg2. -pub const MSG2_WIRE_SIZE: usize = COMMON_PREFIX_SIZE + 4 + 4 + HANDSHAKE_MSG2_SIZE; // 69 bytes +/// Minimum size of handshake msg2 wire packet: prefix + sender_idx + receiver_idx + noise_msg2. +/// Actual size may be larger due to optional negotiation payload. +pub const MSG2_WIRE_SIZE: usize = COMMON_PREFIX_SIZE + 4 + 4 + XX_HANDSHAKE_MSG2_SIZE; // 118 bytes + +/// Minimum size of handshake msg3 wire packet: prefix + sender_idx + receiver_idx + noise_msg3. +/// Actual size may be larger due to optional negotiation payload. +pub const MSG3_WIRE_SIZE: usize = COMMON_PREFIX_SIZE + 4 + 4 + XX_HANDSHAKE_MSG3_SIZE; // 85 bytes /// Minimum size for encrypted frame: header + tag (no plaintext). pub const ENCRYPTED_MIN_SIZE: usize = ESTABLISHED_HEADER_SIZE + TAG_SIZE; // 32 bytes @@ -164,7 +173,8 @@ impl EncryptedHeader { let payload_len = u16::from_le_bytes([data[2], data[3]]); let receiver_idx = SessionIndex::from_le_bytes([data[4], data[5], data[6], data[7]]); let counter = u64::from_le_bytes([ - data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15], + data[8], data[9], data[10], data[11], + data[12], data[13], data[14], data[15], ]); let mut header_bytes = [0u8; ESTABLISHED_HEADER_SIZE]; @@ -195,11 +205,11 @@ impl EncryptedHeader { // Msg1 Header // ============================================================================ -/// Parsed Noise IK message 1 header (phase 0x1). +/// Parsed handshake message 1 header (phase 0x1). /// -/// Wire format (114 bytes): +/// Wire format (41 bytes, Noise XX): /// ```text -/// [0x01][0x00][payload_len:2 LE][sender_idx:4 LE][noise_msg1:106] +/// [0x11][0x00][payload_len:2 LE][sender_idx:4 LE][noise_msg1:33] /// ``` #[derive(Clone, Debug)] pub struct Msg1Header { @@ -249,12 +259,13 @@ impl Msg1Header { // Msg2 Header // ============================================================================ -/// Parsed Noise IK message 2 header (phase 0x2). +/// Parsed handshake message 2 header (phase 0x2). /// -/// Wire format (69 bytes): +/// Wire format (118+ bytes, Noise XX): /// ```text -/// [0x02][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg2:57] +/// [0x12][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg2:106+] /// ``` +/// Size is variable due to optional negotiation payload appended after base XX msg2. #[derive(Clone, Debug)] pub struct Msg2Header { /// Session index chosen by the responder. @@ -268,9 +279,10 @@ pub struct Msg2Header { impl Msg2Header { /// Parse a msg2 header from packet data. /// - /// Returns None if the packet has wrong size or version/phase. + /// Returns None if the packet is too short or has wrong version/phase. + /// Accepts variable size (base + optional negotiation payload). pub fn parse(data: &[u8]) -> Option { - if data.len() != MSG2_WIRE_SIZE { + if data.len() < MSG2_WIRE_SIZE { return None; } @@ -296,11 +308,83 @@ impl Msg2Header { }) } - /// Get the Noise msg2 payload from the original packet. + /// Get the Noise msg2 payload from the original packet (variable length). #[cfg(test)] pub fn noise_msg2<'a>(&self, data: &'a [u8]) -> &'a [u8] { &data[self.noise_msg2_offset..] } + + /// Get the total noise payload length (base + optional negotiation). + #[allow(dead_code)] + pub fn noise_payload_len(&self, data: &[u8]) -> usize { + data.len() - self.noise_msg2_offset + } +} + +// ============================================================================ +// Msg3 Header +// ============================================================================ + +/// Parsed handshake message 3 header (phase 0x3, XX pattern). +/// +/// Wire format (85+ bytes, Noise XX): +/// ```text +/// [0x13][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg3:73+] +/// ``` +/// Size is variable due to optional negotiation payload appended after base XX msg3. +#[derive(Clone, Debug)] +pub struct Msg3Header { + /// Session index chosen by the initiator (echo of msg1 sender_idx). + pub sender_idx: SessionIndex, + /// Echo of the responder's sender_idx from msg2. + pub receiver_idx: SessionIndex, + /// Offset where Noise msg3 payload begins. + pub noise_msg3_offset: usize, +} + +impl Msg3Header { + /// Parse a msg3 header from packet data. + /// + /// Returns None if the packet is too short or has wrong version/phase. + /// Accepts variable size (base + optional negotiation payload). + pub fn parse(data: &[u8]) -> Option { + if data.len() < MSG3_WIRE_SIZE { + return None; + } + + let version = data[0] >> 4; + let phase = data[0] & 0x0F; + + if version != FMP_VERSION || phase != PHASE_MSG3 { + return None; + } + + // flags must be zero during handshake + if data[1] != 0 { + return None; + } + + let sender_idx = SessionIndex::from_le_bytes([data[4], data[5], data[6], data[7]]); + let receiver_idx = SessionIndex::from_le_bytes([data[8], data[9], data[10], data[11]]); + + Some(Self { + sender_idx, + receiver_idx, + noise_msg3_offset: COMMON_PREFIX_SIZE + 4 + 4, // 12 + }) + } + + /// Get the Noise msg3 payload from the original packet (variable length). + #[cfg(test)] + pub fn noise_msg3<'a>(&self, data: &'a [u8]) -> &'a [u8] { + &data[self.noise_msg3_offset..] + } + + /// Get the total noise payload length (base + optional negotiation). + #[allow(dead_code)] + pub fn noise_payload_len(&self, data: &[u8]) -> usize { + data.len() - self.noise_msg3_offset + } } // ============================================================================ @@ -309,9 +393,9 @@ impl Msg2Header { /// Build a wire-format msg1 packet. /// -/// Format: `[0x01][0x00][payload_len:2 LE][sender_idx:4 LE][noise_msg1:106]` +/// Format: `[0x11][0x00][payload_len:2 LE][sender_idx:4 LE][noise_msg1:33]` pub fn build_msg1(sender_idx: SessionIndex, noise_msg1: &[u8]) -> Vec { - debug_assert_eq!(noise_msg1.len(), HANDSHAKE_MSG1_SIZE); + debug_assert_eq!(noise_msg1.len(), XX_HANDSHAKE_MSG1_SIZE); let payload_len = (4 + noise_msg1.len()) as u16; // sender_idx + noise_msg1 @@ -326,17 +410,15 @@ pub fn build_msg1(sender_idx: SessionIndex, noise_msg1: &[u8]) -> Vec { /// Build a wire-format msg2 packet. /// -/// Format: `[0x02][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg2:57]` -pub fn build_msg2( - sender_idx: SessionIndex, - receiver_idx: SessionIndex, - noise_msg2: &[u8], -) -> Vec { - debug_assert_eq!(noise_msg2.len(), HANDSHAKE_MSG2_SIZE); +/// Format: `[0x12][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg2:106+]` +/// The noise_msg2 may include an optional negotiation payload beyond the base XX msg2. +pub fn build_msg2(sender_idx: SessionIndex, receiver_idx: SessionIndex, noise_msg2: &[u8]) -> Vec { + debug_assert!(noise_msg2.len() >= XX_HANDSHAKE_MSG2_SIZE); let payload_len = (4 + 4 + noise_msg2.len()) as u16; // sender + receiver + noise + let total = COMMON_PREFIX_SIZE + 4 + 4 + noise_msg2.len(); - let mut packet = Vec::with_capacity(MSG2_WIRE_SIZE); + let mut packet = Vec::with_capacity(total); packet.push(CommonPrefix::ver_phase_byte(FMP_VERSION, PHASE_MSG2)); packet.push(0x00); // flags must be zero packet.extend_from_slice(&payload_len.to_le_bytes()); @@ -346,6 +428,26 @@ pub fn build_msg2( packet } +/// Build a wire-format msg3 packet (XX handshake completion). +/// +/// Format: `[0x13][0x00][payload_len:2 LE][sender_idx:4 LE][receiver_idx:4 LE][noise_msg3:73+]` +/// The noise_msg3 may include an optional negotiation payload beyond the base XX msg3. +pub fn build_msg3(sender_idx: SessionIndex, receiver_idx: SessionIndex, noise_msg3: &[u8]) -> Vec { + debug_assert!(noise_msg3.len() >= XX_HANDSHAKE_MSG3_SIZE); + + let payload_len = (4 + 4 + noise_msg3.len()) as u16; // sender + receiver + noise + let total = COMMON_PREFIX_SIZE + 4 + 4 + noise_msg3.len(); + + let mut packet = Vec::with_capacity(total); + packet.push(CommonPrefix::ver_phase_byte(FMP_VERSION, PHASE_MSG3)); + packet.push(0x00); // flags must be zero + packet.extend_from_slice(&payload_len.to_le_bytes()); + packet.extend_from_slice(&sender_idx.to_le_bytes()); + packet.extend_from_slice(&receiver_idx.to_le_bytes()); + packet.extend_from_slice(noise_msg3); + packet +} + /// Build the 16-byte outer header for an established frame. /// /// Returns the header bytes (for use as AAD) separately from the construction. @@ -414,9 +516,9 @@ mod tests { #[test] fn test_common_prefix_parse() { - let data = [0x00, 0x04, 0x20, 0x00]; // ver=0, phase=0, flags=SP, payload_len=32 + let data = [0x10, 0x04, 0x20, 0x00]; // ver=1, phase=0, flags=SP, payload_len=32 let prefix = CommonPrefix::parse(&data).unwrap(); - assert_eq!(prefix.version, 0); + assert_eq!(prefix.version, 1); assert_eq!(prefix.phase, 0); assert_eq!(prefix.flags, FLAG_SP); assert_eq!(prefix.payload_len, 32); @@ -439,7 +541,7 @@ mod tests { let packet = build_encrypted(&header, &ciphertext); assert_eq!(packet.len(), ESTABLISHED_HEADER_SIZE + 48); - assert_eq!(packet[0], 0x00); // ver=0, phase=0 + assert_eq!(packet[0], 0x10); // ver=1, phase=0 let parsed = EncryptedHeader::parse(&packet).expect("should parse"); assert_eq!(parsed.receiver_idx, receiver_idx); @@ -459,26 +561,26 @@ mod tests { #[test] fn test_encrypted_header_wrong_phase() { let mut packet = vec![0x00; ENCRYPTED_MIN_SIZE]; - packet[0] = 0x01; // phase 1 (msg1), not established + packet[0] = 0x11; // ver=1, phase 1 (msg1), not established assert!(EncryptedHeader::parse(&packet).is_none()); } #[test] fn test_encrypted_header_wrong_version() { let mut packet = vec![0x00; ENCRYPTED_MIN_SIZE]; - packet[0] = 0x10; // version 1, phase 0 + packet[0] = 0x00; // version 0 (old), phase 0 assert!(EncryptedHeader::parse(&packet).is_none()); } #[test] fn test_msg1_header_parse() { let sender_idx = SessionIndex::new(0xABCDEF01); - let noise_msg1 = vec![0xbb; HANDSHAKE_MSG1_SIZE]; + let noise_msg1 = vec![0xbb; XX_HANDSHAKE_MSG1_SIZE]; let packet = build_msg1(sender_idx, &noise_msg1); assert_eq!(packet.len(), MSG1_WIRE_SIZE); - assert_eq!(packet[0], 0x01); // ver=0, phase=1 + assert_eq!(packet[0], 0x11); // ver=1, phase=1 let header = Msg1Header::parse(&packet).expect("should parse"); assert_eq!(header.sender_idx, sender_idx); @@ -488,23 +590,23 @@ mod tests { #[test] fn test_msg1_header_wrong_size() { - let packet = vec![0x01; MSG1_WIRE_SIZE - 1]; + let packet = vec![0x11; MSG1_WIRE_SIZE - 1]; assert!(Msg1Header::parse(&packet).is_none()); - let packet = vec![0x01; MSG1_WIRE_SIZE + 1]; + let packet = vec![0x11; MSG1_WIRE_SIZE + 1]; assert!(Msg1Header::parse(&packet).is_none()); } #[test] fn test_msg1_header_wrong_phase() { let mut packet = vec![0x00; MSG1_WIRE_SIZE]; - packet[0] = 0x02; // phase 2, not phase 1 + packet[0] = 0x12; // ver=1, phase 2, not phase 1 assert!(Msg1Header::parse(&packet).is_none()); } #[test] fn test_msg1_header_nonzero_flags() { - let mut packet = build_msg1(SessionIndex::new(1), &[0u8; HANDSHAKE_MSG1_SIZE]); + let mut packet = build_msg1(SessionIndex::new(1), &[0u8; XX_HANDSHAKE_MSG1_SIZE]); packet[1] = 0x01; // flags must be zero during handshake assert!(Msg1Header::parse(&packet).is_none()); } @@ -513,12 +615,12 @@ mod tests { fn test_msg2_header_parse() { let sender_idx = SessionIndex::new(0x11223344); let receiver_idx = SessionIndex::new(0x55667788); - let noise_msg2 = vec![0xcc; HANDSHAKE_MSG2_SIZE]; + let noise_msg2 = vec![0xcc; XX_HANDSHAKE_MSG2_SIZE]; let packet = build_msg2(sender_idx, receiver_idx, &noise_msg2); assert_eq!(packet.len(), MSG2_WIRE_SIZE); - assert_eq!(packet[0], 0x02); // ver=0, phase=2 + assert_eq!(packet[0], 0x12); // ver=1, phase=2 let header = Msg2Header::parse(&packet).expect("should parse"); assert_eq!(header.sender_idx, sender_idx); @@ -529,24 +631,29 @@ mod tests { #[test] fn test_msg2_header_wrong_size() { - let packet = vec![0x02; MSG2_WIRE_SIZE - 1]; + let packet = vec![0x12; MSG2_WIRE_SIZE - 1]; assert!(Msg2Header::parse(&packet).is_none()); - let packet = vec![0x02; MSG2_WIRE_SIZE + 1]; - assert!(Msg2Header::parse(&packet).is_none()); + // Larger than minimum is now accepted (variable-length negotiation payload) + let mut packet = vec![0x12; MSG2_WIRE_SIZE + 10]; + packet[0] = 0x12; // ver=1, phase=2 + packet[1] = 0x00; + let header = Msg2Header::parse(&packet); + assert!(header.is_some()); } #[test] fn test_msg2_header_wrong_phase() { let mut packet = vec![0x00; MSG2_WIRE_SIZE]; - packet[0] = 0x00; // phase 0, not phase 2 + packet[0] = 0x10; // ver=1, phase 0, not phase 2 assert!(Msg2Header::parse(&packet).is_none()); } #[test] fn test_wire_sizes() { - assert_eq!(MSG1_WIRE_SIZE, 114); // 4 + 4 + 106 - assert_eq!(MSG2_WIRE_SIZE, 69); // 4 + 4 + 4 + 57 + assert_eq!(MSG1_WIRE_SIZE, 41); // 4 + 4 + 33 (XX msg1) + assert_eq!(MSG2_WIRE_SIZE, 118); // 4 + 4 + 4 + 106 (XX msg2 minimum) + assert_eq!(MSG3_WIRE_SIZE, 85); // 4 + 4 + 4 + 73 (XX msg3 minimum) assert_eq!(ENCRYPTED_MIN_SIZE, 32); // 16 + 16 assert_eq!(COMMON_PREFIX_SIZE, 4); assert_eq!(ESTABLISHED_HEADER_SIZE, 16); @@ -557,7 +664,7 @@ mod tests { fn test_roundtrip_indices() { let idx = SessionIndex::new(0xDEADBEEF); - let msg1 = build_msg1(idx, &[0u8; HANDSHAKE_MSG1_SIZE]); + let msg1 = build_msg1(idx, &[0u8; XX_HANDSHAKE_MSG1_SIZE]); let parsed = Msg1Header::parse(&msg1).unwrap(); assert_eq!(parsed.sender_idx.as_u32(), 0xDEADBEEF); @@ -588,17 +695,22 @@ mod tests { #[test] fn test_flags_byte() { - let header = - build_established_header(SessionIndex::new(1), 0, FLAG_KEY_EPOCH | FLAG_SP, 100); + let header = build_established_header( + SessionIndex::new(1), + 0, + FLAG_KEY_EPOCH | FLAG_SP, + 100, + ); assert_eq!(header[1], 0x05); // bits 0 and 2 set let parsed = EncryptedHeader::parse(&[ - header[0], header[1], header[2], header[3], header[4], header[5], header[6], header[7], - header[8], header[9], header[10], header[11], header[12], header[13], header[14], - header[15], // minimum: TAG_SIZE bytes of ciphertext + header[0], header[1], header[2], header[3], + header[4], header[5], header[6], header[7], + header[8], header[9], header[10], header[11], + header[12], header[13], header[14], header[15], + // minimum: TAG_SIZE bytes of ciphertext 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - ]) - .unwrap(); + ]).unwrap(); assert_eq!(parsed.flags & FLAG_KEY_EPOCH, FLAG_KEY_EPOCH); assert_eq!(parsed.flags & FLAG_CE, 0); assert_eq!(parsed.flags & FLAG_SP, FLAG_SP); @@ -606,10 +718,10 @@ mod tests { #[test] fn test_payload_len_in_msg1() { - let packet = build_msg1(SessionIndex::new(1), &[0u8; HANDSHAKE_MSG1_SIZE]); + let packet = build_msg1(SessionIndex::new(1), &[0u8; XX_HANDSHAKE_MSG1_SIZE]); let prefix = CommonPrefix::parse(&packet).unwrap(); - // payload_len = sender_idx(4) + noise_msg1(106) = 110 - assert_eq!(prefix.payload_len, 110); + // payload_len = sender_idx(4) + noise_msg1(33) = 37 + assert_eq!(prefix.payload_len, 37); } #[test] @@ -617,10 +729,71 @@ mod tests { let packet = build_msg2( SessionIndex::new(1), SessionIndex::new(2), - &[0u8; HANDSHAKE_MSG2_SIZE], + &[0u8; XX_HANDSHAKE_MSG2_SIZE], ); let prefix = CommonPrefix::parse(&packet).unwrap(); - // payload_len = sender_idx(4) + receiver_idx(4) + noise_msg2(57) = 65 - assert_eq!(prefix.payload_len, 65); + // payload_len = sender_idx(4) + receiver_idx(4) + noise_msg2(106) = 114 + assert_eq!(prefix.payload_len, 114); + } + + #[test] + fn test_msg3_header_parse() { + let sender_idx = SessionIndex::new(0xAABBCCDD); + let receiver_idx = SessionIndex::new(0x11223344); + let noise_msg3 = vec![0xdd; XX_HANDSHAKE_MSG3_SIZE]; + + let packet = build_msg3(sender_idx, receiver_idx, &noise_msg3); + + assert_eq!(packet.len(), MSG3_WIRE_SIZE); + assert_eq!(packet[0], 0x13); // ver=1, phase=3 + + let header = Msg3Header::parse(&packet).expect("should parse"); + assert_eq!(header.sender_idx, sender_idx); + assert_eq!(header.receiver_idx, receiver_idx); + assert_eq!(header.noise_msg3_offset, 12); + assert_eq!(header.noise_msg3(&packet), &noise_msg3[..]); + } + + #[test] + fn test_msg3_header_wrong_size() { + let packet = vec![0x13; MSG3_WIRE_SIZE - 1]; + assert!(Msg3Header::parse(&packet).is_none()); + + // Larger than minimum is now accepted (variable-length negotiation payload) + let mut packet = vec![0x13; MSG3_WIRE_SIZE + 10]; + packet[0] = 0x13; // ver=1, phase=3 + packet[1] = 0x00; + let header = Msg3Header::parse(&packet); + assert!(header.is_some()); + } + + #[test] + fn test_msg3_header_wrong_phase() { + let mut packet = vec![0x00; MSG3_WIRE_SIZE]; + packet[0] = 0x12; // ver=1, phase 2, not phase 3 + assert!(Msg3Header::parse(&packet).is_none()); + } + + #[test] + fn test_msg3_header_nonzero_flags() { + let mut packet = build_msg3( + SessionIndex::new(1), + SessionIndex::new(2), + &[0u8; XX_HANDSHAKE_MSG3_SIZE], + ); + packet[1] = 0x01; // flags must be zero during handshake + assert!(Msg3Header::parse(&packet).is_none()); + } + + #[test] + fn test_payload_len_in_msg3() { + let packet = build_msg3( + SessionIndex::new(1), + SessionIndex::new(2), + &[0u8; XX_HANDSHAKE_MSG3_SIZE], + ); + let prefix = CommonPrefix::parse(&packet).unwrap(); + // payload_len = sender_idx(4) + receiver_idx(4) + noise_msg3(73) = 81 + assert_eq!(prefix.payload_len, 81); } } diff --git a/src/noise/handshake.rs b/src/noise/handshake.rs index 737c0f0..ab43b60 100644 --- a/src/noise/handshake.rs +++ b/src/noise/handshake.rs @@ -1,12 +1,13 @@ use super::{ - CipherState, EPOCH_ENCRYPTED_SIZE, EPOCH_SIZE, HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE, - HandshakeProgress, HandshakeRole, NoiseError, NoisePattern, NoiseSession, PROTOCOL_NAME_IK, - PROTOCOL_NAME_XK, PUBKEY_SIZE, XK_HANDSHAKE_MSG1_SIZE, XK_HANDSHAKE_MSG2_SIZE, - XK_HANDSHAKE_MSG3_SIZE, + CipherState, HandshakeProgress, HandshakeRole, NoiseError, NoisePattern, NoiseSession, + EPOCH_ENCRYPTED_SIZE, EPOCH_SIZE, HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE, + PROTOCOL_NAME_IK, PROTOCOL_NAME_XK, PROTOCOL_NAME_XX, PUBKEY_SIZE, + XK_HANDSHAKE_MSG1_SIZE, XK_HANDSHAKE_MSG2_SIZE, XK_HANDSHAKE_MSG3_SIZE, + XX_HANDSHAKE_MSG1_SIZE, XX_HANDSHAKE_MSG2_SIZE, XX_HANDSHAKE_MSG3_SIZE, }; use hkdf::Hkdf; use rand::Rng; -use secp256k1::{Keypair, PublicKey, Secp256k1, SecretKey, ecdh::shared_secret_point}; +use secp256k1::{ecdh::shared_secret_point, Keypair, PublicKey, Secp256k1, SecretKey}; use sha2::{Digest, Sha256}; use std::fmt; @@ -102,7 +103,7 @@ impl SymmetricState { } } -/// Handshake state for Noise IK and XK patterns. +/// Handshake state for Noise IK, XK, and XX patterns. pub struct HandshakeState { /// Which Noise pattern is being used. pattern: NoisePattern, @@ -121,6 +122,8 @@ pub struct HandshakeState { /// For IK responder: learned from message 1. /// For XK initiator: known before handshake (from config). /// For XK responder: learned from message 3. + /// For XX initiator: learned from message 2. + /// For XX responder: learned from message 3. remote_static: Option, /// Remote ephemeral public key (learned during handshake). remote_ephemeral: Option, @@ -343,12 +346,8 @@ impl HandshakeState { }); } - let remote_static = self - .remote_static - .expect("initiator must have remote static"); - let epoch = self - .local_epoch - .expect("local epoch must be set before write_message_1"); + let remote_static = self.remote_static.expect("initiator must have remote static"); + let epoch = self.local_epoch.expect("local epoch must be set before write_message_1"); // Generate ephemeral keypair self.generate_ephemeral(); @@ -466,9 +465,7 @@ impl HandshakeState { } let re = self.remote_ephemeral.expect("should have remote ephemeral"); - let epoch = self - .local_epoch - .expect("local epoch must be set before write_message_2"); + let epoch = self.local_epoch.expect("local epoch must be set before write_message_2"); // Generate ephemeral keypair self.generate_ephemeral(); @@ -578,9 +575,7 @@ impl HandshakeState { }); } - let remote_static = self - .remote_static - .expect("initiator must have remote static"); + let remote_static = self.remote_static.expect("initiator must have remote static"); // Generate ephemeral keypair self.generate_ephemeral(); @@ -665,9 +660,7 @@ impl HandshakeState { } let re = self.remote_ephemeral.expect("should have remote ephemeral"); - let epoch = self - .local_epoch - .expect("local epoch must be set before write_xk_message_2"); + let epoch = self.local_epoch.expect("local epoch must be set before write_xk_message_2"); // Generate ephemeral keypair self.generate_ephemeral(); @@ -765,12 +758,8 @@ impl HandshakeState { }); } - let re = self - .remote_ephemeral - .expect("should have remote ephemeral after msg2"); - let epoch = self - .local_epoch - .expect("local epoch must be set before write_xk_message_3"); + let re = self.remote_ephemeral.expect("should have remote ephemeral after msg2"); + let epoch = self.local_epoch.expect("local epoch must be set before write_xk_message_3"); let mut message = Vec::with_capacity(XK_HANDSHAKE_MSG3_SIZE); @@ -827,10 +816,7 @@ impl HandshakeState { // -> se: DH(e, rs), mix into key // (responder uses their ephemeral with initiator's now-known static) - let ephemeral = self - .ephemeral_keypair - .as_ref() - .expect("should have ephemeral after msg2"); + let ephemeral = self.ephemeral_keypair.as_ref().expect("should have ephemeral after msg2"); let se = self.ecdh(&ephemeral.secret_key(), &rs); self.symmetric.mix_key(&se); @@ -848,6 +834,367 @@ impl HandshakeState { Ok(()) } + // ======================================================================== + // XX Pattern Methods + // ======================================================================== + + /// Create a new XX handshake as initiator. + /// + /// XX: neither side knows the other's static key. No pre-message. + pub fn new_xx_initiator(static_keypair: Keypair) -> Self { + let secp = Secp256k1::new(); + Self { + pattern: NoisePattern::Xx, + role: HandshakeRole::Initiator, + progress: HandshakeProgress::Initial, + symmetric: SymmetricState::initialize(PROTOCOL_NAME_XX), + static_keypair, + ephemeral_keypair: None, + remote_static: None, + remote_ephemeral: None, + secp, + local_epoch: None, + remote_epoch: None, + } + // No pre-message: neither side's static is mixed into hash. + } + + /// Create a new XX handshake as responder. + /// + /// XX: neither side knows the other's static key. No pre-message. + pub fn new_xx_responder(static_keypair: Keypair) -> Self { + let secp = Secp256k1::new(); + Self { + pattern: NoisePattern::Xx, + role: HandshakeRole::Responder, + progress: HandshakeProgress::Initial, + symmetric: SymmetricState::initialize(PROTOCOL_NAME_XX), + static_keypair, + ephemeral_keypair: None, + remote_static: None, + remote_ephemeral: None, + secp, + local_epoch: None, + remote_epoch: None, + } + // No pre-message: neither side's static is mixed into hash. + } + + /// Write XX message 1 (initiator only). + /// + /// XX msg1: `-> e` + /// - e: ephemeral public key (33 bytes) + /// - No DH operations (responder's static is unknown) + /// + /// Total: 33 bytes + pub fn write_xx_message_1(&mut self) -> Result, NoiseError> { + if self.role != HandshakeRole::Initiator { + return Err(NoiseError::WrongState { + expected: "initiator".to_string(), + got: "responder".to_string(), + }); + } + if self.progress != HandshakeProgress::Initial { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Initial.to_string(), + got: self.progress.to_string(), + }); + } + + // Generate ephemeral keypair + self.generate_ephemeral(); + let ephemeral = self.ephemeral_keypair.as_ref().unwrap(); + let e_pub = ephemeral.public_key().serialize(); + + let mut message = Vec::with_capacity(XX_HANDSHAKE_MSG1_SIZE); + + // -> e: send ephemeral, mix into hash + message.extend_from_slice(&e_pub); + self.symmetric.mix_hash(&e_pub); + + // No DH here — responder's static is unknown in XX + + self.progress = HandshakeProgress::Message1Done; + + Ok(message) + } + + /// Read XX message 1 (responder only). + /// + /// Parses the initiator's ephemeral key. No identity learned. + pub fn read_xx_message_1(&mut self, message: &[u8]) -> Result<(), NoiseError> { + if self.role != HandshakeRole::Responder { + return Err(NoiseError::WrongState { + expected: "responder".to_string(), + got: "initiator".to_string(), + }); + } + if self.progress != HandshakeProgress::Initial { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Initial.to_string(), + got: self.progress.to_string(), + }); + } + if message.len() != XX_HANDSHAKE_MSG1_SIZE { + return Err(NoiseError::MessageTooShort { + expected: XX_HANDSHAKE_MSG1_SIZE, + got: message.len(), + }); + } + + // -> e: parse remote ephemeral, mix into hash + let re = PublicKey::from_slice(&message[..PUBKEY_SIZE]) + .map_err(|_| NoiseError::InvalidPublicKey)?; + self.remote_ephemeral = Some(re); + self.symmetric.mix_hash(&message[..PUBKEY_SIZE]); + + // No DH here — responder's static is not in pre-message for XX + + self.progress = HandshakeProgress::Message1Done; + + Ok(()) + } + + /// Write XX message 2 (responder only). + /// + /// XX msg2: `<- e, ee, s, es` + encrypted epoch + /// - e: ephemeral public key (33 bytes) + /// - ee: DH(e_priv, re_pub), mix_key + /// - s: encrypt_and_hash(s_pub) — encrypted static (49 bytes) + /// - es: DH(s_priv, re_pub), mix_key + /// - encrypted epoch (24 bytes) + /// + /// Total: 106 bytes + pub fn write_xx_message_2(&mut self) -> Result, NoiseError> { + if self.role != HandshakeRole::Responder { + return Err(NoiseError::WrongState { + expected: "responder".to_string(), + got: "initiator".to_string(), + }); + } + if self.progress != HandshakeProgress::Message1Done { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Message1Done.to_string(), + got: self.progress.to_string(), + }); + } + + let re = self.remote_ephemeral.expect("should have remote ephemeral"); + let epoch = self.local_epoch.expect("local epoch must be set before write_xx_message_2"); + + // Generate ephemeral keypair + self.generate_ephemeral(); + let ephemeral = self.ephemeral_keypair.as_ref().unwrap(); + let e_pub = ephemeral.public_key().serialize(); + + let mut message = Vec::with_capacity(XX_HANDSHAKE_MSG2_SIZE); + + // <- e: send ephemeral, mix into hash + message.extend_from_slice(&e_pub); + self.symmetric.mix_hash(&e_pub); + + // <- ee: DH(e, re), mix into key + let ee = self.ecdh(&ephemeral.secret_key(), &re); + self.symmetric.mix_key(&ee); + + // <- s: encrypt our static and send + let our_static = self.static_keypair.public_key().serialize(); + let encrypted_static = self.symmetric.encrypt_and_hash(&our_static)?; + message.extend_from_slice(&encrypted_static); + + // <- es: DH(s, re), mix into key + let es = self.ecdh(&self.static_keypair.secret_key(), &re); + self.symmetric.mix_key(&es); + + // <- epoch: encrypt startup epoch for restart detection + let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?; + debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE); + message.extend_from_slice(&encrypted_epoch); + + self.progress = HandshakeProgress::Message2Done; + + Ok(message) + } + + /// Read XX message 2 (initiator only). + /// + /// Processes the responder's ephemeral and encrypted static key. + /// After this, the initiator learns the responder's identity. + pub fn read_xx_message_2(&mut self, message: &[u8]) -> Result<(), NoiseError> { + if self.role != HandshakeRole::Initiator { + return Err(NoiseError::WrongState { + expected: "initiator".to_string(), + got: "responder".to_string(), + }); + } + if self.progress != HandshakeProgress::Message1Done { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Message1Done.to_string(), + got: self.progress.to_string(), + }); + } + if message.len() != XX_HANDSHAKE_MSG2_SIZE { + return Err(NoiseError::MessageTooShort { + expected: XX_HANDSHAKE_MSG2_SIZE, + got: message.len(), + }); + } + + // <- e: parse remote ephemeral, mix into hash + let e_pub = &message[..PUBKEY_SIZE]; + let re = PublicKey::from_slice(e_pub).map_err(|_| NoiseError::InvalidPublicKey)?; + self.remote_ephemeral = Some(re); + self.symmetric.mix_hash(e_pub); + + // <- ee: DH(e, re), mix into key + let ephemeral = self.ephemeral_keypair.as_ref().unwrap(); + let ee = self.ecdh(&ephemeral.secret_key(), &re); + self.symmetric.mix_key(&ee); + + // <- s: decrypt responder's static + let encrypted_static_end = PUBKEY_SIZE + PUBKEY_SIZE + super::TAG_SIZE; + let encrypted_static = &message[PUBKEY_SIZE..encrypted_static_end]; + let decrypted_static = self.symmetric.decrypt_and_hash(encrypted_static)?; + let rs = + PublicKey::from_slice(&decrypted_static).map_err(|_| NoiseError::InvalidPublicKey)?; + self.remote_static = Some(rs); + + // <- es: DH(e, rs), mix into key + let es = self.ecdh(&ephemeral.secret_key(), &rs); + self.symmetric.mix_key(&es); + + // <- epoch: decrypt responder's startup epoch + let encrypted_epoch = &message[encrypted_static_end..]; + debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE); + let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?; + debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE); + let mut epoch = [0u8; EPOCH_SIZE]; + epoch.copy_from_slice(&decrypted_epoch); + self.remote_epoch = Some(epoch); + + self.progress = HandshakeProgress::Message2Done; + + Ok(()) + } + + /// Write XX message 3 (initiator only). + /// + /// XX msg3: `-> s, se` + encrypted epoch + /// - s: encrypt_and_hash(s_pub) — encrypted static (49 bytes) + /// - se: DH(s_priv, re_pub), mix_key + /// - encrypted epoch (24 bytes) + /// + /// Total: 73 bytes + pub fn write_xx_message_3(&mut self) -> Result, NoiseError> { + if self.role != HandshakeRole::Initiator { + return Err(NoiseError::WrongState { + expected: "initiator".to_string(), + got: "responder".to_string(), + }); + } + if self.progress != HandshakeProgress::Message2Done { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Message2Done.to_string(), + got: self.progress.to_string(), + }); + } + + let re = self.remote_ephemeral.expect("should have remote ephemeral after msg2"); + let epoch = self.local_epoch.expect("local epoch must be set before write_xx_message_3"); + + let mut message = Vec::with_capacity(XX_HANDSHAKE_MSG3_SIZE); + + // -> s: encrypt our static and send + let our_static = self.static_keypair.public_key().serialize(); + let encrypted_static = self.symmetric.encrypt_and_hash(&our_static)?; + message.extend_from_slice(&encrypted_static); + + // -> se: DH(s, re), mix into key + let se = self.ecdh(&self.static_keypair.secret_key(), &re); + self.symmetric.mix_key(&se); + + // -> epoch: encrypt startup epoch for restart detection + let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?; + debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE); + message.extend_from_slice(&encrypted_epoch); + + self.progress = HandshakeProgress::Complete; + + Ok(message) + } + + /// Read XX message 3 (responder only). + /// + /// Processes the initiator's encrypted static key and epoch. + /// After this, the responder learns the initiator's identity. + pub fn read_xx_message_3(&mut self, message: &[u8]) -> Result<(), NoiseError> { + if self.role != HandshakeRole::Responder { + return Err(NoiseError::WrongState { + expected: "responder".to_string(), + got: "initiator".to_string(), + }); + } + if self.progress != HandshakeProgress::Message2Done { + return Err(NoiseError::WrongState { + expected: HandshakeProgress::Message2Done.to_string(), + got: self.progress.to_string(), + }); + } + if message.len() != XX_HANDSHAKE_MSG3_SIZE { + return Err(NoiseError::MessageTooShort { + expected: XX_HANDSHAKE_MSG3_SIZE, + got: message.len(), + }); + } + + // -> s: decrypt initiator's static + let encrypted_static_end = PUBKEY_SIZE + super::TAG_SIZE; + let encrypted_static = &message[..encrypted_static_end]; + let decrypted_static = self.symmetric.decrypt_and_hash(encrypted_static)?; + let rs = + PublicKey::from_slice(&decrypted_static).map_err(|_| NoiseError::InvalidPublicKey)?; + self.remote_static = Some(rs); + + // -> se: DH(e, rs), mix into key + // (responder uses their ephemeral with initiator's now-known static) + let ephemeral = self.ephemeral_keypair.as_ref().expect("should have ephemeral after msg2"); + let se = self.ecdh(&ephemeral.secret_key(), &rs); + self.symmetric.mix_key(&se); + + // -> epoch: decrypt initiator's startup epoch + let encrypted_epoch = &message[encrypted_static_end..]; + debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE); + let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?; + debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE); + let mut epoch = [0u8; EPOCH_SIZE]; + epoch.copy_from_slice(&decrypted_epoch); + self.remote_epoch = Some(epoch); + + self.progress = HandshakeProgress::Complete; + + Ok(()) + } + + // ======================================================================== + // Payload Encryption (for negotiation payload in XX msg2/msg3) + // ======================================================================== + + /// Encrypt additional payload and mix into the handshake hash. + /// + /// Called after write_xx_message_2/3 to append negotiation payload. + /// Must be called before `into_session()`. + pub fn encrypt_payload(&mut self, plaintext: &[u8]) -> Result, NoiseError> { + self.symmetric.encrypt_and_hash(plaintext) + } + + /// Decrypt additional payload and mix ciphertext into the handshake hash. + /// + /// Called after read_xx_message_2/3 to extract negotiation payload. + /// Must be called before `into_session()`. + pub fn decrypt_payload(&mut self, ciphertext: &[u8]) -> Result, NoiseError> { + self.symmetric.decrypt_and_hash(ciphertext) + } + /// Complete the handshake and return a NoiseSession. /// /// Must be called after the handshake is complete. diff --git a/src/noise/mod.rs b/src/noise/mod.rs index 6c71c45..f5c40ab 100644 --- a/src/noise/mod.rs +++ b/src/noise/mod.rs @@ -11,6 +11,10 @@ //! own identity until msg3, providing stronger identity hiding. Three-message //! handshake. //! +//! - **XX pattern**: Neither side knows the other's static key. Both identities +//! are revealed during the handshake: responder in msg2, initiator in msg3. +//! Three-message handshake. Will replace IK/XK for both layers. +//! //! ## IK Handshake Pattern (Link Layer) //! //! ```text @@ -28,6 +32,14 @@ //! -> s, se (msg3: encrypted static + DH) //! ``` //! +//! ## XX Handshake Pattern +//! +//! ```text +//! -> e (msg1: ephemeral only, no DH) +//! <- e, ee, s, es (msg2: ephemeral + encrypted static) +//! -> s, se (msg3: encrypted static) +//! ``` +//! //! ## Separation of Concerns //! //! The IK pattern handles **link-layer peer authentication** — securing the @@ -58,6 +70,10 @@ pub(crate) const PROTOCOL_NAME_IK: &[u8] = b"Noise_IK_secp256k1_ChaChaPoly_SHA25 /// Format: Noise_XK_secp256k1_ChaChaPoly_SHA256 pub(crate) const PROTOCOL_NAME_XK: &[u8] = b"Noise_XK_secp256k1_ChaChaPoly_SHA256"; +/// Protocol name for Noise XX with secp256k1. +/// Format: Noise_XX_secp256k1_ChaChaPoly_SHA256 +pub(crate) const PROTOCOL_NAME_XX: &[u8] = b"Noise_XX_secp256k1_ChaChaPoly_SHA256"; + /// Maximum message size for noise transport messages. pub const MAX_MESSAGE_SIZE: usize = 65535; @@ -88,6 +104,15 @@ pub const XK_HANDSHAKE_MSG2_SIZE: usize = PUBKEY_SIZE + EPOCH_ENCRYPTED_SIZE; /// XK msg3: encrypted static (33 + 16 tag) + encrypted epoch (8 + 16 tag) = 73 bytes. pub const XK_HANDSHAKE_MSG3_SIZE: usize = PUBKEY_SIZE + TAG_SIZE + EPOCH_ENCRYPTED_SIZE; +/// XX msg1: ephemeral only (33 bytes). No DH, no encryption. +pub const XX_HANDSHAKE_MSG1_SIZE: usize = PUBKEY_SIZE; + +/// XX msg2: ephemeral (33) + encrypted static (33 + 16 tag) + encrypted epoch (8 + 16 tag) = 106 bytes. +pub const XX_HANDSHAKE_MSG2_SIZE: usize = PUBKEY_SIZE + PUBKEY_SIZE + TAG_SIZE + EPOCH_ENCRYPTED_SIZE; + +/// XX msg3: encrypted static (33 + 16 tag) + encrypted epoch (8 + 16 tag) = 73 bytes. +pub const XX_HANDSHAKE_MSG3_SIZE: usize = PUBKEY_SIZE + TAG_SIZE + EPOCH_ENCRYPTED_SIZE; + /// Replay window size in packets (matching WireGuard). pub const REPLAY_WINDOW_SIZE: usize = 2048; @@ -153,6 +178,8 @@ pub enum NoisePattern { Ik, /// Noise XK: three-message handshake (session layer). Xk, + /// Noise XX: three-message handshake, no prior key knowledge. + Xx, } /// Handshake state machine states. @@ -162,7 +189,7 @@ pub enum HandshakeProgress { Initial, /// Message 1 sent/received, ready for message 2. Message1Done, - /// Message 2 sent/received, ready for message 3 (XK only). + /// Message 2 sent/received, ready for message 3 (XK/XX only). Message2Done, /// Handshake complete, ready for transport. Complete, diff --git a/src/noise/tests.rs b/src/noise/tests.rs index f452f8b..16df6f9 100644 --- a/src/noise/tests.rs +++ b/src/noise/tests.rs @@ -795,3 +795,317 @@ fn test_xk_invalid_msg3_size() { .is_err() ); } + +// ===== XX Handshake Tests ===== + +#[test] +fn test_xx_full_handshake() { + let initiator_keypair = generate_keypair(); + let responder_keypair = generate_keypair(); + let initiator_epoch = generate_epoch(); + let responder_epoch = generate_epoch(); + + // XX: neither side knows the other's static key + let mut initiator = HandshakeState::new_xx_initiator(initiator_keypair); + initiator.set_local_epoch(initiator_epoch); + let mut responder = HandshakeState::new_xx_responder(responder_keypair); + responder.set_local_epoch(responder_epoch); + + assert_eq!(initiator.role(), HandshakeRole::Initiator); + assert_eq!(responder.role(), HandshakeRole::Responder); + + // Neither side knows the other's identity + assert!(initiator.remote_static().is_none()); + assert!(responder.remote_static().is_none()); + + // Message 1: Initiator -> Responder (e only) + let msg1 = initiator.write_xx_message_1().unwrap(); + assert_eq!(msg1.len(), XX_HANDSHAKE_MSG1_SIZE); + assert_eq!(msg1.len(), 33); + + responder.read_xx_message_1(&msg1).unwrap(); + + // After msg1: still no identities known + assert!(initiator.remote_static().is_none()); + assert!(responder.remote_static().is_none()); + + // Message 2: Responder -> Initiator (e, ee, s, es + epoch) + let msg2 = responder.write_xx_message_2().unwrap(); + assert_eq!(msg2.len(), XX_HANDSHAKE_MSG2_SIZE); + assert_eq!(msg2.len(), 106); + + initiator.read_xx_message_2(&msg2).unwrap(); + + // After msg2: initiator knows responder's identity + assert!(initiator.remote_static().is_some()); + assert_eq!( + initiator.remote_static().unwrap(), + &responder_keypair.public_key() + ); + assert_eq!(initiator.remote_epoch(), Some(responder_epoch)); + // Responder still doesn't know initiator + assert!(responder.remote_static().is_none()); + + // Neither side is complete yet + assert!(!initiator.is_complete()); + assert!(!responder.is_complete()); + + // Message 3: Initiator -> Responder (s, se + epoch) + let msg3 = initiator.write_xx_message_3().unwrap(); + assert_eq!(msg3.len(), XX_HANDSHAKE_MSG3_SIZE); + assert_eq!(msg3.len(), 73); + + responder.read_xx_message_3(&msg3).unwrap(); + + // Both should be complete now + assert!(initiator.is_complete()); + assert!(responder.is_complete()); + + // After msg3: responder knows initiator's identity + assert!(responder.remote_static().is_some()); + assert_eq!( + responder.remote_static().unwrap(), + &initiator_keypair.public_key() + ); + assert_eq!(responder.remote_epoch(), Some(initiator_epoch)); + + // Handshake hashes should match + assert_eq!(initiator.handshake_hash(), responder.handshake_hash()); + + // Convert to sessions + let mut initiator_session = initiator.into_session().unwrap(); + let mut responder_session = responder.into_session().unwrap(); + + // Test bidirectional encryption + let plaintext = b"Hello via XX!"; + let ciphertext = initiator_session.encrypt(plaintext).unwrap(); + let decrypted = responder_session.decrypt(&ciphertext).unwrap(); + assert_eq!(decrypted, plaintext); + + let plaintext2 = b"XX reply!"; + let ciphertext2 = responder_session.encrypt(plaintext2).unwrap(); + let decrypted2 = initiator_session.decrypt(&ciphertext2).unwrap(); + assert_eq!(decrypted2, plaintext2); +} + +#[test] +fn test_xx_message_sizes() { + assert_eq!(XX_HANDSHAKE_MSG1_SIZE, 33); // ephemeral only + assert_eq!(XX_HANDSHAKE_MSG2_SIZE, 33 + 33 + 16 + 24); // e + encrypted static + encrypted epoch + assert_eq!(XX_HANDSHAKE_MSG3_SIZE, 33 + 16 + 24); // encrypted static + encrypted epoch +} + +#[test] +fn test_xx_identity_timing() { + // XX property: initiator learns responder in msg2, responder learns initiator in msg3 + let initiator_keypair = generate_keypair(); + let responder_keypair = generate_keypair(); + + let mut initiator = HandshakeState::new_xx_initiator(initiator_keypair); + initiator.set_local_epoch(generate_epoch()); + let mut responder = HandshakeState::new_xx_responder(responder_keypair); + responder.set_local_epoch(generate_epoch()); + + // Before any messages: neither side knows + assert!(initiator.remote_static().is_none()); + assert!(responder.remote_static().is_none()); + + // After msg1 + let msg1 = initiator.write_xx_message_1().unwrap(); + responder.read_xx_message_1(&msg1).unwrap(); + assert!(initiator.remote_static().is_none(), "XX: initiator should NOT know identity after msg1"); + assert!(responder.remote_static().is_none(), "XX: responder should NOT know identity after msg1"); + + // After msg2: initiator learns responder + let msg2 = responder.write_xx_message_2().unwrap(); + initiator.read_xx_message_2(&msg2).unwrap(); + assert!(initiator.remote_static().is_some(), "XX: initiator should know responder after msg2"); + assert_eq!(initiator.remote_static().unwrap(), &responder_keypair.public_key()); + assert!(responder.remote_static().is_none(), "XX: responder should NOT know initiator after msg2"); + + // After msg3: responder learns initiator + let msg3 = initiator.write_xx_message_3().unwrap(); + responder.read_xx_message_3(&msg3).unwrap(); + assert!(responder.remote_static().is_some(), "XX: responder should know initiator after msg3"); + assert_eq!(responder.remote_static().unwrap(), &initiator_keypair.public_key()); +} + +#[test] +fn test_xx_wrong_state_errors() { + let keypair1 = generate_keypair(); + let keypair2 = generate_keypair(); + + // Initiator can't read XX msg1 + let mut initiator = HandshakeState::new_xx_initiator(keypair1); + initiator.set_local_epoch(generate_epoch()); + assert!(initiator.read_xx_message_1(&[0u8; XX_HANDSHAKE_MSG1_SIZE]).is_err()); + + // Initiator can't write msg2 + assert!(initiator.write_xx_message_2().is_err()); + + // Initiator can't write msg3 before msg2 + assert!(initiator.write_xx_message_3().is_err()); + + // Responder can't write msg1 + let mut responder = HandshakeState::new_xx_responder(keypair2); + responder.set_local_epoch(generate_epoch()); + assert!(responder.write_xx_message_1().is_err()); + + // Responder can't read msg3 before msg2 + assert!(responder.read_xx_message_3(&[0u8; XX_HANDSHAKE_MSG3_SIZE]).is_err()); +} + +#[test] +fn test_xx_handshake_hash_differs() { + // XX should produce different handshake hashes from both IK and XK + let keypair1 = generate_keypair(); + let keypair2 = generate_keypair(); + let epoch1 = generate_epoch(); + let epoch2 = generate_epoch(); + + // Complete an IK handshake + let mut ik_init = HandshakeState::new_initiator(keypair1, keypair2.public_key()); + ik_init.set_local_epoch(epoch1); + let mut ik_resp = HandshakeState::new_responder(keypair2); + ik_resp.set_local_epoch(epoch2); + let msg1 = ik_init.write_message_1().unwrap(); + ik_resp.read_message_1(&msg1).unwrap(); + let msg2 = ik_resp.write_message_2().unwrap(); + ik_init.read_message_2(&msg2).unwrap(); + let ik_hash = ik_init.handshake_hash(); + + // Complete an XK handshake with the same keys + let mut xk_init = HandshakeState::new_xk_initiator(keypair1, keypair2.public_key()); + xk_init.set_local_epoch(epoch1); + let mut xk_resp = HandshakeState::new_xk_responder(keypair2); + xk_resp.set_local_epoch(epoch2); + let msg1 = xk_init.write_xk_message_1().unwrap(); + xk_resp.read_xk_message_1(&msg1).unwrap(); + let msg2 = xk_resp.write_xk_message_2().unwrap(); + xk_init.read_xk_message_2(&msg2).unwrap(); + let msg3 = xk_init.write_xk_message_3().unwrap(); + xk_resp.read_xk_message_3(&msg3).unwrap(); + let xk_hash = xk_init.handshake_hash(); + + // Complete an XX handshake with the same keys + let mut xx_init = HandshakeState::new_xx_initiator(keypair1); + xx_init.set_local_epoch(epoch1); + let mut xx_resp = HandshakeState::new_xx_responder(keypair2); + xx_resp.set_local_epoch(epoch2); + let msg1 = xx_init.write_xx_message_1().unwrap(); + xx_resp.read_xx_message_1(&msg1).unwrap(); + let msg2 = xx_resp.write_xx_message_2().unwrap(); + xx_init.read_xx_message_2(&msg2).unwrap(); + let msg3 = xx_init.write_xx_message_3().unwrap(); + xx_resp.read_xx_message_3(&msg3).unwrap(); + let xx_hash = xx_init.handshake_hash(); + + assert_ne!(ik_hash, xx_hash, "IK and XX should produce different handshake hashes"); + assert_ne!(xk_hash, xx_hash, "XK and XX should produce different handshake hashes"); +} + +#[test] +fn test_xx_multiple_messages_after_handshake() { + let keypair1 = generate_keypair(); + let keypair2 = generate_keypair(); + + let mut initiator = HandshakeState::new_xx_initiator(keypair1); + initiator.set_local_epoch(generate_epoch()); + let mut responder = HandshakeState::new_xx_responder(keypair2); + responder.set_local_epoch(generate_epoch()); + + let msg1 = initiator.write_xx_message_1().unwrap(); + responder.read_xx_message_1(&msg1).unwrap(); + let msg2 = responder.write_xx_message_2().unwrap(); + initiator.read_xx_message_2(&msg2).unwrap(); + let msg3 = initiator.write_xx_message_3().unwrap(); + responder.read_xx_message_3(&msg3).unwrap(); + + let mut init_session = initiator.into_session().unwrap(); + let mut resp_session = responder.into_session().unwrap(); + + // Send many messages + for i in 0..100 { + let msg = format!("XX message {}", i); + let ct = init_session.encrypt(msg.as_bytes()).unwrap(); + let pt = resp_session.decrypt(&ct).unwrap(); + assert_eq!(pt, msg.as_bytes()); + } + + assert_eq!(init_session.send_nonce(), 100); + assert_eq!(resp_session.recv_nonce(), 100); +} + +#[test] +fn test_xx_with_odd_parity() { + // XX: no pre-message normalization needed, but ECDH x-only hashing + // must still produce matching shared secrets regardless of parity. + let secp = secp256k1::Secp256k1::new(); + + // Node A (initiator) - even parity key + let sk_a = secp256k1::SecretKey::from_slice( + &hex::decode("0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20") + .unwrap(), + ) + .unwrap(); + let kp_a = secp256k1::Keypair::from_secret_key(&secp, &sk_a); + + // Node B (responder) - odd parity key + let sk_b = secp256k1::SecretKey::from_slice( + &hex::decode("b102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1fb0") + .unwrap(), + ) + .unwrap(); + let kp_b = secp256k1::Keypair::from_secret_key(&secp, &sk_b); + let (_, parity_b) = kp_b.public_key().x_only_public_key(); + assert_eq!(parity_b, Parity::Odd, "Test requires odd-parity responder key"); + + let mut initiator = HandshakeState::new_xx_initiator(kp_a); + initiator.set_local_epoch(generate_epoch()); + let mut responder = HandshakeState::new_xx_responder(kp_b); + responder.set_local_epoch(generate_epoch()); + + let msg1 = initiator.write_xx_message_1().unwrap(); + responder.read_xx_message_1(&msg1).unwrap(); + let msg2 = responder.write_xx_message_2().unwrap(); + initiator.read_xx_message_2(&msg2).unwrap(); + let msg3 = initiator.write_xx_message_3().unwrap(); + responder.read_xx_message_3(&msg3).unwrap(); + + assert!(initiator.is_complete()); + assert!(responder.is_complete()); + + let mut sender = initiator.into_session().unwrap(); + let mut receiver = responder.into_session().unwrap(); + + let counter = sender.current_send_counter(); + let ciphertext = sender.encrypt(b"xx parity test").unwrap(); + let plaintext = receiver.decrypt_with_replay_check(&ciphertext, counter).unwrap(); + assert_eq!(plaintext, b"xx parity test"); +} + +#[test] +fn test_xx_invalid_msg_sizes() { + let keypair1 = generate_keypair(); + let keypair2 = generate_keypair(); + + // Wrong size for msg1 + let mut responder = HandshakeState::new_xx_responder(keypair1); + responder.set_local_epoch(generate_epoch()); + assert!(responder.read_xx_message_1(&[0u8; HANDSHAKE_MSG1_SIZE]).is_err()); // IK msg1 size + assert!(responder.read_xx_message_1(&[0u8; 10]).is_err()); + + // Wrong size for msg3 + let mut initiator = HandshakeState::new_xx_initiator(keypair1); + initiator.set_local_epoch(generate_epoch()); + let mut responder = HandshakeState::new_xx_responder(keypair2); + responder.set_local_epoch(generate_epoch()); + + let msg1 = initiator.write_xx_message_1().unwrap(); + responder.read_xx_message_1(&msg1).unwrap(); + let _msg2 = responder.write_xx_message_2().unwrap(); + + // Responder is now in Message2Done, try wrong-size msg3 + assert!(responder.read_xx_message_3(&[0u8; 10]).is_err()); + assert!(responder.read_xx_message_3(&[0u8; XX_HANDSHAKE_MSG3_SIZE + 1]).is_err()); +} diff --git a/src/peer/active.rs b/src/peer/active.rs index ce708c5..ca42030 100644 --- a/src/peer/active.rs +++ b/src/peer/active.rs @@ -5,10 +5,10 @@ use crate::bloom::BloomFilter; use crate::mmp::{MmpConfig, MmpPeerState}; +use crate::utils::index::SessionIndex; use crate::noise::{HandshakeState as NoiseHandshakeState, NoiseError, NoiseSession}; use crate::transport::{LinkId, LinkStats, TransportAddr, TransportId}; use crate::tree::{ParentDeclaration, TreeCoordinate}; -use crate::utils::index::SessionIndex; use crate::{FipsAddress, NodeAddr, PeerIdentity}; use secp256k1::XOnlyPublicKey; use std::fmt; @@ -32,10 +32,7 @@ pub enum ConnectivityState { impl ConnectivityState { /// Check if the peer is usable for sending traffic. pub fn can_send(&self) -> bool { - matches!( - self, - ConnectivityState::Connected | ConnectivityState::Stale - ) + matches!(self, ConnectivityState::Connected | ConnectivityState::Stale) } /// Check if this is a terminal state requiring cleanup. @@ -182,6 +179,11 @@ pub struct ActivePeer { rekey_msg1: Option>, /// In-progress rekey: next resend timestamp (Unix ms). rekey_msg1_next_resend: u64, + // === Rekey Responder State (XX pattern) === + /// In-progress rekey responder: Noise handshake state awaiting msg3. + rekey_responder_handshake: Option, + /// In-progress rekey responder: our new session index. + rekey_responder_our_index: Option, } impl ActivePeer { @@ -233,6 +235,8 @@ impl ActivePeer { rekey_our_index: None, rekey_msg1: None, rekey_msg1_next_resend: 0, + rekey_responder_handshake: None, + rekey_responder_our_index: None, } } @@ -313,6 +317,8 @@ impl ActivePeer { rekey_our_index: None, rekey_msg1: None, rekey_msg1_next_resend: 0, + rekey_responder_handshake: None, + rekey_responder_our_index: None, } } @@ -747,7 +753,12 @@ impl ActivePeer { // === Filter Updates === /// Update peer's inbound filter. - pub fn update_filter(&mut self, filter: BloomFilter, sequence: u64, current_time_ms: u64) { + pub fn update_filter( + &mut self, + filter: BloomFilter, + sequence: u64, + current_time_ms: u64, + ) { self.inbound_filter = Some(filter); self.filter_sequence = sequence; self.filter_received_at = current_time_ms; @@ -962,11 +973,12 @@ impl ActivePeer { self.rekey_msg1_next_resend = 0; self.rekey_in_progress = false; // Return whichever index needs freeing - self.rekey_our_index.take().or_else(|| { - self.pending_new_session = None; - self.pending_their_index = None; - self.pending_our_index.take() - }) + self.rekey_our_index.take() + .or_else(|| { + self.pending_new_session = None; + self.pending_their_index = None; + self.pending_our_index.take() + }) } // === Rekey Handshake State (Initiator) === @@ -991,33 +1003,92 @@ impl ActivePeer { self.rekey_our_index } - /// Complete the rekey by processing msg2 (initiator side). + /// Complete the rekey by processing msg2 (initiator side, XX pattern). /// - /// Takes the stored handshake state, reads msg2, and returns the - /// completed NoiseSession. Clears the handshake-related fields but - /// leaves rekey_our_index for set_pending_session to use. - pub fn complete_rekey_msg2(&mut self, msg2_bytes: &[u8]) -> Result { - let mut hs = self - .rekey_handshake + /// Takes the stored handshake state, reads XX msg2, generates XX msg3, + /// and returns (msg3_bytes, completed NoiseSession). Clears the + /// handshake-related fields but leaves rekey_our_index for + /// set_pending_session to use. + pub fn complete_rekey_msg2( + &mut self, + msg2_bytes: &[u8], + ) -> Result<(Vec, NoiseSession), NoiseError> { + let mut hs = self.rekey_handshake .take() .ok_or_else(|| NoiseError::WrongState { expected: "rekey handshake in progress".to_string(), got: "no handshake state".to_string(), })?; - hs.read_message_2(msg2_bytes)?; + // Split msg2 into base XX part and any extra (negotiation payload) + let base_size = crate::noise::XX_HANDSHAKE_MSG2_SIZE; + let (base_msg2, extra) = if msg2_bytes.len() > base_size { + (&msg2_bytes[..base_size], Some(&msg2_bytes[base_size..])) + } else { + (msg2_bytes, None) + }; + + hs.read_xx_message_2(base_msg2)?; + + // Must decrypt negotiation payload (if present) to keep hash chain + // in sync, even though rekey doesn't use the negotiation result. + if let Some(encrypted_neg) = extra { + let _ = hs.decrypt_payload(encrypted_neg)?; + } + + let msg3 = hs.write_xx_message_3()?; let session = hs.into_session()?; // Clear msg1 resend state self.rekey_msg1 = None; self.rekey_msg1_next_resend = 0; + Ok((msg3, session)) + } + + /// Complete the rekey by processing msg3 (responder side, XX pattern). + /// + /// Takes the stored responder handshake state, reads XX msg3, and returns + /// the completed NoiseSession. + pub fn complete_rekey_msg3( + &mut self, + msg3_bytes: &[u8], + ) -> Result { + let mut hs = self.rekey_responder_handshake + .take() + .ok_or_else(|| NoiseError::WrongState { + expected: "rekey responder handshake awaiting msg3".to_string(), + got: "no responder handshake state".to_string(), + })?; + + // Split msg3 into base XX part and any extra (negotiation payload) + let base_size = crate::noise::XX_HANDSHAKE_MSG3_SIZE; + let (base_msg3, extra) = if msg3_bytes.len() > base_size { + (&msg3_bytes[..base_size], Some(&msg3_bytes[base_size..])) + } else { + (msg3_bytes, None) + }; + + hs.read_xx_message_3(base_msg3)?; + + // Must decrypt negotiation payload (if present) to keep hash chain + // in sync, even though rekey doesn't use the negotiation result. + if let Some(encrypted_neg) = extra { + let _ = hs.decrypt_payload(encrypted_neg)?; + } + + let session = hs.into_session()?; + + self.rekey_responder_our_index = None; + Ok(session) } /// Check if msg1 needs resending. pub fn needs_msg1_resend(&self, now_ms: u64) -> bool { - self.rekey_in_progress && self.rekey_msg1.is_some() && now_ms >= self.rekey_msg1_next_resend + self.rekey_in_progress + && self.rekey_msg1.is_some() + && now_ms >= self.rekey_msg1_next_resend } /// Get msg1 bytes for resend (without consuming). @@ -1029,6 +1100,37 @@ impl ActivePeer { pub fn set_msg1_next_resend(&mut self, next_ms: u64) { self.rekey_msg1_next_resend = next_ms; } + + // === Rekey Responder State (XX pattern) === + + /// Whether this peer has a rekey responder handshake awaiting msg3. + pub fn has_rekey_responder_handshake(&self) -> bool { + self.rekey_responder_handshake.is_some() + } + + /// Get the rekey responder our_index. + pub fn rekey_responder_our_index(&self) -> Option { + self.rekey_responder_our_index + } + + /// Store rekey responder handshake state after sending msg2. + /// + /// Called when processing a rekey msg1 from the peer. The handshake + /// state is held here until msg3 arrives to complete the rekey. + pub fn set_rekey_responder_state( + &mut self, + handshake: NoiseHandshakeState, + our_index: SessionIndex, + ) { + self.rekey_responder_handshake = Some(handshake); + self.rekey_responder_our_index = Some(our_index); + } + + /// Clear rekey responder state (on failure or abandonment). + pub fn clear_rekey_responder(&mut self) { + self.rekey_responder_handshake = None; + self.rekey_responder_our_index = None; + } } #[cfg(test)] diff --git a/src/peer/connection.rs b/src/peer/connection.rs index bc6c6e2..58ed315 100644 --- a/src/peer/connection.rs +++ b/src/peer/connection.rs @@ -1,21 +1,21 @@ //! Peer Connection (Handshake Phase) //! //! Represents an in-progress connection before authentication completes. -//! PeerConnection tracks the Noise IK handshake state and transitions to +//! PeerConnection tracks the Noise XX handshake state and transitions to //! ActivePeer upon successful authentication. -use crate::PeerIdentity; +use crate::utils::index::SessionIndex; use crate::noise::{self, NoiseError, NoiseSession}; use crate::transport::{LinkDirection, LinkId, LinkStats, TransportAddr, TransportId}; -use crate::utils::index::SessionIndex; +use crate::PeerIdentity; use secp256k1::Keypair; use std::fmt; /// Handshake protocol state machine. /// -/// For Noise IK pattern: -/// - Initiator: Initial → SentMsg1 → Complete -/// - Responder: Initial → ReceivedMsg1 → Complete +/// For Noise XX pattern: +/// - Initiator: Initial → SentMsg1 → Complete (after processing msg2 + sending msg3) +/// - Responder: Initial → ReceivedMsg1 → Complete (after processing msg3) #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum HandshakeState { /// Initial state, ready to start handshake. @@ -398,8 +398,8 @@ impl PeerConnection { /// 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. + /// For outbound connections only. Returns the Noise XX msg1 bytes. + /// XX msg1 is ephemeral-only (33 bytes) — no identity or epoch. pub fn start_handshake( &mut self, our_keypair: Keypair, @@ -420,15 +420,10 @@ impl PeerConnection { }); } - let remote_static = self - .expected_identity - .as_ref() - .expect("outbound must have expected identity") - .pubkey_full(); - - let mut hs = noise::HandshakeState::new_initiator(our_keypair, remote_static); + // XX initiator: no remote static needed upfront + let mut hs = noise::HandshakeState::new_xx_initiator(our_keypair); hs.set_local_epoch(epoch); - let msg1 = hs.write_message_1()?; + let msg1 = hs.write_xx_message_1()?; self.noise_handshake = Some(hs); self.handshake_state = HandshakeState::SentMsg1; @@ -439,13 +434,19 @@ impl PeerConnection { /// 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. + /// For inbound connections only. Returns the Noise XX msg2 bytes. + /// XX: identity is NOT learned from msg1 (only ephemeral exchange). + /// The responder learns the initiator's identity from msg3. + /// The handshake remains in ReceivedMsg1 state (not Complete). + /// + /// If `negotiation_payload` is provided, it is encrypted and appended + /// to the returned msg2 bytes. pub fn receive_handshake_init( &mut self, our_keypair: Keypair, epoch: [u8; 8], message: &[u8], + negotiation_payload: Option<&[u8]>, current_time_ms: u64, ) -> Result, NoiseError> { if self.direction != LinkDirection::Inbound { @@ -462,41 +463,45 @@ impl PeerConnection { }); } - let mut hs = noise::HandshakeState::new_responder(our_keypair); + let mut hs = noise::HandshakeState::new_xx_responder(our_keypair); hs.set_local_epoch(epoch); - // Process message 1 (this reveals the initiator's identity and epoch) - hs.read_message_1(message)?; + // Process XX message 1 (ephemeral only — no identity learned) + hs.read_xx_message_1(message)?; - // Extract the discovered identity - let remote_static = *hs - .remote_static() - .expect("remote static available after msg1"); - self.expected_identity = Some(PeerIdentity::from_pubkey_full(remote_static)); + // Generate XX message 2 (sends our static + epoch) + let mut msg2 = hs.write_xx_message_2()?; - // Capture remote epoch from msg1 - self.remote_epoch = hs.remote_epoch(); + // Append encrypted negotiation payload if provided + if let Some(payload) = negotiation_payload { + let encrypted = hs.encrypt_payload(payload)?; + msg2.extend_from_slice(&encrypted); + } - // 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.handshake_state = HandshakeState::Complete; + // XX: handshake NOT complete yet — need msg3. + // Keep the handshake state for complete_handshake_msg3(). + self.noise_handshake = Some(hs); + self.handshake_state = HandshakeState::ReceivedMsg1; self.last_activity = current_time_ms; Ok(msg2) } - /// Complete the handshake by processing message 2. + /// Complete the handshake by processing message 2 and generating message 3. /// - /// For outbound connections only (initiator completing handshake). + /// For outbound connections only (initiator). Processes the responder's + /// msg2 (learning their identity and epoch), then generates msg3. + /// Returns the Noise XX msg3 bytes to send. + /// + /// If `negotiation_payload` is provided, it is encrypted and appended + /// to the returned msg3 bytes. If the received msg2 contains a negotiation + /// payload (bytes beyond the base XX msg2), it is decrypted and returned. pub fn complete_handshake( &mut self, message: &[u8], + negotiation_payload: Option<&[u8]>, current_time_ms: u64, - ) -> Result<(), NoiseError> { + ) -> Result<(Vec, Option>), NoiseError> { if self.handshake_state != HandshakeState::SentMsg1 { return Err(NoiseError::WrongState { expected: "sent_msg1 state".to_string(), @@ -509,17 +514,109 @@ impl PeerConnection { .take() .expect("noise handshake must exist in SentMsg1 state"); - hs.read_message_2(message)?; + // Split msg2 into base XX part and optional negotiation + let base_size = noise::XX_HANDSHAKE_MSG2_SIZE; + let (base_msg2, extra) = if message.len() > base_size { + (&message[..base_size], Some(&message[base_size..])) + } else { + (message, None) + }; + + // Process XX msg2 (learns responder identity + epoch) + hs.read_xx_message_2(base_msg2)?; + + // Decrypt negotiation payload from msg2 if present + let received_negotiation = if let Some(encrypted) = extra { + Some(hs.decrypt_payload(encrypted)?) + } else { + None + }; + + // Learn responder identity from msg2 + let remote_static = *hs + .remote_static() + .expect("remote static available after XX msg2"); + self.expected_identity = Some(PeerIdentity::from_pubkey_full(remote_static)); // Capture remote epoch from msg2 self.remote_epoch = hs.remote_epoch(); + // Generate XX msg3 + let mut msg3 = hs.write_xx_message_3()?; + + // Append encrypted negotiation payload if provided + if let Some(payload) = negotiation_payload { + let encrypted = hs.encrypt_payload(payload)?; + msg3.extend_from_slice(&encrypted); + } + + // Handshake complete for initiator let session = hs.into_session()?; self.noise_session = Some(session); self.handshake_state = HandshakeState::Complete; self.last_activity = current_time_ms; - Ok(()) + Ok((msg3, received_negotiation)) + } + + /// Complete the responder handshake by processing message 3. + /// + /// For inbound connections only (responder). Processes the initiator's + /// msg3, learning their identity and epoch. + /// + /// If the msg3 contains a negotiation payload (bytes beyond base XX msg3), + /// it is decrypted and returned. + pub fn complete_handshake_msg3( + &mut self, + message: &[u8], + current_time_ms: u64, + ) -> Result>, NoiseError> { + if self.handshake_state != HandshakeState::ReceivedMsg1 { + return Err(NoiseError::WrongState { + expected: "received_msg1 state".to_string(), + got: self.handshake_state.to_string(), + }); + } + + let mut hs = self + .noise_handshake + .take() + .expect("noise handshake must exist in ReceivedMsg1 state"); + + // Split msg3 into base XX part and optional negotiation + let base_size = noise::XX_HANDSHAKE_MSG3_SIZE; + let (base_msg3, extra) = if message.len() > base_size { + (&message[..base_size], Some(&message[base_size..])) + } else { + (message, None) + }; + + // Process XX msg3 (learns initiator identity + epoch) + hs.read_xx_message_3(base_msg3)?; + + // Decrypt negotiation payload from msg3 if present + let received_negotiation = if let Some(encrypted) = extra { + Some(hs.decrypt_payload(encrypted)?) + } else { + None + }; + + // Learn initiator identity from msg3 + let remote_static = *hs + .remote_static() + .expect("remote static available after XX msg3"); + self.expected_identity = Some(PeerIdentity::from_pubkey_full(remote_static)); + + // Capture remote epoch from msg3 + self.remote_epoch = hs.remote_epoch(); + + // Handshake complete for responder + let session = hs.into_session()?; + self.noise_session = Some(session); + self.handshake_state = HandshakeState::Complete; + self.last_activity = current_time_ms; + + Ok(received_negotiation) } /// Take the completed Noise session. @@ -558,6 +655,7 @@ impl PeerConnection { pub fn is_timed_out(&self, current_time_ms: u64, timeout_ms: u64) -> bool { self.idle_time(current_time_ms) > timeout_ms } + } impl fmt::Debug for PeerConnection { @@ -650,35 +748,40 @@ mod tests { 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 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(); + // Initiator starts XX handshake + let msg1 = initiator_conn.start_handshake(initiator_keypair, initiator_epoch, 1100).unwrap(); assert_eq!(initiator_conn.handshake_state(), HandshakeState::SentMsg1); - // Responder processes msg1 and sends msg2 + // Responder processes msg1 and sends msg2 (XX: does NOT complete yet) let msg2 = responder_conn - .receive_handshake_init(responder_keypair, responder_epoch, &msg1, 1200) + .receive_handshake_init(responder_keypair, responder_epoch, &msg1, None, 1200) .unwrap(); - assert_eq!(responder_conn.handshake_state(), HandshakeState::Complete); + assert_eq!(responder_conn.handshake_state(), HandshakeState::ReceivedMsg1); + // Responder does NOT know initiator's identity yet (XX property) + assert!(responder_conn.expected_identity().is_none()); - // 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(); + // Initiator processes msg2 and generates msg3 + let (msg3, _neg) = initiator_conn.complete_handshake(&msg2, None, 1300).unwrap(); assert_eq!(initiator_conn.handshake_state(), HandshakeState::Complete); - // Initiator learned responder's epoch + // Initiator learned responder's identity from msg2 + let discovered = initiator_conn.expected_identity().unwrap(); + assert_eq!(discovered.pubkey(), responder_identity.pubkey()); assert_eq!(initiator_conn.remote_epoch(), Some(responder_epoch)); + // Responder processes msg3 (completes handshake) + let _neg = responder_conn.complete_handshake_msg3(&msg3, 1400).unwrap(); + assert_eq!(responder_conn.handshake_state(), HandshakeState::Complete); + + // Responder learned initiator's identity from msg3 + let discovered = responder_conn.expected_identity().unwrap(); + assert_eq!(discovered.pubkey(), initiator_identity.pubkey()); + assert_eq!(responder_conn.remote_epoch(), Some(initiator_epoch)); + // Both have sessions assert!(initiator_conn.has_session()); assert!(responder_conn.has_session()); @@ -723,18 +826,12 @@ mod tests { // 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() - ); + assert!(outbound + .receive_handshake_init(keypair, make_epoch(), &[0u8; 33], None, 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() - ); + assert!(inbound.start_handshake(keypair, make_epoch(), 1100).is_err()); } } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 1e90519..f0496c1 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -24,25 +24,27 @@ mod discovery; mod error; mod filter; mod link; +mod negotiation; mod session; mod tree; // Re-export all public types at protocol:: level -pub use discovery::{LookupRequest, LookupResponse}; pub use error::ProtocolError; -pub use filter::FilterAnnounce; pub use link::{ - Disconnect, DisconnectReason, HandshakeMessageType, LinkMessageType, - SESSION_DATAGRAM_HEADER_SIZE, SessionDatagram, + Disconnect, DisconnectReason, HandshakeMessageType, LinkMessageType, SessionDatagram, + SESSION_DATAGRAM_HEADER_SIZE, }; +pub use tree::TreeAnnounce; +pub use filter::FilterAnnounce; +pub use discovery::{LookupRequest, LookupResponse}; +pub use negotiation::{NegotiationPayload, TlvEntry, NEGOTIATION_HEADER_SIZE}; pub use session::{ - COORDS_REQUIRED_SIZE, CoordsRequired, FspFlags, FspInnerFlags, MTU_EXCEEDED_SIZE, MtuExceeded, - PATH_MTU_NOTIFICATION_SIZE, PathBroken, PathMtuNotification, SESSION_RECEIVER_REPORT_SIZE, - SESSION_SENDER_REPORT_SIZE, SessionAck, SessionFlags, SessionMessageType, SessionMsg3, - SessionReceiverReport, SessionSenderReport, SessionSetup, + CoordsRequired, FspFlags, FspInnerFlags, MtuExceeded, PathBroken, PathMtuNotification, + SessionAck, SessionFlags, SessionMessageType, SessionMsg3, SessionReceiverReport, + SessionSenderReport, SessionSetup, COORDS_REQUIRED_SIZE, MTU_EXCEEDED_SIZE, + PATH_MTU_NOTIFICATION_SIZE, SESSION_RECEIVER_REPORT_SIZE, SESSION_SENDER_REPORT_SIZE, }; pub(crate) use session::{coords_wire_size, decode_optional_coords, encode_coords}; -pub use tree::TreeAnnounce; /// Protocol version for message compatibility. pub const PROTOCOL_VERSION: u8 = 1; diff --git a/src/protocol/negotiation.rs b/src/protocol/negotiation.rs new file mode 100644 index 0000000..63a122e --- /dev/null +++ b/src/protocol/negotiation.rs @@ -0,0 +1,289 @@ +//! Protocol negotiation payload codec. +//! +//! Encodes/decodes the negotiation payload embedded in XX handshake +//! messages (msg2/msg3). Each layer (FMP, FSP) uses the same wire +//! format with layer-specific version ranges and feature catalogs. +//! +//! ## Wire Format +//! +//! ```text +//! Byte 0: format (must be 0) +//! Byte 1: [version_min:4 high][version_max:4 low] +//! Bytes 2-9: feature bitfield (64 bits, LE) +//! Bytes 10+: TLV entries, each: +//! [field_num:2 LE][length:2 LE][value:N] +//! ``` + +use super::ProtocolError; + +/// Size of the fixed negotiation header (format + version + features). +pub const NEGOTIATION_HEADER_SIZE: usize = 10; + +/// Format byte value for the initial negotiation format. +const NEGOTIATION_FORMAT_V0: u8 = 0; + +/// A TLV entry in the negotiation payload. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TlvEntry { + /// Field number identifying this TLV. + pub field_num: u16, + /// Raw value bytes. + pub value: Vec, +} + +/// Protocol negotiation payload. +/// +/// Carried in XX msg2/msg3 encrypted payloads. Shared codec for both +/// FMP and FSP layers, with layer-specific version ranges and feature +/// bit assignments. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NegotiationPayload { + /// Format byte (must be 0). + pub format: u8, + /// Minimum protocol version supported (4-bit, 0-15). + pub version_min: u8, + /// Maximum protocol version supported (4-bit, 0-15). + pub version_max: u8, + /// Feature bitfield (64 bits, LE). + pub features: u64, + /// Optional TLV extension entries. + pub tlv_entries: Vec, +} + +impl NegotiationPayload { + /// Create a new negotiation payload. + pub fn new(version_min: u8, version_max: u8, features: u64) -> Self { + Self { + format: NEGOTIATION_FORMAT_V0, + version_min, + version_max, + features, + tlv_entries: Vec::new(), + } + } + + /// Add a TLV entry. + pub fn with_tlv(mut self, field_num: u16, value: Vec) -> Self { + self.tlv_entries.push(TlvEntry { field_num, value }); + self + } + + /// Encode to wire format. + pub fn encode(&self) -> Vec { + let mut buf = Vec::with_capacity(NEGOTIATION_HEADER_SIZE); + + buf.push(self.format); + buf.push((self.version_min << 4) | (self.version_max & 0x0F)); + buf.extend_from_slice(&self.features.to_le_bytes()); + + for entry in &self.tlv_entries { + buf.extend_from_slice(&entry.field_num.to_le_bytes()); + let len = entry.value.len() as u16; + buf.extend_from_slice(&len.to_le_bytes()); + buf.extend_from_slice(&entry.value); + } + + buf + } + + /// Decode from wire format. + pub fn decode(data: &[u8]) -> Result { + if data.len() < NEGOTIATION_HEADER_SIZE { + return Err(ProtocolError::MessageTooShort { + expected: NEGOTIATION_HEADER_SIZE, + got: data.len(), + }); + } + + let format = data[0]; + if format != NEGOTIATION_FORMAT_V0 { + return Err(ProtocolError::Malformed(format!( + "unknown negotiation format: {format}" + ))); + } + + let version_min = data[1] >> 4; + let version_max = data[1] & 0x0F; + if version_min > version_max { + return Err(ProtocolError::Malformed(format!( + "version_min ({version_min}) > version_max ({version_max})" + ))); + } + + let features = u64::from_le_bytes(data[2..10].try_into().unwrap()); + + let mut tlv_entries = Vec::new(); + let mut offset = NEGOTIATION_HEADER_SIZE; + while offset < data.len() { + // Need at least 4 bytes for field_num + length + if offset + 4 > data.len() { + return Err(ProtocolError::Malformed( + "truncated TLV header".to_string(), + )); + } + + let field_num = u16::from_le_bytes(data[offset..offset + 2].try_into().unwrap()); + let length = + u16::from_le_bytes(data[offset + 2..offset + 4].try_into().unwrap()) as usize; + offset += 4; + + if offset + length > data.len() { + return Err(ProtocolError::Malformed(format!( + "TLV field {field_num}: declared length {length} exceeds remaining data {}", + data.len() - offset + ))); + } + + let value = data[offset..offset + length].to_vec(); + offset += length; + + tlv_entries.push(TlvEntry { field_num, value }); + } + + Ok(Self { + format, + version_min, + version_max, + features, + tlv_entries, + }) + } + + /// Agree on a protocol version with a peer's negotiation payload. + /// + /// Returns `min(our_max, their_max)`, rejecting if the agreed version + /// is below either side's minimum. + pub fn agree_version(&self, other: &Self) -> Result { + let agreed = self.version_max.min(other.version_max); + if agreed < self.version_min || agreed < other.version_min { + return Err(ProtocolError::Malformed(format!( + "version mismatch: ours [{},{}] theirs [{},{}]", + self.version_min, self.version_max, other.version_min, other.version_max + ))); + } + Ok(agreed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_encode_decode_roundtrip() { + let payload = NegotiationPayload::new(1, 3, 0x00000000_0000002A); + let encoded = payload.encode(); + assert_eq!(encoded.len(), NEGOTIATION_HEADER_SIZE); + + let decoded = NegotiationPayload::decode(&encoded).unwrap(); + assert_eq!(decoded, payload); + } + + #[test] + fn test_encode_decode_with_tlv() { + let payload = NegotiationPayload::new(0, 1, 0) + .with_tlv(1, vec![0xAA, 0xBB]) + .with_tlv(256, vec![0x01, 0x02, 0x03, 0x04]); + + let encoded = payload.encode(); + // 10 header + (2+2+2) + (2+2+4) = 10 + 6 + 8 = 24 + assert_eq!(encoded.len(), 24); + + let decoded = NegotiationPayload::decode(&encoded).unwrap(); + assert_eq!(decoded, payload); + assert_eq!(decoded.tlv_entries.len(), 2); + assert_eq!(decoded.tlv_entries[0].field_num, 1); + assert_eq!(decoded.tlv_entries[0].value, vec![0xAA, 0xBB]); + assert_eq!(decoded.tlv_entries[1].field_num, 256); + assert_eq!(decoded.tlv_entries[1].value, vec![0x01, 0x02, 0x03, 0x04]); + } + + #[test] + fn test_version_agreement_basic() { + let ours = NegotiationPayload::new(1, 3, 0); + let theirs = NegotiationPayload::new(1, 2, 0); + assert_eq!(ours.agree_version(&theirs).unwrap(), 2); + } + + #[test] + fn test_version_agreement_mismatch() { + let ours = NegotiationPayload::new(3, 5, 0); + let theirs = NegotiationPayload::new(1, 2, 0); + assert!(ours.agree_version(&theirs).is_err()); + } + + #[test] + fn test_version_agreement_asymmetric() { + // Ours: [2,5], theirs: [1,4] → agreed = min(5,4) = 4, 4 >= 2 and 4 >= 1 → ok + let ours = NegotiationPayload::new(2, 5, 0); + let theirs = NegotiationPayload::new(1, 4, 0); + assert_eq!(ours.agree_version(&theirs).unwrap(), 4); + + // Ours: [1,4], theirs: [2,5] → agreed = min(4,5) = 4, 4 >= 1 and 4 >= 2 → ok + assert_eq!(theirs.agree_version(&ours).unwrap(), 4); + } + + #[test] + fn test_unknown_format_rejected() { + let mut data = NegotiationPayload::new(0, 0, 0).encode(); + data[0] = 1; // Set format to 1 + assert!(NegotiationPayload::decode(&data).is_err()); + } + + #[test] + fn test_invalid_version_range() { + let mut data = NegotiationPayload::new(0, 0, 0).encode(); + // Set version_min=5, version_max=3 (invalid: min > max) + data[1] = (5 << 4) | 3; + assert!(NegotiationPayload::decode(&data).is_err()); + } + + #[test] + fn test_unknown_tlv_forward_compat() { + // Unknown field_nums should be preserved through encode/decode + let payload = NegotiationPayload::new(0, 1, 0) + .with_tlv(9999, vec![0xFF, 0xFE, 0xFD]); + + let encoded = payload.encode(); + let decoded = NegotiationPayload::decode(&encoded).unwrap(); + assert_eq!(decoded.tlv_entries.len(), 1); + assert_eq!(decoded.tlv_entries[0].field_num, 9999); + assert_eq!(decoded.tlv_entries[0].value, vec![0xFF, 0xFE, 0xFD]); + } + + #[test] + fn test_empty_payload() { + let payload = NegotiationPayload::new(0, 0, 0); + let encoded = payload.encode(); + assert_eq!(encoded.len(), NEGOTIATION_HEADER_SIZE); + + let decoded = NegotiationPayload::decode(&encoded).unwrap(); + assert_eq!(decoded.version_min, 0); + assert_eq!(decoded.version_max, 0); + assert_eq!(decoded.features, 0); + assert!(decoded.tlv_entries.is_empty()); + } + + #[test] + fn test_truncated_payload() { + // Less than header size + assert!(NegotiationPayload::decode(&[0u8; 5]).is_err()); + assert!(NegotiationPayload::decode(&[]).is_err()); + } + + #[test] + fn test_truncated_tlv() { + let payload = NegotiationPayload::new(0, 1, 0) + .with_tlv(1, vec![0xAA, 0xBB, 0xCC]); + let mut encoded = payload.encode(); + + // Truncate the TLV value (remove last byte) + encoded.pop(); + assert!(NegotiationPayload::decode(&encoded).is_err()); + + // Truncate to just partial TLV header (only 2 of 4 header bytes) + let mut partial = NegotiationPayload::new(0, 1, 0).encode(); + partial.extend_from_slice(&[0x01, 0x00]); // Only field_num, no length + assert!(NegotiationPayload::decode(&partial).is_err()); + } +} diff --git a/src/transport/tcp/mod.rs b/src/transport/tcp/mod.rs index 9a325cf..659d94b 100644 --- a/src/transport/tcp/mod.rs +++ b/src/transport/tcp/mod.rs @@ -1153,7 +1153,7 @@ mod tests { let payload_len = 4u16; let total = 4 + 12 + payload_len as usize + 16; let mut frame = vec![0u8; total]; - frame[0] = 0x00; // ver=0, phase=0 (established) + frame[0] = 0x10; // ver=1, phase=0 (established) frame[1] = 0x00; // flags frame[2..4].copy_from_slice(&payload_len.to_le_bytes()); // Fill the rest with a recognizable pattern @@ -1194,10 +1194,10 @@ mod tests { let addr2 = t2.local_addr().unwrap(); // Build valid FMP msg1 frame (114 bytes) - let mut msg1_frame = vec![0xAA; 114]; - msg1_frame[0] = 0x01; // phase=msg1 + let mut msg1_frame = vec![0xAA; 41]; + msg1_frame[0] = 0x11; // ver=1, phase=msg1 msg1_frame[1] = 0x00; - msg1_frame[2..4].copy_from_slice(&110u16.to_le_bytes()); // payload_len = 110 + msg1_frame[2..4].copy_from_slice(&37u16.to_le_bytes()); // payload_len = 37 // Send from t1 to t2 t1.send_async(&TransportAddr::from_string(&addr2.to_string()), &msg1_frame) @@ -1211,10 +1211,10 @@ mod tests { assert_eq!(packet.data, msg1_frame); // Build valid FMP msg2 frame (69 bytes) - let mut msg2_frame = vec![0xBB; 69]; - msg2_frame[0] = 0x02; // phase=msg2 + let mut msg2_frame = vec![0xBB; 118]; + msg2_frame[0] = 0x12; // ver=1, phase=msg2 msg2_frame[1] = 0x00; - msg2_frame[2..4].copy_from_slice(&65u16.to_le_bytes()); // payload_len = 65 + msg2_frame[2..4].copy_from_slice(&114u16.to_le_bytes()); // payload_len = 114 // Send from t2 to t1 t2.send_async(&TransportAddr::from_string(&addr1.to_string()), &msg2_frame) @@ -1270,10 +1270,10 @@ mod tests { let remote = TransportAddr::from_string(&addr2.to_string()); // Build valid msg1 frame to establish connection - let mut msg1 = vec![0xAA; 114]; - msg1[0] = 0x01; + let mut msg1 = vec![0xAA; 41]; + msg1[0] = 0x11; msg1[1] = 0x00; - msg1[2..4].copy_from_slice(&110u16.to_le_bytes()); + msg1[2..4].copy_from_slice(&37u16.to_le_bytes()); t1.send_async(&remote, &msg1).await.unwrap(); @@ -1371,10 +1371,10 @@ mod tests { let remote = TransportAddr::from_string(&addr2.to_string()); // Build valid msg1 frame - let mut msg1 = vec![0xAA; 114]; - msg1[0] = 0x01; + let mut msg1 = vec![0xAA; 41]; + msg1[0] = 0x11; msg1[1] = 0x00; - msg1[2..4].copy_from_slice(&110u16.to_le_bytes()); + msg1[2..4].copy_from_slice(&37u16.to_le_bytes()); // First send establishes connection t1.send_async(&remote, &msg1).await.unwrap(); @@ -1424,10 +1424,10 @@ mod tests { assert_eq!(state, ConnectionState::Connected); // Now send should work (connection already established) - let mut msg1 = vec![0xAA; 114]; - msg1[0] = 0x01; + let mut msg1 = vec![0xAA; 41]; + msg1[0] = 0x11; msg1[1] = 0x00; - msg1[2..4].copy_from_slice(&110u16.to_le_bytes()); + msg1[2..4].copy_from_slice(&37u16.to_le_bytes()); t1.send_async(&remote, &msg1).await.unwrap(); @@ -1527,10 +1527,10 @@ mod tests { ); // Build valid FMP msg1 frame - let mut msg1 = vec![0xAA; 114]; - msg1[0] = 0x01; + let mut msg1 = vec![0xAA; 41]; + msg1[0] = 0x11; msg1[1] = 0x00; - msg1[2..4].copy_from_slice(&110u16.to_le_bytes()); + msg1[2..4].copy_from_slice(&37u16.to_le_bytes()); // Send using the pre-established connection t1.send_async(&remote, &msg1).await.unwrap(); @@ -1577,10 +1577,10 @@ mod tests { // Connect using IP string — build a valid FMP frame (114 bytes) let addr = TransportAddr::from_string(&format!("127.0.0.1:{}", port2)); - let mut frame = vec![0xAA; 114]; - frame[0] = 0x01; // ver=0, phase=1 + let mut frame = vec![0xAA; 41]; + frame[0] = 0x11; // ver=1, phase=1 frame[1] = 0x00; // flags - frame[2..4].copy_from_slice(&110u16.to_le_bytes()); // payload_len + frame[2..4].copy_from_slice(&37u16.to_le_bytes()); // payload_len t1.send_async(&addr, &frame).await.unwrap(); // Receive on t2 diff --git a/src/transport/tcp/stream.rs b/src/transport/tcp/stream.rs index 4323fdc..27e801d 100644 --- a/src/transport/tcp/stream.rs +++ b/src/transport/tcp/stream.rs @@ -12,6 +12,7 @@ use tokio::io::{AsyncRead, AsyncReadExt}; const PHASE_ESTABLISHED: u8 = 0x0; const PHASE_MSG1: u8 = 0x1; const PHASE_MSG2: u8 = 0x2; +const PHASE_MSG3: u8 = 0x3; /// Size of the FMP common prefix. const PREFIX_SIZE: usize = 4; @@ -36,6 +37,7 @@ pub enum StreamError { max_payload_len: u16, }, /// Handshake packet has unexpected payload_len for its phase. + /// For msg1, expected is exact; for msg2/msg3, expected is the minimum. HandshakeSizeMismatch { phase: u8, expected: u16, got: u16 }, /// I/O error (including EOF). Io(std::io::Error), @@ -80,17 +82,22 @@ impl From for StreamError { } } -/// Known wire sizes for handshake messages. -/// msg1: 4 (prefix) + 4 (sender_idx) + 106 (noise_msg1) = 114 bytes -/// msg2: 4 (prefix) + 4 (sender_idx) + 4 (receiver_idx) + 57 (noise_msg2) = 69 bytes -const MSG1_WIRE_SIZE: usize = 114; -const MSG2_WIRE_SIZE: usize = 69; +/// Known wire sizes for handshake messages (Noise XX). +/// msg1: 4 (prefix) + 4 (sender_idx) + 33 (noise_msg1) = 41 bytes (exact) +/// msg2: 4 (prefix) + 4 (sender_idx) + 4 (receiver_idx) + 106+ (noise_msg2) = 118+ bytes (minimum) +/// msg3: 4 (prefix) + 4 (sender_idx) + 4 (receiver_idx) + 73+ (noise_msg3) = 85+ bytes (minimum) +const MSG1_WIRE_SIZE: usize = 41; +const MSG2_MIN_WIRE_SIZE: usize = 118; +const MSG3_MIN_WIRE_SIZE: usize = 85; -/// Expected payload_len for msg1: sender_idx(4) + noise_msg1(106) = 110. +/// Expected payload_len for msg1: sender_idx(4) + noise_msg1(33) = 37. const MSG1_PAYLOAD_LEN: u16 = (MSG1_WIRE_SIZE - PREFIX_SIZE) as u16; -/// Expected payload_len for msg2: sender_idx(4) + receiver_idx(4) + noise_msg2(57) = 65. -const MSG2_PAYLOAD_LEN: u16 = (MSG2_WIRE_SIZE - PREFIX_SIZE) as u16; +/// Minimum payload_len for msg2: sender_idx(4) + receiver_idx(4) + noise_msg2(106) = 114. +const MSG2_MIN_PAYLOAD_LEN: u16 = (MSG2_MIN_WIRE_SIZE - PREFIX_SIZE) as u16; + +/// Minimum payload_len for msg3: sender_idx(4) + receiver_idx(4) + noise_msg3(73) = 81. +const MSG3_MIN_PAYLOAD_LEN: u16 = (MSG3_MIN_WIRE_SIZE - PREFIX_SIZE) as u16; /// Read one complete FMP packet from an async reader. /// @@ -120,7 +127,7 @@ pub async fn read_fmp_packet( let version = prefix[0] >> 4; let phase = prefix[0] & 0x0F; - if version != 0 { + if version != 1 { return Err(StreamError::UnknownVersion(version)); } @@ -155,10 +162,20 @@ pub async fn read_fmp_packet( payload_len as usize } PHASE_MSG2 => { - if payload_len != MSG2_PAYLOAD_LEN { + if payload_len < MSG2_MIN_PAYLOAD_LEN { return Err(StreamError::HandshakeSizeMismatch { phase, - expected: MSG2_PAYLOAD_LEN, + expected: MSG2_MIN_PAYLOAD_LEN, + got: payload_len, + }); + } + payload_len as usize + } + PHASE_MSG3 => { + if payload_len < MSG3_MIN_PAYLOAD_LEN { + return Err(StreamError::HandshakeSizeMismatch { + phase, + expected: MSG3_MIN_PAYLOAD_LEN, got: payload_len, }); } @@ -193,7 +210,7 @@ mod tests { let total = PREFIX_SIZE + ESTABLISHED_REMAINING_HEADER + payload_len as usize + AEAD_TAG_SIZE; let mut frame = vec![0u8; total]; - frame[0] = 0x00; // ver=0, phase=0 (established) + frame[0] = 0x10; // ver=1, phase=0 (established) frame[1] = 0x00; // flags frame[2..4].copy_from_slice(&payload_len.to_le_bytes()); // Fill remaining with pattern for verification @@ -203,21 +220,30 @@ mod tests { frame } - /// Build a msg1 frame (114 bytes total). + /// Build a msg1 frame (41 bytes total). fn build_msg1_frame() -> Vec { let mut frame = vec![0xAA; MSG1_WIRE_SIZE]; - frame[0] = 0x01; // ver=0, phase=1 + frame[0] = 0x11; // ver=1, phase=1 frame[1] = 0x00; // flags frame[2..4].copy_from_slice(&MSG1_PAYLOAD_LEN.to_le_bytes()); frame } - /// Build a msg2 frame (69 bytes total). + /// Build a msg2 frame (minimum 118 bytes). fn build_msg2_frame() -> Vec { - let mut frame = vec![0xBB; MSG2_WIRE_SIZE]; - frame[0] = 0x02; // ver=0, phase=2 + let mut frame = vec![0xBB; MSG2_MIN_WIRE_SIZE]; + frame[0] = 0x12; // ver=1, phase=2 frame[1] = 0x00; // flags - frame[2..4].copy_from_slice(&MSG2_PAYLOAD_LEN.to_le_bytes()); + frame[2..4].copy_from_slice(&MSG2_MIN_PAYLOAD_LEN.to_le_bytes()); + frame + } + + /// Build a msg3 frame (minimum 85 bytes). + fn build_msg3_frame() -> Vec { + let mut frame = vec![0xCC; MSG3_MIN_WIRE_SIZE]; + frame[0] = 0x13; // ver=1, phase=3 + frame[1] = 0x00; // flags + frame[2..4].copy_from_slice(&MSG3_MIN_PAYLOAD_LEN.to_le_bytes()); frame } @@ -250,7 +276,7 @@ mod tests { let mut cursor = Cursor::new(frame); let packet = read_fmp_packet(&mut cursor, 1400).await.unwrap(); - assert_eq!(packet.len(), MSG2_WIRE_SIZE); + assert_eq!(packet.len(), MSG2_MIN_WIRE_SIZE); assert_eq!(packet, expected); } @@ -272,24 +298,48 @@ mod tests { assert_eq!(p2, est); let p3 = read_fmp_packet(&mut cursor, 1400).await.unwrap(); - assert_eq!(p3.len(), MSG2_WIRE_SIZE); + assert_eq!(p3.len(), MSG2_MIN_WIRE_SIZE); + } + + #[tokio::test] + async fn test_read_msg3_frame() { + let frame = build_msg3_frame(); + let expected = frame.clone(); + + let mut cursor = Cursor::new(frame); + let packet = read_fmp_packet(&mut cursor, 1400).await.unwrap(); + assert_eq!(packet.len(), MSG3_MIN_WIRE_SIZE); + assert_eq!(packet, expected); } #[tokio::test] async fn test_unknown_version_error() { // TLS ClientHello starts with 0x16 (record type "Handshake"), - // which parses as FMP version=1, phase=6. + // which parses as FMP version=1, phase=6. But now FMP version IS 1, + // so test with version=0 (old protocol) instead. let mut frame = vec![0u8; 100]; - frame[0] = 0x16; + frame[0] = 0x00; // version 0 (old), phase 0 let mut cursor = Cursor::new(frame); let err = read_fmp_packet(&mut cursor, 1400).await.unwrap_err(); - assert!(matches!(err, StreamError::UnknownVersion(1))); + assert!(matches!(err, StreamError::UnknownVersion(0))); + } + + #[tokio::test] + async fn test_unknown_version_tls() { + // TLS ClientHello: 0x16 → version=1, phase=6. + // Version 1 is now valid, so this triggers UnknownPhase, not UnknownVersion. + let mut frame = vec![0u8; 100]; + frame[0] = 0x16; + frame[2..4].copy_from_slice(&10u16.to_le_bytes()); + let mut cursor = Cursor::new(frame); + let err = read_fmp_packet(&mut cursor, 1400).await.unwrap_err(); + assert!(matches!(err, StreamError::UnknownPhase(0x6))); } #[tokio::test] async fn test_unknown_phase_error() { let mut frame = vec![0u8; 100]; - frame[0] = 0x05; // unknown phase + frame[0] = 0x15; // ver=1, unknown phase 5 frame[2..4].copy_from_slice(&10u16.to_le_bytes()); let mut cursor = Cursor::new(frame); @@ -302,7 +352,7 @@ mod tests { // mtu=100, max_payload_len = 100 - 32 = 68 let payload_len = 100u16; // exceeds max of 68 let mut prefix = [0u8; 4]; - prefix[0] = 0x00; // established + prefix[0] = 0x10; // ver=1, established prefix[2..4].copy_from_slice(&payload_len.to_le_bytes()); // Provide enough bytes for the reader to read prefix @@ -317,8 +367,8 @@ mod tests { #[tokio::test] async fn test_handshake_size_mismatch_msg1() { let mut frame = vec![0u8; 200]; - frame[0] = 0x01; // msg1 - // Wrong payload_len (should be 110) + frame[0] = 0x11; // ver=1, msg1 + // Wrong payload_len (should be 37) frame[2..4].copy_from_slice(&50u16.to_le_bytes()); let mut cursor = Cursor::new(frame); @@ -332,8 +382,8 @@ mod tests { #[tokio::test] async fn test_handshake_size_mismatch_msg2() { let mut frame = vec![0u8; 200]; - frame[0] = 0x02; // msg2 - // Wrong payload_len (should be 65) + frame[0] = 0x12; // ver=1, msg2 + // Wrong payload_len (should be >= 114) frame[2..4].copy_from_slice(&50u16.to_le_bytes()); let mut cursor = Cursor::new(frame); @@ -344,6 +394,18 @@ mod tests { )); } + #[tokio::test] + async fn test_handshake_size_mismatch_msg3() { + let mut frame = vec![0u8; 200]; + frame[0] = 0x13; // ver=1, msg3 + // Wrong payload_len (should be >= 81) + frame[2..4].copy_from_slice(&50u16.to_le_bytes()); + + let mut cursor = Cursor::new(frame); + let err = read_fmp_packet(&mut cursor, 1400).await.unwrap_err(); + assert!(matches!(err, StreamError::HandshakeSizeMismatch { phase: 0x3, .. })); + } + #[tokio::test] async fn test_eof_on_prefix() { // Only 2 bytes available (need 4 for prefix) @@ -356,8 +418,8 @@ mod tests { #[tokio::test] async fn test_eof_on_body() { // Valid msg1 prefix but truncated body - let mut data = vec![0u8; 10]; // need 114 total - data[0] = 0x01; // msg1 + let mut data = vec![0u8; 10]; // need 41 total + data[0] = 0x11; // ver=1, msg1 data[2..4].copy_from_slice(&MSG1_PAYLOAD_LEN.to_le_bytes()); let mut cursor = Cursor::new(data); @@ -394,7 +456,7 @@ mod tests { // mtu=1400, max_payload_len = 1368, try 1369 let over = 1400u16 - 32 + 1; let mut prefix = [0u8; 4]; - prefix[0] = 0x00; // established + prefix[0] = 0x10; // ver=1, established prefix[2..4].copy_from_slice(&over.to_le_bytes()); let mut data = prefix.to_vec(); diff --git a/src/transport/tor/mod.rs b/src/transport/tor/mod.rs index f475850..4fbf0c3 100644 --- a/src/transport/tor/mod.rs +++ b/src/transport/tor/mod.rs @@ -1557,14 +1557,14 @@ mod tests { use mock_socks5::MockSocks5Server; /// msg1 wire size: 4 prefix + 4 sender_idx + 106 noise_msg1 = 114 bytes. - const MSG1_WIRE_SIZE: usize = 114; - /// msg1 payload_len: sender_idx(4) + noise_msg1(106) = 110. + const MSG1_WIRE_SIZE: usize = 41; + /// msg1 payload_len: sender_idx(4) + noise_msg1(33) = 37. const MSG1_PAYLOAD_LEN: u16 = (MSG1_WIRE_SIZE - 4) as u16; - /// Build a msg1 frame (114 bytes) for testing. + /// Build a msg1 frame (41 bytes) for testing. fn build_msg1_frame() -> Vec { let mut frame = vec![0xAA; MSG1_WIRE_SIZE]; - frame[0] = 0x01; // ver=0, phase=1 + frame[0] = 0x11; // ver=1, phase=1 frame[1] = 0x00; // flags frame[2..4].copy_from_slice(&MSG1_PAYLOAD_LEN.to_le_bytes()); frame