waf_detection/
rate_limit.rs

1// SPDX-FileCopyrightText: 2026 0x00spor3
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use tracing::warn;
7use waf_core::{
8    BucketParams, Clock, Config, Decision, Phase, RateLimitAction, RateLimitKey, RateLimitState,
9    RequestContext, WafModule,
10};
11
12// The token-bucket store, the `StateStore` seam, and the clock live in `waf-core`
13// (`waf_core::state`) so the enterprise can inject a distributed store without a
14// fork. This module is just the `WafModule` that drives that store from config.
15
16pub struct RateLimitModule {
17    enabled: bool,
18    params: BucketParams,
19    action: RateLimitAction,
20    score: u32,
21    key: RateLimitKey,
22    state: RateLimitState,
23}
24
25impl Default for RateLimitModule {
26    fn default() -> Self {
27        Self::with_state(RateLimitState::new())
28    }
29}
30
31impl RateLimitModule {
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Construct with a custom clock (used by tests): an in-memory store driven by
37    /// that clock, default tracked-key cap.
38    pub fn with_clock(clock: Arc<dyn Clock>) -> Self {
39        Self::with_state(RateLimitState::in_memory_with_clock(clock, DEFAULT_MAX_TRACKED_KEYS))
40    }
41
42    /// Construct with a **shared** store, reinjected across reloads so the throttle
43    /// state persists when only config parameters change.
44    pub fn with_state(state: RateLimitState) -> Self {
45        Self {
46            enabled: false,
47            params: BucketParams { capacity: 0.0, refill_per_sec: 0.0 },
48            action: RateLimitAction::Block,
49            score: 0,
50            key: RateLimitKey::ClientIp,
51            state,
52        }
53    }
54
55    fn key_for(&self, ctx: &RequestContext) -> String {
56        match self.key {
57            // Peer socket address. Behind an LB/CDN this is the proxy IP — see
58            // ARCHITECTURE §8 (X-Forwarded-For + trusted hops is the extension).
59            RateLimitKey::ClientIp => ctx.client_ip.to_string(),
60        }
61    }
62
63    /// Seconds until at least one token is available again (>= 1).
64    fn retry_after(&self, tokens: f64) -> u64 {
65        if self.params.refill_per_sec <= 0.0 {
66            return 1;
67        }
68        ((1.0 - tokens) / self.params.refill_per_sec).ceil().max(1.0) as u64
69    }
70}
71
72/// Default tracked-key cap used when constructing a test store via `with_clock`.
73/// Mirrors `RateLimitConfig`'s default; the authoritative cap reaching production
74/// comes from config through `RateLimitState::in_memory` at bind time.
75const DEFAULT_MAX_TRACKED_KEYS: usize = 100_000;
76
77impl WafModule for RateLimitModule {
78    fn id(&self) -> &str {
79        "rate_limit"
80    }
81
82    fn phase(&self) -> Phase {
83        Phase::Connection
84    }
85
86    fn init(&mut self, cfg: &Config) {
87        let rl = &cfg.rate_limit;
88        self.enabled = rl.enabled;
89        let window = rl.window_seconds.max(1) as f64;
90        self.params = BucketParams {
91            capacity: rl.burst.unwrap_or(rl.requests).max(1) as f64,
92            refill_per_sec: rl.requests as f64 / window,
93        };
94        self.action = rl.action;
95        self.score = rl.score;
96        self.key = rl.key;
97    }
98
99    fn inspect(&self, ctx: &RequestContext) -> Decision {
100        if !self.enabled {
101            return Decision::Allow;
102        }
103
104        let key = self.key_for(ctx);
105
106        // One atomic refill-then-consume against the (possibly distributed) store.
107        let outcome = self.state.try_acquire(&key, 1.0, self.params);
108        if outcome.allowed {
109            return Decision::Allow;
110        }
111        let tokens_after = outcome.tokens_remaining;
112
113        // Over budget.
114        match self.action {
115            RateLimitAction::Block => {
116                let retry = self.retry_after(tokens_after);
117                warn!(
118                    request_id = %ctx.request_id,
119                    key = %key,
120                    retry_after = retry,
121                    "rate limit exceeded (block)"
122                );
123                Decision::Reject {
124                    rule_id: "rate-limit".to_string(),
125                    reason: "rate limit exceeded".to_string(),
126                    status: 429,
127                    retry_after: Some(retry),
128                }
129            }
130            RateLimitAction::Score => {
131                warn!(
132                    request_id = %ctx.request_id,
133                    key = %key,
134                    points = self.score,
135                    "rate limit exceeded (score)"
136                );
137                Decision::Score {
138                    rule_id: "rate-limit-exceeded".to_string(),
139                    points: self.score,
140                }
141            }
142        }
143    }
144}
waf_detection/rate_limit.rs

waf_detection/
rate_limit.rs
