mirror of
https://github.com/jmcorgan/fips.git
synced 2026-08-09 16:24:45 +00:00
Meter established-link msg1 separately from stranger admission
The msg1 rate limiter ran before the established-peer carve-out, so a rekey or restart msg1 arriving on an existing link was refused on exactly the same terms as a stranger's first packet and the carve-out below it never applied to the traffic it was written for. On a node with many peers this refuses a large share of ordinary maintenance 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. Classify the source before metering it, and give established-link msg1 its own token bucket instead of a bypass. A bypass was rejected deliberately: the limiter is global precisely because UDP sources are spoofable, and an established-peer exemption is by construction keyed on source address, so metering the exempted class keeps that property where bypassing discards it. The bucket is derived from settings the operator already sets, burst from max_peers and rate from max_peers, the rekey interval and the resend budget, so raising the peer limit sizes it automatically rather than leaving a constant nobody revisits. Both parameters can be set explicitly; an explicit zero burst or non-positive rate is rejected at config validation, because it would refuse every rekey msg1 from an established peer rather than disabling the limit. Split out the established-link test as its own predicate so the rate-limit classifier and the accept_connections gate cannot drift, and convert the limiter's pending slot to a guard released on drop. The slot was previously acquired in one place and released explicitly at eighteen exit paths; once some paths stop acquiring one, any path that still released one would have freed a slot belonging to a different in-flight handshake, lifting effective concurrency above the configured maximum with no counter moving and no log firing. The "Msg1 rate limited" line now reports which limb refused, the pending count or the token bucket, which it did not distinguish before. Classification costs an O(peers) scan on every inbound msg1 including refused ones, where the previous order refused at O(1). The scan is only needed because addr_to_link is keyed on the unresolved dial address; correcting that keying reduces this to a single O(1) lookup. Recorded at the predicate.
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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<u32>,
|
||||
/// 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<f64>,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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).
|
||||
|
||||
+28
-10
@@ -611,6 +611,32 @@ pub struct Node {
|
||||
tokio::sync::mpsc::UnboundedSender<decrypt_worker::DecryptWorkerEvent>,
|
||||
}
|
||||
|
||||
/// 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<Self, NodeError> {
|
||||
@@ -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;
|
||||
|
||||
+364
-55
@@ -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<Cell<usize>>` 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<AtomicUsize>,
|
||||
}
|
||||
|
||||
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<AtomicUsize>,
|
||||
/// 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<PendingHandshake, Msg1Refusal> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
+361
-5
@@ -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<u8>| 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<usize> {
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user