featherbit/plugins/native/
limit_count.rs

1//! Fixed-window request-count limiting plugin (`limit-count`).
2//!
3//! The featherbit port of Apache APISIX's `limit-count`: counts requests per
4//! resolved key within a fixed time window using a shared [`CounterStore`]
5//! backend (see [`crate::ratelimit`]). Requests that exceed `count` within
6//! `time_window` seconds are rejected through the node's error port with the
7//! configured status. Unlike the token-bucket `rate-limit` plugin (smooth
8//! refill), this enforces a hard cap per discrete window, matching APISIX
9//! semantics.
10
11use async_trait::async_trait;
12use bytes::Bytes;
13use std::collections::HashMap;
14use std::sync::Arc;
15use std::time::Duration;
16
17use crate::context::{Context, GatewayError};
18use crate::plugins::resources::PluginResources;
19use crate::plugins::{Plugin, PluginExecutionError, PluginOutput, PluginResult};
20use crate::ratelimit::CounterStore;
21use crate::vars::interpolate;
22
23/// Enforces a per-key request count within a fixed time window.
24///
25/// On each request the resolved key is counted against a [`CounterStore`]
26/// (currently only the in-memory `local` backend). Within the limit the
27/// request passes through the `success` port; over the limit it is rejected
28/// with `rejected_code` and a JSON `{"error_msg": ...}` body through the
29/// `error` port (code `RATE_LIMITED`). When `show_limit_quota_header` is set
30/// the `X-RateLimit-Limit`/`-Remaining`/`-Reset` headers are written onto
31/// `context.response` in both cases.
32pub struct LimitCountPlugin {
33    /// Maximum number of requests allowed per window.
34    count: u64,
35    /// Length of the fixed window.
36    window: Duration,
37    /// Key template (`$var` interpolated); empty result falls back to the
38    /// client remote address.
39    key_template: String,
40    /// Optional counter-key prefix so multiple nodes share one counter.
41    group: Option<String>,
42    /// Status returned when a request is rejected.
43    rejected_code: u16,
44    /// Optional custom message used in the rejection body.
45    rejected_msg: Option<String>,
46    /// Whether to emit the `X-RateLimit-*` quota headers.
47    show_limit_quota_header: bool,
48    /// When true, a counter-backend error lets the request through instead of
49    /// failing it.
50    allow_degradation: bool,
51    /// Resolved counter backend (from `policy`).
52    store: Arc<dyn CounterStore>,
53}
54
55impl LimitCountPlugin {
56    /// Builds the plugin from node config.
57    ///
58    /// Accepted keys:
59    /// - `count` (integer > 0, **required**): requests allowed per window.
60    /// - `time_window` (integer > 0, **required**): window length in seconds.
61    /// - `key` (string, default `"$remote_addr"`): a `$var` template resolved
62    ///   per request (e.g. `$remote_addr`, `$consumer_name`,
63    ///   `$http_x_api_key`). An empty resolved value falls back to the client
64    ///   remote address.
65    /// - `policy` (string, default `"local"`): counter backend. Only `local`
66    ///   is available; `redis`/others are rejected at config load.
67    /// - `group` (string, optional): prefixes the counter key so multiple
68    ///   nodes share one counter.
69    /// - `rejected_code` (integer 200-599, default `503`): status for
70    ///   over-limit requests.
71    /// - `rejected_msg` (string, optional): message placed in the rejection
72    ///   body (`{"error_msg": ...}`).
73    /// - `show_limit_quota_header` (bool, default `true`): emit the
74    ///   `X-RateLimit-*` headers onto the response.
75    /// - `allow_degradation` (bool, default `false`): on a counter-backend
76    ///   error, allow the request through instead of failing it.
77    ///
78    /// ```yaml
79    /// type: limit-count
80    /// config:
81    ///   count: 100
82    ///   time_window: 60
83    ///   key: "$remote_addr"
84    ///   policy: local
85    ///   rejected_code: 429
86    ///   show_limit_quota_header: true
87    /// ```
88    pub fn from_config(
89        config: &HashMap<String, serde_json::Value>,
90        resources: &Arc<PluginResources>,
91    ) -> Result<Self, String> {
92        let count = config
93            .get("count")
94            .and_then(|v| v.as_u64())
95            .filter(|n| *n > 0)
96            .ok_or("limit-count requires 'count' as an integer greater than 0")?;
97
98        let time_window = config
99            .get("time_window")
100            .and_then(|v| v.as_u64())
101            .filter(|n| *n > 0)
102            .ok_or("limit-count requires 'time_window' as an integer greater than 0 (seconds)")?;
103
104        let key_template = config
105            .get("key")
106            .and_then(|v| v.as_str())
107            .filter(|s| !s.is_empty())
108            .unwrap_or("$remote_addr")
109            .to_string();
110
111        let policy = config
112            .get("policy")
113            .and_then(|v| v.as_str())
114            .unwrap_or("local");
115        // Resolves the backend (and surfaces the supported-policy list on an
116        // unknown/not-yet-implemented policy such as `redis`).
117        let store = resources.counters.get(policy)?;
118
119        let group = config
120            .get("group")
121            .and_then(|v| v.as_str())
122            .filter(|s| !s.is_empty())
123            .map(String::from);
124
125        let rejected_code = config
126            .get("rejected_code")
127            .and_then(|v| v.as_u64())
128            .filter(|n| (200..=599).contains(n))
129            .map(|n| n as u16)
130            .unwrap_or(503);
131
132        let rejected_msg = config
133            .get("rejected_msg")
134            .and_then(|v| v.as_str())
135            .filter(|s| !s.is_empty())
136            .map(String::from);
137
138        let show_limit_quota_header = config
139            .get("show_limit_quota_header")
140            .and_then(|v| v.as_bool())
141            .unwrap_or(true);
142
143        let allow_degradation = config
144            .get("allow_degradation")
145            .and_then(|v| v.as_bool())
146            .unwrap_or(false);
147
148        Ok(Self {
149            count,
150            window: Duration::from_secs(time_window),
151            key_template,
152            group,
153            rejected_code,
154            rejected_msg,
155            show_limit_quota_header,
156            allow_degradation,
157            store,
158        })
159    }
160
161    /// Resolves the per-request counter key: interpolates the `key` template
162    /// and, when it resolves to empty, falls back to the remote address
163    /// (matching APISIX). The `group` prefix, when set, is prepended so nodes
164    /// in the same group share one counter.
165    fn resolve_key(&self, ctx: &Context) -> String {
166        let mut key = interpolate(ctx, &self.key_template);
167        if key.is_empty() {
168            key = interpolate(ctx, "$remote_addr");
169        }
170        match &self.group {
171            Some(group) => format!("{}:{}", group, key),
172            None => key,
173        }
174    }
175
176    /// Writes the `X-RateLimit-*` quota headers onto the response.
177    fn set_quota_headers(&self, ctx: &mut Context, remaining: u64, reset: Duration) {
178        if !self.show_limit_quota_header {
179            return;
180        }
181        ctx.response.headers.insert(
182            "x-ratelimit-limit".to_string(),
183            vec![self.count.to_string()],
184        );
185        ctx.response.headers.insert(
186            "x-ratelimit-remaining".to_string(),
187            vec![remaining.to_string()],
188        );
189        // Whole seconds until the window resets (rounded up so a partial
190        // second is never reported as 0 while the window is still open).
191        let reset_secs = reset.as_secs_f64().ceil() as u64;
192        ctx.response.headers.insert(
193            "x-ratelimit-reset".to_string(),
194            vec![reset_secs.to_string()],
195        );
196    }
197}
198
199#[async_trait]
200impl Plugin for LimitCountPlugin {
201    fn plugin_type(&self) -> &str {
202        "limit-count"
203    }
204
205    async fn execute(
206        &self,
207        mut ctx: Context,
208        _named_inputs: &HashMap<String, serde_json::Value>,
209    ) -> PluginResult {
210        let key = self.resolve_key(&ctx);
211
212        let result = match self
213            .store
214            .incr_fixed_window(&key, self.count, self.window)
215            .await
216        {
217            Ok(r) => r,
218            Err(_) => {
219                // Counter backend failed. Fail open when configured to
220                // degrade, otherwise reject with a 500 through the error port.
221                if self.allow_degradation {
222                    return Ok(PluginOutput {
223                        context: ctx,
224                        named_outputs: HashMap::new(),
225                    });
226                }
227                ctx.response.status_code = 500;
228                ctx.response.body = Bytes::from(r#"{"error_msg": "failed to limit count"}"#);
229                ctx.response.headers.insert(
230                    "content-type".to_string(),
231                    vec!["application/json".to_string()],
232                );
233                return Err(PluginExecutionError {
234                    context: ctx,
235                    error: GatewayError {
236                        node_id: String::new(),
237                        code: "RATE_LIMIT_UNAVAILABLE".to_string(),
238                        message: "failed to limit count".to_string(),
239                        metadata: HashMap::new(),
240                    },
241                });
242            }
243        };
244
245        if result.allowed {
246            self.set_quota_headers(&mut ctx, result.remaining, result.reset);
247            return Ok(PluginOutput {
248                context: ctx,
249                named_outputs: HashMap::new(),
250            });
251        }
252
253        // Rejected: quota headers show 0 remaining, then a JSON error body.
254        self.set_quota_headers(&mut ctx, 0, result.reset);
255        let msg = self
256            .rejected_msg
257            .clone()
258            .unwrap_or_else(|| "Requests over the limit".to_string());
259        let body = serde_json::json!({ "error_msg": msg }).to_string();
260        ctx.response.status_code = self.rejected_code;
261        ctx.response.body = Bytes::from(body);
262        ctx.response.headers.insert(
263            "content-type".to_string(),
264            vec!["application/json".to_string()],
265        );
266
267        Err(PluginExecutionError {
268            context: ctx,
269            error: GatewayError {
270                node_id: String::new(),
271                code: "RATE_LIMITED".to_string(),
272                message: msg,
273                metadata: HashMap::new(),
274            },
275        })
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::context::{GatewayRequest, GatewayResponse, Protocol};
283
284    fn test_ctx() -> Context {
285        let mut headers = HashMap::new();
286        headers.insert("x-api-key".to_string(), vec!["abc123".to_string()]);
287        Context {
288            request: GatewayRequest {
289                method: "GET".to_string(),
290                path: "/api".to_string(),
291                host: "example.com".to_string(),
292                scheme: "http".to_string(),
293                headers,
294                query_params: HashMap::new(),
295                body: Bytes::new(),
296                remote_addr: "10.1.2.3:44321".to_string(),
297                protocol: Protocol::Http1,
298            },
299            response: GatewayResponse {
300                status_code: 0,
301                headers: HashMap::new(),
302                body: Bytes::new(),
303            },
304            message: HashMap::new(),
305            errors: Vec::new(),
306        }
307    }
308
309    fn plugin(config: serde_json::Value) -> Result<LimitCountPlugin, String> {
310        let map: HashMap<String, serde_json::Value> = serde_json::from_value(config).unwrap();
311        LimitCountPlugin::from_config(&map, &PluginResources::empty())
312    }
313
314    #[test]
315    fn test_config_requires_count_and_window() {
316        assert!(plugin(serde_json::json!({ "time_window": 60 })).is_err());
317        assert!(plugin(serde_json::json!({ "count": 10 })).is_err());
318        assert!(plugin(serde_json::json!({ "count": 0, "time_window": 60 })).is_err());
319        assert!(plugin(serde_json::json!({ "count": 10, "time_window": 0 })).is_err());
320        assert!(plugin(serde_json::json!({ "count": 10, "time_window": 60 })).is_ok());
321    }
322
323    #[test]
324    fn test_unknown_policy_rejected() {
325        let err = match plugin(serde_json::json!({
326            "count": 10,
327            "time_window": 60,
328            "policy": "redis"
329        })) {
330            Ok(_) => panic!("'redis' policy should not resolve yet"),
331            Err(e) => e,
332        };
333        assert!(err.contains("policy"), "{err}");
334    }
335
336    #[test]
337    fn test_key_interpolation() {
338        let ctx = test_ctx();
339
340        // default: remote_addr (port stripped by the var resolver)
341        let p = plugin(serde_json::json!({ "count": 10, "time_window": 60 })).unwrap();
342        assert_eq!(p.resolve_key(&ctx), "10.1.2.3");
343
344        // custom var template
345        let p = plugin(serde_json::json!({
346            "count": 10, "time_window": 60, "key": "$http_x_api_key"
347        }))
348        .unwrap();
349        assert_eq!(p.resolve_key(&ctx), "abc123");
350
351        // empty resolution falls back to remote_addr
352        let p = plugin(serde_json::json!({
353            "count": 10, "time_window": 60, "key": "$http_missing"
354        }))
355        .unwrap();
356        assert_eq!(p.resolve_key(&ctx), "10.1.2.3");
357
358        // group prefixes the key
359        let p = plugin(serde_json::json!({
360            "count": 10, "time_window": 60, "key": "$remote_addr", "group": "svc"
361        }))
362        .unwrap();
363        assert_eq!(p.resolve_key(&ctx), "svc:10.1.2.3");
364    }
365
366    #[tokio::test]
367    async fn test_rejects_after_count_requests() {
368        let count = 3u64;
369        let p = plugin(serde_json::json!({
370            "count": count, "time_window": 60, "rejected_code": 429
371        }))
372        .unwrap();
373
374        // First `count` requests pass.
375        for i in 0..count {
376            let out = p.execute(test_ctx(), &HashMap::new()).await;
377            assert!(out.is_ok(), "request {i} should pass");
378            let ctx = out.unwrap().context;
379            assert_eq!(
380                ctx.response.headers.get("x-ratelimit-limit"),
381                Some(&vec!["3".to_string()])
382            );
383            assert_eq!(
384                ctx.response.headers.get("x-ratelimit-remaining"),
385                Some(&vec![(count - 1 - i).to_string()])
386            );
387        }
388
389        // The next one is rejected.
390        let err = p.execute(test_ctx(), &HashMap::new()).await.unwrap_err();
391        assert_eq!(err.error.code, "RATE_LIMITED");
392        assert_eq!(err.context.response.status_code, 429);
393        assert_eq!(
394            err.context.response.headers.get("x-ratelimit-remaining"),
395            Some(&vec!["0".to_string()])
396        );
397        let body = String::from_utf8(err.context.response.body.to_vec()).unwrap();
398        assert!(body.contains("error_msg"), "{body}");
399    }
400
401    #[tokio::test]
402    async fn test_rejected_msg_used() {
403        let p = plugin(serde_json::json!({
404            "count": 1, "time_window": 60, "rejected_msg": "slow down"
405        }))
406        .unwrap();
407        assert!(p.execute(test_ctx(), &HashMap::new()).await.is_ok());
408        let err = p.execute(test_ctx(), &HashMap::new()).await.unwrap_err();
409        let body = String::from_utf8(err.context.response.body.to_vec()).unwrap();
410        assert!(body.contains("slow down"), "{body}");
411        assert_eq!(err.error.message, "slow down");
412    }
413
414    #[tokio::test]
415    async fn test_show_limit_quota_header_false_omits_headers() {
416        let p = plugin(serde_json::json!({
417            "count": 5, "time_window": 60, "show_limit_quota_header": false
418        }))
419        .unwrap();
420        let ctx = p
421            .execute(test_ctx(), &HashMap::new())
422            .await
423            .unwrap()
424            .context;
425        assert!(!ctx.response.headers.contains_key("x-ratelimit-limit"));
426    }
427}
featherbit/plugins/native/limit_count.rs

featherbit/plugins/native/
limit_count.rs
