archivist_core/
ratelimit.rs

1//! Cross-platform rate limiting.
2//!
3//! Buckets key on `archivist:ratelimit:<fichub_user_id>` when the caller is
4//! linked to a FicHub account, so a spammer on Telegram gets throttled on
5//! Discord/Matrix/Web simultaneously. When the user is not linked, the bucket
6//! keys on `archivist:ratelimit:<platform>:<ext_id>` (per-platform fallback).
7//!
8//! Sliding-window semantics are approximated with a simple fixed window:
9//! `INCR` + `EXPIRE`; when the counter exceeds `limit` inside the window the
10//! call is rejected.
11
12use redis::aio::ConnectionManager;
13
14use crate::config::BotConfig;
15use crate::error::{BotError, Result};
16
17/// Default max commands per window per user.
18pub const DEFAULT_LIMIT: u64 = 30;
19/// Default window length (seconds).
20pub const DEFAULT_WINDOW_SECS: u64 = 60;
21
22/// Rate limiter over the shared Redis instance.
23#[derive(Clone)]
24pub struct RateLimiter {
25    conn: ConnectionManager,
26    limit: u64,
27    window_secs: u64,
28}
29
30impl RateLimiter {
31    /// Connect to Redis.
32    pub async fn connect(config: &BotConfig) -> Result<Self> {
33        let client = redis::Client::open(config.redis_url.as_str())?;
34        let conn = ConnectionManager::new(client).await?;
35        Ok(Self {
36            conn,
37            limit: DEFAULT_LIMIT,
38            window_secs: DEFAULT_WINDOW_SECS,
39        })
40    }
41
42    /// Key for a linked FicHub user (cross-platform shared bucket).
43    fn user_key(fichub_user_id: i64) -> String {
44        format!("archivist:ratelimit:{fichub_user_id}")
45    }
46
47    /// Key for an unlinked platform identity.
48    fn platform_key(platform: &str, ext_id: &str) -> String {
49        format!("archivist:ratelimit:{platform}:{ext_id}")
50    }
51
52    /// Check whether a call is allowed; increments the bucket.
53    ///
54    /// Returns `Err(BotError::RateLimited { retry_after_secs })` when the
55    /// window budget is exhausted.
56    pub async fn check(&self, platform: &str, ext_id: &str) -> Result<()> {
57        let key = Self::platform_key(platform, ext_id);
58        self.check_key(&key).await
59    }
60
61    /// Check against the cross-platform bucket for a linked FicHub user.
62    pub async fn check_user(&self, platform: &str, ext_id: &str, fichub_user_id: i64) -> Result<()> {
63        let key = Self::user_key(fichub_user_id);
64        self.check_key(&key).await?;
65        // Also count toward the per-platform bucket so unlinked fallback
66        // still has a sane budget.
67        self.check_key(&Self::platform_key(platform, ext_id)).await
68    }
69
70    async fn check_key(&self, key: &str) -> Result<()> {
71        let count: u64 = redis::cmd("INCR")
72            .arg(key)
73            .query_async(&mut self.conn.clone())
74            .await?;
75        if count == 1 {
76            // First hit in this window — set expiry.
77            let _: () = redis::cmd("EXPIRE")
78                .arg(key)
79                .arg(self.window_secs)
80                .query_async(&mut self.conn.clone())
81                .await?;
82        }
83        if count > self.limit {
84            // Compute rough retry-after from the fixed window.
85            let ttl: i64 = redis::cmd("TTL")
86                .arg(key)
87                .query_async(&mut self.conn.clone())
88                .await
89                .unwrap_or(self.window_secs as i64);
90            return Err(BotError::RateLimited {
91                retry_after_secs: ttl.max(1) as u64,
92            });
93        }
94        Ok(())
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn user_key_format() {
104        assert_eq!(
105            RateLimiter::user_key(42),
106            "archivist:ratelimit:42"
107        );
108    }
109
110    #[test]
111    fn platform_key_format() {
112        assert_eq!(
113            RateLimiter::platform_key("matrix", "@u:server"),
114            "archivist:ratelimit:matrix:@u:server"
115        );
116        assert_eq!(
117            RateLimiter::platform_key("telegram", "123"),
118            "archivist:ratelimit:telegram:123"
119        );
120    }
121}
archivist_core/ratelimit.rs

archivist_core/
ratelimit.rs
