skp_ratelimit/key/
mod.rs

1//! Key extraction for rate limiting.
2//!
3//! This module provides the `Key` trait for extracting rate limiting keys from
4//! HTTP requests, along with pre-built extractors for common patterns.
5//!
6//! # Overview
7//!
8//! Rate limiting keys determine how requests are grouped together. For example:
9//! - Limit by IP address: all requests from the same IP share a quota
10//! - Limit by user ID: all requests from the same user share a quota
11//! - Limit by route: different quotas for different endpoints
12//!
13//! # Example
14//!
15//! ```ignore
16//! use skp_ratelimit::key::{Key, IpKey, CompositeKey};
17//!
18//! // Simple IP-based key
19//! let ip_key = IpKey::new();
20//!
21//! // Composite key: IP + path
22//! let composite = CompositeKey::new(IpKey::new(), PathKey::new());
23//! ```
24
25mod composite;
26mod extractors;
27
28pub use composite::{CompositeKey, CompositeKey3, EitherKey, OptionalKey};
29pub use extractors::*;
30
31/// Trait for extracting rate limiting keys from requests.
32///
33/// The key determines how requests are grouped for rate limiting purposes.
34/// Return `None` if the key cannot be extracted (e.g., missing header),
35/// which typically results in the request being allowed.
36///
37/// # Type Parameters
38///
39/// - `R`: The request type (e.g., `axum::extract::Request`, `actix_web::HttpRequest`)
40pub trait Key<R>: Send + Sync + 'static {
41    /// Extract a rate limiting key from the request.
42    ///
43    /// Returns `None` if the key cannot be extracted, which typically
44    /// means the request should be allowed (fail open).
45    fn extract(&self, request: &R) -> Option<String>;
46
47    /// Get the key name for logging/metrics.
48    fn name(&self) -> &'static str;
49}
50
51/// A constant key that applies the same limit to all requests.
52#[derive(Debug, Clone, Default)]
53pub struct GlobalKey;
54
55impl GlobalKey {
56    /// Create a new global key.
57    pub fn new() -> Self {
58        Self
59    }
60}
61
62impl<R> Key<R> for GlobalKey {
63    fn extract(&self, _request: &R) -> Option<String> {
64        Some("global".to_string())
65    }
66
67    fn name(&self) -> &'static str {
68        "global"
69    }
70}
71
72/// A key that extracts a specific field from the request.
73///
74/// This is a generic extractor that can be configured with a closure.
75#[derive(Clone)]
76pub struct FnKey<F> {
77    extractor: F,
78    name: &'static str,
79}
80
81impl<F> std::fmt::Debug for FnKey<F> {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("FnKey").field("name", &self.name).finish()
84    }
85}
86
87impl<F> FnKey<F> {
88    /// Create a new function-based key extractor.
89    pub fn new(name: &'static str, extractor: F) -> Self {
90        Self { extractor, name }
91    }
92}
93
94impl<R, F> Key<R> for FnKey<F>
95where
96    F: Fn(&R) -> Option<String> + Send + Sync + 'static,
97{
98    fn extract(&self, request: &R) -> Option<String> {
99        (self.extractor)(request)
100    }
101
102    fn name(&self) -> &'static str {
103        self.name
104    }
105}
106
107/// A key that always returns a static value.
108#[derive(Debug, Clone)]
109pub struct StaticKey {
110    key: String,
111}
112
113impl StaticKey {
114    /// Create a new static key.
115    pub fn new(key: impl Into<String>) -> Self {
116        Self { key: key.into() }
117    }
118}
119
120impl<R> Key<R> for StaticKey {
121    fn extract(&self, _request: &R) -> Option<String> {
122        Some(self.key.clone())
123    }
124
125    fn name(&self) -> &'static str {
126        "static"
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn test_global_key() {
136        let key = GlobalKey::new();
137        let request = ();
138        assert_eq!(Key::<()>::extract(&key, &request), Some("global".to_string()));
139        assert_eq!(Key::<()>::name(&key), "global");
140    }
141
142    #[test]
143    fn test_static_key() {
144        let key = StaticKey::new("my-key");
145        let request = ();
146        assert_eq!(key.extract(&request), Some("my-key".to_string()));
147    }
148
149    #[test]
150    fn test_fn_key() {
151        let key: FnKey<fn(&i32) -> Option<String>> = FnKey::new("custom", |_: &i32| Some("from-fn".to_string()));
152        assert_eq!(key.extract(&42), Some("from-fn".to_string()));
153        assert_eq!(key.name(), "custom");
154    }
155}
skp_ratelimit/key/mod.rs

skp_ratelimit/key/
mod.rs
