mirror of
https://github.com/jmcorgan/fips.git
synced 2026-08-09 08:14:42 +00:00
Fix secp256k1 parity in Noise IK, add disconnect protocol, cross-connection handling, timeout cleanup
Noise IK parity fix: - Pre-message hash normalizes responder static key to even parity (0x02) so initiator and responder hash chains match regardless of actual parity - ECDH uses shared_secret_point() + SHA-256(x-only) instead of SharedSecret::new() which includes a parity-dependent version byte - Fixes handshake failure for ~50% of keys when initiator has only npub Graceful disconnect protocol (link message 0x50): - DisconnectReason enum with 8 reason codes - Disconnect struct with encode/decode - send_encrypted_link_message() reusable helper - handle_disconnect() with immediate peer removal - send_disconnect_to_all_peers() called during Node::stop() Cross-connection fix in handle_msg1(): - addr_to_link check now distinguishes inbound duplicates (reject) from outbound links (cross-connection, allow and resolve via tie-breaker) - remove_link() only clears addr_to_link if entry maps to same link_id - Link cleanup and addr_to_link restoration in cross-connection branches Handshake timeout cleanup: - RX loop uses tokio::select! with 1-second interval tick - check_timeouts() scans for stale (>30s) and failed connections - cleanup_stale_connection() removes all associated state Tests: 279 passing (4 new: cross-connection, stale cleanup, failed cleanup, odd-parity handshake)
This commit is contained in:
Johnathan Corgan
+199
-20
@@ -1,6 +1,7 @@
|
||||
//! RX event loop and message handlers.
|
||||
|
||||
use super::*;
|
||||
use crate::rate_limit::HANDSHAKE_TIMEOUT_SECS;
|
||||
|
||||
impl Node {
|
||||
// === RX Event Loop ===
|
||||
@@ -13,16 +14,32 @@ impl Node {
|
||||
/// - 0x01: Handshake message 1 (initiator -> responder)
|
||||
/// - 0x02: Handshake message 2 (responder -> initiator)
|
||||
///
|
||||
/// Also runs a periodic tick (1s) to clean up stale handshake connections
|
||||
/// that never received a response. This prevents resource leaks when peers
|
||||
/// are unreachable.
|
||||
///
|
||||
/// 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 tick = tokio::time::interval(Duration::from_secs(1));
|
||||
|
||||
info!("RX event loop started");
|
||||
|
||||
while let Some(packet) = packet_rx.recv().await {
|
||||
self.process_packet(packet).await;
|
||||
loop {
|
||||
tokio::select! {
|
||||
packet = packet_rx.recv() => {
|
||||
match packet {
|
||||
Some(p) => self.process_packet(p).await,
|
||||
None => break, // channel closed
|
||||
}
|
||||
}
|
||||
_ = tick.tick() => {
|
||||
self.check_timeouts();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!("RX event loop stopped (channel closed)");
|
||||
@@ -159,16 +176,32 @@ impl Node {
|
||||
}
|
||||
};
|
||||
|
||||
// Check for existing connection from this address
|
||||
// Check for existing connection from this address.
|
||||
//
|
||||
// If we already have an *inbound* link from this address, drop the msg1
|
||||
// (duplicate or replay). But 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 to proceed — promote_connection() will resolve via tie-breaker.
|
||||
let addr_key = (packet.transport_id, packet.remote_addr.clone());
|
||||
if self.addr_to_link.contains_key(&addr_key) {
|
||||
self.msg1_rate_limiter.complete_handshake();
|
||||
debug!(
|
||||
transport_id = %packet.transport_id,
|
||||
remote_addr = %packet.remote_addr,
|
||||
"Already have connection from this address"
|
||||
);
|
||||
return;
|
||||
if let Some(&existing_link_id) = self.addr_to_link.get(&addr_key) {
|
||||
if let Some(link) = self.links.get(&existing_link_id) {
|
||||
if link.direction() == LinkDirection::Inbound {
|
||||
self.msg1_rate_limiter.complete_handshake();
|
||||
debug!(
|
||||
transport_id = %packet.transport_id,
|
||||
remote_addr = %packet.remote_addr,
|
||||
"Already have inbound connection from this address"
|
||||
);
|
||||
return;
|
||||
}
|
||||
// Outbound link to this address — cross-connection, allow msg1
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// === CRYPTO COST PAID HERE ===
|
||||
@@ -279,13 +312,22 @@ impl Node {
|
||||
);
|
||||
}
|
||||
PromotionResult::CrossConnectionWon { loser_link_id, node_addr } => {
|
||||
// Clean up the losing connection's link
|
||||
self.remove_link(&loser_link_id);
|
||||
info!(
|
||||
node_addr = %node_addr,
|
||||
loser_link_id = %loser_link_id,
|
||||
"Inbound cross-connection won"
|
||||
"Inbound cross-connection won, loser link cleaned up"
|
||||
);
|
||||
}
|
||||
PromotionResult::CrossConnectionLost { winner_link_id } => {
|
||||
// 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,
|
||||
);
|
||||
info!(
|
||||
winner_link_id = %winner_link_id,
|
||||
"Inbound cross-connection lost, keeping existing"
|
||||
@@ -300,9 +342,7 @@ impl Node {
|
||||
"Failed to promote inbound connection"
|
||||
);
|
||||
// Clean up on promotion failure
|
||||
self.links.remove(&link_id);
|
||||
self.addr_to_link
|
||||
.remove(&(packet.transport_id, packet.remote_addr));
|
||||
self.remove_link(&link_id);
|
||||
let _ = self.index_allocator.free(our_index);
|
||||
}
|
||||
}
|
||||
@@ -392,16 +432,30 @@ impl Node {
|
||||
);
|
||||
}
|
||||
PromotionResult::CrossConnectionWon { loser_link_id, node_addr } => {
|
||||
// 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,
|
||||
);
|
||||
info!(
|
||||
node_addr = %node_addr,
|
||||
loser_link_id = %loser_link_id,
|
||||
"Cross-connection won"
|
||||
"Outbound cross-connection won, loser link cleaned up"
|
||||
);
|
||||
}
|
||||
PromotionResult::CrossConnectionLost { winner_link_id } => {
|
||||
// This connection lost — clean up its link
|
||||
self.remove_link(&link_id);
|
||||
// Ensure addr_to_link points to the winner's link
|
||||
self.addr_to_link.insert(
|
||||
(packet.transport_id, packet.remote_addr.clone()),
|
||||
winner_link_id,
|
||||
);
|
||||
info!(
|
||||
winner_link_id = %winner_link_id,
|
||||
"Cross-connection lost"
|
||||
"Outbound cross-connection lost, keeping existing"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -575,15 +629,15 @@ impl Node {
|
||||
/// Dispatch a decrypted link message to the appropriate handler.
|
||||
///
|
||||
/// Link messages are protocol messages exchanged between authenticated peers.
|
||||
async fn dispatch_link_message(&mut self, _from: &NodeAddr, plaintext: &[u8]) {
|
||||
async fn dispatch_link_message(&mut self, from: &NodeAddr, plaintext: &[u8]) {
|
||||
if plaintext.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let msg_type = plaintext[0];
|
||||
let _payload = &plaintext[1..];
|
||||
let payload = &plaintext[1..];
|
||||
|
||||
// TODO: Implement link message handlers
|
||||
// TODO: Implement remaining link message handlers
|
||||
match msg_type {
|
||||
0x10 => {
|
||||
// TreeAnnounce
|
||||
@@ -605,9 +659,134 @@ impl Node {
|
||||
// SessionDatagram
|
||||
debug!("Received SessionDatagram (not yet implemented)");
|
||||
}
|
||||
0x50 => {
|
||||
// Disconnect
|
||||
self.handle_disconnect(from, payload);
|
||||
}
|
||||
_ => {
|
||||
debug!(msg_type = msg_type, "Unknown link message type");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a Disconnect notification from a peer.
|
||||
///
|
||||
/// The peer is signaling an orderly departure. We immediately remove
|
||||
/// them from all state rather than waiting for timeout detection.
|
||||
fn handle_disconnect(&mut self, from: &NodeAddr, payload: &[u8]) {
|
||||
let disconnect = match crate::protocol::Disconnect::decode(payload) {
|
||||
Ok(msg) => msg,
|
||||
Err(e) => {
|
||||
debug!(from = %from, error = %e, "Malformed disconnect message");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
node_addr = %from,
|
||||
reason = %disconnect.reason,
|
||||
"Peer sent disconnect notification"
|
||||
);
|
||||
|
||||
self.remove_active_peer(from);
|
||||
}
|
||||
|
||||
/// Remove an active peer and clean up all associated state.
|
||||
///
|
||||
/// Frees session index, removes link and address mappings. Used for
|
||||
/// both graceful disconnect and timeout-based eviction.
|
||||
pub(super) fn remove_active_peer(&mut self, node_addr: &NodeAddr) {
|
||||
let peer = match self.peers.remove(node_addr) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
debug!(node_addr = %node_addr, "Peer already removed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let link_id = peer.link_id();
|
||||
|
||||
// Free session index
|
||||
if let (Some(tid), Some(idx)) = (peer.transport_id(), peer.our_index()) {
|
||||
self.peers_by_index.remove(&(tid, idx.as_u32()));
|
||||
let _ = self.index_allocator.free(idx);
|
||||
}
|
||||
|
||||
// Remove link and address mapping
|
||||
self.remove_link(&link_id);
|
||||
|
||||
info!(
|
||||
node_addr = %node_addr,
|
||||
link_id = %link_id,
|
||||
"Peer removed and state cleaned up"
|
||||
);
|
||||
}
|
||||
|
||||
// === Timeout Management ===
|
||||
|
||||
/// Check for timed-out handshake connections and clean them up.
|
||||
///
|
||||
/// Called periodically by the RX event loop. Removes connections that have
|
||||
/// been idle longer than HANDSHAKE_TIMEOUT_SECS or are in Failed state.
|
||||
pub(super) fn check_timeouts(&mut self) {
|
||||
if self.connections.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
let timeout_ms = HANDSHAKE_TIMEOUT_SECS * 1000;
|
||||
|
||||
let stale: Vec<LinkId> = self.connections.iter()
|
||||
.filter(|(_, conn)| conn.is_timed_out(now_ms, timeout_ms) || conn.is_failed())
|
||||
.map(|(link_id, _)| *link_id)
|
||||
.collect();
|
||||
|
||||
for link_id in stale {
|
||||
self.cleanup_stale_connection(link_id, now_ms);
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a stale or failed handshake connection and all associated state.
|
||||
///
|
||||
/// Frees the session index, removes pending_outbound entry, and cleans up
|
||||
/// the link and address mapping.
|
||||
fn cleanup_stale_connection(&mut self, link_id: LinkId, now_ms: u64) {
|
||||
let conn = match self.connections.remove(&link_id) {
|
||||
Some(c) => c,
|
||||
None => return,
|
||||
};
|
||||
|
||||
let direction = conn.direction();
|
||||
let idle_ms = conn.idle_time(now_ms);
|
||||
let is_failed = conn.is_failed();
|
||||
|
||||
// Free session index and pending_outbound if allocated
|
||||
if let Some(idx) = conn.our_index() {
|
||||
if let Some(tid) = conn.transport_id() {
|
||||
self.pending_outbound.remove(&(tid, idx.as_u32()));
|
||||
}
|
||||
let _ = self.index_allocator.free(idx);
|
||||
}
|
||||
|
||||
// Remove link and addr_to_link
|
||||
self.remove_link(&link_id);
|
||||
|
||||
if is_failed {
|
||||
info!(
|
||||
link_id = %link_id,
|
||||
direction = %direction,
|
||||
"Failed handshake connection cleaned up"
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
link_id = %link_id,
|
||||
direction = %direction,
|
||||
idle_secs = idle_ms / 1000,
|
||||
"Stale handshake connection timed out"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+44
-1
@@ -1,6 +1,7 @@
|
||||
//! Node lifecycle management: start, stop, and peer connection initiation.
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::{Disconnect, DisconnectReason};
|
||||
|
||||
impl Node {
|
||||
/// Initiate connections to configured static peers.
|
||||
@@ -316,7 +317,10 @@ impl Node {
|
||||
self.state = NodeState::Stopping;
|
||||
info!(state = %self.state, "Node stopping");
|
||||
|
||||
// Shutdown transports first (they're packet producers)
|
||||
// Send disconnect notifications to all active peers before closing transports
|
||||
self.send_disconnect_to_all_peers(DisconnectReason::Shutdown).await;
|
||||
|
||||
// Shutdown transports (they're packet producers)
|
||||
let transport_ids: Vec<_> = self.transports.keys().cloned().collect();
|
||||
for transport_id in transport_ids {
|
||||
if let Some(mut handle) = self.transports.remove(&transport_id) {
|
||||
@@ -368,4 +372,43 @@ impl Node {
|
||||
info!(state = %self.state, "Node stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send disconnect notifications to all active peers.
|
||||
///
|
||||
/// Best-effort: send failures are logged and ignored since the transport
|
||||
/// may already be degraded. This runs before transports are shut down.
|
||||
async fn send_disconnect_to_all_peers(&mut self, reason: DisconnectReason) {
|
||||
let disconnect = Disconnect::new(reason);
|
||||
let plaintext = disconnect.encode();
|
||||
|
||||
// Collect node_addrs to avoid borrow conflict with send helper
|
||||
let peer_addrs: Vec<NodeAddr> = self.peers.iter()
|
||||
.filter(|(_, peer)| peer.can_send() && peer.has_session())
|
||||
.map(|(addr, _)| *addr)
|
||||
.collect();
|
||||
|
||||
if peer_addrs.is_empty() {
|
||||
debug!(
|
||||
total_peers = self.peers.len(),
|
||||
"No sendable peers for disconnect notification"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let mut sent = 0usize;
|
||||
for node_addr in &peer_addrs {
|
||||
match self.send_encrypted_link_message(node_addr, &plaintext).await {
|
||||
Ok(()) => sent += 1,
|
||||
Err(e) => {
|
||||
debug!(
|
||||
node_addr = %node_addr,
|
||||
error = %e,
|
||||
"Failed to send disconnect (transport may be down)"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!(sent, total = peer_addrs.len(), reason = %reason, "Sent disconnect notifications");
|
||||
}
|
||||
}
|
||||
|
||||
+74
-3
@@ -24,7 +24,7 @@ use crate::transport::udp::UdpTransport;
|
||||
use crate::tree::TreeState;
|
||||
use crate::tun::{run_tun_reader, shutdown_tun_interface, TunDevice, TunError, TunState, TunTx};
|
||||
use crate::wire::{
|
||||
build_msg1, build_msg2, EncryptedHeader, Msg1Header, Msg2Header,
|
||||
build_encrypted, build_msg1, build_msg2, EncryptedHeader, Msg1Header, Msg2Header,
|
||||
DISCRIMINATOR_ENCRYPTED, DISCRIMINATOR_MSG1, DISCRIMINATOR_MSG2,
|
||||
};
|
||||
use crate::{Config, ConfigError, Identity, IdentityError, NodeAddr, PeerIdentity};
|
||||
@@ -89,6 +89,9 @@ pub enum NodeError {
|
||||
#[error("promotion failed for link {link_id}: {reason}")]
|
||||
PromotionFailed { link_id: LinkId, reason: String },
|
||||
|
||||
#[error("send failed to {node_addr}: {reason}")]
|
||||
SendFailed { node_addr: NodeAddr, reason: String },
|
||||
|
||||
#[error("config error: {0}")]
|
||||
Config(#[from] ConfigError),
|
||||
|
||||
@@ -594,11 +597,17 @@ impl Node {
|
||||
}
|
||||
|
||||
/// Remove a link.
|
||||
///
|
||||
/// Only removes the addr_to_link reverse lookup if it still points to this
|
||||
/// link. In cross-connection scenarios, a newer link may have replaced the
|
||||
/// entry for the same address.
|
||||
pub fn remove_link(&mut self, link_id: &LinkId) -> Option<Link> {
|
||||
if let Some(link) = self.links.remove(link_id) {
|
||||
// Clean up reverse lookup
|
||||
// Clean up reverse lookup only if it still maps to this link
|
||||
let key = (link.transport_id(), link.remote_addr().clone());
|
||||
self.addr_to_link.remove(&key);
|
||||
if self.addr_to_link.get(&key) == Some(link_id) {
|
||||
self.addr_to_link.remove(&key);
|
||||
}
|
||||
Some(link)
|
||||
} else {
|
||||
None
|
||||
@@ -708,6 +717,68 @@ impl Node {
|
||||
pub fn tun_tx(&self) -> Option<&TunTx> {
|
||||
self.tun_tx.as_ref()
|
||||
}
|
||||
|
||||
// === Sending ===
|
||||
|
||||
/// Encrypt and send a link-layer message to an authenticated peer.
|
||||
///
|
||||
/// The plaintext should include the message type byte followed by the
|
||||
/// message-specific payload (e.g., `[0x50, reason]` for Disconnect).
|
||||
///
|
||||
/// This is the standard path for sending any link-layer control message
|
||||
/// to a peer over their encrypted Noise session.
|
||||
pub(super) async fn send_encrypted_link_message(
|
||||
&mut self,
|
||||
node_addr: &NodeAddr,
|
||||
plaintext: &[u8],
|
||||
) -> Result<(), NodeError> {
|
||||
let peer = self.peers.get_mut(node_addr)
|
||||
.ok_or(NodeError::PeerNotFound(*node_addr))?;
|
||||
|
||||
let their_index = peer.their_index().ok_or_else(|| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: "no their_index".into(),
|
||||
})?;
|
||||
let transport_id = peer.transport_id().ok_or_else(|| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: "no transport_id".into(),
|
||||
})?;
|
||||
let remote_addr = peer.current_addr().cloned().ok_or_else(|| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: "no current_addr".into(),
|
||||
})?;
|
||||
|
||||
let session = peer.noise_session_mut().ok_or_else(|| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: "no noise session".into(),
|
||||
})?;
|
||||
|
||||
// Get counter before encrypt (encrypt increments it)
|
||||
let counter = session.current_send_counter();
|
||||
let ciphertext = session.encrypt(plaintext).map_err(|e| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: format!("encryption failed: {}", e),
|
||||
})?;
|
||||
|
||||
let wire_packet = build_encrypted(their_index, counter, &ciphertext);
|
||||
|
||||
// Re-borrow peer for stats update after sending
|
||||
let transport = self.transports.get(&transport_id)
|
||||
.ok_or(NodeError::TransportNotFound(transport_id))?;
|
||||
|
||||
let bytes_sent = transport.send(&remote_addr, &wire_packet).await
|
||||
.map_err(|e| NodeError::SendFailed {
|
||||
node_addr: *node_addr,
|
||||
reason: format!("transport send: {}", e),
|
||||
})?;
|
||||
|
||||
// Update send statistics
|
||||
if let Some(peer) = self.peers.get_mut(node_addr) {
|
||||
peer.link_stats_mut().record_sent(bytes_sent);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for Node {
|
||||
|
||||
@@ -896,3 +896,269 @@ async fn test_run_rx_loop_handshake() {
|
||||
t.stop().await.ok();
|
||||
}
|
||||
}
|
||||
|
||||
/// Integration test: simultaneous cross-connection (both nodes initiate).
|
||||
///
|
||||
/// Simulates the live scenario where both nodes have auto_connect to each other.
|
||||
/// Both send msg1 simultaneously, creating a cross-connection that must be
|
||||
/// resolved by the tie-breaker rule. Exercises the addr_to_link fix that allows
|
||||
/// inbound msg1 when an outbound link to the same address already exists.
|
||||
#[tokio::test]
|
||||
async fn test_cross_connection_both_initiate() {
|
||||
use crate::config::UdpConfig;
|
||||
use crate::transport::udp::UdpTransport;
|
||||
use crate::wire::build_msg1;
|
||||
use tokio::time::{timeout, Duration};
|
||||
|
||||
// === Setup: Two nodes with UDP transports on localhost ===
|
||||
|
||||
let mut node_a = make_node();
|
||||
let mut node_b = make_node();
|
||||
|
||||
let transport_id_a = TransportId::new(1);
|
||||
let transport_id_b = TransportId::new(1);
|
||||
|
||||
let udp_config = UdpConfig {
|
||||
bind_addr: Some("127.0.0.1:0".to_string()),
|
||||
mtu: Some(1280),
|
||||
};
|
||||
|
||||
let (packet_tx_a, mut packet_rx_a) = packet_channel(64);
|
||||
let (packet_tx_b, mut packet_rx_b) = packet_channel(64);
|
||||
|
||||
let mut transport_a =
|
||||
UdpTransport::new(transport_id_a, None, udp_config.clone(), packet_tx_a);
|
||||
let mut transport_b =
|
||||
UdpTransport::new(transport_id_b, None, udp_config, packet_tx_b);
|
||||
|
||||
transport_a.start_async().await.unwrap();
|
||||
transport_b.start_async().await.unwrap();
|
||||
|
||||
let addr_a = transport_a.local_addr().unwrap();
|
||||
let addr_b = transport_b.local_addr().unwrap();
|
||||
let remote_addr_b = TransportAddr::from_string(&addr_b.to_string());
|
||||
let remote_addr_a = TransportAddr::from_string(&addr_a.to_string());
|
||||
|
||||
node_a
|
||||
.transports
|
||||
.insert(transport_id_a, TransportHandle::Udp(transport_a));
|
||||
node_b
|
||||
.transports
|
||||
.insert(transport_id_b, TransportHandle::Udp(transport_b));
|
||||
|
||||
// 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_node_addr = *peer_b_identity.node_addr();
|
||||
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) ===
|
||||
|
||||
// Node A initiates to Node B
|
||||
let link_id_a_out = node_a.allocate_link_id();
|
||||
let mut conn_a = PeerConnection::outbound(link_id_a_out, peer_b_identity.clone(), 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, 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());
|
||||
|
||||
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),
|
||||
);
|
||||
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.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 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.clone(), 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, 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());
|
||||
|
||||
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),
|
||||
);
|
||||
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.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);
|
||||
|
||||
// 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");
|
||||
|
||||
let transport = node_b.transports.get(&transport_id_b).unwrap();
|
||||
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.
|
||||
|
||||
// B receives A's msg1
|
||||
let packet_at_b = timeout(Duration::from_secs(1), packet_rx_b.recv())
|
||||
.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");
|
||||
|
||||
// A receives B's msg1
|
||||
let packet_at_a = timeout(Duration::from_secs(1), packet_rx_a.recv())
|
||||
.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");
|
||||
|
||||
// === 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.
|
||||
|
||||
// A receives B's msg2 (response to A's original msg1)
|
||||
let msg2_at_a = timeout(Duration::from_secs(1), packet_rx_a.recv())
|
||||
.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)
|
||||
let msg2_at_b = timeout(Duration::from_secs(1), packet_rx_b.recv())
|
||||
.await.expect("Timeout waiting for msg2 at B").expect("Channel closed");
|
||||
node_b.handle_msg2(msg2_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");
|
||||
|
||||
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");
|
||||
assert!(peer_b_on_a.can_send(), "Peer B on A should be sendable");
|
||||
assert!(peer_a_on_b.can_send(), "Peer A on B should be sendable");
|
||||
|
||||
// Clean up transports
|
||||
for (_, t) in node_a.transports.iter_mut() {
|
||||
t.stop().await.ok();
|
||||
}
|
||||
for (_, t) in node_b.transports.iter_mut() {
|
||||
t.stop().await.ok();
|
||||
}
|
||||
}
|
||||
|
||||
/// Test that stale handshake connections are cleaned up by check_timeouts().
|
||||
///
|
||||
/// Simulates the scenario where a node initiates a handshake to a peer that
|
||||
/// isn't running. The outbound connection should be cleaned up after the
|
||||
/// handshake timeout expires.
|
||||
#[tokio::test]
|
||||
async fn test_stale_connection_cleanup() {
|
||||
let mut node = make_node();
|
||||
let transport_id = TransportId::new(1);
|
||||
|
||||
let peer_identity = make_peer_identity();
|
||||
let remote_addr = TransportAddr::from_string("10.0.0.2:4000");
|
||||
|
||||
// Create outbound connection with a timestamp far in the past
|
||||
let past_time_ms = 1000; // A very early timestamp
|
||||
let link_id = node.allocate_link_id();
|
||||
let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), past_time_ms);
|
||||
|
||||
// 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, 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),
|
||||
);
|
||||
node.links.insert(link_id, link);
|
||||
node.addr_to_link.insert((transport_id, remote_addr.clone()), link_id);
|
||||
node.connections.insert(link_id, conn);
|
||||
node.pending_outbound.insert((transport_id, our_index.as_u32()), link_id);
|
||||
|
||||
// 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_eq!(node.index_allocator.count(), 1);
|
||||
|
||||
// Connection was created at time 1000ms. check_timeouts uses SystemTime::now(),
|
||||
// which is far beyond the 30s timeout. The connection should be cleaned up.
|
||||
node.check_timeouts();
|
||||
|
||||
// Verify everything was cleaned up
|
||||
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.get(&(transport_id, remote_addr)).is_none(),
|
||||
"addr_to_link should be cleaned up");
|
||||
}
|
||||
|
||||
/// Test that failed connections are cleaned up by check_timeouts().
|
||||
#[tokio::test]
|
||||
async fn test_failed_connection_cleanup() {
|
||||
let mut node = make_node();
|
||||
let transport_id = TransportId::new(1);
|
||||
|
||||
let peer_identity = make_peer_identity();
|
||||
let remote_addr = TransportAddr::from_string("10.0.0.2:4000");
|
||||
|
||||
// Create a connection and mark it failed (simulating a send failure)
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
let link_id = node.allocate_link_id();
|
||||
let mut conn = PeerConnection::outbound(link_id, peer_identity.clone(), now_ms);
|
||||
|
||||
let our_index = node.index_allocator.allocate().unwrap();
|
||||
let our_keypair = node.identity.keypair();
|
||||
let _noise_msg1 = conn.start_handshake(our_keypair, 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),
|
||||
);
|
||||
node.links.insert(link_id, link);
|
||||
node.addr_to_link.insert((transport_id, remote_addr.clone()), link_id);
|
||||
node.connections.insert(link_id, conn);
|
||||
node.pending_outbound.insert((transport_id, our_index.as_u32()), link_id);
|
||||
|
||||
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.link_count(), 0, "Failed link should be removed");
|
||||
assert_eq!(node.index_allocator.count(), 0, "Session index should be freed");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user