throttle_kit/
redis.rs

1use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
2
3use redis::Script;
4use redis::aio::MultiplexedConnection;
5
6use crate::backend::RateLimitBackend;
7use crate::metrics::RateLimitResult;
8use crate::quota::Quota;
9
10/// Redis-backed rate limiter using the GCRA algorithm via a Lua script.
11pub struct RedisBackend {
12    conn: MultiplexedConnection,
13}
14
15impl RedisBackend {
16    /// Create a new `RedisBackend` from an existing `MultiplexedConnection`.
17    pub fn new(conn: MultiplexedConnection) -> Self {
18        Self { conn }
19    }
20
21    /// Connect to a Redis instance at `url` (e.g. `"redis://127.0.0.1/"`).
22    pub async fn connect(url: &str) -> Result<Self, crate::error::RateLimitError> {
23        let client = redis::Client::open(url)
24            .map_err(|e| crate::error::RateLimitError::BackendError(e.to_string()))?;
25        let conn = client
26            .get_multiplexed_async_connection()
27            .await
28            .map_err(|e| crate::error::RateLimitError::BackendError(e.to_string()))?;
29        Ok(Self::new(conn))
30    }
31
32    /// Connect using a `redis::Client`.
33    pub async fn from_client(client: redis::Client) -> Result<Self, crate::error::RateLimitError> {
34        let conn = client
35            .get_multiplexed_async_connection()
36            .await
37            .map_err(|e| crate::error::RateLimitError::BackendError(e.to_string()))?;
38        Ok(Self::new(conn))
39    }
40}
41
42impl RedisBackend {
43    /// Execute the GCRA Lua script against Redis.
44    ///
45    /// Returns `(allowed, remaining, retry_after_ms)`.
46    async fn gcra_check(
47        &self,
48        key: &str,
49        emission: u64,
50        burst: u32,
51        cost: u32,
52        now_ms: u64,
53        ttl_ms: u64,
54    ) -> Result<(bool, u32, u64), crate::error::RateLimitError> {
55        let script = Script::new(GCRA_SCRIPT);
56        let mut conn = self.conn.clone();
57
58        let result: Vec<i64> = script
59            .key(key)
60            .arg(emission)
61            .arg(burst)
62            .arg(cost)
63            .arg(now_ms)
64            .arg(ttl_ms)
65            .invoke_async(&mut conn)
66            .await
67            .map_err(|e| crate::error::RateLimitError::BackendError(e.to_string()))?;
68
69        // The script returns a 3-element table; treat a short reply as
70        // "deny" rather than panicking on attacker- or failure-shaped data.
71        let allowed = result.first().is_some_and(|&v| v != 0);
72        let remaining = result.get(1).copied().unwrap_or(0) as u32;
73        let retry_after_ms = result.get(2).copied().unwrap_or(0) as u64;
74
75        Ok((allowed, remaining, retry_after_ms))
76    }
77}
78
79impl RateLimitBackend for RedisBackend {
80    async fn check(&self, key: &str, quota: &Quota) -> RateLimitResult {
81        let emission = quota.interval().as_millis() as u64;
82        let burst = quota.burst;
83        let cost: u32 = 1;
84        let now_ms = SystemTime::now()
85            .duration_since(UNIX_EPOCH)
86            .unwrap_or_default()
87            .as_millis() as u64;
88        let ttl_ms = emission * burst as u64;
89
90        match self
91            .gcra_check(key, emission, burst, cost, now_ms, ttl_ms)
92            .await
93        {
94            Ok((allowed, remaining, retry_after_ms)) => {
95                let reset_at = if retry_after_ms > 0 {
96                    Instant::now() + Duration::from_millis(retry_after_ms)
97                } else {
98                    Instant::now() + quota.interval()
99                };
100
101                RateLimitResult {
102                    allowed,
103                    remaining: remaining as u64,
104                    reset_at,
105                    limit: burst as u64,
106                    retry_after: if allowed {
107                        None
108                    } else {
109                        Some(Duration::from_millis(retry_after_ms))
110                    },
111                }
112            }
113            Err(crate::error::RateLimitError::BackendError(msg)) => {
114                // On Redis errors, fail open — allow the request.
115                eprintln!("redis backend error, failing open: {msg}");
116                RateLimitResult {
117                    allowed: true,
118                    remaining: 0,
119                    reset_at: Instant::now() + quota.interval(),
120                    limit: quota.burst as u64,
121                    retry_after: None,
122                }
123            }
124            Err(e) => {
125                tracing::warn!("unexpected rate limit error, failing open: {e}");
126                RateLimitResult {
127                    allowed: true,
128                    remaining: 0,
129                    reset_at: Instant::now() + quota.interval(),
130                    limit: quota.burst as u64,
131                    retry_after: None,
132                }
133            }
134        }
135    }
136}
137
138const GCRA_SCRIPT: &str = r#"
139-- GCRA Rate Limiter for Redis
140local key = KEYS[1]
141local emission = tonumber(ARGV[1])
142local burst = tonumber(ARGV[2])
143local cost = tonumber(ARGV[3])
144local now = tonumber(ARGV[4])
145local ttl = tonumber(ARGV[5])
146
147local tat = tonumber(redis.call('GET', key))
148if tat == nil then tat = now end
149
150local increment = emission * cost
151local burst_offset = emission * burst
152local new_tat = math.max(tat, now) + increment
153local allow_at = new_tat - burst_offset
154
155if now >= allow_at then
156    redis.call('SET', key, new_tat, 'PX', ttl)
157    local remaining = math.floor((new_tat - now) / emission)
158    return {1, remaining, 0}
159else
160    -- retry_after is in milliseconds, matching the Rust contract
161    -- (`RateLimitResult::retry_after` is a `Duration` built from this
162    -- value via `from_millis`). Dividing by 1000 here made denials
163    -- claim a retry window 1000x too short.
164    local retry_after = math.ceil(allow_at - now)
165    return {0, 0, retry_after}
166end
167"#;
throttle_kit/redis.rs

throttle_kit/
redis.rs
