fips/src/node/rate_limit.rs

//! Rate Limiting for FIPS Protocol
//!
//! Provides token bucket rate limiting for protecting against DoS attacks,
//! particularly on the Noise handshake path where msg1 processing involves
//! expensive cryptographic operations.
//!
//! ## Design
//!
//! - Token bucket algorithm with configurable burst and refill rate
//! - Global rate limit (not per-source, since UDP sources are spoofable)
//! - Applied before expensive DH operations in handshake processing
//!
//! ## Default Parameters
//!
//! - 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).
pub const DEFAULT_BURST_CAPACITY: u32 = 100;

/// Default refill rate (tokens per second).
pub const DEFAULT_REFILL_RATE: f64 = 10.0;

/// Token bucket rate limiter.
///
/// Uses a classic token bucket algorithm where tokens are consumed for each
/// operation and refilled at a constant rate. When tokens are exhausted,
/// operations are rate-limited until tokens refill.
#[derive(Debug, Clone)]
pub struct TokenBucket {
    /// Maximum number of tokens (burst capacity).
    capacity: u32,
    /// Current number of available tokens (may be fractional during refill).
    tokens: f64,
    /// Tokens added per second.
    refill_rate: f64,
    /// Last time tokens were refilled.
    last_refill: Instant,
}

impl TokenBucket {
    /// Create a new token bucket with default parameters.
    ///
    /// - Burst capacity: 100 tokens
    /// - Refill rate: 10 tokens/second
    pub fn new() -> Self {
        Self::with_params(DEFAULT_BURST_CAPACITY, DEFAULT_REFILL_RATE)
    }

    /// Create a token bucket with custom parameters.
    ///
    /// # Arguments
    ///
    /// * `capacity` - Maximum number of tokens (burst capacity)
    /// * `refill_rate` - Tokens added per second
    pub fn with_params(capacity: u32, refill_rate: f64) -> Self {
        Self {
            capacity,
            tokens: capacity as f64,
            refill_rate,
            last_refill: Instant::now(),
        }
    }

    /// Try to consume one token.
    ///
    /// Returns `true` if a token was available and consumed, `false` if
    /// rate limited (no tokens available).
    pub fn try_acquire(&mut self) -> bool {
        self.try_acquire_n(1)
    }

    /// Try to consume n tokens.
    ///
    /// Returns `true` if n tokens were available and consumed, `false` if
    /// rate limited (insufficient tokens).
    pub fn try_acquire_n(&mut self, n: u32) -> bool {
        self.refill();

        if self.tokens >= n as f64 {
            self.tokens -= n as f64;
            true
        } else {
            false
        }
    }

    /// Check if tokens are available without consuming them.
    #[cfg(test)]
    pub fn available(&mut self) -> bool {
        self.refill();
        self.tokens >= 1.0
    }

    /// Get the current number of available tokens.
    #[cfg(test)]
    pub fn tokens(&mut self) -> f64 {
        self.refill();
        self.tokens
    }

    /// Get the capacity (max tokens).
    #[cfg(test)]
    pub fn capacity(&self) -> u32 {
        self.capacity
    }

    /// Refill tokens based on elapsed time.
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);
        let elapsed_secs = elapsed.as_secs_f64();

        // Add tokens based on time elapsed
        self.tokens += elapsed_secs * self.refill_rate;

        // Cap at capacity
        if self.tokens > self.capacity as f64 {
            self.tokens = self.capacity as f64;
        }

        self.last_refill = now;
    }

    /// Reset to full capacity.
    #[cfg(test)]
    pub fn reset(&mut self) {
        self.tokens = self.capacity as f64;
        self.last_refill = Instant::now();
    }

    /// Time until the next token is available.
    ///
    /// Returns `Duration::ZERO` if tokens are available, otherwise the
    /// estimated time until one token will be available.
    #[cfg(test)]
    pub fn time_until_available(&mut self) -> std::time::Duration {
        self.refill();

        if self.tokens >= 1.0 {
            std::time::Duration::ZERO
        } else {
            let needed = 1.0 - self.tokens;
            let secs = needed / self.refill_rate;
            std::time::Duration::from_secs_f64(secs)
        }
    }
}

impl Default for TokenBucket {
    fn default() -> Self {
        Self::new()
    }
}

/// 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. 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 stranger msg1.
    bucket: TokenBucket,
    /// 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, established: TokenBucket, max_pending: usize) -> Self {
        Self {
            bucket,
            established,
            pending: Arc::new(AtomicUsize::new(0)),
            max_pending,
        }
    }

    /// Check if a new handshake of `class` can be started.
    ///
    /// Returns `true` if:
    /// - 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, class: Msg1Class) -> bool {
        self.bucket_for(class).available()
            && self.pending.load(Ordering::Relaxed) < self.max_pending
    }

    /// Start a new handshake, consuming a token and taking a pending slot.
    ///
    /// 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_for(class).try_acquire() {
            return Err(Msg1Refusal::RateLimit);
        }

        self.pending.fetch_add(1, Ordering::Relaxed);
        Ok(PendingHandshake {
            pending: Arc::clone(&self.pending),
        })
    }

    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.load(Ordering::Relaxed)
    }

    /// 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.established.reset();
        self.pending.store(0, Ordering::Relaxed);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_token_bucket_basic() {
        let mut bucket = TokenBucket::with_params(10, 1.0);

        // Should have full capacity
        assert_eq!(bucket.capacity(), 10);
        assert!(bucket.tokens() >= 9.9); // Allow for timing

        // Consume all tokens
        for _ in 0..10 {
            assert!(bucket.try_acquire());
        }

        // Should be empty
        assert!(!bucket.try_acquire());
        assert!(!bucket.available());
    }

    #[test]
    fn test_token_bucket_refill() {
        let mut bucket = TokenBucket::with_params(10, 100.0); // 100 tokens/sec

        // Drain completely
        for _ in 0..10 {
            bucket.try_acquire();
        }
        assert!(!bucket.available());

        // Wait for refill, measuring actual elapsed time to avoid sensitivity
        // to OS scheduler variance (sleep can overshoot by a large margin).
        let before = Instant::now();
        thread::sleep(Duration::from_millis(50));
        let elapsed_secs = before.elapsed().as_secs_f64();

        // Expected tokens = elapsed * rate, capped at capacity.
        // Allow ±20% tolerance around the actual elapsed time.
        let expected = (elapsed_secs * 100.0).min(10.0);
        let lo = (expected * 0.8).min(expected - 0.5).max(0.0);
        let hi = (expected * 1.2).max(expected + 0.5).min(10.0);

        let tokens = bucket.tokens();
        assert!(
            (lo..=hi).contains(&tokens),
            "tokens: {}, expected ~{:.2} (range {:.2}..={:.2})",
            tokens,
            expected,
            lo,
            hi
        );
    }

    #[test]
    fn test_token_bucket_try_acquire_n() {
        let mut bucket = TokenBucket::with_params(10, 1.0);

        // Acquire 5
        assert!(bucket.try_acquire_n(5));
        assert!(bucket.tokens() >= 4.9 && bucket.tokens() <= 5.1);

        // Acquire 5 more
        assert!(bucket.try_acquire_n(5));

        // Can't acquire more
        assert!(!bucket.try_acquire_n(1));
    }

    #[test]
    fn test_token_bucket_reset() {
        let mut bucket = TokenBucket::with_params(10, 1.0);

        // Drain
        for _ in 0..10 {
            bucket.try_acquire();
        }

        // Reset
        bucket.reset();

        // Should be full again
        assert!(bucket.tokens() >= 9.9);
    }

    #[test]
    fn test_token_bucket_time_until_available() {
        let mut bucket = TokenBucket::with_params(10, 10.0); // 10 tokens/sec

        // When full, should be zero
        assert_eq!(bucket.time_until_available(), Duration::ZERO);

        // Drain completely
        for _ in 0..10 {
            bucket.try_acquire();
        }

        // Should need ~100ms for one token at 10/sec
        let wait = bucket.time_until_available();
        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 = test_limiter(TokenBucket::new(), 100);

        assert!(limiter.can_start_handshake(Msg1Class::Stranger));
        assert_eq!(limiter.pending_count(), 0);

        // Start a handshake
        let slot = limiter.start_handshake(Msg1Class::Stranger).unwrap();
        assert_eq!(limiter.pending_count(), 1);

        // Complete it
        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 = test_limiter(bucket, 3);

        // Start 3 handshakes
        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(Msg1Class::Stranger));
        assert_eq!(
            limiter.start_handshake(Msg1Class::Stranger).unwrap_err(),
            Msg1Refusal::PendingLimit
        );

        // Complete one
        drop(a);

        // Now should be able to start another
        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 = test_limiter(bucket, 100);

        // Start 5 handshakes (exhausts tokens), releasing each immediately
        for _ in 0..5 {
            let slot = limiter.start_handshake(Msg1Class::Stranger).unwrap();
            drop(slot);
        }

        // Tokens exhausted, even though pending is 0
        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 = test_limiter(TokenBucket::new(), 100);

        // Start some handshakes
        let _a = limiter.start_handshake(Msg1Class::Stranger).unwrap();
        let _b = limiter.start_handshake(Msg1Class::Stranger).unwrap();
        assert_eq!(limiter.pending_count(), 2);

        // Reset
        limiter.reset();

        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);
    }
}