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}