caelix_core/
throttle.rs

1use crate::{
2    BoxFuture, Container, HttpResponse, Injectable, InternalServerErrorException, Module,
3    ModuleMetadata, ProviderDependency, RequestContext, Result, TooManyRequestsException,
4};
5use ipnet::IpNet;
6use std::{
7    cmp::Reverse,
8    collections::{BinaryHeap, HashMap, VecDeque},
9    marker::PhantomData,
10    net::IpAddr,
11    sync::{Arc, Mutex},
12    time::{Duration, Instant},
13};
14
15/// A fixed-window request policy.
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub struct ThrottlePolicy {
18    /// Maximum attempts accepted during the window.
19    pub limit: u64,
20    /// Window duration.
21    pub window: Duration,
22}
23
24impl ThrottlePolicy {
25    /// Creates a policy from a positive request limit and window length.
26    pub const fn new(limit: u64, window_seconds: u64) -> Self {
27        Self {
28            limit,
29            window: Duration::from_secs(window_seconds),
30        }
31    }
32
33    /// Validates that this policy can be represented by the runtime clock.
34    pub fn validate(self) -> Result<()> {
35        validate_policy(self)
36    }
37}
38
39/// The result of an atomic store increment.
40#[derive(Clone, Copy, Debug)]
41pub struct ThrottleStoreRecord {
42    /// Count after the increment.
43    pub count: u64,
44    /// Time remaining in the current window.
45    pub retry_after: Duration,
46}
47
48/// Atomic storage used by the throttler.
49pub trait ThrottleStore: Send + Sync + 'static {
50    /// Atomically increments `key` in its first-hit fixed window.
51    fn increment<'a>(
52        &'a self,
53        key: &'a str,
54        window: Duration,
55    ) -> BoxFuture<'a, Result<ThrottleStoreRecord>>;
56}
57
58struct MemoryBucket {
59    count: u64,
60    expires_at: Instant,
61    generation: u64,
62}
63
64#[derive(Default)]
65struct MemoryThrottleState {
66    buckets: HashMap<String, MemoryBucket>,
67    expirations: BinaryHeap<Reverse<(Instant, u64, String)>>,
68    insertion_order: VecDeque<(u64, String)>,
69    next_generation: u64,
70}
71
72/// A bounded, process-local atomic throttle store.
73pub struct MemoryThrottleStore {
74    state: Mutex<MemoryThrottleState>,
75    capacity: usize,
76}
77
78impl MemoryThrottleStore {
79    /// Creates a store with the default capacity of 100,000 active buckets.
80    pub fn new() -> Self {
81        Self::with_capacity(100_000)
82    }
83
84    /// Creates a store with a configurable active-bucket capacity.
85    pub fn with_capacity(capacity: usize) -> Self {
86        Self {
87            state: Mutex::new(MemoryThrottleState::default()),
88            capacity: capacity.max(1),
89        }
90    }
91}
92
93impl Default for MemoryThrottleStore {
94    fn default() -> Self {
95        Self::new()
96    }
97}
98
99impl ThrottleStore for MemoryThrottleStore {
100    fn increment<'a>(
101        &'a self,
102        key: &'a str,
103        window: Duration,
104    ) -> BoxFuture<'a, Result<ThrottleStoreRecord>> {
105        Box::pin(async move {
106            let now = Instant::now();
107            let mut state = self.state.lock().map_err(|_| {
108                InternalServerErrorException::new(std::io::Error::other(
109                    "throttle store lock poisoned",
110                ))
111            })?;
112            while let Some(Reverse((expires_at, generation, expired_key))) =
113                state.expirations.peek().cloned()
114            {
115                if expires_at > now {
116                    break;
117                }
118                state.expirations.pop();
119                if state.buckets.get(&expired_key).is_some_and(|bucket| {
120                    bucket.generation == generation && bucket.expires_at <= now
121                }) {
122                    state.buckets.remove(&expired_key);
123                }
124            }
125            let new_expires_at = if state.buckets.contains_key(key) {
126                None
127            } else {
128                Some(now.checked_add(window).ok_or_else(|| {
129                    crate::exception::startup_error("throttle window exceeds the clock range")
130                })?)
131            };
132            if !state.buckets.contains_key(key) && state.buckets.len() >= self.capacity {
133                while let Some((generation, oldest)) = state.insertion_order.pop_front() {
134                    if state
135                        .buckets
136                        .get(&oldest)
137                        .is_some_and(|bucket| bucket.generation == generation)
138                    {
139                        state.buckets.remove(&oldest);
140                        break;
141                    }
142                }
143            }
144            if !state.buckets.contains_key(key) {
145                let owned_key = key.to_owned();
146                let expires_at =
147                    new_expires_at.expect("new throttle bucket must have an expiration");
148                state.next_generation = state.next_generation.wrapping_add(1);
149                let generation = state.next_generation;
150                state.buckets.insert(
151                    owned_key.clone(),
152                    MemoryBucket {
153                        count: 0,
154                        expires_at,
155                        generation,
156                    },
157                );
158                state
159                    .expirations
160                    .push(Reverse((expires_at, generation, owned_key.clone())));
161                state.insertion_order.push_back((generation, owned_key));
162            }
163            if state.expirations.len() > self.capacity.saturating_mul(2) {
164                state.expirations = state
165                    .buckets
166                    .iter()
167                    .map(|(key, bucket)| {
168                        Reverse((bucket.expires_at, bucket.generation, key.clone()))
169                    })
170                    .collect();
171            }
172            if state.insertion_order.len() > self.capacity.saturating_mul(2) {
173                let mut active = state
174                    .buckets
175                    .iter()
176                    .map(|(key, bucket)| (bucket.generation, key.clone()))
177                    .collect::<Vec<_>>();
178                active.sort_unstable_by_key(|(generation, _)| *generation);
179                state.insertion_order = active.into();
180            }
181            let bucket = state
182                .buckets
183                .get_mut(key)
184                .expect("throttle bucket was just inserted");
185            bucket.count = bucket.count.saturating_add(1);
186            Ok(ThrottleStoreRecord {
187                count: bucket.count,
188                retry_after: bucket.expires_at.saturating_duration_since(now),
189            })
190        })
191    }
192}
193
194/// Derives the client identity used in throttle keys.
195pub trait ThrottleTracker: Send + Sync + 'static {
196    /// Returns a stable identity for this request.
197    fn track<'a>(&'a self, context: &'a RequestContext) -> BoxFuture<'a, Result<String>>;
198}
199
200/// Tracks clients by IP, with explicit trusted-proxy support.
201pub struct IpThrottleTracker {
202    trusted_proxies: Vec<IpNet>,
203}
204
205impl IpThrottleTracker {
206    /// Creates an IP tracker that trusts no forwarding headers.
207    pub fn new() -> Self {
208        Self {
209            trusted_proxies: Vec::new(),
210        }
211    }
212
213    /// Creates a tracker with trusted proxy IP addresses or CIDR ranges.
214    pub fn with_trusted_proxies<I, S>(proxies: I) -> Result<Self>
215    where
216        I: IntoIterator<Item = S>,
217        S: AsRef<str>,
218    {
219        let trusted_proxies = proxies
220            .into_iter()
221            .map(|value| {
222                value.as_ref().parse::<IpNet>().map_err(|error| {
223                    crate::exception::startup_error(format!(
224                        "invalid trusted proxy '{}': {error}",
225                        value.as_ref()
226                    ))
227                })
228            })
229            .collect::<Result<Vec<_>>>()?;
230        Ok(Self { trusted_proxies })
231    }
232
233    fn trusted(&self, address: IpAddr) -> bool {
234        self.trusted_proxies
235            .iter()
236            .any(|network| network.contains(&address))
237    }
238}
239
240impl Default for IpThrottleTracker {
241    fn default() -> Self {
242        Self::new()
243    }
244}
245
246impl ThrottleTracker for IpThrottleTracker {
247    fn track<'a>(&'a self, context: &'a RequestContext) -> BoxFuture<'a, Result<String>> {
248        Box::pin(async move {
249            let peer = context.peer_addr().ok_or_else(|| {
250                InternalServerErrorException::new(std::io::Error::other(
251                    "request peer address is unavailable",
252                ))
253            })?;
254            if !self.trusted(peer.ip()) {
255                return Ok(peer.ip().to_string());
256            }
257            let Some(forwarded) = context.header("x-forwarded-for") else {
258                return Ok(peer.ip().to_string());
259            };
260            let chain = forwarded
261                .split(',')
262                .map(|value| value.trim().parse::<IpAddr>())
263                .collect::<std::result::Result<Vec<_>, _>>()
264                .map_err(|_| {
265                    InternalServerErrorException::new(std::io::Error::other(
266                        "malformed X-Forwarded-For header",
267                    ))
268                })?;
269            if chain.is_empty() {
270                return Err(InternalServerErrorException::new(std::io::Error::other(
271                    "malformed X-Forwarded-For header",
272                )));
273            }
274            Ok(chain
275                .iter()
276                .rev()
277                .find(|address| !self.trusted(**address))
278                .copied()
279                .unwrap_or(peer.ip())
280                .to_string())
281        })
282    }
283}
284
285/// Runtime throttling configuration.
286#[derive(Clone)]
287pub struct ThrottleOptions {
288    /// Global policy applied to unannotated controller routes.
289    pub policy: ThrottlePolicy,
290    /// Whether rejected responses include `Retry-After`.
291    pub retry_after_header: bool,
292    /// Atomic counter storage.
293    pub store: Arc<dyn ThrottleStore>,
294    /// Client identity tracker.
295    pub tracker: Arc<dyn ThrottleTracker>,
296}
297
298impl Default for ThrottleOptions {
299    fn default() -> Self {
300        Self {
301            policy: ThrottlePolicy::new(60, 60),
302            retry_after_header: true,
303            store: Arc::new(MemoryThrottleStore::new()),
304            tracker: Arc::new(IpThrottleTracker::new()),
305        }
306    }
307}
308
309impl ThrottleOptions {
310    /// Replaces the atomic store.
311    pub fn with_store(mut self, store: Arc<dyn ThrottleStore>) -> Self {
312        self.store = store;
313        self
314    }
315
316    /// Replaces the client tracker.
317    pub fn with_tracker(mut self, tracker: Arc<dyn ThrottleTracker>) -> Self {
318        self.tracker = tracker;
319        self
320    }
321
322    /// Controls whether rejected responses carry `Retry-After`.
323    pub fn with_retry_after_header(mut self, enabled: bool) -> Self {
324        self.retry_after_header = enabled;
325        self
326    }
327}
328
329/// Supplies options for [`ThrottleModule`].
330pub trait ThrottleConfig: Send + Sync + 'static {
331    /// Modules whose exported providers are used while building options.
332    fn imports() -> Vec<crate::ModuleDef> {
333        vec![]
334    }
335
336    /// Providers that may be resolved while building throttle options.
337    fn dependencies() -> Vec<ProviderDependency> {
338        vec![]
339    }
340
341    /// Builds the application throttle options.
342    fn options(container: &Container) -> Result<ThrottleOptions>;
343}
344
345/// Default 60 requests per 60 seconds throttle configuration.
346pub struct DefaultThrottleConfig;
347
348impl ThrottleConfig for DefaultThrottleConfig {
349    fn options(_: &Container) -> Result<ThrottleOptions> {
350        Ok(ThrottleOptions::default())
351    }
352}
353
354/// The global request-throttling service.
355pub struct Throttle {
356    options: ThrottleOptions,
357}
358
359impl Throttle {
360    /// Creates a throttle service from application options.
361    pub fn new(options: ThrottleOptions) -> Result<Self> {
362        validate_policy(options.policy)?;
363        Ok(Self { options })
364    }
365
366    /// Checks and increments the quota, returning a response only when rejected.
367    pub async fn check(
368        &self,
369        context: &RequestContext,
370        method: &str,
371        route: &str,
372        policy: ThrottlePolicy,
373    ) -> Result<Option<HttpResponse>> {
374        validate_policy(policy)?;
375        let client = self.options.tracker.track(context).await?;
376        let key = format!(
377            "{}:{client}{}:{method}{}:{route}",
378            client.len(),
379            method.len(),
380            route.len()
381        );
382        let record = self.options.store.increment(&key, policy.window).await?;
383        if record.count <= policy.limit {
384            return Ok(None);
385        }
386        let mut response = crate::IntoCaelixResponse::into_response(TooManyRequestsException::new(
387            "Rate limit exceeded",
388        ));
389        if self.options.retry_after_header {
390            let seconds = (record.retry_after.as_secs()
391                + u64::from(record.retry_after.subsec_nanos() > 0))
392            .max(1);
393            response.insert_header("Retry-After", seconds.to_string());
394        }
395        Ok(Some(response))
396    }
397
398    /// Returns the configured global policy.
399    pub fn policy(&self) -> ThrottlePolicy {
400        self.options.policy
401    }
402}
403
404/// Global module enabling throttling for macro-generated controller routes.
405pub struct ThrottleModule<C = DefaultThrottleConfig>(PhantomData<C>);
406
407impl<C: ThrottleConfig> Module for ThrottleModule<C> {
408    fn register() -> ModuleMetadata {
409        let mut metadata = ModuleMetadata::global();
410        metadata.imports.extend(C::imports());
411        metadata
412            .provider_async_factory::<Throttle, _, crate::HttpException>(
413                C::dependencies(),
414                |container| async move {
415                    let options = C::options(&container)?;
416                    Throttle::new(options)
417                },
418            )
419            .export::<Throttle>()
420    }
421}
422
423fn validate_policy(policy: ThrottlePolicy) -> Result<()> {
424    if policy.limit == 0 || policy.window.is_zero() {
425        return Err(crate::exception::startup_error(
426            "throttle limit and window must be greater than zero",
427        ));
428    }
429    Instant::now().checked_add(policy.window).ok_or_else(|| {
430        crate::exception::startup_error("throttle window exceeds the clock range")
431    })?;
432    Ok(())
433}
434
435impl Injectable for Throttle {
436    fn create(_: &Container) -> BoxFuture<'_, Result<Self>> {
437        Box::pin(async { Throttle::new(ThrottleOptions::default()) })
438    }
439
440    fn dependencies() -> Vec<ProviderDependency> {
441        vec![]
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::*;
448    use std::{
449        collections::HashMap,
450        net::{Ipv4Addr, Ipv6Addr, SocketAddr},
451    };
452
453    #[tokio::test]
454    async fn increments_atomically_and_isolates_keys() {
455        let store = Arc::new(MemoryThrottleStore::new());
456        let mut tasks = Vec::new();
457        for _ in 0..100 {
458            let store = store.clone();
459            tasks.push(tokio::spawn(async move {
460                store
461                    .increment("a", Duration::from_secs(60))
462                    .await
463                    .unwrap()
464                    .count
465            }));
466        }
467        let mut counts = Vec::new();
468        for task in tasks {
469            counts.push(task.await.unwrap());
470        }
471        counts.sort_unstable();
472        assert_eq!(counts, (1..=100).collect::<Vec<_>>());
473        assert_eq!(
474            store
475                .increment("b", Duration::from_secs(60))
476                .await
477                .unwrap()
478                .count,
479            1
480        );
481    }
482
483    #[tokio::test]
484    async fn evicts_oldest_bucket_at_capacity() {
485        let store = MemoryThrottleStore::with_capacity(1);
486        store
487            .increment("old", Duration::from_secs(60))
488            .await
489            .unwrap();
490        store
491            .increment("new", Duration::from_secs(60))
492            .await
493            .unwrap();
494        assert_eq!(
495            store
496                .increment("old", Duration::from_secs(60))
497                .await
498                .unwrap()
499                .count,
500            1
501        );
502    }
503
504    #[tokio::test]
505    async fn cleanup_uses_each_buckets_own_window() {
506        let store = MemoryThrottleStore::new();
507        store
508            .increment("long", Duration::from_secs(1))
509            .await
510            .unwrap();
511        store
512            .increment("short", Duration::from_millis(10))
513            .await
514            .unwrap();
515        tokio::time::sleep(Duration::from_millis(20)).await;
516        store
517            .increment("cleanup", Duration::from_millis(10))
518            .await
519            .unwrap();
520        assert_eq!(
521            store
522                .increment("long", Duration::from_secs(1))
523                .await
524                .unwrap()
525                .count,
526            2
527        );
528    }
529
530    #[tokio::test]
531    async fn stale_generations_do_not_change_oldest_active_eviction() {
532        let store = MemoryThrottleStore::with_capacity(2);
533        store
534            .increment("a", Duration::from_millis(10))
535            .await
536            .unwrap();
537        store.increment("b", Duration::from_secs(60)).await.unwrap();
538        tokio::time::sleep(Duration::from_millis(20)).await;
539        store.increment("c", Duration::from_secs(60)).await.unwrap();
540        store.increment("a", Duration::from_secs(60)).await.unwrap();
541        assert_eq!(
542            store
543                .increment("b", Duration::from_secs(60))
544                .await
545                .unwrap()
546                .count,
547            1
548        );
549        assert_eq!(
550            store
551                .increment("a", Duration::from_secs(60))
552                .await
553                .unwrap()
554                .count,
555            2
556        );
557    }
558
559    #[tokio::test]
560    async fn auxiliary_indexes_remain_bounded_during_churn() {
561        let store = MemoryThrottleStore::with_capacity(4);
562        for index in 0..100 {
563            store
564                .increment(&format!("key-{index}"), Duration::from_secs(60))
565                .await
566                .unwrap();
567        }
568        let state = store.state.lock().unwrap();
569        assert!(state.buckets.len() <= 4);
570        assert!(state.expirations.len() <= 8);
571        assert!(state.insertion_order.len() <= 8);
572    }
573
574    #[tokio::test]
575    async fn failed_new_bucket_does_not_evict_an_active_bucket() {
576        let store = MemoryThrottleStore::with_capacity(1);
577        store
578            .increment("active", Duration::from_secs(60))
579            .await
580            .unwrap();
581        assert!(
582            store
583                .increment("overflow", Duration::from_secs(u64::MAX))
584                .await
585                .is_err()
586        );
587        assert_eq!(
588            store
589                .increment("active", Duration::from_secs(60))
590                .await
591                .unwrap()
592                .count,
593            2
594        );
595    }
596
597    #[test]
598    fn rejects_zero_programmatic_policies() {
599        let mut options = ThrottleOptions::default();
600        options.policy = ThrottlePolicy::new(0, 60);
601        assert!(Throttle::new(options).is_err());
602
603        let mut options = ThrottleOptions::default();
604        options.policy = ThrottlePolicy::new(1, 0);
605        assert!(Throttle::new(options).is_err());
606
607        let mut options = ThrottleOptions::default();
608        options.policy = ThrottlePolicy::new(1, u64::MAX);
609        assert!(Throttle::new(options).is_err());
610    }
611
612    #[tokio::test]
613    async fn ip_tracker_ignores_forwarding_from_untrusted_peers() {
614        let tracker = IpThrottleTracker::new();
615        let context = RequestContext::new(
616            "GET",
617            "/",
618            HashMap::from([("X-Forwarded-For".into(), "203.0.113.9".into())]),
619        )
620        .with_peer_addr(SocketAddr::new(Ipv4Addr::new(192, 0, 2, 4).into(), 80));
621        assert_eq!(
622            tracker.track(&context).await.unwrap(),
623            Ipv4Addr::new(192, 0, 2, 4).to_string()
624        );
625    }
626
627    #[tokio::test]
628    async fn ip_tracker_walks_trusted_proxy_chain_right_to_left() {
629        let tracker =
630            IpThrottleTracker::with_trusted_proxies(["10.0.0.0/8", "192.0.2.0/24"]).unwrap();
631        let context = RequestContext::new(
632            "GET",
633            "/",
634            HashMap::from([(
635                "X-Forwarded-For".into(),
636                "2001:db8::7, 192.0.2.8, 10.1.1.3".into(),
637            )]),
638        )
639        .with_peer_addr(SocketAddr::new(Ipv4Addr::new(10, 0, 0, 2).into(), 80));
640        assert_eq!(tracker.track(&context).await.unwrap(), "2001:db8::7");
641    }
642
643    #[tokio::test]
644    async fn ip_tracker_fails_closed_for_malformed_trusted_chain_and_missing_peer() {
645        let tracker = IpThrottleTracker::with_trusted_proxies(["::1/128"]).unwrap();
646        let malformed = RequestContext::new(
647            "GET",
648            "/",
649            HashMap::from([("X-Forwarded-For".into(), "not-an-ip".into())]),
650        )
651        .with_peer_addr(SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 80));
652        assert!(tracker.track(&malformed).await.is_err());
653        assert!(
654            tracker
655                .track(&RequestContext::new("GET", "/", HashMap::new()))
656                .await
657                .is_err()
658        );
659    }
660}
caelix_core/throttle.rs

caelix_core/
throttle.rs
