openauth_core/options/
rate_limit.rs

1use std::fmt;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::time::Duration;
6
7use http::Request;
8
9use crate::error::OpenAuthError;
10
11/// Rate limiting defaults.
12#[derive(Clone)]
13pub struct RateLimitOptions {
14    pub enabled: Option<bool>,
15    pub window: u64,
16    pub max: u64,
17    pub storage: RateLimitStorageOption,
18    pub custom_rules: Vec<RateLimitPathRule>,
19    pub dynamic_rules: Vec<DynamicRateLimitPathRule>,
20    pub custom_store: Option<Arc<dyn RateLimitStore>>,
21    pub custom_storage: Option<Arc<dyn RateLimitStorage>>,
22    pub hybrid: HybridRateLimitOptions,
23    pub memory_cleanup_interval: Option<Duration>,
24}
25
26impl Default for RateLimitOptions {
27    fn default() -> Self {
28        Self {
29            enabled: None,
30            window: 10,
31            max: 100,
32            storage: RateLimitStorageOption::Memory,
33            custom_rules: Vec::new(),
34            dynamic_rules: Vec::new(),
35            custom_store: None,
36            custom_storage: None,
37            hybrid: HybridRateLimitOptions::default(),
38            memory_cleanup_interval: Some(Duration::from_secs(60 * 60)),
39        }
40    }
41}
42
43impl RateLimitOptions {
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    pub fn builder() -> Self {
49        Self::new()
50    }
51
52    pub fn memory() -> Self {
53        Self {
54            storage: RateLimitStorageOption::Memory,
55            ..Self::default()
56        }
57    }
58
59    pub fn database<S>(store: S) -> Self
60    where
61        S: RateLimitStore,
62    {
63        Self::database_arc(Arc::new(store))
64    }
65
66    pub fn database_arc(store: Arc<dyn RateLimitStore>) -> Self {
67        Self {
68            storage: RateLimitStorageOption::Database,
69            custom_store: Some(store),
70            ..Self::default()
71        }
72    }
73
74    pub fn secondary_storage<S>(store: S) -> Self
75    where
76        S: RateLimitStore,
77    {
78        Self::secondary_storage_arc(Arc::new(store))
79    }
80
81    pub fn secondary_storage_arc(store: Arc<dyn RateLimitStore>) -> Self {
82        Self {
83            storage: RateLimitStorageOption::SecondaryStorage,
84            custom_store: Some(store),
85            ..Self::default()
86        }
87    }
88
89    #[must_use]
90    pub fn enabled(mut self, enabled: bool) -> Self {
91        self.enabled = Some(enabled);
92        self
93    }
94
95    #[must_use]
96    pub fn window(mut self, window: u64) -> Self {
97        self.window = window;
98        self
99    }
100
101    #[must_use]
102    pub fn max(mut self, max: u64) -> Self {
103        self.max = max;
104        self
105    }
106
107    #[must_use]
108    pub fn storage(mut self, storage: RateLimitStorageOption) -> Self {
109        self.storage = storage;
110        self
111    }
112
113    #[must_use]
114    pub fn custom_store<S>(mut self, store: S) -> Self
115    where
116        S: RateLimitStore,
117    {
118        self.custom_store = Some(Arc::new(store));
119        self
120    }
121
122    #[must_use]
123    pub fn custom_store_arc(mut self, store: Arc<dyn RateLimitStore>) -> Self {
124        self.custom_store = Some(store);
125        self
126    }
127
128    #[must_use]
129    pub fn custom_storage(mut self, storage: Arc<dyn RateLimitStorage>) -> Self {
130        self.custom_storage = Some(storage);
131        self
132    }
133
134    #[must_use]
135    pub fn custom_rule(mut self, path: impl Into<String>, rule: RateLimitRule) -> Self {
136        self.custom_rules.push(RateLimitPathRule {
137            path: path.into(),
138            rule: Some(rule),
139        });
140        self
141    }
142
143    #[must_use]
144    pub fn disabled_path(mut self, path: impl Into<String>) -> Self {
145        self.custom_rules.push(RateLimitPathRule {
146            path: path.into(),
147            rule: None,
148        });
149        self
150    }
151
152    #[must_use]
153    pub fn dynamic_rule<P>(mut self, path: impl Into<String>, provider: P) -> Self
154    where
155        P: RateLimitRuleProvider,
156    {
157        self.dynamic_rules
158            .push(DynamicRateLimitPathRule::new(path, provider));
159        self
160    }
161
162    #[must_use]
163    pub fn hybrid(mut self, hybrid: HybridRateLimitOptions) -> Self {
164        self.hybrid = hybrid;
165        self
166    }
167
168    #[must_use]
169    pub fn memory_cleanup_interval(mut self, interval: Option<Duration>) -> Self {
170        self.memory_cleanup_interval = interval;
171        self
172    }
173}
174
175impl fmt::Debug for RateLimitOptions {
176    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
177        formatter
178            .debug_struct("RateLimitOptions")
179            .field("enabled", &self.enabled)
180            .field("window", &self.window)
181            .field("max", &self.max)
182            .field("storage", &self.storage)
183            .field("custom_rules", &self.custom_rules)
184            .field("dynamic_rules", &self.dynamic_rules)
185            .field(
186                "custom_store",
187                &self.custom_store.as_ref().map(|_| "<custom-store>"),
188            )
189            .field(
190                "custom_storage",
191                &self.custom_storage.as_ref().map(|_| "<custom-storage>"),
192            )
193            .field("hybrid", &self.hybrid)
194            .field("memory_cleanup_interval", &self.memory_cleanup_interval)
195            .finish()
196    }
197}
198
199#[derive(Debug, Clone, PartialEq, Eq)]
200pub struct RateLimitRule {
201    pub window: u64,
202    pub max: u64,
203}
204
205impl RateLimitRule {
206    pub fn new(window: u64, max: u64) -> Self {
207        Self { window, max }
208    }
209}
210
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub struct HybridRateLimitOptions {
213    pub enabled: bool,
214    pub local_multiplier: u64,
215}
216
217impl Default for HybridRateLimitOptions {
218    fn default() -> Self {
219        Self {
220            enabled: false,
221            local_multiplier: 2,
222        }
223    }
224}
225
226impl HybridRateLimitOptions {
227    pub fn new() -> Self {
228        Self::default()
229    }
230
231    pub fn builder() -> Self {
232        Self::new()
233    }
234
235    pub fn enabled() -> Self {
236        Self {
237            enabled: true,
238            ..Self::default()
239        }
240    }
241
242    pub fn disabled() -> Self {
243        Self::default()
244    }
245
246    #[must_use]
247    pub fn set_enabled(mut self, enabled: bool) -> Self {
248        self.enabled = enabled;
249        self
250    }
251
252    #[must_use]
253    pub fn local_multiplier(mut self, multiplier: u64) -> Self {
254        self.local_multiplier = multiplier;
255        self
256    }
257}
258
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub struct RateLimitPathRule {
261    pub path: String,
262    pub rule: Option<RateLimitRule>,
263}
264
265pub trait RateLimitRuleProvider: Send + Sync + 'static {
266    fn resolve(
267        &self,
268        request: &Request<Vec<u8>>,
269        current_rule: &RateLimitRule,
270    ) -> Result<Option<RateLimitRule>, OpenAuthError>;
271}
272
273impl<F> RateLimitRuleProvider for F
274where
275    F: Fn(&Request<Vec<u8>>, &RateLimitRule) -> Result<Option<RateLimitRule>, OpenAuthError>
276        + Send
277        + Sync
278        + 'static,
279{
280    fn resolve(
281        &self,
282        request: &Request<Vec<u8>>,
283        current_rule: &RateLimitRule,
284    ) -> Result<Option<RateLimitRule>, OpenAuthError> {
285        self(request, current_rule)
286    }
287}
288
289#[derive(Clone)]
290pub struct DynamicRateLimitPathRule {
291    pub path: String,
292    pub provider: Arc<dyn RateLimitRuleProvider>,
293}
294
295impl DynamicRateLimitPathRule {
296    pub fn new<P>(path: impl Into<String>, provider: P) -> Self
297    where
298        P: RateLimitRuleProvider,
299    {
300        Self {
301            path: path.into(),
302            provider: Arc::new(provider),
303        }
304    }
305}
306
307impl fmt::Debug for DynamicRateLimitPathRule {
308    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
309        formatter
310            .debug_struct("DynamicRateLimitPathRule")
311            .field("path", &self.path)
312            .field("provider", &"<request-aware>")
313            .finish()
314    }
315}
316
317/// Rate limit storage record.
318#[derive(Debug, Clone, PartialEq, Eq)]
319pub struct RateLimitRecord {
320    pub key: String,
321    pub count: u64,
322    pub last_request: i64,
323}
324
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub struct RateLimitConsumeInput {
327    pub key: String,
328    pub rule: RateLimitRule,
329    pub now_ms: i64,
330}
331
332#[derive(Debug, Clone, PartialEq, Eq)]
333pub struct RateLimitDecision {
334    pub permitted: bool,
335    pub retry_after: u64,
336    pub limit: u64,
337    pub remaining: u64,
338    pub reset_after: u64,
339}
340
341pub type RateLimitFuture<'a> =
342    Pin<Box<dyn Future<Output = Result<RateLimitDecision, OpenAuthError>> + Send + 'a>>;
343
344/// Atomic rate limit storage contract.
345///
346/// Implementations must make the check-and-increment decision in one atomic
347/// operation when used for cross-process or distributed enforcement.
348pub trait RateLimitStore: Send + Sync + 'static {
349    fn consume<'a>(&'a self, input: RateLimitConsumeInput) -> RateLimitFuture<'a>;
350}
351
352/// Synchronous storage contract for router-level rate limiting.
353///
354/// This legacy contract is preserved for compatibility. It is not atomic across
355/// multiple processes unless the implementation makes `get`/`set` externally
356/// serializable.
357pub trait RateLimitStorage: Send + Sync + 'static {
358    fn get(&self, key: &str) -> Result<Option<RateLimitRecord>, OpenAuthError>;
359    fn set(
360        &self,
361        key: &str,
362        value: RateLimitRecord,
363        ttl_seconds: u64,
364        update: bool,
365    ) -> Result<(), OpenAuthError>;
366}
367
368/// Rate limit storage selector.
369#[derive(Debug, Clone, Copy, PartialEq, Eq)]
370pub enum RateLimitStorageOption {
371    Memory,
372    Database,
373    SecondaryStorage,
374}
openauth_core/options/rate_limit.rs

openauth_core/options/
rate_limit.rs
