throttle_ro/
lib.rs

1//! A throttling service that limits the number of attempts (hits) from an IP address
2//! within a specified time period.
3//!
4//! # Example: Basic Rate Limiting
5//!
6//! ```
7//! use std::time::Duration;
8//! use cache_ro::Cache;
9//! use throttle_ro::ThrottlesService;
10//!
11//! // Create a cache instance (in-memory for this example)
12//! let cache = Cache::new(cache_ro::CacheConfig {
13//!     persistent: false,
14//!     ..Default::default()
15//! }).unwrap();
16//!
17//! // Create a throttling service for IP "127.0.0.1"
18//! // allowing maximum 5 attempts per minute
19//! let ip = "127.0.0.1".to_string();
20//! let mut service = ThrottlesService::new(
21//!     ip,
22//!     5, // max attempts
23//!     Duration::from_secs(60), // time window
24//!     "api_rate_limit_" // cache key prefix
25//! );
26//!
27//! // Check if the IP is allowed to proceed
28//! if service.can_go(&cache) {
29//!     // Record the attempt
30//!     service.hit(&cache);
31//!     println!("Request allowed");
32//!     // Process the request...
33//! } else {
34//!     println!("Rate limit exceeded - please try again later");
35//!     // Return error or wait...
36//! }
37//!
38//! // You can also manually clear the throttle if needed
39//! // service.remove(&cache);
40//! ```
41//!
42
43
44use cache_ro::Cache;
45use std::time::Duration;
46
47/// A service for throttling attempts from an IP address.
48///
49/// Tracks the number of attempts (hits) from a given IP address and determines
50/// whether further attempts should be allowed based on configured limits.
51///
52/// # Examples
53///
54/// ```
55/// use std::time::Duration;
56/// use cache_ro::Cache;
57/// use throttle_ro::ThrottlesService;
58///
59/// let cache = Cache::new(Default::default()); // In real usage, configure properly
60/// let ip = "127.0.0.1".to_string();
61/// let mut service = ThrottlesService::new(
62///     ip,
63///     5, // max attempts
64///     Duration::from_secs(60), // time window
65///     "rate_limit_"
66/// );
67///
68/// if service.can_go(&cache) {
69///     service.hit(&cache);
70///     // Process the request
71/// } else {
72///     // Reject the request - rate limit exceeded
73/// }
74/// ```
75pub struct ThrottlesService {
76    ip: String,
77    max_attempts: u32,
78    period: Duration,
79    prefix: String,
80}
81
82impl ThrottlesService {
83    /// Creates a new `ThrottlesService` instance.
84    ///
85    /// # Arguments
86    ///
87    /// * `ip` - The IP address to track
88    /// * `max_attempts` - Maximum number of allowed attempts in the time period
89    /// * `period` - Duration of the throttling window
90    /// * `prefix` - Prefix for cache keys to avoid collisions
91    pub fn new(ip: String, max_attempts: u32, period: Duration, prefix: &str) -> Self {
92        Self {
93            ip,
94            max_attempts,
95            period,
96            prefix: prefix.to_string(),
97        }
98    }
99
100    /// Checks whether the IP is allowed to make another attempt.
101    ///
102    /// Returns `true` if the current attempt count is below the maximum allowed.
103    pub fn can_go(&mut self, cache: &Cache) -> bool {
104        let v = self.get_value(cache).unwrap_or(0);
105        v < self.max_attempts
106    }
107
108    fn get_value(&mut self, cache: &Cache) -> Option<u32> {
109        cache.get::<u32>(&self.key())
110    }
111
112    /// Generates the cache key for this IP.
113    pub fn key(&self) -> String {
114        format!("{}{}", self.prefix, self.ip)
115    }
116
117    /// Gets the remaining duration for the current throttling window.
118    ///
119    /// Returns the configured period if no expiration is set in the cache.
120    pub fn get_expire(&mut self, cache: &Cache) -> Duration {
121        let ex = cache.expire(&self.key());
122        match ex {
123            None => self.period,
124            Some(a) => a,
125        }
126    }
127
128    /// Records an attempt (hit) from the IP.
129    ///
130    /// Increments the attempt count and resets the expiration time.
131    pub fn hit(&mut self, cache: &Cache) {
132        let key = self.key();
133        let expire = self.get_expire(cache);
134
135        match self.get_value(cache) {
136            None => cache.set::<u32>(&key, 1, expire).unwrap(),
137            Some(v) => cache.set::<u32>(&key, v + 1, expire).unwrap(),
138        }
139    }
140
141    /// Clears the attempt count for the IP.
142    pub fn remove(&self, cache: &Cache) {
143        cache.remove(&self.key()).unwrap();
144    }
145}
throttle_ro/lib.rs

throttle_ro/
lib.rs
