diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f29d05..1b6c018 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -101,6 +101,15 @@ with v0.4.x or earlier peers. ### 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. + - The receive-path `RejectReason` classification (shipped in 0.4.0) is additionally wired into the Noise XX handshake cluster (msg1/msg2/msg3) and the rekey-initiator outbound sites on `next`. @@ -130,6 +139,21 @@ with v0.4.x or earlier peers. ### Changed +- Inbound msg1 is classified before it is rate limited, and rekey or restart + msg1 arriving on a link belonging to a promoted peer 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. On the XX handshake path the classifier keys on + promotion state rather than on the presence of an address-map entry, because + msg1 creates such an entry for a still-pending inbound connection before any + identity is known; a still-handshaking stranger therefore stays in the + stranger class for its whole lifetime, retransmits included. 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. + - Rekey timer jitter is enabled on next's XX FMP rekey path (`REKEY_JITTER_SECS = 15` at `src/node/mod.rs`), matching the IK-line behavior on maint/master. It had been temporarily set to diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 2b7cf83..89a00bb 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -135,6 +135,28 @@ Handshake rate limiting protects against DoS on the Noise XX 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. + +"Established link" here means a **promoted** peer, which is stricter than +it sounds on the XX handshake path. An inbound handshake that has sent +msg1 but not yet completed msg3 is not promoted, so it draws on the +stranger bucket for its whole lifetime, including every msg1 retransmit. +Only traffic from a source already matching a promoted peer reaches the +established bucket. ### Retry / Backoff (`node.retry.*`) diff --git a/src/config/mod.rs b/src/config/mod.rs index 90ef742..dd2fb8a 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -777,6 +777,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(()) } @@ -1689,6 +1715,48 @@ node: .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 634d4d1..6cd9a53 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 42600f9..020137e 100644 --- a/src/node/handlers/handshake.rs +++ b/src/node/handlers/handshake.rs @@ -9,6 +9,7 @@ use crate::NodeAddr; use crate::PeerIdentity; use crate::node::acl::PeerAclContext; use crate::node::dataplane::PeerActionCtx; +use crate::node::rate_limit::Msg1Class; use crate::node::reject::{HandshakeReject, RejectReason}; use crate::node::{Node, NodeError}; use crate::peer::ActivePeer; @@ -96,6 +97,54 @@ impl Node { } } + /// Returns true if an inbound msg1's source matches a link belonging to a + /// **promoted** peer, i.e. it is rekey/restart maintenance traffic rather + /// than a stranger's fresh handshake. + /// + /// Deliberately *not* expressed in terms of `should_admit_msg1`, and + /// deliberately not its building block, which is how the `master` lineage + /// arranges the same pair. On XX, `handle_msg1` inserts into + /// `addr_to_link` for a still-pending inbound connection before any + /// identity is known, and `initiate_connection` inserts for an outbound + /// dial in flight. A bare `addr_to_link` hit therefore does not mean + /// "established" here, and the two predicates answer different questions: + /// `should_admit_msg1` asks whether the `accept_connections` gate applies + /// (a pending outbound dial must bypass it, or the dual-init tie-breaker + /// deadlocks), while this asks whether the source is already a promoted + /// peer, which is the only safe basis for exempting traffic from + /// stranger-class metering. + /// + /// Two ways to be a promoted peer at this `(transport_id, addr)`: + /// + /// 1. `addr_to_link` maps the tuple to a link that some entry in `peers` + /// owns. Catches the peer whose registered `TransportAddr` form matches + /// the form inbound packets carry. + /// 2. A peer's `current_addr()` matches the tuple. `current_addr` is + /// seeded at promotion from the handshake's source address and updated + /// from inbound encrypted frames, so it is always numeric + /// `SocketAddr`-form; this catches the peer whose `addr_to_link` key is + /// hostname-form because `initiate_connection` populated it from a + /// hostname-bearing peer config. + /// + /// Cost: one O(1) map lookup plus one O(peers) scan covering both limbs, + /// run on every inbound msg1 including those about to be refused. The scan + /// exists only because `addr_to_link` is keyed on the *unresolved* dial + /// address; correcting that keying reduces this to O(1). + pub(in crate::node) fn is_established_link_msg1( + &self, + transport_id: crate::transport::TransportId, + remote_addr: &crate::transport::TransportAddr, + ) -> bool { + let link_at_addr = self + .addr_to_link + .get(&(transport_id, remote_addr.clone())) + .copied(); + self.peers.values().any(|p| { + Some(p.link_id()) == link_at_addr + || (p.transport_id() == Some(transport_id) && p.current_addr() == Some(remote_addr)) + }) + } + /// Returns true if an inbound msg1 should be admitted past the /// `accept_connections` gate. /// @@ -122,6 +171,18 @@ impl Node { /// /// Otherwise the transport's `accept_connections` config decides; /// absence of a registered transport admits (no gate to apply). + /// + /// Intentionally independent of `is_established_link_msg1` on this + /// branch, and not built from it. The two are deliberately allowed to + /// drift because they answer different questions: predicate 1 above + /// admits on a bare `addr_to_link` hit, which at XX also covers a + /// *pending* inbound connection and an outbound dial still in flight. + /// That breadth is required here — it is what admits the peer's inbound + /// msg1 when the larger-`NodeAddr` side has `accept_connections: false`, + /// without which the dual-init tie-breaker deadlocks — and is exactly + /// what disqualifies it as a metering classifier, since an unpromoted + /// stranger would then draw on the established-link bucket. Do not + /// collapse the two into one predicate. pub(in crate::node) fn should_admit_msg1( &self, transport_id: crate::transport::TransportId, @@ -150,22 +211,48 @@ impl Node { /// (revealing its own identity), and stores the connection in /// pending_inbound to await msg3. pub(in crate::node) async fn handle_msg1(&mut self, packet: ReceivedPacket) { - // === RATE LIMITING (before any processing) === - if !self.msg1_rate_limiter.start_handshake() { - 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 one map lookup plus an O(peers) scan, and now runs + // on every inbound msg1 including refused ones. See + // `is_established_link_msg1` for why the scan is still needed. + // + // `_slot` is an RAII guard: it releases the limiter's pending slot on + // drop, which is every return path below and the end of the function. + // The binding name matters. Renaming it to a bare `_` drops the guard + // right here instead, releasing the slot at acquire time — silently, + // with no test and no clippy lint catching the difference. Do not + // "tidy" this binding. + // + // Known coverage gap: `handle_msg1`'s slot is held across its `.await` + // points by binding alone, and an early drop is unobserved by any + // test. It is structurally unobservable at node level — `handle_msg1` + // takes `&mut self`, so no second msg1 can be in flight to notice the + // slot missing while this one awaits, and the pending count at return + // is identical either way. The `#[must_use]` on `PendingHandshake` and + // this comment are the only defences. + let class = if self.is_established_link_msg1(packet.transport_id, &packet.remote_addr) { + 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(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -175,7 +262,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)); @@ -225,7 +311,6 @@ impl Node { HandshakeReject::UnknownConnection, )); } - self.msg1_rate_limiter.complete_handshake(); return; } // Active peer on this address — allow the new handshake. @@ -276,7 +361,6 @@ impl Node { ) { Ok(m) => m, Err(e) => { - self.msg1_rate_limiter.complete_handshake(); debug!( error = %e, "Failed to process msg1" @@ -301,7 +385,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)); @@ -363,7 +446,6 @@ impl Node { .remove(&(packet.transport_id, packet.remote_addr)); let _ = self.index_allocator.free(our_index); self.remove_peer_machine(link_id); - self.msg1_rate_limiter.complete_handshake(); self.stats_mut() .record_reject(RejectReason::Handshake(HandshakeReject::BadState)); return; @@ -375,8 +457,6 @@ impl Node { // Store in pending_inbound for msg3 dispatch. self.pending_inbound .insert((packet.transport_id, our_index.as_u32()), link_id); - - self.msg1_rate_limiter.complete_handshake(); } /// 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 f1f0158..7cb8860 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -580,6 +580,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 { @@ -626,11 +652,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; @@ -784,11 +806,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 69a4b29..8638ddb 100644 --- a/src/node/tests/handshake.rs +++ b/src/node/tests/handshake.rs @@ -1461,6 +1461,181 @@ async fn test_should_admit_msg1_admits_rekey_when_addr_form_differs() { ); } +/// `is_established_link_msg1` and `should_admit_msg1` are deliberately +/// independent on this branch, and this asserts the three cases where they +/// must disagree. +/// +/// The `master` lineage defines `should_admit_msg1` as +/// `is_established_link_msg1() || accept_connections()`, which is sound at IK +/// but not at XX: here a bare `addr_to_link` hit also covers a *pending* +/// inbound connection and an outbound dial still in flight. The gate needs +/// that breadth (it is what breaks the dual-init deadlock), and the metering +/// classifier must not have it, or an unpromoted stranger draws on the +/// established-link bucket. +/// +/// Every case below registers a transport whose `accept_connections()` is +/// false. That is not incidental: with no transport registered the gate's +/// fallback admits unconditionally, and a collapsed `should_admit_msg1` would +/// still read true, so the disagreement this test exists to pin would vanish. +#[tokio::test] +async fn established_predicate_is_independent_of_should_admit_msg1() { + use crate::config::UdpConfig; + use crate::peer::ActivePeer; + use crate::transport::Link; + use crate::transport::udp::UdpTransport; + use std::time::Duration; + + let transport_id = TransportId::new(1); + + // --- Case 1: no transport registered at all. -------------------------- + // Documentation only, NOT a discriminator: under the collapsed-predicate + // break this case still passes, because the gate's no-transport fallback + // admits regardless of what the classifier says. + { + let node = make_node(); + let addr = TransportAddr::from_string("10.0.0.2:2121"); + assert!( + node.should_admit_msg1(transport_id, &addr), + "no registered transport means no gate to apply" + ); + assert!( + !node.is_established_link_msg1(transport_id, &addr), + "an address with no peer behind it is never established-link" + ); + } + + // --- Case 2: pending INBOUND link, no promoted peer. ------------------ + // This is the state `handle_msg1` leaves behind between msg1 and msg3. + { + let mut node = make_node(); + let cfg = UdpConfig { + bind_addr: Some("127.0.0.1:0".to_string()), + accept_connections: Some(false), + ..Default::default() + }; + let (tx, _rx) = packet_channel(64); + let udp = UdpTransport::new(transport_id, None, cfg, tx); + node.transports + .insert(transport_id, TransportHandle::Udp(udp)); + + let addr = TransportAddr::from_string("10.0.0.2:2121"); + let link_id = node.allocate_link_id(); + node.links.insert( + link_id, + Link::connectionless( + link_id, + transport_id, + addr.clone(), + LinkDirection::Inbound, + Duration::from_millis(100), + ), + ); + node.addr_to_link + .insert((transport_id, addr.clone()), link_id); + + assert!( + node.peers.is_empty(), + "precondition: the link is pending, nothing is promoted" + ); + assert!( + node.should_admit_msg1(transport_id, &addr), + "the gate admits on a bare addr_to_link hit" + ); + assert!( + !node.is_established_link_msg1(transport_id, &addr), + "a pending inbound handshake is a stranger's, and must keep \ + drawing on the stranger bucket for its whole lifetime" + ); + } + + // --- Case 3: pending OUTBOUND dial, no promoted peer. ----------------- + // This is the dual-init carve-out. The assertion is what stops a later + // "simplification" collapsing the two predicates. + { + let mut node = make_node(); + let cfg = UdpConfig { + outbound_only: Some(true), + ..Default::default() + }; + let (tx, _rx) = packet_channel(64); + let udp = UdpTransport::new(transport_id, None, cfg, tx); + node.transports + .insert(transport_id, TransportHandle::Udp(udp)); + + let addr = TransportAddr::from_string("10.0.0.3:2121"); + let link_id = node.allocate_link_id(); + node.links.insert( + link_id, + Link::connectionless( + link_id, + transport_id, + addr.clone(), + LinkDirection::Outbound, + Duration::from_millis(100), + ), + ); + node.addr_to_link + .insert((transport_id, addr.clone()), link_id); + + assert!( + node.should_admit_msg1(transport_id, &addr), + "an outbound dial in flight must admit the peer's inbound msg1, \ + or the dual-init tie-breaker deadlocks" + ); + assert!( + !node.is_established_link_msg1(transport_id, &addr), + "an outbound dial is not a promoted peer" + ); + } + + // --- Agreement case: a genuinely promoted peer, both addr forms. ------ + { + let mut node = make_node(); + let cfg = UdpConfig { + outbound_only: Some(true), + ..Default::default() + }; + let (tx, _rx) = packet_channel(64); + let udp = UdpTransport::new(transport_id, None, cfg, tx); + node.transports + .insert(transport_id, TransportHandle::Udp(udp)); + + let hostname_addr = TransportAddr::from_string("core-vm.example:2121"); + let link_id = node.allocate_link_id(); + node.addr_to_link + .insert((transport_id, hostname_addr.clone()), link_id); + + let peer_full = crate::Identity::generate(); + let peer_identity = PeerIdentity::from_pubkey(peer_full.pubkey()); + let peer_node_addr = *peer_identity.node_addr(); + let mut peer = ActivePeer::new(peer_identity, link_id, 1000); + let numeric_addr = TransportAddr::from_string("100.64.0.5:2121"); + peer.set_current_addr(transport_id, numeric_addr.clone()); + node.peers.insert(peer_node_addr, peer); + + // Limb 1: addr_to_link maps the hostname form to a link a peer owns. + assert!(node.should_admit_msg1(transport_id, &hostname_addr)); + assert!( + node.is_established_link_msg1(transport_id, &hostname_addr), + "hostname-keyed promoted peer must class as established-link" + ); + + // Limb 2: the numeric form matches the peer's current_addr. + assert!(node.should_admit_msg1(transport_id, &numeric_addr)); + assert!( + node.is_established_link_msg1(transport_id, &numeric_addr), + "rekey msg1 arriving in numeric form from a promoted peer must \ + class as established-link even though addr_to_link is keyed on \ + the hostname form" + ); + + // Negative: a stranger elsewhere is neither admitted nor established. + let stranger_addr = TransportAddr::from_string("198.51.100.1:2121"); + assert!(!node.should_admit_msg1(transport_id, &stranger_addr)); + assert!(!node.is_established_link_msg1(transport_id, &stranger_addr)); + } +} + // =========================================================================== // Regression: `handle_msg3` must return the msg1-allocated session index to // the allocator on the two inbound-establish arms that abandon the pending diff --git a/src/node/tests/unit.rs b/src/node/tests/unit.rs index 28ae440..4b1e98a 100644 --- a/src/node/tests/unit.rs +++ b/src/node/tests/unit.rs @@ -1,8 +1,9 @@ use super::*; +use crate::node::rate_limit::{Msg1Class, Msg1Refusal}; use crate::nostr::{BootstrapEvent, NostrRendezvous}; use crate::proto::fmp::PromotionResult; use crate::transport::udp::UdpTransport; -use crate::transport::{TransportHandle, packet_channel}; +use crate::transport::{Link, TransportHandle, packet_channel}; use std::sync::Arc; #[test] @@ -761,14 +762,20 @@ fn test_rate_limiter_initialized() { 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("a fresh limiter admits the first stranger msg1"); assert_eq!(node.msg1_rate_limiter.pending_count(), 1); - // Complete it - node.msg1_rate_limiter.complete_handshake(); + // Complete it (the guard releases the slot on drop) + drop(slot); assert_eq!(node.msg1_rate_limiter.pending_count(), 0); } @@ -2917,3 +2924,496 @@ fn test_peer_display_name_tracks_alias_change() { peer_identity.short_npub() ); } + +// =========================================================================== +// msg1 rate-limiter metering: established-link traffic vs stranger admission +// =========================================================================== + +/// Drive exactly one msg1 from `node_a` into `node_b` and stop there. +/// +/// This is the first leg of [`drive_xx_handshake`] on its own. Metering tests +/// need a *single* msg1 rather than a completed three-message handshake, and +/// it has to carry a real Noise payload: the `build_msg1`-with-garbage idiom +/// used elsewhere in this file lands in the processing-failure arm, which +/// returns before the `addr_to_link` insert that the pending-inbound state +/// depends on. +/// +/// The msg1 goes out over `node_a`'s registered transport, so `node_b` +/// observes `node_a`'s real socket address as the source. That matters: the +/// classifier keys on the source address, so a synthetic sender socket would +/// not exercise the same lookup. +async fn pump_one_msg1( + node_a: &mut Node, + node_b: &mut Node, + transport_id: TransportId, + addr_b: std::net::SocketAddr, + packet_rx_b: &mut crate::transport::PacketRx, +) { + use crate::proto::fmp::wire::build_msg1; + use tokio::time::{Duration, timeout}; + + let remote_addr_b = TransportAddr::from_string(&addr_b.to_string()); + let peer_b_identity = PeerIdentity::from_pubkey_full(node_b.identity().pubkey_full()); + + let link_id_a = node_a.allocate_link_id(); + let our_index_a = node_a.index_allocator.allocate().unwrap(); + node_a + .seed_handshake_machine( + HandshakeSeed::outbound(link_id_a, peer_b_identity, 1000) + .with_our_index(our_index_a) + .with_transport_id(transport_id) + .with_source_addr(remote_addr_b.clone()), + ) + .unwrap(); + let our_keypair_a = node_a.identity().keypair(); + let startup_epoch_a = node_a.startup_epoch(); + let noise_msg1 = node_a + .peer_machines + .get_mut(&link_id_a) + .unwrap() + .start_handshake(our_keypair_a, startup_epoch_a, 1000) + .unwrap(); + let wire_msg1 = build_msg1(our_index_a, &noise_msg1); + + node_a + .transports + .get(&transport_id) + .unwrap() + .send(&remote_addr_b, &wire_msg1) + .await + .expect("Failed to send msg1"); + + let packet_b = timeout(Duration::from_secs(1), packet_rx_b.recv()) + .await + .expect("Timeout waiting for msg1") + .expect("Channel closed"); + node_b.handle_msg1(packet_b).await; +} + +/// Drain a packet channel to empty, returning how many packets were taken, +/// and assert that none of them was handshake traffic. +/// +/// `drive_xx_handshake` does not leave the channels empty: the +/// `PromoteToActive` executor arm sends a tree announce for *both* +/// directions, after each channel's last read, so on return both channels +/// hold a queued encrypted `TreeAnnounce`. A metering test that reads a +/// channel expecting a msg2 would otherwise match that announce and pass +/// under its own break-check. +/// +/// `try_recv` cannot be used: the transport's reader task fills the channel +/// asynchronously, so an empty `try_recv` proves nothing about what is still +/// in flight. The inner assertion is what keeps this from becoming a blanket +/// swallow. +async fn drain_handshake_free(rx: &mut crate::transport::PacketRx) -> usize { + use crate::proto::fmp::wire::{Msg1Header, Msg2Header, Msg3Header}; + use tokio::time::{Duration, timeout}; + + let mut n = 0; + while let Ok(Some(p)) = timeout(Duration::from_millis(200), rx.recv()).await { + assert!( + Msg1Header::parse(&p.data).is_none() + && Msg2Header::parse(&p.data).is_none() + && Msg3Header::parse(&p.data).is_none(), + "handshake traffic left queued after drive_xx_handshake; the \ + helper's read sequence changed and this test's observation \ + point is no longer clean" + ); + n += 1; + } + n +} + +/// The acceptance criterion: the two msg1 classes are metered separately. +/// +/// With the stranger bucket drained to empty and never refilling, a msg1 +/// from a genuinely promoted peer must still be admitted and answered with +/// msg2, because it draws on the established-link bucket instead. +#[tokio::test] +async fn established_link_msg1_admitted_when_stranger_bucket_drained() { + let transport_id = TransportId::new(1); + + let mut config_b = Config::new(); + // Stranger bucket never refills. The burst must be greater than 1: the + // setup's own promotion runs a msg1 through the limiter at a moment when + // node_b holds no promoted peer at node_a's address, so that msg1 is + // correctly classed Stranger and spends a token itself. + config_b.node.rate_limit.handshake_burst = 4; + config_b.node.rate_limit.handshake_rate = 0.0; + let mut node_b = make_node_with(config_b); + let mut node_a = make_node(); + + let (addr_a, mut packet_rx_a) = register_udp_transport(&mut node_a).await; + let (addr_b, mut packet_rx_b) = register_udp_transport(&mut node_b).await; + + drive_xx_handshake( + &mut node_a, + &mut node_b, + transport_id, + addr_b, + &mut packet_rx_a, + &mut packet_rx_b, + ) + .await; + + assert_eq!( + node_b.peer_count(), + 1, + "setup precondition: node_b promoted node_a" + ); + + // Both channels hold a post-promotion tree announce. Take them out, or + // the msg2 assertion below would match one of them. + drain_handshake_free(&mut packet_rx_a).await; + drain_handshake_free(&mut packet_rx_b).await; + + // The classification precondition. Without this, a setup that promoted + // at some other address would send the second msg1 down the Stranger + // path for a reason that has nothing to do with the code under test. + let addr_a_transport = TransportAddr::from_string(&addr_a.to_string()); + assert!( + node_b.is_established_link_msg1(transport_id, &addr_a_transport), + "setup did not produce a promoted peer at node_a's source address, \ + so the second msg1 would be classed Stranger for the wrong reason" + ); + + // Drain the stranger bucket by counting, rather than assuming how many + // tokens the promotion above spent. + while let Ok(slot) = node_b + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + { + drop(slot); + } + // This probe is non-destructive: start_handshake checks pending first and + // takes the slot only after try_acquire succeeds, so a RateLimit refusal + // consumes no token and leaves no slot behind. + assert_eq!( + node_b + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .unwrap_err(), + Msg1Refusal::RateLimit, + "the established-link msg1 below must be the only thing that could answer" + ); + + // A second msg1 from the promoted peer's own socket. + pump_one_msg1( + &mut node_a, + &mut node_b, + transport_id, + addr_b, + &mut packet_rx_b, + ) + .await; + + let reply = tokio::time::timeout(std::time::Duration::from_secs(1), packet_rx_a.recv()) + .await + .expect("established-link msg1 was refused: no reply at all") + .expect("channel closed"); + assert!( + crate::proto::fmp::wire::Msg2Header::parse(&reply.data).is_some(), + "the reply to an established-link msg1 must be a msg2; the stranger \ + bucket is empty, so answering at all requires the second bucket" + ); +} + +/// The test that distinguishes this port from a verbatim copy of the +/// `master` predicate, and the one with no `master` equivalent. +/// +/// A stranger's *pending* inbound handshake leaves an `addr_to_link` entry +/// behind at XX, before any identity is known. A predicate that reads a bare +/// `addr_to_link` hit as "established" would therefore let a stranger promote +/// itself into the established-link bucket by sending one msg1 first, which +/// is exactly the population that bucket exists to protect against. +#[tokio::test] +async fn pending_inbound_stranger_msg1_stays_stranger_class() { + let transport_id = TransportId::new(1); + + let mut config_b = Config::new(); + // Exactly one stranger token, never refilled. + config_b.node.rate_limit.handshake_burst = 1; + config_b.node.rate_limit.handshake_rate = 0.0; + // The established bucket set explicitly small so the drain count below is + // exact. The rate must be finite and strictly positive or validate() + // rejects it; at 0.001/s no whole token refills over this test's life. + config_b.node.rate_limit.established_handshake_burst = Some(2); + config_b.node.rate_limit.established_handshake_rate = Some(0.001); + let mut node_b = make_node_with(config_b); + let mut node_a = make_node(); + + let (_addr_a, mut packet_rx_a) = register_udp_transport(&mut node_a).await; + let (addr_b, mut packet_rx_b) = register_udp_transport(&mut node_b).await; + + // First msg1: admitted on the single stranger token, msg2 sent, and an + // addr_to_link entry created for a still-pending inbound connection. + pump_one_msg1( + &mut node_a, + &mut node_b, + transport_id, + addr_b, + &mut packet_rx_b, + ) + .await; + + // Prove the contaminated state was actually reached, so the assertions + // below cannot pass vacuously. + assert_eq!( + node_b.peer_count(), + 0, + "XX promotes nobody at msg1: the link must still be pending" + ); + assert_eq!( + node_b.addr_to_link.len(), + 1, + "msg1 must have left an addr_to_link entry for the pending link" + ); + + // The first msg1's msg2 is on the wire; take it out of the way. + let first_reply = tokio::time::timeout(std::time::Duration::from_secs(1), packet_rx_a.recv()) + .await + .expect("first msg1 should have been admitted on the one stranger token") + .expect("channel closed"); + assert!( + crate::proto::fmp::wire::Msg2Header::parse(&first_reply.data).is_some(), + "the first msg1 must be answered with msg2" + ); + + // Second msg1 from the same source address. + pump_one_msg1( + &mut node_a, + &mut node_b, + transport_id, + addr_b, + &mut packet_rx_b, + ) + .await; + + // Discriminator 1: the established bucket is untouched. Counted by + // draining rather than read directly — established_bucket() hands back a + // shared reference and TokenBucket::tokens() needs a mutable one. + let mut drained = 0; + while let Ok(slot) = node_b + .msg1_rate_limiter + .start_handshake(Msg1Class::EstablishedLink) + { + drop(slot); + drained += 1; + } + assert_eq!( + drained, 2, + "a pending inbound stranger's msg1 must not have drawn on the \ + established bucket" + ); + + // Discriminator 2: no msg2 came back. Correct behaviour classes the + // second msg1 Stranger, finds that bucket empty, and refuses it at the + // limiter before any crypto. + let second_reply = + tokio::time::timeout(std::time::Duration::from_millis(300), packet_rx_a.recv()).await; + assert!( + second_reply.is_err(), + "the second msg1 from a still-pending stranger must be refused by \ + the stranger bucket, not answered" + ); +} + +/// Every reject arm of `handle_msg1` must release its own pending slot and +/// no one else's. A foreign slot is held for the whole test; if an arm +/// released a slot it never took, the count would fall below 1. +/// +/// Coverage gap, recorded rather than discharged: the index-allocation arm +/// and the msg2-send-failure arm of `handle_msg1` are not driven here, so +/// their slot release is unexercised. Both need the handler to get most of +/// the way through a successful msg1 and then fail on a resource the test +/// cannot withhold without contorting the setup. +#[tokio::test] +async fn msg1_reject_arms_do_not_release_another_handshakes_slot() { + use crate::config::UdpConfig; + use crate::transport::udp::UdpTransport; + + let transport_id = TransportId::new(1); + + // --- Arm: accept_connections gate ------------------------------------ + { + let mut node = make_node(); + let cfg = UdpConfig { + outbound_only: Some(true), + ..Default::default() + }; + let (tx, _rx) = packet_channel(64); + node.transports.insert( + transport_id, + TransportHandle::Udp(UdpTransport::new(transport_id, None, cfg, tx)), + ); + + let seed = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("foreign slot"); + assert_eq!(node.msg1_rate_limiter.pending_count(), 1); + + node.handle_msg1(ReceivedPacket::with_timestamp( + transport_id, + TransportAddr::from_string("198.51.100.7:2121"), + crate::proto::fmp::wire::build_msg1( + SessionIndex::new(7), + &[0u8; crate::noise::HANDSHAKE_MSG1_SIZE], + ), + 1000, + )) + .await; + + assert_eq!( + node.stats().handshake.bad_state, + 1, + "setup must actually reach the accept_connections arm" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "accept_connections arm must not release the foreign slot" + ); + drop(seed); + assert_eq!(node.msg1_rate_limiter.pending_count(), 0); + } + + // --- Arm: invalid msg1 header ---------------------------------------- + { + let mut node = make_node(); + let seed = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("foreign slot"); + + node.handle_msg1(ReceivedPacket::with_timestamp( + transport_id, + TransportAddr::from_string("198.51.100.8:2121"), + vec![0u8; 4], + 1000, + )) + .await; + + assert_eq!( + node.stats().handshake.bad_state, + 1, + "setup must actually reach the invalid-header arm" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "invalid-header arm must not release the foreign slot" + ); + drop(seed); + assert_eq!(node.msg1_rate_limiter.pending_count(), 0); + } + + // --- Arm: receive_handshake_init failure ------------------------------ + { + let mut node = make_node(); + let seed = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("foreign slot"); + + // Well-formed framing, garbage Noise payload. + node.handle_msg1(ReceivedPacket::with_timestamp( + transport_id, + TransportAddr::from_string("198.51.100.9:2121"), + crate::proto::fmp::wire::build_msg1( + SessionIndex::new(7), + &[0u8; crate::noise::HANDSHAKE_MSG1_SIZE], + ), + 1000, + )) + .await; + + assert_eq!( + node.stats().handshake.bad_state, + 1, + "setup must actually reach the Noise-processing arm" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "Noise-processing arm must not release the foreign slot" + ); + drop(seed); + assert_eq!(node.msg1_rate_limiter.pending_count(), 0); + } + + // --- Arm: duplicate msg1 with no stored msg2 -------------------------- + { + let mut node = make_node(); + let addr = TransportAddr::from_string("198.51.100.10:2121"); + + // An inbound link at this address with no control machine behind it, + // so find_stored_msg2 misses and the arm fires. + let link_id = node.allocate_link_id(); + node.links.insert( + link_id, + Link::connectionless( + link_id, + transport_id, + addr.clone(), + LinkDirection::Inbound, + Duration::from_millis(100), + ), + ); + node.addr_to_link + .insert((transport_id, addr.clone()), link_id); + + let seed = node + .msg1_rate_limiter + .start_handshake(Msg1Class::Stranger) + .expect("foreign slot"); + + node.handle_msg1(ReceivedPacket::with_timestamp( + transport_id, + addr, + crate::proto::fmp::wire::build_msg1( + SessionIndex::new(7), + &[0u8; crate::noise::HANDSHAKE_MSG1_SIZE], + ), + 1000, + )) + .await; + + assert_eq!( + node.stats().handshake.unknown_connection, + 1, + "setup must actually reach the duplicate-msg1 arm" + ); + assert_eq!( + node.msg1_rate_limiter.pending_count(), + 1, + "duplicate-msg1 arm must not release the foreign slot" + ); + 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" + ); +}