archivist-core 0.2.0

//! Cross-platform rate limiting.
//!
//! Buckets key on `archivist:ratelimit:<fichub_user_id>` when the caller is
//! linked to a FicHub account, so a spammer on Telegram gets throttled on
//! Discord/Matrix/Web simultaneously. When the user is not linked, the bucket
//! keys on `archivist:ratelimit:<platform>:<ext_id>` (per-platform fallback).
//!
//! Sliding-window semantics are approximated with a simple fixed window:
//! `INCR` + `EXPIRE`; when the counter exceeds `limit` inside the window the
//! call is rejected.

use redis::aio::ConnectionManager;

use crate::config::BotConfig;
use crate::error::{BotError, Result};

/// Default max commands per window per user.
pub const DEFAULT_LIMIT: u64 = 30;
/// Default window length (seconds).
pub const DEFAULT_WINDOW_SECS: u64 = 60;

/// Rate limiter over the shared Redis instance.
#[derive(Clone)]
pub struct RateLimiter {
    conn: ConnectionManager,
    limit: u64,
    window_secs: u64,
}

impl RateLimiter {
    /// Connect to Redis.
    pub async fn connect(config: &BotConfig) -> Result<Self> {
        let client = redis::Client::open(config.redis_url.as_str())?;
        let conn = ConnectionManager::new(client).await?;
        Ok(Self {
            conn,
            limit: DEFAULT_LIMIT,
            window_secs: DEFAULT_WINDOW_SECS,
        })
    }

    /// Key for a linked FicHub user (cross-platform shared bucket).
    fn user_key(fichub_user_id: i64) -> String {
        format!("archivist:ratelimit:{fichub_user_id}")
    }

    /// Key for an unlinked platform identity.
    fn platform_key(platform: &str, ext_id: &str) -> String {
        format!("archivist:ratelimit:{platform}:{ext_id}")
    }

    /// Check whether a call is allowed; increments the bucket.
    ///
    /// Returns `Err(BotError::RateLimited { retry_after_secs })` when the
    /// window budget is exhausted.
    pub async fn check(&self, platform: &str, ext_id: &str) -> Result<()> {
        let key = Self::platform_key(platform, ext_id);
        self.check_key(&key).await
    }

    /// Check against the cross-platform bucket for a linked FicHub user.
    pub async fn check_user(&self, platform: &str, ext_id: &str, fichub_user_id: i64) -> Result<()> {
        let key = Self::user_key(fichub_user_id);
        self.check_key(&key).await?;
        // Also count toward the per-platform bucket so unlinked fallback
        // still has a sane budget.
        self.check_key(&Self::platform_key(platform, ext_id)).await
    }

    async fn check_key(&self, key: &str) -> Result<()> {
        let count: u64 = redis::cmd("INCR")
            .arg(key)
            .query_async(&mut self.conn.clone())
            .await?;
        if count == 1 {
            // First hit in this window — set expiry.
            let _: () = redis::cmd("EXPIRE")
                .arg(key)
                .arg(self.window_secs)
                .query_async(&mut self.conn.clone())
                .await?;
        }
        if count > self.limit {
            // Compute rough retry-after from the fixed window.
            let ttl: i64 = redis::cmd("TTL")
                .arg(key)
                .query_async(&mut self.conn.clone())
                .await
                .unwrap_or(self.window_secs as i64);
            return Err(BotError::RateLimited {
                retry_after_secs: ttl.max(1) as u64,
            });
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn user_key_format() {
        assert_eq!(
            RateLimiter::user_key(42),
            "archivist:ratelimit:42"
        );
    }

    #[test]
    fn platform_key_format() {
        assert_eq!(
            RateLimiter::platform_key("matrix", "@u:server"),
            "archivist:ratelimit:matrix:@u:server"
        );
        assert_eq!(
            RateLimiter::platform_key("telegram", "123"),
            "archivist:ratelimit:telegram:123"
        );
    }
}