diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e3b502..e34fb39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,8 +9,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- `node.rate_limit.established_handshake_burst` and + `node.rate_limit.established_handshake_rate`, the parameters of the new + established-link msg1 token bucket. Both are optional; omitting them (the + normal case) derives the bucket from `node.limits.max_peers`, + `node.rekey.after_secs` and `node.rate_limit.handshake_max_resends`, so + raising the peer limit sizes the bucket automatically. An explicit zero + burst or a non-positive rate is rejected at config validation rather than + silently refusing all rekey traffic. + ### Changed +- Inbound msg1 is classified before it is rate limited, and rekey or restart + msg1 arriving on an established link now draws on its own token bucket + instead of competing with stranger admission for a single shared one. On a + node with many peers the shared bucket refused a large share of ordinary + rekey traffic: a field node at roughly 245 peers refused 8753 msg1 in 25 + minutes, and 159 of the 201 distinct sources were peers it already held + sessions with. Nodes upgrade with no config change. The `Msg1 rate limited` + log line now reports which limb refused, the pending count or the token + bucket, which it previously did not distinguish. + - `SessionDatagram::decrement_ttl` and `SessionDatagram::can_forward` now match the forwarder's IP hop-limit semantics: `decrement_ttl` decrements first and reports false when the result is zero, and `can_forward` is true only at a diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 35d1f4b..2db3775 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -134,6 +134,21 @@ Handshake rate limiting protects against DoS on the Noise IK handshake path. | `node.rate_limit.handshake_resend_interval_ms` | u64 | `1000` | Initial handshake message resend interval | | `node.rate_limit.handshake_resend_backoff` | f64 | `2.0` | Resend backoff multiplier (1s, 2s, 4s, 8s, 16s with defaults) | | `node.rate_limit.handshake_max_resends` | u32 | `5` | Max resends per handshake attempt | +| `node.rate_limit.established_handshake_burst` | u32 | derived | Burst capacity of the established-link bucket. Derived default is `node.limits.max_peers` (128) | +| `node.rate_limit.established_handshake_rate` | f64 | derived | Refill rate of that bucket. Derived default is `(max_peers / max(node.rekey.after_secs, 1)) * (1 + handshake_max_resends)`, floored at 1.0/s — 6.4/s at shipped defaults | + +Msg1 whose source matches an established link (rekey and restart +maintenance traffic) draws on a second bucket rather than competing with +stranger admission. Both keys are optional; leaving them unset keeps the +derived sizing, which tracks `max_peers` and the rekey period +automatically instead of becoming a constant nobody revisits. +`max_peers: 0` (unlimited) has no peer-count-derived size, so the +derivation falls back to `handshake_burst` / `handshake_rate`. + +The node's total admitted msg1 rate is the **sum** of the two buckets: 228 +burst and 16.4/s at shipped defaults, of which the established half is +reachable only by a source that already matches a live link. Size against +the sum when budgeting handshake crypto load for a host. ### Retry / Backoff (`node.retry.*`) diff --git a/src/config/mod.rs b/src/config/mod.rs index 2bc0ebf..210e175 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -713,6 +713,32 @@ impl Config { ))); } + // The established-link msg1 bucket. Both keys are `Option`, and + // absent means "derive from max_peers", which is the intended path. + // An explicit zero is the dangerous state: it does not disable the + // bucket, it refuses every rekey and restart msg1 from an already + // established peer, which is a worse failure than the shared-bucket + // starvation the second bucket exists to prevent. + let rl = &self.node.rate_limit; + + if rl.established_handshake_burst == Some(0) { + return Err(ConfigError::Validation( + "`node.rate_limit.established_handshake_burst` is 0, which refuses every rekey and restart msg1 from an established peer rather than disabling the limit. \ + Omit the key to derive it from `node.limits.max_peers`, or set a positive burst." + .to_string(), + )); + } + + if let Some(rate) = rl.established_handshake_rate + && !(rate.is_finite() && rate > 0.0) + { + return Err(ConfigError::Validation(format!( + "`node.rate_limit.established_handshake_rate` is {rate}, but must be a finite value greater than 0; \ + a non-positive or non-finite refill rate never replenishes the established-link bucket, so rekey msg1 stops being admitted once the initial burst is spent. \ + Omit the key to derive it from `node.limits.max_peers` and `node.rekey.after_secs`." + ))); + } + Ok(()) } @@ -1493,6 +1519,48 @@ peers: .expect("shipped default rekey settings must validate"); } + #[test] + fn test_validate_established_burst_zero_rejected() { + let mut config = Config::default(); + config.node.rate_limit.established_handshake_burst = Some(0); + + let err = config.validate().expect_err("validation should fail"); + let msg = err.to_string(); + assert!(msg.contains("established_handshake_burst"), "got: {msg}"); + } + + #[test] + fn test_validate_established_rate_non_positive_rejected() { + for bad in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let mut config = Config::default(); + config.node.rate_limit.established_handshake_rate = Some(bad); + + let err = config + .validate() + .expect_err(&format!("validation should fail for {bad}")); + let msg = err.to_string(); + assert!( + msg.contains("established_handshake_rate"), + "for {bad}, got: {msg}" + ); + } + } + + #[test] + fn test_validate_established_bucket_absent_and_positive_accepted() { + // Absent is the normal path (derived from max_peers) and must validate. + Config::default() + .validate() + .expect("omitted established-bucket keys must validate"); + + let mut config = Config::default(); + config.node.rate_limit.established_handshake_burst = Some(1); + config.node.rate_limit.established_handshake_rate = Some(0.5); + config + .validate() + .expect("positive established-bucket values must validate"); + } + #[test] fn test_validate_rekey_after_messages_zero_rejected() { let mut config = Config::default(); diff --git a/src/config/node.rs b/src/config/node.rs index f791071..061c66d 100644 --- a/src/config/node.rs +++ b/src/config/node.rs @@ -78,6 +78,19 @@ pub struct RateLimitConfig { /// Max handshake resends per attempt (`node.rate_limit.handshake_max_resends`). #[serde(default = "RateLimitConfig::default_handshake_max_resends")] pub handshake_max_resends: u32, + /// Burst capacity of the established-link msg1 bucket + /// (`node.rate_limit.established_handshake_burst`). + /// + /// Absent (the normal case) derives it from `node.limits.max_peers`. + #[serde(default)] + pub established_handshake_burst: Option, + /// Tokens/sec refill rate of the established-link msg1 bucket + /// (`node.rate_limit.established_handshake_rate`). + /// + /// Absent (the normal case) derives it from `node.limits.max_peers`, + /// `node.rekey.after_secs` and `handshake_max_resends`. + #[serde(default)] + pub established_handshake_rate: Option, } impl Default for RateLimitConfig { @@ -89,6 +102,8 @@ impl Default for RateLimitConfig { handshake_resend_interval_ms: 1000, handshake_resend_backoff: 2.0, handshake_max_resends: 5, + established_handshake_burst: None, + established_handshake_rate: None, } } } diff --git a/src/node/handlers/handshake.rs b/src/node/handlers/handshake.rs index df46000..4d5d938 100644 --- a/src/node/handlers/handshake.rs +++ b/src/node/handlers/handshake.rs @@ -2,6 +2,7 @@ use crate::PeerIdentity; use crate::node::acl::PeerAclContext; +use crate::node::rate_limit::Msg1Class; use crate::node::reject::{HandshakeReject, RejectReason}; use crate::node::wire::{Msg1Header, Msg2Header, build_msg2}; use crate::node::{Node, NodeError}; @@ -11,13 +12,14 @@ use std::time::Duration; use tracing::{debug, info, warn}; impl Node { - /// Returns true if an inbound msg1 should be admitted past the - /// `accept_connections` gate. + /// Returns true if an inbound msg1's source matches an established + /// link, i.e. it is rekey/restart maintenance traffic rather than a + /// stranger's fresh handshake. /// - /// Rekey/restart msg1 from an established peer is always admitted (the - /// gate is meant to filter fresh handshakes from strangers, not - /// maintenance traffic on established sessions). Two predicates cover - /// "established peer at this transport+addr": + /// This is deliberately separate from the `accept_connections` gate: + /// it is the only half of `should_admit_msg1` that is a safe basis + /// for exempting traffic from stranger-class treatment. Two + /// predicates cover "established peer at this transport+addr": /// /// 1. `addr_to_link` has an entry for `(transport_id, remote_addr)`. /// This is the fast path and matches when the peer registered with @@ -35,9 +37,14 @@ impl Node { /// with `udp.accept_connections: false` or `udp.outbound_only: true` /// (the production trigger for the 2026-04-30 bug). /// - /// Otherwise the transport's `accept_connections` config decides; - /// absence of a registered transport admits (no gate to apply). - pub(in crate::node) fn should_admit_msg1( + /// Cost: predicate 1 is O(1), predicate 2 is O(peers). Because + /// `handle_msg1` classifies before rate limiting, predicate 2 runs on + /// every inbound msg1 including those about to be refused, so a msg1 + /// flood costs O(peers) per dropped packet rather than O(1). Predicate 2 + /// exists only because `addr_to_link` is keyed on the *unresolved* dial + /// address; if that keying is corrected, this becomes a single O(1) + /// lookup and the flood cost returns to O(1). + pub(in crate::node) fn is_established_link_msg1( &self, transport_id: crate::transport::TransportId, remote_addr: &crate::transport::TransportAddr, @@ -53,6 +60,26 @@ impl Node { }) { return true; } + false + } + + /// Returns true if an inbound msg1 should be admitted past the + /// `accept_connections` gate. + /// + /// Rekey/restart msg1 from an established peer is always admitted (the + /// gate is meant to filter fresh handshakes from strangers, not + /// maintenance traffic on established sessions). + /// + /// Otherwise the transport's `accept_connections` config decides; + /// absence of a registered transport admits (no gate to apply). + pub(in crate::node) fn should_admit_msg1( + &self, + transport_id: crate::transport::TransportId, + remote_addr: &crate::transport::TransportAddr, + ) -> bool { + if self.is_established_link_msg1(transport_id, remote_addr) { + return true; + } self.transports .get(&transport_id) .is_none_or(|t| t.accept_connections()) @@ -62,23 +89,47 @@ impl Node { /// /// This creates a new inbound connection. Rate limiting is applied /// before any expensive crypto operations. + /// + /// Classifying the source costs no crypto (two map/scan lookups), so it + /// happens first and selects which bucket the msg1 draws on: rekey and + /// restart traffic from an established link stops competing with + /// stranger admission, while still being metered. pub(in crate::node) async fn handle_msg1(&mut self, packet: ReceivedPacket) { - // === RATE LIMITING (before any processing) === - if !self.msg1_rate_limiter.start_handshake() { - debug!( - transport_id = %packet.transport_id, - remote_addr = %packet.remote_addr, - "Msg1 rate limited" - ); - return; - } + // === CLASSIFY, THEN RATE LIMIT (both before any crypto) === + // Classification is two map lookups; the second is O(peers) and now + // runs on every inbound msg1, including refused ones. See the + // `is_established_link_msg1` doc comment for why the scan is still + // needed and what would retire it. + let established = self.is_established_link_msg1(packet.transport_id, &packet.remote_addr); + let class = if established { + Msg1Class::EstablishedLink + } else { + Msg1Class::Stranger + }; + let _slot = match self.msg1_rate_limiter.start_handshake(class) { + Ok(slot) => slot, + Err(reason) => { + debug!( + transport_id = %packet.transport_id, + remote_addr = %packet.remote_addr, + refused_by = %reason, + "Msg1 rate limited" + ); + return; + } + }; // accept_connections gate. Rekey/restart msg1 on an existing link // is always admitted; the gate only filters truly-fresh connections // from strangers. Without this carve-out, the dual-init tie-breaker // deadlocks when the larger-NodeAddr side has accept_connections=false. - if !self.should_admit_msg1(packet.transport_id, &packet.remote_addr) { - self.msg1_rate_limiter.complete_handshake(); + // + // `!established &&` is not a behaviour change: `should_admit_msg1` + // is `is_established_link_msg1() || accept_connections()`, so the + // short-circuit only skips a second evaluation of the `peers` scan + // on the hot path. The call is left in place so the two predicates + // cannot drift apart. + if !established && !self.should_admit_msg1(packet.transport_id, &packet.remote_addr) { self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -88,7 +139,6 @@ impl Node { let header = match Msg1Header::parse(&packet.data) { Some(h) => h, None => { - self.msg1_rate_limiter.complete_handshake(); debug!("Invalid msg1 header"); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); @@ -146,7 +196,6 @@ impl Node { HandshakeReject::UnknownConnection, )); } - self.msg1_rate_limiter.complete_handshake(); return; } } else { @@ -187,7 +236,6 @@ impl Node { ) { Ok(m) => m, Err(e) => { - self.msg1_rate_limiter.complete_handshake(); debug!( error = %e, "Failed to process msg1" @@ -202,7 +250,6 @@ impl Node { let peer_identity = match conn.expected_identity() { Some(id) => *id, None => { - self.msg1_rate_limiter.complete_handshake(); warn!("Identity not learned from msg1"); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); @@ -243,7 +290,6 @@ impl Node { // `link_id` was allocated above but `conn` is still a local // (not yet inserted into self.connections / self.links / // self.addr_to_link), so the local drop suffices. - self.msg1_rate_limiter.complete_handshake(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -300,7 +346,6 @@ impl Node { ); self.connections.remove(&link_id); self.links.remove(&link_id); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -321,7 +366,6 @@ impl Node { ); self.connections.remove(&link_id); self.links.remove(&link_id); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut().record_reject(RejectReason::Handshake( HandshakeReject::BadState, )); @@ -350,7 +394,6 @@ impl Node { Ok(idx) => idx, Err(e) => { warn!(error = %e, "Failed to allocate index for rekey"); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut().record_reject(RejectReason::Handshake( HandshakeReject::BadState, )); @@ -363,7 +406,6 @@ impl Node { None => { warn!("Rekey msg1: no session from handshake"); let _ = self.index_allocator.free(our_new_index); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut().record_reject(RejectReason::Handshake( HandshakeReject::BadState, )); @@ -390,7 +432,6 @@ impl Node { "Failed to send rekey msg2" ); let _ = self.index_allocator.free(our_new_index); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut().record_reject(RejectReason::Handshake( HandshakeReject::BadState, )); @@ -422,7 +463,6 @@ impl Node { self.connections.remove(&link_id); self.links.remove(&link_id); - self.msg1_rate_limiter.complete_handshake(); return; } @@ -442,7 +482,6 @@ impl Node { ), } } - self.msg1_rate_limiter.complete_handshake(); return; } } @@ -459,7 +498,6 @@ impl Node { ) .is_err() { - self.msg1_rate_limiter.complete_handshake(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -472,7 +510,6 @@ impl Node { let our_index = match self.index_allocator.allocate() { Ok(idx) => idx, Err(e) => { - self.msg1_rate_limiter.complete_handshake(); warn!(error = %e, "Failed to allocate session index for inbound"); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); @@ -525,7 +562,6 @@ impl Node { self.addr_to_link .remove(&(packet.transport_id, packet.remote_addr)); let _ = self.index_allocator.free(our_index); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -618,8 +654,6 @@ impl Node { .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); } } - - self.msg1_rate_limiter.complete_handshake(); } /// Find stored msg2 bytes for a given link (pre- or post-promotion). diff --git a/src/node/mod.rs b/src/node/mod.rs index aebfa73..f5275ef 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -611,6 +611,32 @@ pub struct Node { tokio::sync::mpsc::UnboundedSender, } +/// Build the msg1 limiter's two buckets and shared pending ceiling. +/// +/// The established-link bucket's size is derived from `max_peers`, the +/// rekey period and the resend budget unless the operator overrode it, so +/// raising the peer count moves it automatically. Nothing re-derives on a +/// config reload; that matches how `handshake_burst` already behaves, +/// since both are read once here at construction. +fn build_msg1_rate_limiter(config: &Config) -> HandshakeRateLimiter { + let rl = &config.node.rate_limit; + let (derived_burst, derived_rate) = rate_limit::derive_established_bucket( + config.node.limits.max_peers, + config.node.rekey.after_secs, + rl.handshake_max_resends, + rl.handshake_burst, + rl.handshake_rate, + ); + HandshakeRateLimiter::with_params( + rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate), + rate_limit::TokenBucket::with_params( + rl.established_handshake_burst.unwrap_or(derived_burst), + rl.established_handshake_rate.unwrap_or(derived_rate), + ), + config.node.limits.max_pending_inbound, + ) +} + impl Node { /// Create a new node from configuration. pub fn new(config: Config) -> Result { @@ -652,11 +678,7 @@ impl Node { config.node.cache.coord_size, config.node.cache.coord_ttl_secs * 1000, ); - let rl = &config.node.rate_limit; - let msg1_rate_limiter = HandshakeRateLimiter::with_params( - rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate), - config.node.limits.max_pending_inbound, - ); + let msg1_rate_limiter = build_msg1_rate_limiter(&config); let max_connections = config.node.limits.max_connections; let max_peers = config.node.limits.max_peers; @@ -816,11 +838,7 @@ impl Node { config.node.cache.coord_size, config.node.cache.coord_ttl_secs * 1000, ); - let rl = &config.node.rate_limit; - let msg1_rate_limiter = HandshakeRateLimiter::with_params( - rate_limit::TokenBucket::with_params(rl.handshake_burst, rl.handshake_rate), - config.node.limits.max_pending_inbound, - ); + let msg1_rate_limiter = build_msg1_rate_limiter(&config); let max_connections = config.node.limits.max_connections; let max_peers = config.node.limits.max_peers; diff --git a/src/node/rate_limit.rs b/src/node/rate_limit.rs index 5189333..32a581c 100644 --- a/src/node/rate_limit.rs +++ b/src/node/rate_limit.rs @@ -15,7 +15,30 @@ //! - Burst capacity: 100 tokens (max concurrent handshakes) //! - Refill rate: 10 tokens/second (sustained handshake rate) //! - This allows handling burst traffic while limiting sustained attack impact +//! +//! ## Two buckets, and what the aggregate is +//! +//! Msg1 whose source matches an established link (rekey and restart +//! maintenance traffic) draws on its own bucket rather than competing with +//! stranger admission. It is *metered*, not exempted: an established-peer +//! carve-out is by construction keyed on source address, and the sentence +//! above about spoofable UDP sources is still true, so an off-path attacker +//! who can forge a live peer's `(transport_id, addr)` tuple reaches the +//! second bucket. Metering keeps that exposure bounded. +//! +//! The consequence, stated rather than left implicit: the node's total +//! admitted msg1 rate is the **sum** of the two buckets, not the stranger +//! bucket alone. At shipped defaults that is burst `100 + 128 = 228` and +//! `10.0 + 6.4 = 16.4` msg1/sec, of which the established half is only +//! reachable by a source that already matches a live link. Operators sizing +//! the handshake-crypto ceiling against a host should size against the sum. +//! +//! The concurrency limb (`max_pending`) is deliberately *not* split, so no +//! equivalent inflation happens there: one counter bounds simultaneous +//! in-flight handshake state whoever holds the slot. +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Instant; /// Default burst capacity (max tokens). @@ -156,84 +179,222 @@ impl Default for TokenBucket { } } +/// Floor on the derived established-link refill rate, in tokens/second. +/// +/// Covers the degenerate case of a very long (or effectively disabled) +/// rekey period, where a node must still admit restart msg1 at some rate. +pub const ESTABLISHED_RATE_FLOOR: f64 = 1.0; + +/// Which bucket an inbound msg1 draws on. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Msg1Class { + /// No established link matches the source `(transport_id, addr)`. + Stranger, + /// The source matches an established link: rekey or restart traffic. + EstablishedLink, +} + +/// Why an inbound msg1 was refused by the limiter. +/// +/// `start_handshake` refuses on either limb and the caller's log line is +/// blind to which without this. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Msg1Refusal { + /// The shared in-flight handshake count is at `max_pending`. + PendingLimit, + /// The class's token bucket is empty. + RateLimit, +} + +impl Msg1Refusal { + /// Stable field value for structured logs. + pub fn as_str(self) -> &'static str { + match self { + Msg1Refusal::PendingLimit => "pending_limit", + Msg1Refusal::RateLimit => "rate_limit", + } + } +} + +impl std::fmt::Display for Msg1Refusal { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// An in-flight handshake's pending slot, released on drop. +/// +/// The slot is the limiter's concurrency limb. Releasing it by hand from +/// every exit path of a long handler is the shape that lets one path free +/// a slot belonging to a *different* handshake, silently lifting effective +/// concurrency above `max_pending` with no counter moving and no log +/// firing. Holding the release in `Drop` makes that structurally +/// impossible. +/// +/// Worth knowing when reading `max_pending`: today the count cannot +/// exceed 1 in production. `start_handshake` is reached only from +/// `handle_msg1`, whose sole production caller awaits it to completion +/// inside `process_packet(&mut self)`, itself awaited serially in the rx +/// loop, so no two msg1 handlers are ever in flight at once. The +/// concurrency limb is therefore a structural invariant rather than a +/// limiter that currently binds, and `Msg1Refusal::PendingLimit` does not +/// fire in the field. The guard is what keeps that invariant true as the +/// handler grows exit paths, and what makes it safe for different msg1 +/// classes to take slots on different terms. +/// +/// The atomic is for `Send`-ness, not cross-thread coordination: the guard +/// is held across `.await` points inside `handle_msg1`, and a +/// `Rc>` would make every future containing it non-`Send`. +#[must_use = "binding the slot to a named local is what holds it; \ + dropping it immediately releases it straight away"] +#[derive(Debug)] +pub struct PendingHandshake { + pending: Arc, +} + +impl Drop for PendingHandshake { + fn drop(&mut self) { + let _ = self + .pending + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1)); + } +} + +/// Derive the established-link bucket's parameters from configuration. +/// +/// Every input is a value the operator already sets, so raising `max_peers` +/// moves this automatically rather than leaving a bare constant that nobody +/// revisits. +/// +/// - **burst = `max_peers`.** The worst *legitimate* burst on this bucket is +/// every established peer re-handshaking at once (local restart, partition +/// heal), and that population is bounded by `max_peers` by construction, at +/// one token per msg1. +/// - **rate = `(max_peers / max(rekey_after_secs, 1)) * (1 + max_resends)`,** +/// floored at [`ESTABLISHED_RATE_FLOOR`]. The first factor is the +/// steady-state inbound rekey rate; the second is the worst case where +/// every attempt consumes its full retransmission budget. +/// +/// `max_peers == 0` means unlimited, for which no peer-count-derived size +/// exists; the stranger bucket's own parameters are returned so an +/// unlimited-peers node gets a second bucket the same size as its first +/// rather than a bucket of one token. +/// +/// `rekey.enabled` is deliberately ignored: restart and reconnect msg1 exist +/// regardless, and branching on it would make the disabled-rekey +/// configuration the under-provisioned one. +pub fn derive_established_bucket( + max_peers: usize, + rekey_after_secs: u64, + max_resends: u32, + stranger_burst: u32, + stranger_rate: f64, +) -> (u32, f64) { + if max_peers == 0 { + return (stranger_burst, stranger_rate); + } + let burst = u32::try_from(max_peers).unwrap_or(u32::MAX); + let period = rekey_after_secs.max(1) as f64; + let rate = (max_peers as f64 / period) * (1.0 + f64::from(max_resends)); + (burst, rate.max(ESTABLISHED_RATE_FLOOR)) +} + /// Rate limiter for handshake message 1 processing. /// /// Combines token bucket rate limiting with connection counting to -/// protect against DoS attacks on the handshake path. +/// protect against DoS attacks on the handshake path. The rate limb is +/// split by [`Msg1Class`]; the concurrency limb is shared (see the module +/// doc for why, and for what the aggregate rate becomes). #[derive(Debug)] pub struct HandshakeRateLimiter { - /// Token bucket for rate limiting. + /// Token bucket for stranger msg1. bucket: TokenBucket, - /// Current count of pending inbound connections. - pending_count: usize, + /// Token bucket for established-link msg1 (rekey / restart). + established: TokenBucket, + /// Current count of pending inbound connections, shared with the + /// outstanding [`PendingHandshake`] guards. + pending: Arc, /// Maximum pending inbound connections. max_pending: usize, } impl HandshakeRateLimiter { /// Create a handshake rate limiter with the given parameters. - pub fn with_params(bucket: TokenBucket, max_pending: usize) -> Self { + pub fn with_params(bucket: TokenBucket, established: TokenBucket, max_pending: usize) -> Self { Self { bucket, - pending_count: 0, + established, + pending: Arc::new(AtomicUsize::new(0)), max_pending, } } - /// Check if a new handshake can be started. + /// Check if a new handshake of `class` can be started. /// /// Returns `true` if: - /// - Token bucket has available tokens (rate limit not exceeded) + /// - The class's token bucket has available tokens (rate limit not exceeded) /// - Pending connection count is below maximum /// /// Does NOT consume a token - call `start_handshake` for that. #[cfg(test)] - pub fn can_start_handshake(&mut self) -> bool { - self.bucket.available() && self.pending_count < self.max_pending + pub fn can_start_handshake(&mut self, class: Msg1Class) -> bool { + self.bucket_for(class).available() + && self.pending.load(Ordering::Relaxed) < self.max_pending } - /// Start a new handshake, consuming a token and incrementing pending count. + /// Start a new handshake, consuming a token and taking a pending slot. /// - /// Returns `true` if the handshake was allowed, `false` if rate limited. - pub fn start_handshake(&mut self) -> bool { - if self.pending_count >= self.max_pending { - return false; + /// The returned guard releases the slot when it drops. On refusal the + /// [`Msg1Refusal`] says which limb refused. + pub fn start_handshake(&mut self, class: Msg1Class) -> Result { + if self.pending.load(Ordering::Relaxed) >= self.max_pending { + return Err(Msg1Refusal::PendingLimit); } - if self.bucket.try_acquire() { - self.pending_count += 1; - true - } else { - false + if !self.bucket_for(class).try_acquire() { + return Err(Msg1Refusal::RateLimit); } + + self.pending.fetch_add(1, Ordering::Relaxed); + Ok(PendingHandshake { + pending: Arc::clone(&self.pending), + }) } - /// Mark a handshake as complete (successful or failed). - /// - /// Decrements the pending connection count. - pub fn complete_handshake(&mut self) { - if self.pending_count > 0 { - self.pending_count -= 1; + fn bucket_for(&mut self, class: Msg1Class) -> &mut TokenBucket { + match class { + Msg1Class::Stranger => &mut self.bucket, + Msg1Class::EstablishedLink => &mut self.established, } } /// Get the current pending connection count. #[cfg(test)] pub fn pending_count(&self) -> usize { - self.pending_count + self.pending.load(Ordering::Relaxed) } - /// Get a reference to the token bucket. + /// Get a reference to the stranger token bucket. #[cfg(test)] pub fn bucket(&self) -> &TokenBucket { &self.bucket } + /// Get a reference to the established-link token bucket. + #[cfg(test)] + pub fn established_bucket(&self) -> &TokenBucket { + &self.established + } + /// Reset the rate limiter. + /// + /// Does not affect slots held by live [`PendingHandshake`] guards; they + /// still decrement on drop, saturating at zero. #[cfg(test)] pub fn reset(&mut self) { self.bucket.reset(); - self.pending_count = 0; + self.established.reset(); + self.pending.store(0, Ordering::Relaxed); } } @@ -342,71 +503,82 @@ mod tests { assert!(wait.as_millis() >= 90 && wait.as_millis() <= 110); } + fn test_limiter(bucket: TokenBucket, max_pending: usize) -> HandshakeRateLimiter { + HandshakeRateLimiter::with_params( + bucket, + TokenBucket::with_params(1000, 100.0), + max_pending, + ) + } + #[test] fn test_handshake_rate_limiter_basic() { - let mut limiter = HandshakeRateLimiter::with_params(TokenBucket::new(), 100); + let mut limiter = test_limiter(TokenBucket::new(), 100); - assert!(limiter.can_start_handshake()); + assert!(limiter.can_start_handshake(Msg1Class::Stranger)); assert_eq!(limiter.pending_count(), 0); // Start a handshake - assert!(limiter.start_handshake()); + let slot = limiter.start_handshake(Msg1Class::Stranger).unwrap(); assert_eq!(limiter.pending_count(), 1); // Complete it - limiter.complete_handshake(); + drop(slot); assert_eq!(limiter.pending_count(), 0); } #[test] fn test_handshake_rate_limiter_max_pending() { let bucket = TokenBucket::with_params(1000, 100.0); - let mut limiter = HandshakeRateLimiter::with_params(bucket, 3); + let mut limiter = test_limiter(bucket, 3); // Start 3 handshakes - assert!(limiter.start_handshake()); - assert!(limiter.start_handshake()); - assert!(limiter.start_handshake()); + let a = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + let _b = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + let _c = limiter.start_handshake(Msg1Class::Stranger).unwrap(); // Fourth should fail (max pending) - assert!(!limiter.can_start_handshake()); - assert!(!limiter.start_handshake()); + assert!(!limiter.can_start_handshake(Msg1Class::Stranger)); + assert_eq!( + limiter.start_handshake(Msg1Class::Stranger).unwrap_err(), + Msg1Refusal::PendingLimit + ); // Complete one - limiter.complete_handshake(); + drop(a); // Now should be able to start another - assert!(limiter.can_start_handshake()); - assert!(limiter.start_handshake()); + assert!(limiter.can_start_handshake(Msg1Class::Stranger)); + assert!(limiter.start_handshake(Msg1Class::Stranger).is_ok()); } #[test] fn test_handshake_rate_limiter_token_exhaustion() { let bucket = TokenBucket::with_params(5, 0.0); // No refill - let mut limiter = HandshakeRateLimiter::with_params(bucket, 100); + let mut limiter = test_limiter(bucket, 100); - // Start 5 handshakes (exhausts tokens) + // Start 5 handshakes (exhausts tokens), releasing each immediately for _ in 0..5 { - assert!(limiter.start_handshake()); - } - - // Complete them all - for _ in 0..5 { - limiter.complete_handshake(); + let slot = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + drop(slot); } // Tokens exhausted, even though pending is 0 - assert!(!limiter.can_start_handshake()); - assert!(!limiter.start_handshake()); + assert_eq!(limiter.pending_count(), 0); + assert!(!limiter.can_start_handshake(Msg1Class::Stranger)); + assert_eq!( + limiter.start_handshake(Msg1Class::Stranger).unwrap_err(), + Msg1Refusal::RateLimit + ); } #[test] fn test_handshake_rate_limiter_reset() { - let mut limiter = HandshakeRateLimiter::with_params(TokenBucket::new(), 100); + let mut limiter = test_limiter(TokenBucket::new(), 100); // Start some handshakes - limiter.start_handshake(); - limiter.start_handshake(); + let _a = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + let _b = limiter.start_handshake(Msg1Class::Stranger).unwrap(); assert_eq!(limiter.pending_count(), 2); // Reset @@ -415,4 +587,141 @@ mod tests { assert_eq!(limiter.pending_count(), 0); assert!(limiter.bucket().tokens >= DEFAULT_BURST_CAPACITY as f64 - 0.1); } + + /// The pending slot is released by `Drop`, not by any reachable manual + /// call, and the release saturates at zero rather than underflowing. + #[test] + fn pending_slot_releases_on_drop_and_saturates() { + let mut limiter = test_limiter(TokenBucket::with_params(100, 100.0), 100); + + let outer = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + { + let _inner = limiter.start_handshake(Msg1Class::Stranger).unwrap(); + assert_eq!(limiter.pending_count(), 2); + } + assert_eq!( + limiter.pending_count(), + 1, + "inner guard released its own slot" + ); + + // A reset zeroes the counter while a guard is still live; the + // guard's later drop must not underflow it. + limiter.reset(); + assert_eq!(limiter.pending_count(), 0); + drop(outer); + assert_eq!(limiter.pending_count(), 0, "release saturates at zero"); + } + + /// The refusal reason discriminates the two limbs. This is what the + /// `refused_by` log field in `handle_msg1` carries; the log line itself + /// is not asserted here. + #[test] + fn refusal_reason_names_the_limb() { + // Pending limb: max_pending 1, one live guard, tokens plentiful. + let mut pending_bound = test_limiter(TokenBucket::with_params(1000, 100.0), 1); + let _held = pending_bound.start_handshake(Msg1Class::Stranger).unwrap(); + assert_eq!( + pending_bound + .start_handshake(Msg1Class::Stranger) + .unwrap_err(), + Msg1Refusal::PendingLimit + ); + + // Rate limb: pending headroom, zero-capacity/zero-refill bucket. + let mut rate_bound = test_limiter(TokenBucket::with_params(0, 0.0), 1000); + assert_eq!( + rate_bound.start_handshake(Msg1Class::Stranger).unwrap_err(), + Msg1Refusal::RateLimit + ); + + assert_eq!(Msg1Refusal::PendingLimit.to_string(), "pending_limit"); + assert_eq!(Msg1Refusal::RateLimit.to_string(), "rate_limit"); + } + + /// The two classes draw on separate buckets: draining one leaves the + /// other untouched. + #[test] + fn msg1_classes_draw_on_separate_buckets() { + let mut limiter = HandshakeRateLimiter::with_params( + TokenBucket::with_params(1, 0.0), + TokenBucket::with_params(3, 0.0), + 1000, + ); + + drop(limiter.start_handshake(Msg1Class::Stranger).unwrap()); + assert_eq!( + limiter.start_handshake(Msg1Class::Stranger).unwrap_err(), + Msg1Refusal::RateLimit, + "stranger bucket drained" + ); + + for _ in 0..3 { + drop( + limiter + .start_handshake(Msg1Class::EstablishedLink) + .expect("established bucket is independent of the stranger bucket"), + ); + } + assert_eq!( + limiter + .start_handshake(Msg1Class::EstablishedLink) + .unwrap_err(), + Msg1Refusal::RateLimit, + "established bucket drains on its own terms" + ); + } + + #[test] + fn derive_established_bucket_from_shipped_defaults() { + // max_peers 128, rekey.after_secs 120, handshake_max_resends 5. + let (burst, rate) = derive_established_bucket(128, 120, 5, 100, 10.0); + assert_eq!(burst, 128); + assert!( + (rate - 6.4).abs() < 1e-9, + "expected 6.4 tokens/sec, got {rate}" + ); + + // Scales with max_peers rather than needing a revisit. + let (burst, rate) = derive_established_bucket(512, 120, 5, 100, 10.0); + assert_eq!(burst, 512); + assert!( + (rate - 25.6).abs() < 1e-9, + "expected 25.6 tokens/sec, got {rate}" + ); + } + + /// `after_secs = 0` must not divide by zero. It is clamped to 1s, which + /// yields a *large* rate, not the floor — the floor binds at the other + /// end, for a very long rekey period. Both ends are asserted here + /// because the two are easy to conflate. + #[test] + fn derive_established_bucket_degenerate_rekey_periods() { + let (burst, rate) = derive_established_bucket(128, 0, 5, 100, 10.0); + assert_eq!(burst, 128); + assert!(rate.is_finite(), "after_secs = 0 must not divide by zero"); + assert_eq!( + rate, + derive_established_bucket(128, 1, 5, 100, 10.0).1, + "after_secs = 0 is clamped to 1s" + ); + assert!( + rate > ESTABLISHED_RATE_FLOOR, + "a zero rekey period is the high end, not the floor" + ); + + // The floor binds for a very long / effectively disabled period. + let (_, rate) = derive_established_bucket(1, 100_000, 5, 100, 10.0); + assert_eq!(rate, ESTABLISHED_RATE_FLOOR); + } + + /// `max_peers == 0` means unlimited. Deriving a burst from it would + /// yield a zero- or one-token bucket, which is worse than the bug this + /// second bucket exists to fix; the stranger parameters are reused. + #[test] + fn derive_established_bucket_unlimited_peers_reuses_stranger_params() { + let (burst, rate) = derive_established_bucket(0, 120, 5, 100, 10.0); + assert_eq!(burst, 100); + assert_eq!(rate, 10.0); + } } diff --git a/src/node/tests/handshake.rs b/src/node/tests/handshake.rs index 0e49fe2..822eb3d 100644 --- a/src/node/tests/handshake.rs +++ b/src/node/tests/handshake.rs @@ -954,11 +954,22 @@ async fn test_duplicate_msg2_dropped() { /// `should_admit_msg1` admits when no transport is registered for the id. /// (No gate to apply — the caller's other checks decide the outcome.) +/// +/// This node is also the discriminator for the extraction of +/// `is_established_link_msg1`: with no transport registered the +/// `accept_connections` fallback admits, so the two predicates disagree +/// here and nowhere else. An extraction that dragged the fallback into +/// `is_established_link_msg1` fails the second assertion. #[test] fn test_should_admit_msg1_no_transport() { let node = make_node(); let addr = TransportAddr::from_string("10.0.0.2:2121"); assert!(node.should_admit_msg1(TransportId::new(1), &addr)); + assert!( + !node.is_established_link_msg1(TransportId::new(1), &addr), + "the accept_connections fallback must not be part of the \ + established-link predicate" + ); } /// `should_admit_msg1` rejects a fresh msg1 (no addr_to_link entry) when @@ -1132,4 +1143,10 @@ async fn test_should_admit_msg1_admits_rekey_when_addr_form_differs() { !node.should_admit_msg1(transport_id, &stranger_addr), "fresh msg1 from unknown source must still be rejected" ); + + // The same two predicates read directly: both addr-forms of the + // established peer are established links, the stranger is not. + assert!(node.is_established_link_msg1(transport_id, &hostname_addr)); + assert!(node.is_established_link_msg1(transport_id, &numeric_addr)); + assert!(!node.is_established_link_msg1(transport_id, &stranger_addr)); } diff --git a/src/node/tests/unit.rs b/src/node/tests/unit.rs index 1a82a96..66494a6 100644 --- a/src/node/tests/unit.rs +++ b/src/node/tests/unit.rs @@ -560,17 +560,25 @@ async fn test_node_rx_loop_takes_channel() { #[test] fn test_rate_limiter_initialized() { + use crate::node::rate_limit::Msg1Class; + let mut node = make_node(); // Rate limiter should allow handshakes initially - assert!(node.msg1_rate_limiter.can_start_handshake()); + assert!( + node.msg1_rate_limiter + .can_start_handshake(Msg1Class::Stranger) + ); // Start a handshake - assert!(node.msg1_rate_limiter.start_handshake()); + let slot = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("fresh limiter admits"); assert_eq!(node.msg1_rate_limiter.pending_count(), 1); // Complete it - node.msg1_rate_limiter.complete_handshake(); + drop(slot); assert_eq!(node.msg1_rate_limiter.pending_count(), 0); } @@ -1990,8 +1998,8 @@ async fn handle_msg1_silent_drops_at_cap_for_new_peer() { assert_eq!( node.msg1_rate_limiter.pending_count(), before_pending, - "rate limiter must rebalance: start_handshake() then \ - complete_handshake() before silent-drop return" + "rate limiter must rebalance: the pending slot start_handshake() \ + took is released by its guard before the silent-drop return" ); // Wire-observable discriminator: with the early gate in place, no @@ -2083,6 +2091,354 @@ async fn handle_msg1_admits_existing_peer_at_cap() { ); } +/// Every reject arm of `handle_msg1` releases *its own* pending slot and +/// nothing else. +/// +/// The failure this guards is invisible by construction: a slot released +/// by a path that never took one frees a slot belonging to a **different** +/// in-flight handshake, lifting effective concurrency above `max_pending` +/// with no counter moving, no log firing, and no underflow (the release +/// saturates at zero). The only way to observe it is to hold a *foreign* +/// slot across the calls and watch whether it survives — which is what +/// `seed` is. A borrow-based guard could not be held here at all, since +/// `handle_msg1` needs `&mut node` throughout. +/// +/// Each arm additionally asserts the reject counter it is supposed to +/// bump, so a setup that silently failed to reach the arm (wrong addr, +/// packet rejected earlier) shows up as a red rather than as a +/// vacuously-stable pending count. +#[tokio::test] +async fn msg1_reject_arms_do_not_release_another_handshakes_slot() { + use crate::config::UdpConfig; + use crate::node::rate_limit::Msg1Class; + use crate::node::wire::build_msg1; + use crate::noise::HANDSHAKE_MSG1_SIZE; + use crate::utils::index::SessionIndex; + + // max_peers 2 so the early-cap arm is reachable on the *same* node the + // foreign slot is seeded on. That would otherwise derive a 2-token + // established bucket, which refuses the third arm before it runs, so + // both buckets are overridden to sizes this test never approaches: + // the subject here is slot accounting, not admission. + let mut config = Config::new(); + config.node.limits.max_peers = 2; + config.node.rate_limit.handshake_burst = 100; + config.node.rate_limit.established_handshake_burst = Some(100); + let mut node = make_node_with(config); + + let transport_id = TransportId::new(1); + let udp_config = UdpConfig { + bind_addr: Some("127.0.0.1:0".to_string()), + mtu: Some(1280), + ..Default::default() + }; + let (packet_tx, mut packet_rx) = packet_channel(64); + let mut transport = UdpTransport::new(transport_id, None, udp_config, packet_tx); + transport.start_async().await.unwrap(); + let node_udp_addr = transport.local_addr().unwrap(); + node.transports + .insert(transport_id, TransportHandle::Udp(transport)); + + let socket_a = tokio::net::UdpSocket::bind("127.0.0.1:0") + .await + .expect("bind sender socket"); + let wire_addr = TransportAddr::from_string(&socket_a.local_addr().unwrap().to_string()); + + // A foreign in-flight handshake's slot, held for the whole test. + let seed = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("fresh limiter admits the seed"); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "baseline: exactly one foreign slot outstanding" + ); + + // Both source tuples are established links, so the msg1s below are in + // the exempted class. The link ids are deliberately absent from + // `node.links` for now, so the duplicate-msg1 branch is skipped. + let hand_addr = TransportAddr::from_string("198.51.100.7:2121"); + let hand_link_id = node.allocate_link_id(); + let wire_link_id = node.allocate_link_id(); + node.addr_to_link + .insert((transport_id, hand_addr.clone()), hand_link_id); + node.addr_to_link + .insert((transport_id, wire_addr.clone()), wire_link_id); + + let hand_packet = |data: Vec| ReceivedPacket { + transport_id, + remote_addr: hand_addr.clone(), + data, + timestamp_ms: 1000, + }; + let garbage_msg1 = build_msg1(SessionIndex::new(0x4242), &[0u8; HANDSHAKE_MSG1_SIZE]); + + // Arm 1: invalid header (truncated body). + let before = node.stats().handshake.bad_state; + node.handle_msg1(hand_packet(vec![0u8; 8])).await; + assert_eq!( + node.stats().handshake.bad_state, + before + 1, + "arm 1 must reach the invalid-header reject" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "invalid-header arm must not release the foreign slot" + ); + + // Arm 2: well-formed header, unusable Noise payload. + let before = node.stats().handshake.bad_state; + node.handle_msg1(hand_packet(garbage_msg1.clone())).await; + assert_eq!( + node.stats().handshake.bad_state, + before + 1, + "arm 2 must reach the receive_handshake_init reject" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "handshake-init-failure arm must not release the foreign slot" + ); + + // Arm 3: duplicate msg1 on a pending inbound link with no stored msg2. + node.links.insert( + hand_link_id, + Link::connectionless( + hand_link_id, + transport_id, + hand_addr.clone(), + LinkDirection::Inbound, + Duration::from_millis(100), + ), + ); + let before = node.stats().handshake.unknown_connection; + node.handle_msg1(hand_packet(garbage_msg1)).await; + assert_eq!( + node.stats().handshake.unknown_connection, + before + 1, + "arm 3 must reach the duplicate-msg1-no-stored-msg2 reject" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "duplicate-msg1 arm must not release the foreign slot" + ); + + // Arm 4: the early max_peers cap gate, on a genuine crafted msg1 that + // gets all the way through the Noise step. + inject_dummy_peers(&mut node, 2); + assert_eq!(node.peer_count(), 2, "precondition: at cap"); + let sender = Identity::generate(); + let sender_node_addr = + craft_and_send_msg1(&node, &sender, &socket_a, node_udp_addr, 2000).await; + let before = node.stats().handshake.bad_state; + pump_one_msg1_into_node(&mut node, &mut packet_rx, 1000) + .await + .expect("crafted msg1 must reach packet_rx"); + assert_eq!( + node.stats().handshake.bad_state, + before + 1, + "arm 4 must reach the max_peers cap reject" + ); + assert!( + !node.peers.contains_key(&sender_node_addr), + "arm 4 must not admit the new identity" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "max_peers-cap arm must not release the foreign slot" + ); + + // The foreign slot was still ours to release all along. + drop(seed); + assert_eq!(node.msg1_rate_limiter.pending_count(), 0); +} + +/// The established-link bucket is wired from config at construction: +/// derived from `max_peers` by default, overridden when the operator sets +/// the key. This is the only test covering the config → limiter path. +#[test] +fn node_established_bucket_is_derived_then_overridable() { + let mut config = Config::new(); + config.node.limits.max_peers = 300; + let node = make_node_with(config); + assert_eq!( + node.msg1_rate_limiter.established_bucket().capacity(), + 300, + "derived burst tracks max_peers" + ); + + let mut config = Config::new(); + config.node.limits.max_peers = 300; + config.node.rate_limit.established_handshake_burst = Some(7); + let node = make_node_with(config); + assert_eq!( + node.msg1_rate_limiter.established_bucket().capacity(), + 7, + "an explicit key wins over the derivation" + ); +} + +/// Build a node whose *stranger* msg1 bucket holds exactly one token and +/// never refills, with a UDP transport bound and started. Returns the node, +/// its transport id, its wire address and its packet receiver. +/// +/// The established-link bucket is left at its derived size (128 burst at +/// default `max_peers`), which is the whole point: the two classes are +/// metered separately. +async fn node_with_single_stranger_token() -> ( + Node, + TransportId, + std::net::SocketAddr, + crate::transport::PacketRx, +) { + use crate::config::UdpConfig; + + let mut config = Config::new(); + config.node.rate_limit.handshake_burst = 1; + config.node.rate_limit.handshake_rate = 0.0; + let mut node = make_node_with(config); + + let transport_id = TransportId::new(1); + let udp_config = UdpConfig { + bind_addr: Some("127.0.0.1:0".to_string()), + mtu: Some(1280), + ..Default::default() + }; + let (packet_tx, packet_rx) = packet_channel(64); + let mut transport = UdpTransport::new(transport_id, None, udp_config, packet_tx); + transport.start_async().await.unwrap(); + let wire_addr = transport.local_addr().unwrap(); + node.transports + .insert(transport_id, TransportHandle::Udp(transport)); + + (node, transport_id, wire_addr, packet_rx) +} + +/// Poll a sender socket for a msg2 reply. `None` means nothing arrived +/// within 300 ms, the same wire-observable discriminator the max_peers cap +/// tests use. +async fn poll_for_msg2(socket: &tokio::net::UdpSocket) -> Option { + use tokio::time::{Duration, timeout}; + let mut buf = [0u8; 2048]; + timeout(Duration::from_millis(300), socket.recv_from(&mut buf)) + .await + .ok() + .and_then(|inner| inner.ok()) + .map(|(n, _)| n) +} + +/// An established link's rekey/restart msg1 is admitted even when the +/// stranger bucket is empty, because it draws on its own bucket. +/// +/// This is the symptom the whole change exists to fix: before it, one +/// drained global bucket refused an established peer's maintenance traffic +/// on exactly the same terms as a stranger's first packet. +#[tokio::test] +async fn established_link_msg1_admitted_when_stranger_bucket_drained() { + use crate::node::rate_limit::Msg1Class; + + let (mut node, transport_id, wire_addr, mut packet_rx) = + node_with_single_stranger_token().await; + + // Spend the one stranger token, so any msg1 classed as a stranger is + // refused from here on. + drop( + node.msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("the single stranger token is available at t=0"), + ); + + for attempt in 0..3 { + let socket = tokio::net::UdpSocket::bind("127.0.0.1:0") + .await + .expect("bind sender socket"); + let sender_addr = TransportAddr::from_string(&socket.local_addr().unwrap().to_string()); + + // Mark the source as an established link. The link id is + // deliberately absent from `node.links`, so the duplicate-msg1 + // branch is skipped and the msg1 is processed as a fresh + // connection that answers with msg2. + let link_id = node.allocate_link_id(); + node.addr_to_link + .insert((transport_id, sender_addr), link_id); + + let sender = Identity::generate(); + craft_and_send_msg1(&node, &sender, &socket, wire_addr, 1000 + attempt).await; + pump_one_msg1_into_node(&mut node, &mut packet_rx, 1000) + .await + .expect("msg1 must reach packet_rx"); + + assert!( + poll_for_msg2(&socket).await.is_some(), + "attempt {attempt}: established-link msg1 must be answered with \ + msg2 while the stranger bucket is empty" + ); + } +} + +/// The twin of the test above: a stranger is still refused once the +/// stranger bucket is drained. The second bucket must not become a way in +/// for sources that match no established link. +/// +/// The first sender is a **positive control** on the same node, transport +/// and code path: it proves the setup really delivers a msg1 and really +/// produces a msg2 on the wire, so the second sender's silence is +/// attributable to the drained bucket rather than to a msg1 that never +/// arrived. +#[tokio::test] +async fn stranger_msg1_still_refused_when_bucket_drained() { + let (mut node, _transport_id, wire_addr, mut packet_rx) = + node_with_single_stranger_token().await; + + // Positive control: one token is available, so this stranger is + // admitted and answered. + let socket_ok = tokio::net::UdpSocket::bind("127.0.0.1:0") + .await + .expect("bind sender socket"); + craft_and_send_msg1(&node, &Identity::generate(), &socket_ok, wire_addr, 1000).await; + pump_one_msg1_into_node(&mut node, &mut packet_rx, 1000) + .await + .expect("control msg1 must reach packet_rx"); + assert!( + poll_for_msg2(&socket_ok).await.is_some(), + "positive control: a stranger with a token available must be \ + answered with msg2 — without this the silence below proves nothing" + ); + assert_eq!(node.peer_count(), 1, "control msg1 was fully processed"); + + // The bucket is now empty and never refills. A second stranger, at a + // different source addr and with no established link, gets nothing. + let socket_refused = tokio::net::UdpSocket::bind("127.0.0.1:0") + .await + .expect("bind second sender socket"); + craft_and_send_msg1( + &node, + &Identity::generate(), + &socket_refused, + wire_addr, + 2000, + ) + .await; + pump_one_msg1_into_node(&mut node, &mut packet_rx, 1000) + .await + .expect("refused msg1 must still reach packet_rx"); + let bytes = poll_for_msg2(&socket_refused).await; + assert!( + bytes.is_none(), + "stranger msg1 must stay refused once the stranger bucket is \ + drained; observed {bytes:?} wire bytes" + ); + assert_eq!( + node.peer_count(), + 1, + "the refused stranger must not have been admitted" + ); +} + // ===== Transport kernel-drop detection (sans-IO) ===== // // The drop-detection edge-detector, tested directly. It replaces the