featherbit/plugins/native/
rate_limit.rs

1//! Token-bucket rate limiting plugin (`rate-limit`).
2//!
3//! Maintains one in-memory token bucket per client key (remote address or a
4//! configured header) in a concurrent `DashMap`; requests that find the
5//! bucket empty are rejected with 429 through the node's error port.
6
7use async_trait::async_trait;
8use bytes::Bytes;
9use dashmap::DashMap;
10use std::collections::HashMap;
11use std::sync::Arc;
12use std::time::Instant;
13
14use crate::context::{Context, GatewayError};
15use crate::plugins::{Plugin, PluginExecutionError, PluginOutput, PluginResult};
16
17/// Enforces a per-client request rate using the token bucket algorithm.
18///
19/// Each client key gets its own bucket sized `burst` and refilled at
20/// `requests_per_second`; the key is the remote address by default or the
21/// value of a configured header (falling back to the remote address when the
22/// header is absent). Over-limit requests are rejected with a 429 JSON
23/// response, a `Retry-After: 1` header, and error code `RATE_LIMITED`. Does
24/// not write to `context.message`.
25pub struct RateLimitPlugin {
26    /// Sustained refill rate: tokens added per second.
27    requests_per_second: u64,
28    /// Bucket capacity: maximum requests allowed in a burst.
29    burst: u64,
30    /// Where the per-client key is taken from.
31    key_source: KeySource,
32    /// Per-key token buckets, created lazily on first request.
33    buckets: Arc<DashMap<String, TokenBucket>>,
34}
35
36/// Source of the rate-limiting key for a request.
37#[derive(Debug, Clone)]
38enum KeySource {
39    /// Key on the client's remote address (the default).
40    RemoteAddr,
41    /// Key on the first value of the named request header.
42    Header(String),
43}
44
45/// Classic token bucket: refilled continuously based on elapsed time, capped
46/// at `max_tokens`; each request consumes one token.
47struct TokenBucket {
48    tokens: f64,
49    last_refill: Instant,
50    max_tokens: f64,
51    refill_rate: f64,
52}
53
54impl TokenBucket {
55    fn new(max_tokens: f64, refill_rate: f64) -> Self {
56        Self {
57            tokens: max_tokens,
58            last_refill: Instant::now(),
59            max_tokens,
60            refill_rate,
61        }
62    }
63
64    /// Refills the bucket for the elapsed time, then attempts to take one
65    /// token; returns false when the request should be rate-limited.
66    fn try_consume(&mut self) -> bool {
67        let now = Instant::now();
68        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
69        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
70        self.last_refill = now;
71
72        if self.tokens >= 1.0 {
73            self.tokens -= 1.0;
74            true
75        } else {
76            false
77        }
78    }
79}
80
81impl RateLimitPlugin {
82    /// Builds the plugin from node config. Never fails; every key has a
83    /// default.
84    ///
85    /// Accepted keys:
86    /// - `requests_per_second` (integer, default `100`): sustained rate per
87    ///   client key.
88    /// - `burst` (integer, default = `requests_per_second`): bucket capacity.
89    /// - `key_from` (string, default remote address): use
90    ///   `"header:<name>"` to key on a request header instead; any other
91    ///   value keys on the remote address.
92    ///
93    /// ```yaml
94    /// type: rate-limit
95    /// config:
96    ///   requests_per_second: 10
97    ///   burst: 20
98    ///   key_from: "header:x-api-key"
99    /// ```
100    pub fn from_config(config: &HashMap<String, serde_json::Value>) -> Result<Self, String> {
101        let requests_per_second = config
102            .get("requests_per_second")
103            .and_then(|v| v.as_u64())
104            .unwrap_or(100);
105
106        let burst = config
107            .get("burst")
108            .and_then(|v| v.as_u64())
109            .unwrap_or(requests_per_second);
110
111        let key_source = match config.get("key_from").and_then(|v| v.as_str()) {
112            Some(header) if header.starts_with("header:") => {
113                KeySource::Header(header[7..].to_string())
114            }
115            _ => KeySource::RemoteAddr,
116        };
117
118        Ok(Self {
119            requests_per_second,
120            burst,
121            key_source,
122            buckets: Arc::new(DashMap::new()),
123        })
124    }
125}
126
127#[async_trait]
128impl Plugin for RateLimitPlugin {
129    fn plugin_type(&self) -> &str {
130        "rate-limit"
131    }
132
133    async fn execute(
134        &self,
135        mut ctx: Context,
136        _named_inputs: &HashMap<String, serde_json::Value>,
137    ) -> PluginResult {
138        let key = match &self.key_source {
139            KeySource::RemoteAddr => ctx.request.remote_addr.clone(),
140            KeySource::Header(header) => ctx
141                .request
142                .headers
143                .get(header)
144                .and_then(|v| v.first())
145                .cloned()
146                .unwrap_or_else(|| ctx.request.remote_addr.clone()),
147        };
148
149        let allowed = {
150            let mut bucket = self.buckets.entry(key).or_insert_with(|| {
151                TokenBucket::new(self.burst as f64, self.requests_per_second as f64)
152            });
153            bucket.try_consume()
154        };
155
156        if allowed {
157            Ok(PluginOutput {
158                context: ctx,
159                named_outputs: HashMap::new(),
160            })
161        } else {
162            ctx.response.status_code = 429;
163            ctx.response.body =
164                Bytes::from(r#"{"error": "rate_limited", "message": "Too many requests"}"#);
165            ctx.response.headers.insert(
166                "content-type".to_string(),
167                vec!["application/json".to_string()],
168            );
169            ctx.response
170                .headers
171                .insert("retry-after".to_string(), vec!["1".to_string()]);
172
173            let error = GatewayError {
174                node_id: String::new(),
175                code: "RATE_LIMITED".to_string(),
176                message: "Too many requests".to_string(),
177                metadata: HashMap::new(),
178            };
179            Err(PluginExecutionError {
180                context: ctx,
181                error,
182            })
183        }
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    //! Behavioral tests inspired by Apache APISIX's `t/plugin/limit-req.t`
190    //! (token/leaky bucket rate limiting), adapted to featherbit's token bucket.
191    //!
192    //! The bucket starts full (`burst` tokens) and refills at `requests_per_second`
193    //! per second. Over the sub-millisecond span of a test, refill is negligible
194    //! (< 1 token), so "fire `burst` requests then one more" is deterministic.
195    use super::*;
196    use crate::context::{GatewayRequest, GatewayResponse, Protocol};
197
198    fn ctx(remote_addr: &str, header: Option<(&str, &str)>) -> Context {
199        let mut headers = HashMap::new();
200        if let Some((k, v)) = header {
201            headers.insert(k.to_string(), vec![v.to_string()]);
202        }
203        Context {
204            request: GatewayRequest {
205                method: "GET".to_string(),
206                path: "/hello".to_string(),
207                host: "h".to_string(),
208                scheme: "http".to_string(),
209                headers,
210                query_params: HashMap::new(),
211                body: Bytes::new(),
212                remote_addr: remote_addr.to_string(),
213                protocol: Protocol::Http1,
214            },
215            response: GatewayResponse {
216                status_code: 0,
217                headers: HashMap::new(),
218                body: Bytes::new(),
219            },
220            message: HashMap::new(),
221            errors: Vec::new(),
222        }
223    }
224
225    fn plugin(config: serde_json::Value) -> RateLimitPlugin {
226        let map: HashMap<String, serde_json::Value> =
227            config.as_object().unwrap().clone().into_iter().collect();
228        RateLimitPlugin::from_config(&map).unwrap()
229    }
230
231    /// Requests within the burst allowance pass; the next is throttled with 429.
232    #[tokio::test]
233    async fn test_burst_then_throttle() {
234        let p = plugin(serde_json::json!({ "requests_per_second": 1, "burst": 2 }));
235        // Two tokens available -> two passes.
236        assert!(p
237            .execute(ctx("1.2.3.4:5", None), &HashMap::new())
238            .await
239            .is_ok());
240        assert!(p
241            .execute(ctx("1.2.3.4:5", None), &HashMap::new())
242            .await
243            .is_ok());
244        // Third exhausts the bucket -> 429.
245        let err = p
246            .execute(ctx("1.2.3.4:5", None), &HashMap::new())
247            .await
248            .unwrap_err();
249        assert_eq!(err.error.code, "RATE_LIMITED");
250        assert_eq!(err.context.response.status_code, 429);
251        assert_eq!(
252            err.context
253                .response
254                .headers
255                .get("retry-after")
256                .and_then(|v| v.first())
257                .map(String::as_str),
258            Some("1")
259        );
260    }
261
262    /// Each client key gets its own bucket: one client's exhaustion does not
263    /// throttle another.
264    #[tokio::test]
265    async fn test_buckets_are_per_key() {
266        let p = plugin(serde_json::json!({ "requests_per_second": 1, "burst": 1 }));
267        assert!(p
268            .execute(ctx("10.0.0.1:5", None), &HashMap::new())
269            .await
270            .is_ok());
271        // Same client is now throttled...
272        assert!(p
273            .execute(ctx("10.0.0.1:5", None), &HashMap::new())
274            .await
275            .is_err());
276        // ...but a different client still has a full bucket.
277        assert!(p
278            .execute(ctx("10.0.0.2:5", None), &HashMap::new())
279            .await
280            .is_ok());
281    }
282
283    /// `key_from: header:<name>` keys on a request header instead of the address.
284    #[tokio::test]
285    async fn test_key_from_header() {
286        let p = plugin(serde_json::json!({
287            "requests_per_second": 1, "burst": 1, "key_from": "header:x-api-key"
288        }));
289        // Two different remote addrs but the SAME api key share one bucket.
290        assert!(p
291            .execute(
292                ctx("10.0.0.1:5", Some(("x-api-key", "k1"))),
293                &HashMap::new()
294            )
295            .await
296            .is_ok());
297        assert!(p
298            .execute(
299                ctx("10.0.0.2:5", Some(("x-api-key", "k1"))),
300                &HashMap::new()
301            )
302            .await
303            .is_err());
304        // A different key has its own bucket.
305        assert!(p
306            .execute(
307                ctx("10.0.0.3:5", Some(("x-api-key", "k2"))),
308                &HashMap::new()
309            )
310            .await
311            .is_ok());
312    }
313
314    #[test]
315    fn test_burst_defaults_to_rps() {
316        let p = plugin(serde_json::json!({ "requests_per_second": 7 }));
317        assert_eq!(p.burst, 7);
318    }
319}
featherbit/plugins/native/rate_limit.rs

featherbit/plugins/native/
rate_limit.rs
