axum_limit/backend/
mod.rs

1mod memory;
2
3#[cfg(feature = "redis")]
4mod redis;
5
6use crate::codec::CodecError;
7use crate::policy::{PolicyState, RateLimitPolicy};
8use crate::quota::Quota;
9use crate::snapshot::RateLimitSnapshot;
10use async_trait::async_trait;
11use std::error::Error;
12use std::fmt::{Display, Formatter, Result as FmtResult};
13
14pub use memory::MemoryBackend;
15
16#[cfg(feature = "redis")]
17pub use redis::RedisBackend;
18
19/// Errors returned by storage backends.
20#[derive(Debug)]
21pub enum BackendError {
22    /// Failed to encode or decode policy state.
23    Codec(CodecError),
24    /// Optimistic transaction failed after too many retries.
25    Contention,
26    /// Redis returned an error.
27    #[cfg(feature = "redis")]
28    Redis(::redis::RedisError),
29}
30
31impl Display for BackendError {
32    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
33        match self {
34            BackendError::Codec(error) => write!(f, "{error}"),
35            BackendError::Contention => write!(f, "rate limit storage contention"),
36            #[cfg(feature = "redis")]
37            BackendError::Redis(error) => write!(f, "redis error: {error}"),
38        }
39    }
40}
41
42impl Error for BackendError {
43    fn source(&self) -> Option<&(dyn Error + 'static)> {
44        match self {
45            BackendError::Codec(error) => Some(error),
46            BackendError::Contention => None,
47            #[cfg(feature = "redis")]
48            BackendError::Redis(error) => Some(error),
49        }
50    }
51}
52
53impl From<CodecError> for BackendError {
54    fn from(value: CodecError) -> Self {
55        Self::Codec(value)
56    }
57}
58
59/// Encodes a subject key into a stable storage representation.
60pub trait StorageKey: crate::Key {
61    /// Returns a stable string representation for storage backends.
62    fn storage_key(&self) -> String;
63}
64
65/// Storage backend for rate limit policy state.
66///
67/// Implement this trait to provide custom backends such as database or Consul storage.
68#[async_trait]
69pub trait RateLimitBackend: Send + Sync + Clone + 'static {
70    /// Error type returned by this backend.
71    type Error: Error + Send + Sync + 'static;
72
73    /// Namespace used to isolate keys for this backend instance.
74    fn namespace(&self) -> &str;
75
76    /// Atomically loads policy state, applies the rate limit algorithm, and persists the result.
77    async fn transact<P>(
78        &self,
79        storage_key: &str,
80        quota: Quota,
81        now_ms: u64,
82    ) -> Result<RateLimitSnapshot, Self::Error>
83    where
84        P: RateLimitPolicy;
85}
86
87/// Builds a storage key for a subject and quota fingerprint.
88pub fn build_storage_key<P>(namespace: &str, subject: &str, quota: Quota) -> String
89where
90    P: RateLimitPolicy,
91{
92    let fingerprint = quota.fingerprint();
93    let burst = fingerprint
94        .burst
95        .map(|value| value.to_string())
96        .unwrap_or_else(|| "-".to_string());
97
98    format!(
99        "{namespace}:{}:{}:{}:{}:{subject}",
100        P::STATE_ID,
101        fingerprint.max,
102        fingerprint.per_ms,
103        burst,
104    )
105}
106
107/// Applies a rate limit policy to optional encoded state and returns updated bytes and snapshot.
108///
109/// Custom backends can use this helper inside [`RateLimitBackend::transact`].
110pub fn apply_policy<P>(
111    bytes: Option<&[u8]>,
112    quota: Quota,
113    now_ms: u64,
114) -> Result<(Vec<u8>, RateLimitSnapshot), CodecError>
115where
116    P: RateLimitPolicy,
117{
118    let mut state = match bytes {
119        Some(payload) => P::State::decode(payload, quota)?,
120        None => P::State::create(quota, now_ms),
121    };
122    let snapshot = state.try_acquire(now_ms);
123    Ok((state.encode()?, snapshot))
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use crate::policy::{FixedWindowPolicy, TokenBucketPolicy};
130    use crate::quota::Quota;
131
132    #[test]
133    fn build_storage_key_includes_policy_quota_and_subject() {
134        let quota = Quota::with_burst(5, 1000, 10);
135        let key = build_storage_key::<TokenBucketPolicy>("ns", "user-1", quota);
136        assert_eq!(key, "ns:token_bucket:5:1000:10:user-1");
137    }
138
139    #[test]
140    fn different_quotas_produce_different_storage_keys() {
141        let subject = "shared";
142        let key_a = build_storage_key::<FixedWindowPolicy>("ns", subject, Quota::new(1, 1000));
143        let key_b = build_storage_key::<FixedWindowPolicy>("ns", subject, Quota::new(5, 1000));
144        assert_ne!(key_a, key_b);
145    }
146}
axum_limit/backend/mod.rs

axum_limit/backend/
mod.rs
