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 `limited` 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;
15use crate::plugins::{Plugin, 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 and a `Retry-After: 1` header, exiting through the `limited`
24/// port. Does 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(&self, mut ctx: Context) -> PluginResult {
134        let key = match &self.key_source {
135            KeySource::RemoteAddr => ctx.request.remote_addr.clone(),
136            KeySource::Header(header) => ctx
137                .request
138                .headers
139                .get(header)
140                .and_then(|v| v.first())
141                .cloned()
142                .unwrap_or_else(|| ctx.request.remote_addr.clone()),
143        };
144
145        let allowed = {
146            let mut bucket = self.buckets.entry(key).or_insert_with(|| {
147                TokenBucket::new(self.burst as f64, self.requests_per_second as f64)
148            });
149            bucket.try_consume()
150        };
151
152        if allowed {
153            Ok(PluginOutput::success(ctx))
154        } else {
155            ctx.response.status_code = 429;
156            ctx.response.body =
157                Bytes::from(r#"{"error": "rate_limited", "message": "Too many requests"}"#);
158            ctx.response.headers.insert(
159                "content-type".to_string(),
160                vec!["application/json".to_string()],
161            );
162            ctx.response
163                .headers
164                .insert("retry-after".to_string(), vec!["1".to_string()]);
165
166            Ok(PluginOutput::on_port(ctx, "limited"))
167        }
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    //! Behavioral tests inspired by Apache APISIX's `t/plugin/limit-req.t`
174    //! (token/leaky bucket rate limiting), adapted to featherbit's token bucket.
175    //!
176    //! The bucket starts full (`burst` tokens) and refills at `requests_per_second`
177    //! per second. Over the sub-millisecond span of a test, refill is negligible
178    //! (< 1 token), so "fire `burst` requests then one more" is deterministic.
179    use super::*;
180    use crate::context::{GatewayRequest, GatewayResponse, Protocol};
181
182    fn ctx(remote_addr: &str, header: Option<(&str, &str)>) -> Context {
183        let mut headers = HashMap::new();
184        if let Some((k, v)) = header {
185            headers.insert(k.to_string(), vec![v.to_string()]);
186        }
187        Context {
188            request: GatewayRequest {
189                method: "GET".to_string(),
190                path: "/hello".to_string(),
191                host: "h".to_string(),
192                scheme: "http".to_string(),
193                headers,
194                query_params: HashMap::new(),
195                body: Bytes::new(),
196                remote_addr: remote_addr.to_string(),
197                protocol: Protocol::Http1,
198            },
199            response: GatewayResponse {
200                status_code: 0,
201                headers: HashMap::new(),
202                body: Bytes::new(),
203                stream: None,
204            },
205            message: HashMap::new(),
206            errors: Vec::new(),
207        }
208    }
209
210    fn plugin(config: serde_json::Value) -> RateLimitPlugin {
211        let map: HashMap<String, serde_json::Value> =
212            config.as_object().unwrap().clone().into_iter().collect();
213        RateLimitPlugin::from_config(&map).unwrap()
214    }
215
216    /// Requests within the burst allowance pass; the next is throttled with 429.
217    #[tokio::test]
218    async fn test_burst_then_throttle() {
219        let p = plugin(serde_json::json!({ "requests_per_second": 1, "burst": 2 }));
220        // Two tokens available -> two passes.
221        assert!(p
222            .execute(ctx("1.2.3.4:5", None))
223            .await
224            .unwrap()
225            .port
226            .is_none());
227        assert!(p
228            .execute(ctx("1.2.3.4:5", None))
229            .await
230            .unwrap()
231            .port
232            .is_none());
233        // Third exhausts the bucket -> 429 on `limited`.
234        let out = p.execute(ctx("1.2.3.4:5", None)).await.unwrap();
235        assert_eq!(out.port, Some("limited"));
236        assert_eq!(out.context.response.status_code, 429);
237        assert_eq!(
238            out.context
239                .response
240                .headers
241                .get("retry-after")
242                .and_then(|v| v.first())
243                .map(String::as_str),
244            Some("1")
245        );
246    }
247
248    /// Each client key gets its own bucket: one client's exhaustion does not
249    /// throttle another.
250    #[tokio::test]
251    async fn test_buckets_are_per_key() {
252        let p = plugin(serde_json::json!({ "requests_per_second": 1, "burst": 1 }));
253        assert!(p
254            .execute(ctx("10.0.0.1:5", None))
255            .await
256            .unwrap()
257            .port
258            .is_none());
259        // Same client is now throttled...
260        assert_eq!(
261            p.execute(ctx("10.0.0.1:5", None)).await.unwrap().port,
262            Some("limited")
263        );
264        // ...but a different client still has a full bucket.
265        assert!(p
266            .execute(ctx("10.0.0.2:5", None))
267            .await
268            .unwrap()
269            .port
270            .is_none());
271    }
272
273    /// `key_from: header:<name>` keys on a request header instead of the address.
274    #[tokio::test]
275    async fn test_key_from_header() {
276        let p = plugin(serde_json::json!({
277            "requests_per_second": 1, "burst": 1, "key_from": "header:x-api-key"
278        }));
279        // Two different remote addrs but the SAME api key share one bucket.
280        assert!(p
281            .execute(ctx("10.0.0.1:5", Some(("x-api-key", "k1"))))
282            .await
283            .unwrap()
284            .port
285            .is_none());
286        assert_eq!(
287            p.execute(ctx("10.0.0.2:5", Some(("x-api-key", "k1"))))
288                .await
289                .unwrap()
290                .port,
291            Some("limited")
292        );
293        // A different key has its own bucket.
294        assert!(p
295            .execute(ctx("10.0.0.3:5", Some(("x-api-key", "k2"))))
296            .await
297            .unwrap()
298            .port
299            .is_none());
300    }
301
302    #[test]
303    fn test_burst_defaults_to_rps() {
304        let p = plugin(serde_json::json!({ "requests_per_second": 7 }));
305        assert_eq!(p.burst, 7);
306    }
307}
featherbit/plugins/native/rate_limit.rs

featherbit/plugins/native/
rate_limit.rs
