throttle_kit/
lib.rs

1#![forbid(unsafe_code)]
2#![deny(missing_docs)]
3//! Rate limiting for Rust.
4//!
5//! Implements the **GCRA** (Generic Cell Rate Algorithm) with pluggable
6//! backends (in-memory via `DashMap`, optional Redis) and optional Tower
7//! layer integration.
8//!
9//! # Quick Start
10//!
11//! ```no_run
12//! use ratelimit::{RateLimiter, Quota, InMemoryBackend};
13//!
14//! #[tokio::main]
15//! async fn main() {
16//!     let backend = InMemoryBackend::new();
17//!     let limiter = RateLimiter::new(
18//!         Quota::per_second(10),
19//!         backend,
20//!     );
21//!
22//!     let result = limiter.check("user-123").await;
23//!     assert!(result.allowed);
24//! }
25//! ```
26
27mod backend;
28mod error;
29mod metrics;
30mod quota;
31
32#[cfg(feature = "redis")]
33mod redis;
34
35#[cfg(feature = "sliding-window")]
36mod sliding_window;
37
38#[cfg(feature = "tower")]
39mod tower_layer;
40
41pub use backend::{InMemoryBackend, RateLimitBackend};
42
43#[cfg(feature = "sliding-window")]
44pub use sliding_window::SlidingWindowBackend;
45pub use error::RateLimitError;
46pub use metrics::RateLimitResult;
47pub use quota::Quota;
48
49#[cfg(feature = "redis")]
50pub use redis::RedisBackend;
51
52use std::sync::Arc;
53
54/// Core rate limiter that delegates to a [`RateLimitBackend`].
55#[derive(Clone)]
56pub struct RateLimiter<B: RateLimitBackend> {
57    quota: Quota,
58    backend: Arc<B>,
59}
60
61impl<B: RateLimitBackend> RateLimiter<B> {
62    /// Create a new rate limiter with the given quota and backend.
63    pub fn new(quota: Quota, backend: B) -> Self {
64        Self {
65            quota,
66            backend: Arc::new(backend),
67        }
68    }
69
70    /// Check whether the caller identified by `key` is allowed to proceed.
71    pub async fn check(&self, key: &str) -> RateLimitResult {
72        self.backend.check(key, &self.quota).await
73    }
74
75    /// Synchronous version of [`check`](RateLimiter::check).
76    ///
77    /// Blocks the current thread until the check completes.
78    /// Useful for synchronous frameworks (e.g. Actix, Rocket).
79    pub fn check_sync(&self, key: &str) -> RateLimitResult {
80        let rt = tokio::runtime::Builder::new_current_thread()
81            .build()
82            .expect("failed to create tokio runtime for sync check");
83        rt.block_on(self.check(key))
84    }
85}
86
87#[cfg(feature = "in-memory")]
88pub use keyed::KeyedRateLimiter;
89
90#[cfg(feature = "in-memory")]
91mod keyed {
92    use dashmap::DashMap;
93
94    use crate::backend::RateLimitBackend;
95    use crate::metrics::RateLimitResult;
96    use crate::quota::Quota;
97    use crate::RateLimiter;
98
99    /// Per-key rate limiter that tracks separate rate limits for each key.
100    ///
101    /// Each key gets its own [`RateLimiter`] instance with either a
102    /// custom quota (set via [`with_quota_for_key`](KeyedRateLimiter::with_quota_for_key))
103    /// or the default quota.
104    ///
105    /// # Example
106    ///
107    /// ```no_run
108    /// use ratelimit::{KeyedRateLimiter, Quota, InMemoryBackend};
109    ///
110    /// #[tokio::main]
111    /// async fn main() {
112    ///     let backend = InMemoryBackend::new();
113    ///     let limiter = KeyedRateLimiter::new(Quota::per_second(100), backend);
114    ///
115    ///     // Tighter limit for login endpoint
116    ///     limiter.with_quota_for_key("login", Quota::per_second(10));
117    ///
118    ///     let result = limiter.check("192.168.1.1").await;
119    ///     assert!(result.allowed);
120    /// }
121    /// ```
122    pub struct KeyedRateLimiter<B: RateLimitBackend + Clone> {
123        limiters: DashMap<String, RateLimiter<B>>,
124        default_quota: Quota,
125        backend: B,
126    }
127
128    impl<B: RateLimitBackend + Clone> KeyedRateLimiter<B> {
129        /// Create a new keyed rate limiter with the given default quota and backend.
130        pub fn new(default_quota: Quota, backend: B) -> Self {
131            Self {
132                limiters: DashMap::new(),
133                default_quota,
134                backend,
135            }
136        }
137
138        /// Override the rate limit quota for a specific key.
139        ///
140        /// This replaces any existing limiter for the key. If the key
141        /// has not been seen yet, the next call to [`check`](KeyedRateLimiter::check)
142        /// will use this quota instead of the default.
143        pub fn with_quota_for_key(&self, key: &str, quota: Quota) {
144            self.limiters.insert(
145                key.to_string(),
146                RateLimiter::new(quota, self.backend.clone()),
147            );
148        }
149
150        /// Check whether the caller identified by `key` is allowed to proceed.
151        pub async fn check(&self, key: &str) -> RateLimitResult {
152            let limiter = self
153                .limiters
154                .entry(key.to_string())
155                .or_insert_with(|| {
156                    RateLimiter::new(self.default_quota.clone(), self.backend.clone())
157                })
158                .value()
159                .clone();
160            limiter.check(key).await
161        }
162
163        /// Synchronous version of [`check`](KeyedRateLimiter::check).
164        ///
165        /// Blocks the current thread until the check completes.
166        /// Useful for synchronous frameworks (e.g. Actix, Rocket).
167        pub fn check_sync(&self, key: &str) -> RateLimitResult {
168            let rt = tokio::runtime::Builder::new_current_thread()
169                .build()
170                .expect("failed to create tokio runtime for sync check");
171            rt.block_on(self.check(key))
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use std::time::Duration;
180
181    #[test]
182    fn quota_per_second() {
183        let q = Quota::per_second(10);
184        assert_eq!(q.burst, 10);
185        assert_eq!(q.interval(), Duration::from_millis(100));
186    }
187
188    #[test]
189    fn quota_per_minute() {
190        let q = Quota::per_minute(100);
191        assert_eq!(q.burst, 100);
192        assert_eq!(q.interval(), Duration::from_millis(600));
193    }
194
195    #[test]
196    fn quota_per_hour() {
197        let q = Quota::per_hour(3600);
198        assert_eq!(q.burst, 3600);
199        assert_eq!(q.interval(), Duration::from_secs(1));
200    }
201
202    #[test]
203    fn quota_allow_burst() {
204        let q = Quota::per_second(10).allow_burst(20);
205        assert_eq!(q.burst, 20);
206    }
207
208    #[test]
209    fn rate_limit_result_fields() {
210        let r = RateLimitResult {
211            allowed: true,
212            remaining: 5,
213            reset_at: std::time::Instant::now(),
214            limit: 10,
215            retry_after: None,
216        };
217        assert!(r.allowed);
218        assert_eq!(r.remaining, 5);
219        assert_eq!(r.limit, 10);
220        assert!(r.retry_after.is_none());
221    }
222
223    #[test]
224    fn rate_limit_result_headers() {
225        let r = RateLimitResult {
226            allowed: true,
227            remaining: 7,
228            reset_at: std::time::Instant::now() + Duration::from_secs(30),
229            limit: 10,
230            retry_after: None,
231        };
232        let headers = r.headers();
233        assert_eq!(headers.len(), 3);
234        assert_eq!(headers[0].0, "X-RateLimit-Limit");
235        assert_eq!(headers[0].1, "10");
236        assert_eq!(headers[1].0, "X-RateLimit-Remaining");
237        assert_eq!(headers[1].1, "7");
238        assert_eq!(headers[2].0, "X-RateLimit-Reset");
239    }
240
241    #[test]
242    fn rate_limit_result_retry_after() {
243        let r = RateLimitResult {
244            allowed: false,
245            remaining: 0,
246            reset_at: std::time::Instant::now(),
247            limit: 10,
248            retry_after: Some(Duration::from_millis(100)),
249        };
250        assert!(!r.allowed);
251        assert_eq!(r.retry_after, Some(Duration::from_millis(100)));
252    }
253
254    #[test]
255    fn in_memory_backend_new() {
256        let _ = InMemoryBackend::new();
257    }
258
259    #[test]
260    fn in_memory_backend_default() {
261        let _ = InMemoryBackend::default();
262    }
263
264    #[test]
265    fn rate_limit_error_display() {
266        let e = RateLimitError::RateLimited;
267        assert_eq!(e.to_string(), "rate limit exceeded");
268
269        let e = RateLimitError::BackendError("connection refused".into());
270        assert_eq!(
271            e.to_string(),
272            "rate limit backend error: connection refused"
273        );
274    }
275
276    #[tokio::test]
277    async fn first_request_is_allowed() {
278        let backend = InMemoryBackend::new();
279        let limiter = RateLimiter::new(Quota::per_second(10), backend);
280        let result = limiter.check("user-1").await;
281        assert!(result.allowed);
282        assert_eq!(result.remaining, 9);
283        assert_eq!(result.limit, 10);
284        assert!(result.retry_after.is_none());
285    }
286
287    #[tokio::test]
288    async fn rate_limit_exhaustion() {
289        let backend = InMemoryBackend::new();
290        let limiter = RateLimiter::new(Quota::per_second(2), backend);
291        let r1 = limiter.check("user-1").await;
292        assert!(r1.allowed);
293        let r2 = limiter.check("user-1").await;
294        assert!(r2.allowed);
295        let r3 = limiter.check("user-1").await;
296        assert!(!r3.allowed);
297        assert!(r3.retry_after.is_some());
298    }
299
300    #[tokio::test]
301    async fn keyed_limiter_default_quota() {
302        use crate::KeyedRateLimiter;
303        let backend = InMemoryBackend::new();
304        let limiter = KeyedRateLimiter::new(Quota::per_second(5), backend);
305
306        let r = limiter.check("client-1").await;
307        assert!(r.allowed);
308        assert_eq!(r.limit, 5);
309    }
310
311    #[tokio::test]
312    async fn keyed_limiter_custom_quota_per_key() {
313        use crate::KeyedRateLimiter;
314        let backend = InMemoryBackend::new();
315        let limiter = KeyedRateLimiter::new(Quota::per_second(100), backend);
316
317        // Tighter limit for login
318        limiter.with_quota_for_key("login", Quota::per_second(2));
319
320        let r_login = limiter.check("login").await;
321        assert!(r_login.allowed);
322        assert_eq!(r_login.limit, 2);
323
324        let r_api = limiter.check("api").await;
325        assert!(r_api.allowed);
326        assert_eq!(r_api.limit, 100);
327    }
328
329    #[tokio::test]
330    async fn keyed_limiter_isolation() {
331        use crate::KeyedRateLimiter;
332        let backend = InMemoryBackend::new();
333        let limiter = KeyedRateLimiter::new(Quota::per_second(1), backend);
334
335        let r1 = limiter.check("a").await;
336        assert!(r1.allowed);
337
338        // "a" is exhausted but "b" should be fine
339        let r1_again = limiter.check("a").await;
340        assert!(!r1_again.allowed);
341
342        let r2 = limiter.check("b").await;
343        assert!(r2.allowed);
344    }
345
346    #[test]
347    fn keyed_limiter_sync_check() {
348        use crate::KeyedRateLimiter;
349        let backend = InMemoryBackend::new();
350        let limiter = KeyedRateLimiter::new(Quota::per_second(10), backend);
351
352        let r = limiter.check_sync("client-1");
353        assert!(r.allowed);
354        assert_eq!(r.limit, 10);
355    }
356
357    #[tokio::test]
358    async fn keyed_limiter_headers() {
359        use crate::KeyedRateLimiter;
360        let backend = InMemoryBackend::new();
361        let limiter = KeyedRateLimiter::new(Quota::per_second(10), backend);
362
363        let r = limiter.check("client-1").await;
364        let headers = r.headers();
365        assert!(headers.iter().any(|(k, _)| *k == "X-RateLimit-Limit"));
366        assert!(headers
367            .iter()
368            .any(|(k, v)| *k == "X-RateLimit-Remaining" && v == "9"));
369    }
370
371    #[test]
372    fn sync_check_basic() {
373        let backend = InMemoryBackend::new();
374        let limiter = RateLimiter::new(Quota::per_second(10), backend);
375        let r = limiter.check_sync("user-1");
376        assert!(r.allowed);
377    }
378}
throttle_kit/lib.rs

throttle_kit/
lib.rs
