ri/gateway/rate_limiter.rs
1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//! http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Rate Limiter Module
21//!
22//! This module provides rate limiting functionality for the Ri gateway, allowing for
23//! controlling the rate of requests from clients to prevent abuse and ensure fair usage.
24//!
25//! ## Key Components
26//!
27//! - **RiRateLimitConfig**: Configuration for rate limiting behavior
28//! - **RiRateLimiter**: Token bucket based rate limiter implementation
29//! - **RiSlidingWindowRateLimiter**: Sliding window based rate limiter for fine-grained control
30//! - **RiRateLimitStats**: Metrics for monitoring rate limiter performance
31//!
32//! ## Design Principles
33//!
34//! 1. **Token Bucket Algorithm**: Implements the token bucket algorithm for smooth rate limiting
35//! 2. **Sliding Window**: Provides a sliding window implementation for more precise control
36//! 3. **Thread Safe**: Uses Arc and RwLock for safe operation in multi-threaded environments
37//! 4. **Configurable**: Allows fine-tuning of requests per second, burst size, and window duration
38//! 5. **Metrics Collection**: Tracks and reports rate limiter statistics
39//! 6. **Async Compatibility**: Built with async/await patterns for modern Rust applications
40//! 7. **Burst Support**: Allows for temporary bursts of requests beyond the steady rate
41//! 8. **Key-Based Limiting**: Supports rate limiting by client IP or custom keys
42//!
43//! ## Usage
44//!
45//! ```rust
46//! use ri::prelude::*;
47//!
48//! async fn example() {
49//! // Create a rate limiter with default configuration
50//! let mut limiter = RiRateLimiter::new(RiRateLimitConfig::default());
51//!
52//! // Check if a request should be allowed
53//! let client_ip = "192.168.1.1";
54//! if limiter.check_rate_limit(client_ip, 1).await {
55//! println!("Request allowed");
56//! } else {
57//! println!("Request rate limited");
58//! }
59//!
60//! // Get rate limit stats for a client
61//! if let Some(stats) = limiter.get_stats(client_ip).await {
62//! println!("Current tokens: {}, Total requests: {}",
63//! stats.current_tokens, stats.total_requests);
64//! }
65//!
66//! // Create a sliding window rate limiter
67//! let sliding_limiter = RiSlidingWindowRateLimiter::new(100, 60);
68//! if sliding_limiter.allow_request().await {
69//! println!("Sliding window request allowed");
70//! }
71//! }
72//! ```
73
74use std::collections::HashMap as FxHashMap;
75use std::sync::Arc;
76use std::sync::atomic::{AtomicUsize, Ordering};
77use tokio::sync::RwLock;
78use std::time::{Duration, Instant};
79
80/// Configuration for rate limiting behavior.
81///
82/// This struct defines the parameters that control how the rate limiter behaves,
83/// including the steady rate, burst capacity, and window duration.
84#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
85#[derive(Debug, Clone)]
86pub struct RiRateLimitConfig {
87 /// Maximum number of requests allowed per second in steady state
88 pub requests_per_second: u32,
89
90 /// Maximum number of requests allowed in a burst (temporary spike)
91 pub burst_size: u32,
92
93 /// Duration of the rate limiting window in seconds
94 pub window_seconds: u64,
95
96 /// Maximum number of unique keys (clients) to track
97 /// This prevents memory exhaustion from too many unique clients
98 pub max_keys: usize,
99}
100
101/// Maximum length of a rate limit key
102const MAX_KEY_LENGTH: usize = 256;
103
104#[cfg(feature = "pyo3")]
105#[pyo3::prelude::pymethods]
106impl RiRateLimitConfig {
107 #[new]
108 fn py_new() -> Self {
109 Self::default()
110 }
111
112 #[staticmethod]
113 fn py_new_with_values(requests_per_second: u32, burst_size: u32, window_seconds: u64) -> Self {
114 Self {
115 requests_per_second,
116 burst_size,
117 window_seconds,
118 max_keys: 10000,
119 }
120 }
121
122 fn get_requests_per_second(&self) -> u32 {
123 self.requests_per_second
124 }
125
126 fn set_requests_per_second(&mut self, value: u32) {
127 self.requests_per_second = value;
128 }
129
130 fn get_burst_size(&self) -> u32 {
131 self.burst_size
132 }
133
134 fn set_burst_size(&mut self, value: u32) {
135 self.burst_size = value;
136 }
137
138 fn get_window_seconds(&self) -> u64 {
139 self.window_seconds
140 }
141
142 fn set_window_seconds(&mut self, value: u64) {
143 self.window_seconds = value;
144 }
145}
146
147impl Default for RiRateLimitConfig {
148 /// Creates a default rate limit configuration.
149 ///
150 /// Default values:
151 /// - requests_per_second: 10 requests per second
152 /// - burst_size: 20 requests (temporary burst capacity)
153 /// - window_seconds: 60 seconds window duration
154 /// - max_keys: 10000 unique clients
155 fn default() -> Self {
156 Self {
157 requests_per_second: 10,
158 burst_size: 20,
159 window_seconds: 60,
160 max_keys: 10000,
161 }
162 }
163}
164
165/// Internal token bucket for rate limiting.
166///
167/// This struct implements the token bucket algorithm for rate limiting, tracking
168/// available tokens, last update time, and request count.
169#[derive(Debug)]
170struct RateLimitBucket {
171 /// Current number of available tokens in the bucket
172 tokens: AtomicUsize,
173
174 /// Timestamp of the last token refill
175 last_update: RwLock<Instant>,
176
177 /// Total number of requests processed by this bucket
178 request_count: AtomicUsize,
179}
180
181impl RateLimitBucket {
182 /// Creates a new token bucket with the specified initial tokens.
183 ///
184 /// # Parameters
185 ///
186 /// - `tokens`: Initial number of tokens in the bucket
187 ///
188 /// # Returns
189 ///
190 /// A new `RateLimitBucket` instance
191 fn new(tokens: usize) -> Self {
192 Self {
193 tokens: AtomicUsize::new(tokens),
194 last_update: RwLock::new(Instant::now()),
195 request_count: AtomicUsize::new(0),
196 }
197 }
198
199 /// Attempts to consume tokens from the bucket.
200 ///
201 /// This method refills tokens based on time elapsed since the last update,
202 /// then attempts to consume the requested number of tokens.
203 ///
204 /// # Parameters
205 ///
206 /// - `tokens`: Number of tokens to consume
207 /// - `config`: Rate limit configuration for token refill
208 ///
209 /// # Returns
210 ///
211 /// `true` if tokens were successfully consumed, `false` otherwise
212 async fn try_consume(&self, tokens: usize, config: &RiRateLimitConfig) -> bool {
213 let now = Instant::now();
214 let mut last_update = self.last_update.write().await;
215
216 // Refill tokens based on time elapsed
217 let elapsed = now.duration_since(*last_update).as_secs_f64();
218 let tokens_to_add = (elapsed * config.requests_per_second as f64) as usize;
219
220 if tokens_to_add > 0 {
221 let current_tokens = self.tokens.load(Ordering::Relaxed);
222 let new_tokens = std::cmp::min(current_tokens + tokens_to_add, config.burst_size as usize);
223 self.tokens.store(new_tokens, Ordering::Relaxed);
224 *last_update = now;
225 }
226
227 // Try to consume tokens
228 let current_tokens = self.tokens.load(Ordering::Relaxed);
229 if current_tokens >= tokens {
230 self.tokens.fetch_sub(tokens, Ordering::Relaxed);
231 self.request_count.fetch_add(1, Ordering::Relaxed);
232 true
233 } else {
234 false
235 }
236 }
237
238 /// Gets the current statistics for this bucket.
239 ///
240 /// # Returns
241 ///
242 /// A `RiRateLimitStats` struct containing current tokens and total requests
243 fn get_stats(&self) -> RiRateLimitStats {
244 RiRateLimitStats {
245 current_tokens: self.tokens.load(Ordering::Relaxed),
246 total_requests: self.request_count.load(Ordering::Relaxed),
247 }
248 }
249}
250
251/// Statistics for rate limiting monitoring.
252///
253/// This struct contains metrics about a rate limiter bucket, including the current
254/// number of available tokens and the total number of requests processed.
255#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass(get_all, set_all))]
256#[derive(Debug, Clone)]
257pub struct RiRateLimitStats {
258 /// Current number of available tokens in the bucket
259 pub current_tokens: usize,
260
261 /// Total number of requests processed by the bucket
262 pub total_requests: usize,
263}
264
265#[cfg(feature = "pyo3")]
266#[pyo3::prelude::pymethods]
267impl RiRateLimitStats {
268 #[new]
269 fn py_new(current_tokens: usize, total_requests: usize) -> Self {
270 Self {
271 current_tokens,
272 total_requests,
273 }
274 }
275
276 fn get_current_tokens(&self) -> usize {
277 self.current_tokens
278 }
279
280 fn get_total_requests(&self) -> usize {
281 self.total_requests
282 }
283}
284
285/// Token bucket based rate limiter implementation.
286///
287/// This struct implements the token bucket algorithm for rate limiting, allowing
288/// for both steady-state rate limiting and temporary bursts of requests.
289#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
290pub struct RiRateLimiter {
291 /// Configuration for rate limiting behavior
292 config: RiRateLimitConfig,
293
294 /// Map of key to token bucket instances
295 buckets: RwLock<FxHashMap<String, Arc<RateLimitBucket>>>,
296}
297
298impl RiRateLimiter {
299 /// Creates a new rate limiter with the specified configuration.
300 ///
301 /// # Parameters
302 ///
303 /// - `config`: The configuration for rate limiting behavior
304 ///
305 /// # Returns
306 ///
307 /// A new `RiRateLimiter` instance
308 pub fn new(config: RiRateLimitConfig) -> Self {
309 Self {
310 config,
311 buckets: RwLock::new(FxHashMap::default()),
312 }
313 }
314
315 /// Checks if a gateway request should be allowed based on rate limiting.
316 ///
317 /// This method uses the client IP address as the key for rate limiting.
318 ///
319 /// # Parameters
320 ///
321 /// - `request`: The gateway request to check
322 ///
323 /// # Returns
324 ///
325 /// `true` if the request should be allowed, `false` otherwise
326 pub async fn check_request(&self, request: &crate::gateway::RiGatewayRequest) -> bool {
327 // Use client IP as the key for rate limiting
328 let key = request.remote_addr.clone();
329 self.check_rate_limit(&key, 1)
330 }
331
332 /// Checks if a request with a custom key should be allowed based on rate limiting.
333 ///
334 /// This method attempts to consume tokens from the bucket associated with the given key.
335 /// If no bucket exists for the key, a new one is created.
336 ///
337 /// # Security
338 ///
339 /// This method validates:
340 /// - Key length (max 256 characters)
341 /// - Maximum number of unique keys (prevents memory exhaustion)
342 ///
343 /// # Parameters
344 ///
345 /// - `key`: The key to use for rate limiting (e.g., client IP, API key)
346 /// - `tokens`: Number of tokens to consume for this request
347 ///
348 /// # Returns
349 ///
350 /// `true` if the request should be allowed, `false` otherwise
351 pub fn check_rate_limit(&self, key: &str, tokens: usize) -> bool {
352 // Security: Validate key length
353 if key.is_empty() || key.len() > MAX_KEY_LENGTH {
354 log::warn!(
355 "[Ri.RateLimiter] Invalid key length: {} chars (max {})",
356 key.len(), MAX_KEY_LENGTH
357 );
358 return false;
359 }
360
361 futures::executor::block_on(async {
362 let buckets = self.buckets.read().await;
363
364 if let Some(bucket) = buckets.get(key) {
365 bucket.try_consume(tokens, &self.config).await
366 } else {
367 drop(buckets);
368 let mut buckets = self.buckets.write().await;
369
370 // Security: Check maximum keys limit to prevent memory exhaustion
371 if buckets.len() >= self.config.max_keys {
372 log::warn!(
373 "[Ri.RateLimiter] Maximum keys limit reached: {} (max {})",
374 buckets.len(), self.config.max_keys
375 );
376 // Reject new clients when limit is reached
377 return false;
378 }
379
380 if let Some(bucket) = buckets.get(key) {
381 bucket.try_consume(tokens, &self.config).await
382 } else {
383 let bucket = Arc::new(RateLimitBucket::new(self.config.burst_size as usize));
384 let result = bucket.try_consume(tokens, &self.config).await;
385 buckets.insert(key.to_string(), bucket);
386 result
387 }
388 }
389 })
390 }
391
392 /// Gets rate limit statistics for a specific key.
393 ///
394 /// # Parameters
395 ///
396 /// - `key`: The key to get statistics for
397 ///
398 /// # Returns
399 ///
400 /// An `Option<RiRateLimitStats>` with the statistics, or `None` if no bucket exists for the key
401 pub fn get_stats(&self, key: &str) -> Option<RiRateLimitStats> {
402 futures::executor::block_on(async {
403 let buckets = self.buckets.read().await;
404 buckets.get(key).map(|bucket| bucket.get_stats())
405 })
406 }
407
408 /// Gets the remaining tokens for a specific key.
409 pub fn get_remaining(&self, key: &str) -> Option<f64> {
410 futures::executor::block_on(async {
411 let buckets = self.buckets.read().await;
412 buckets.get(key).map(|bucket| {
413 let stats = bucket.get_stats();
414 stats.current_tokens as f64
415 })
416 })
417 }
418
419 /// Gets rate limit statistics for all keys.
420 ///
421 /// # Returns
422 ///
423 /// A `FxHashMap<String, RiRateLimitStats>` with statistics for all keys
424 pub fn get_all_stats(&self) -> FxHashMap<String, RiRateLimitStats> {
425 futures::executor::block_on(async {
426 let buckets = self.buckets.read().await;
427 let mut stats = FxHashMap::default();
428
429 for (key, bucket) in buckets.iter() {
430 stats.insert(key.clone(), bucket.get_stats());
431 }
432
433 stats
434 })
435 }
436
437 /// Resets the rate limit bucket for a specific key.
438 ///
439 /// This method removes the bucket for the given key, effectively resetting the rate limit.
440 ///
441 /// # Parameters
442 ///
443 /// - `key`: The key to reset the bucket for
444 pub fn reset_bucket(&self, key: &str) {
445 futures::executor::block_on(async {
446 let mut buckets = self.buckets.write().await;
447 buckets.remove(key);
448 })
449 }
450
451 /// Clears all rate limit buckets.
452 ///
453 /// This method removes all buckets, effectively resetting rate limits for all keys.
454 pub fn clear_all_buckets(&self) {
455 futures::executor::block_on(async {
456 let mut buckets = self.buckets.write().await;
457 buckets.clear();
458 })
459 }
460
461 /// Gets the current rate limit configuration.
462 ///
463 /// # Returns
464 ///
465 /// A reference to the current `RiRateLimitConfig`
466 pub fn get_config(&self) -> RiRateLimitConfig {
467 self.config.clone()
468 }
469
470 /// Updates the rate limit configuration.
471 ///
472 /// This method updates the configuration and resets all buckets with the new settings.
473 ///
474 /// # Parameters
475 ///
476 /// - `config`: The new rate limit configuration
477 pub async fn update_config(&mut self, config: RiRateLimitConfig) {
478 self.config = config;
479
480 let mut buckets = self.buckets.write().await;
481 buckets.clear();
482 }
483
484 pub async fn check_multi(&self, keys: &[String], tokens: usize) -> Vec<bool> {
485 let mut results = Vec::with_capacity(keys.len());
486 for key in keys {
487 results.push(self.check_rate_limit(key, tokens));
488 }
489 results
490 }
491
492 pub async fn get_keys(&self) -> Vec<String> {
493 let buckets = self.buckets.read().await;
494 buckets.keys().cloned().collect()
495 }
496
497 pub fn bucket_count(&self) -> usize {
498 futures::executor::block_on(async {
499 let buckets = self.buckets.read().await;
500 buckets.len()
501 })
502 }
503}
504
505/// Sliding window rate limiter for fine-grained control.
506///
507/// This struct implements a sliding window rate limiter, which provides more precise
508/// rate limiting by tracking all requests within a sliding time window.
509#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
510pub struct RiSlidingWindowRateLimiter {
511 /// Maximum number of requests allowed within the window
512 max_requests: u32,
513 /// Duration of the sliding window
514 window_duration: Duration,
515 /// Vector of request timestamps within the window
516 requests: RwLock<Vec<Instant>>,
517}
518
519impl RiSlidingWindowRateLimiter {
520 /// Creates a new sliding window rate limiter.
521 ///
522 /// # Parameters
523 ///
524 /// - `max_requests`: Maximum number of requests allowed within the window
525 /// - `window_seconds`: Duration of the window in seconds
526 ///
527 /// # Returns
528 ///
529 /// A new `RiSlidingWindowRateLimiter` instance
530 pub fn new(max_requests: u32, window_seconds: u64) -> Self {
531 Self {
532 max_requests,
533 window_duration: Duration::from_secs(window_seconds),
534 requests: RwLock::new(Vec::new()),
535 }
536 }
537
538 /// Checks if a request should be allowed based on the sliding window.
539 ///
540 /// This method removes old requests outside the window, then checks if the number
541 /// of remaining requests is below the maximum allowed.
542 ///
543 /// # Returns
544 ///
545 /// `true` if the request should be allowed, `false` otherwise
546 pub fn allow_request(&self) -> bool {
547 futures::executor::block_on(async {
548 let mut requests = self.requests.write().await;
549 let now = Instant::now();
550
551 requests.retain(|×tamp| now.duration_since(timestamp) < self.window_duration);
552
553 if requests.len() < self.max_requests as usize {
554 requests.push(now);
555 true
556 } else {
557 false
558 }
559 })
560 }
561
562 /// Gets the current number of requests within the sliding window.
563 ///
564 /// This method removes old requests outside the window, then returns the count
565 /// of remaining requests.
566 ///
567 /// # Returns
568 ///
569 /// The number of requests within the current window
570 pub fn get_current_count(&self) -> usize {
571 futures::executor::block_on(async {
572 let mut requests = self.requests.write().await;
573 let now = Instant::now();
574
575 requests.retain(|×tamp| now.duration_since(timestamp) < self.window_duration);
576
577 requests.len()
578 })
579 }
580
581 /// Resets the sliding window by clearing all request timestamps.
582 pub fn reset(&self) {
583 futures::executor::block_on(async {
584 let mut requests = self.requests.write().await;
585 requests.clear();
586 })
587 }
588
589 pub fn get_max_requests(&self) -> u32 {
590 self.max_requests
591 }
592
593 pub fn get_window_seconds(&self) -> u64 {
594 self.window_duration.as_secs()
595 }
596}