diff --git a/src/config/node.rs b/src/config/node.rs index c809923..f7a765d 100644 --- a/src/config/node.rs +++ b/src/config/node.rs @@ -231,6 +231,11 @@ pub struct SessionConfig { /// Established sessions with no activity for this duration are removed. #[serde(default = "SessionConfig::default_idle_timeout_secs")] pub idle_timeout_secs: u64, + /// Number of initial DataPackets per session that include COORDS_PRESENT + /// for transit cache warmup (`node.session.coords_warmup_packets`). + /// Also used as the reset count on CoordsRequired receipt. + #[serde(default = "SessionConfig::default_coords_warmup_packets")] + pub coords_warmup_packets: u8, } impl Default for SessionConfig { @@ -240,6 +245,7 @@ impl Default for SessionConfig { pending_packets_per_dest: 16, pending_max_destinations: 256, idle_timeout_secs: 90, + coords_warmup_packets: 5, } } } @@ -249,6 +255,7 @@ impl SessionConfig { fn default_pending_packets_per_dest() -> usize { 16 } fn default_pending_max_destinations() -> usize { 256 } fn default_idle_timeout_secs() -> u64 { 90 } + fn default_coords_warmup_packets() -> u8 { 5 } } /// Internal buffers (`node.buffers.*`). diff --git a/src/node/handlers/session.rs b/src/node/handlers/session.rs index c6e84a0..d614cfe 100644 --- a/src/node/handlers/session.rs +++ b/src/node/handlers/session.rs @@ -211,6 +211,7 @@ impl Node { }; entry.set_state(EndToEndState::Established(session)); + entry.set_coords_warmup_remaining(self.config.node.session.coords_warmup_packets); entry.touch(Self::now_ms()); self.sessions.insert(*src_addr, entry); @@ -278,6 +279,7 @@ impl Node { } }; entry.set_state(EndToEndState::Established(noise_session)); + entry.set_coords_warmup_remaining(self.config.node.session.coords_warmup_packets); debug!(src = %src_addr, "Session established (responder, on first data)"); } @@ -341,6 +343,18 @@ impl Node { ); self.maybe_initiate_lookup(&msg.dest_addr).await; + + // Reset coords warmup counter so the next N packets include + // COORDS_PRESENT, re-warming transit caches along the path. + if let Some(entry) = self.sessions.get_mut(&msg.dest_addr) { + let n = self.config.node.session.coords_warmup_packets; + entry.set_coords_warmup_remaining(n); + debug!( + dest = %msg.dest_addr, + warmup_packets = n, + "Reset coords warmup counter after CoordsRequired" + ); + } } /// Handle a PathBroken error signal from a transit router. @@ -449,8 +463,20 @@ impl Node { reason: format!("session encrypt failed: {}", e), })?; - // Build DataPacket and wrap in SessionDatagram - let data_packet = DataPacket::new(ciphertext); + // Check warmup counter and decrement (while entry is still borrowed) + let include_coords = entry.coords_warmup_remaining() > 0; + if include_coords { + entry.set_coords_warmup_remaining(entry.coords_warmup_remaining() - 1); + } + + // Build DataPacket, conditionally with coordinates + let mut data_packet = DataPacket::new(ciphertext); + if include_coords { + let my_coords = self.tree_state.my_coords().clone(); + let dest_coords = self.get_dest_coords(dest_addr); + data_packet = data_packet.with_coords(my_coords, dest_coords); + } + let my_addr = *self.node_addr(); let datagram = SessionDatagram::new(my_addr, *dest_addr, data_packet.encode()) .with_hop_limit(self.config.node.session.default_hop_limit); diff --git a/src/node/session.rs b/src/node/session.rs index 31c9b0c..903d7ec 100644 --- a/src/node/session.rs +++ b/src/node/session.rs @@ -54,6 +54,10 @@ pub(crate) struct SessionEntry { created_at: u64, /// Last activity timestamp (Unix milliseconds). last_activity: u64, + /// Remaining DataPackets that should include COORDS_PRESENT. + /// Initialized from config when session becomes Established; + /// reset on CoordsRequired receipt. + coords_warmup_remaining: u8, } impl SessionEntry { @@ -70,6 +74,7 @@ impl SessionEntry { state: Some(state), created_at: now_ms, last_activity: now_ms, + coords_warmup_remaining: 0, } } @@ -116,4 +121,15 @@ impl SessionEntry { pub(crate) fn last_activity(&self) -> u64 { self.last_activity } + + /// Remaining DataPackets that should include COORDS_PRESENT. + pub(crate) fn coords_warmup_remaining(&self) -> u8 { + self.coords_warmup_remaining + } + + /// Set the coords warmup counter (used on Established transition + /// and CoordsRequired reset). + pub(crate) fn set_coords_warmup_remaining(&mut self, value: u8) { + self.coords_warmup_remaining = value; + } } diff --git a/src/node/tests/session.rs b/src/node/tests/session.rs index 106e6e5..df1020d 100644 --- a/src/node/tests/session.rs +++ b/src/node/tests/session.rs @@ -1299,6 +1299,90 @@ fn test_purge_idle_sessions_disabled_when_zero() { assert_eq!(node.session_count(), 1, "Sessions should not be purged when idle timeout is disabled"); } +// ============================================================================ +// Unit tests: COORDS_PRESENT warmup counter +// ============================================================================ + +#[test] +fn test_coords_warmup_counter_default_zero_on_new() { + use crate::noise::HandshakeState; + + let identity_a = Identity::generate(); + let identity_b = Identity::generate(); + + let handshake = HandshakeState::new_initiator( + identity_a.keypair(), + identity_b.pubkey_full(), + ); + + let entry = crate::node::session::SessionEntry::new( + *identity_b.node_addr(), + identity_b.pubkey_full(), + EndToEndState::Initiating(handshake), + 1000, + ); + + assert_eq!(entry.coords_warmup_remaining(), 0, + "Counter should be 0 for non-Established sessions"); +} + +#[test] +fn test_coords_warmup_counter_set_and_get() { + let node = make_node(); + let remote = Identity::generate(); + let remote_addr = *remote.node_addr(); + + let session = make_noise_session(node.identity(), &remote); + let mut entry = crate::node::session::SessionEntry::new( + remote_addr, + remote.pubkey_full(), + EndToEndState::Established(session), + 1000, + ); + + assert_eq!(entry.coords_warmup_remaining(), 0); + + entry.set_coords_warmup_remaining(5); + assert_eq!(entry.coords_warmup_remaining(), 5); + + entry.set_coords_warmup_remaining(0); + assert_eq!(entry.coords_warmup_remaining(), 0); +} + +#[test] +fn test_coords_warmup_counter_decrement() { + let node = make_node(); + let remote = Identity::generate(); + let remote_addr = *remote.node_addr(); + + let session = make_noise_session(node.identity(), &remote); + let mut entry = crate::node::session::SessionEntry::new( + remote_addr, + remote.pubkey_full(), + EndToEndState::Established(session), + 1000, + ); + + entry.set_coords_warmup_remaining(3); + + // Simulate the decrement pattern used in send_session_data + for expected in (0..3).rev() { + assert!(entry.coords_warmup_remaining() > 0); + entry.set_coords_warmup_remaining(entry.coords_warmup_remaining() - 1); + assert_eq!(entry.coords_warmup_remaining(), expected); + } + + assert_eq!(entry.coords_warmup_remaining(), 0, + "Counter should reach 0 after N decrements"); +} + +#[test] +fn test_coords_warmup_config_default() { + let config = crate::config::Config::new(); + assert_eq!(config.node.session.coords_warmup_packets, 5, + "Default coords_warmup_packets should be 5"); +} + // ============================================================================ // Unit tests: Identity cache // ============================================================================