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

throttle_kit/
lib.rs
