Skip to main content

rust_supervisor/ipc/security/
limits.rs

1//! Request size limit (C5) and rate limit (C6).
2//!
3//! C5: Rejects requests exceeding `max_bytes` before JSON deserialization.
4//! C6: Token bucket rate limiter, per connection, with configurable
5//!     refill rate and burst capacity.
6
7use crate::config::ipc_security::{RateLimitConfig, RequestSizeLimitConfig};
8use crate::dashboard::error::DashboardError;
9use std::time::Instant;
10
11// ---------------------------------------------------------------------------
12// C5: Request size limit
13// ---------------------------------------------------------------------------
14
15/// Checks request body size against the configured limit (C5).
16///
17/// # Arguments
18///
19/// - `actual_bytes`: Raw byte length of the request body.
20/// - `config`: Size limit configuration.
21///
22/// # Returns
23///
24/// Returns `Ok(())` when within limit, or `Err(DashboardError)` with
25/// code `request_too_large`.
26pub fn check_request_size(
27    actual_bytes: usize,
28    config: &RequestSizeLimitConfig,
29) -> Result<(), DashboardError> {
30    if !config.enabled {
31        return Ok(());
32    }
33    if actual_bytes > config.max_bytes {
34        return Err(DashboardError::request_too_large(
35            actual_bytes,
36            config.max_bytes,
37        ));
38    }
39    Ok(())
40}
41
42// ---------------------------------------------------------------------------
43// C6: Token bucket rate limiter
44// ---------------------------------------------------------------------------
45
46/// Token bucket rate limiter.
47///
48/// Maintains a token pool that refills at a constant rate and depletes by
49/// one token per request. Burst capacity limits the maximum token
50/// accumulation.
51pub struct TokenBucket {
52    /// Current token count.
53    tokens: f64,
54    /// Maximum token capacity.
55    max_tokens: f64,
56    /// Token refill rate in tokens per second.
57    refill_rate: f64,
58    /// Timestamp of last refill.
59    last_refill: Instant,
60}
61
62impl TokenBucket {
63    /// Creates a new token bucket.
64    ///
65    /// # Arguments
66    ///
67    /// - `refill_rate`: Tokens added per second.
68    /// - `burst_capacity`: Maximum tokens the bucket can hold.
69    ///
70    /// # Returns
71    ///
72    /// Returns a fully filled [`TokenBucket`].
73    pub fn new(refill_rate: f64, burst_capacity: u32) -> Self {
74        let max_tokens = burst_capacity as f64;
75        Self {
76            tokens: max_tokens,
77            max_tokens,
78            refill_rate,
79            last_refill: Instant::now(),
80        }
81    }
82
83    /// Creates a token bucket from rate limit configuration.
84    ///
85    /// # Arguments
86    ///
87    /// - `config`: Rate limit configuration.
88    ///
89    /// # Returns
90    ///
91    /// Returns a configured [`TokenBucket`].
92    pub fn from_config(config: &RateLimitConfig) -> Self {
93        Self::new(config.refill_rate, config.burst_capacity)
94    }
95
96    /// Refills tokens based on elapsed time since last refill.
97    fn refill(&mut self) {
98        let now = Instant::now();
99        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
100        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
101        self.last_refill = now;
102    }
103
104    /// Attempts to consume one token.
105    ///
106    /// # Returns
107    ///
108    /// Returns `true` if a token was available and consumed, `false` if
109    /// the bucket is empty.
110    pub fn try_consume(&mut self) -> bool {
111        self.refill();
112        if self.tokens >= 1.0 {
113            self.tokens -= 1.0;
114            true
115        } else {
116            false
117        }
118    }
119
120    /// Checks the rate limit and returns an error if exceeded (C6).
121    ///
122    /// # Arguments
123    ///
124    /// - `config`: Rate limit configuration (for enabled check).
125    ///
126    /// # Returns
127    ///
128    /// Returns `Ok(())` when allowed, or `Err(DashboardError)` with
129    /// code `rate_limit_exceeded`.
130    pub fn check_rate_limit(&mut self, config: &RateLimitConfig) -> Result<(), DashboardError> {
131        if !config.enabled {
132            return Ok(());
133        }
134        if !self.try_consume() {
135            return Err(DashboardError::rate_limit_exceeded());
136        }
137        Ok(())
138    }
139}