siderite_core/middleware/
rate_limit.rs

1//! Hard-bounded token buckets with incremental conservative expiry.
2
3use super::TrustedProxies;
4use super::adapt::{BoxService, Next, from_fn, impl_layer};
5use crate::error::ApiError;
6use crate::extract::Request;
7use crate::response::IntoResponse;
8use http::{HeaderValue, StatusCode, header};
9use std::collections::{HashMap, VecDeque};
10use std::net::IpAddr;
11use std::sync::{Arc, Mutex, PoisonError};
12use std::time::{Duration, Instant};
13use tower::Layer;
14
15const DEFAULT_MAX_CLIENTS: usize = 10_000;
16const EXPIRY_SCAN_BUDGET: usize = 16;
17type Key = Option<IpAddr>;
18type Verdict = Result<(), Duration>;
19
20struct Bucket {
21    tokens: f64,
22    last: Instant,
23}
24
25#[derive(Default)]
26struct Buckets {
27    entries: HashMap<Key, Bucket>,
28    expiry: VecDeque<Key>,
29}
30
31/// Token-bucket rate limiter keyed by validated client IP.
32///
33/// State is process-local with a hard client bound. At saturation new
34/// identities are denied; existing depleted buckets are never evicted.
35/// Only fully replenished buckets expire, with bounded scanning per request.
36/// Proxy headers are ignored unless explicitly trusted. Unknown peers share
37/// one bucket; counters reset on process restart, so replicas multiply limits.
38#[derive(Clone)]
39pub struct RateLimit {
40    capacity: f64,
41    refill_per_sec: f64,
42    max_clients: usize,
43    proxies: TrustedProxies,
44    buckets: Arc<Mutex<Buckets>>,
45}
46
47/// Invalid rate-limiter configuration.
48#[derive(Debug, thiserror::Error, PartialEq, Eq)]
49#[error("rate limits require positive capacity and finite positive refill")]
50pub struct RateLimitError;
51
52impl std::fmt::Debug for RateLimit {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("RateLimit")
55            .field("capacity", &self.capacity)
56            .field("refill_per_sec", &self.refill_per_sec)
57            .field("max_clients", &self.max_clients)
58            .field("proxies", &self.proxies)
59            .finish_non_exhaustive()
60    }
61}
62
63impl RateLimit {
64    /// Configure burst and refill; invalid values fail App validation.
65    ///
66    /// Args:
67    ///     burst: Positive request-token capacity.
68    ///     per_second: Finite, positive replenishment rate.
69    ///
70    /// Returns:
71    ///     Limiter; use [`Self::try_new`] to check configuration immediately.
72    pub fn new(burst: u32, per_second: f64) -> Self {
73        Self {
74            capacity: f64::from(burst),
75            refill_per_sec: per_second,
76            max_clients: DEFAULT_MAX_CLIENTS,
77            proxies: TrustedProxies::default(),
78            buckets: Arc::default(),
79        }
80    }
81
82    /// Validate configuration before constructing a limiter.
83    ///
84    /// Args:
85    ///     burst: Positive request-token capacity.
86    ///     rate: Finite, positive replenishment rate.
87    ///
88    /// Returns:
89    ///     A valid limiter or a configuration error.
90    pub fn try_new(burst: u32, rate: f64) -> Result<Self, RateLimitError> {
91        let limiter = Self::new(burst, rate);
92        if limiter.valid() {
93            Ok(limiter)
94        } else {
95            Err(RateLimitError)
96        }
97    }
98
99    /// Configure a hard client-state bound with fresh counters.
100    ///
101    /// Args:
102    ///     max: Positive tracked-client capacity; zero is invalid.
103    ///
104    /// Returns:
105    ///     A newly configured limiter; existing clones keep their counters.
106    #[must_use]
107    pub fn max_clients(mut self, max: usize) -> Self {
108        self.max_clients = max;
109        self.buckets = Arc::default();
110        self
111    }
112
113    /// Trust forwarded client IP only from configured immediate peers.
114    ///
115    /// Args:
116    ///     proxies: Policy requiring proxies to replace client headers.
117    ///
118    /// Returns:
119    ///     Limiter using validated, normalized client IP addresses.
120    #[must_use]
121    pub fn trusted_proxies(mut self, proxies: TrustedProxies) -> Self {
122        self.proxies = proxies;
123        self
124    }
125
126    /// Explicit compatibility trust for the first X-Forwarded-For hop.
127    ///
128    /// Args:
129    ///     trust: Blanket ingress trust; false ignores forwarded client IP.
130    ///
131    /// Returns:
132    ///     Limiter with the selected policy; malformed IPs are rejected.
133    #[must_use]
134    pub fn trust_forwarded_for(self, trust: bool) -> Self {
135        self.trusted_proxies(if trust {
136            TrustedProxies::dangerously_trust_all()
137        } else {
138            TrustedProxies::default()
139        })
140    }
141
142    pub(super) fn valid(&self) -> bool {
143        self.capacity > 0.0
144            && self.refill_per_sec.is_finite()
145            && self.refill_per_sec > 0.0
146            && self.max_clients > 0
147    }
148
149    fn take(&self, key: Key, now: Instant) -> Verdict {
150        let guard = self.buckets.lock();
151        let mut state = guard.unwrap_or_else(PoisonError::into_inner);
152        if !state.entries.contains_key(&key) {
153            if state.entries.len() >= self.max_clients {
154                self.expire(&mut state, now);
155            }
156            if state.entries.len() >= self.max_clients {
157                return Err(Duration::from_secs(1));
158            }
159            state.expiry.push_back(key);
160        }
161        let bucket = state.entries.entry(key).or_insert(Bucket {
162            tokens: self.capacity,
163            last: now,
164        });
165        let elapsed = now.duration_since(bucket.last).as_secs_f64();
166        let refilled = bucket.tokens + elapsed * self.refill_per_sec;
167        bucket.tokens = refilled.min(self.capacity);
168        bucket.last = now;
169        if bucket.tokens >= 1.0 {
170            bucket.tokens -= 1.0;
171            Ok(())
172        } else {
173            let wait = (1.0 - bucket.tokens) / self.refill_per_sec;
174            let maximum = Duration::from_secs(3600);
175            Err(Duration::try_from_secs_f64(wait).unwrap_or(maximum))
176        }
177    }
178
179    fn expire(&self, state: &mut Buckets, now: Instant) {
180        for _ in 0..state.expiry.len().min(EXPIRY_SCAN_BUDGET) {
181            let Some(key) = state.expiry.pop_front() else {
182                break;
183            };
184            let expired = state.entries.get(&key).is_some_and(|bucket| {
185                let elapsed = now.duration_since(bucket.last).as_secs_f64();
186                bucket.tokens + elapsed * self.refill_per_sec >= self.capacity
187            });
188            if expired {
189                state.entries.remove(&key);
190            } else {
191                state.expiry.push_back(key);
192            }
193        }
194    }
195
196    fn wrap(&self, inner: BoxService) -> BoxService {
197        let this = self.clone();
198        from_fn(move |req: Request, next: Next| {
199            let valid = this.valid();
200            let verdict = if valid {
201                this.proxies
202                    .resolve(&req)
203                    .map(|context| {
204                        let now = Instant::now();
205                        this.take(context.client_ip, now)
206                    })
207                    .map_err(|_| ())
208            } else {
209                Err(())
210            };
211            async move {
212                if !valid {
213                    let message = "invalid rate limiter configuration";
214                    return ApiError::internal(message).into_response();
215                }
216                match verdict {
217                    Err(()) => {
218                        let message = "Invalid forwarded request context.";
219                        ApiError::bad_request(message).into_response()
220                    }
221                    Ok(Ok(())) => next.run(req).await,
222                    Ok(Err(wait)) => {
223                        let status = StatusCode::TOO_MANY_REQUESTS;
224                        let message = "Rate limit exceeded.";
225                        let error = ApiError::new(status, message);
226                        let mut response = error.into_response();
227                        let round_up = u64::from(wait.subsec_nanos() > 0);
228                        let secs = wait.as_secs().saturating_add(round_up);
229                        let seconds = secs.max(1).to_string();
230                        let value = HeaderValue::from_str(&seconds);
231                        if let Ok(value) = value {
232                            let headers = response.headers_mut();
233                            headers.insert(header::RETRY_AFTER, value);
234                        }
235                        response
236                    }
237                }
238            }
239        })
240        .layer(inner)
241    }
242}
243
244impl_layer!(RateLimit);
245
246#[cfg(test)]
247#[path = "rate_limit_tests.rs"]
248mod tests;
siderite_core/middleware/rate_limit.rs

siderite_core/middleware/
rate_limit.rs
