archivist_core/

store.rs

//! Redis-backed auth token store.
//!
//! Maps Discord user ids → FicHub JWT tokens. Tokens live in the same Redis
//! instance FicHub uses (`archivist:token:<discord_id>`), so the bot can be
//! restarted without losing linked accounts. TTL is 30 days by default.

use redis::aio::ConnectionManager;
use serde::{Deserialize, Serialize};

use crate::config::BotConfig;
use crate::error::Result;

/// Default TTL for stored tokens (30 days).
pub const TOKEN_TTL_SECS: u64 = 60 * 60 * 24 * 30;

/// A stored token + the user info it maps to.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredToken {
    /// The FicHub JWT.
    pub token: String,
    /// FicHub user id this token belongs to.
    pub fichub_user_id: i64,
    /// Username on FicHub.
    pub username: String,
    /// ISO timestamp when the link was created.
    pub linked_at: String,
}

/// Redis-backed token store.
#[derive(Clone)]
pub struct TokenStore {
    conn: ConnectionManager,
}

impl TokenStore {
    /// 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 })
    }

    /// Legacy Discord-only key (backward compatible with existing linked
    /// Discord users; new code should use `key_for`).
    fn key(discord_id: &str) -> String {
        format!("archivist:token:{discord_id}")
    }

    /// Platform-scoped key: `archivist:token:<platform>:<ext_id>`.
    fn key_for(platform: &str, ext_id: &str) -> String {
        format!("archivist:token:{platform}:{ext_id}")
    }

    /// Index key mapping a FicHub user id → set of `platform:ext_id` strings
    /// (for the linked-accounts page).
    fn user_index_key(fichub_user_id: i64) -> String {
        format!("archivist:user_tokens:{fichub_user_id}")
    }

    /// Store a token for a Discord user.
    pub async fn set(&self, discord_id: u64, stored: &StoredToken) -> Result<()> {
        let key = Self::key(&discord_id.to_string());
        let data = serde_json::to_vec(stored)?;
        redis::cmd("SETEX")
            .arg(&key)
            .arg(TOKEN_TTL_SECS)
            .arg(data)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// Fetch the stored token for a Discord user.
    pub async fn get(&self, discord_id: &str) -> Result<Option<StoredToken>> {
        let key = Self::key(discord_id);
        let raw: Option<Vec<u8>> = redis::cmd("GET")
            .arg(&key)
            .query_async(&mut self.conn.clone())
            .await?;
        match raw {
            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
            None => Ok(None),
        }
    }

    /// Delete a stored token (unlink).
    pub async fn delete(&self, discord_id: u64) -> Result<()> {
        let key = Self::key(&discord_id.to_string());
        redis::cmd("DEL")
            .arg(&key)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// Refresh a token's TTL (called on each successful authed call).
    pub async fn touch(&self, discord_id: &str) -> Result<()> {
        let key = Self::key(discord_id);
        redis::cmd("EXPIRE")
            .arg(&key)
            .arg(TOKEN_TTL_SECS)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// Store a token for a platform-scoped user id (`matrix:@user:server`,
    /// `telegram:123456`, `slack:U123`, `irc:nick`, `mastodon:user@inst`).
    /// Also maintains the per-FicHub-user index for the linked-accounts page.
    pub async fn set_for(
        &self,
        platform: &str,
        ext_id: &str,
        stored: &StoredToken,
    ) -> Result<()> {
        let key = Self::key_for(platform, ext_id);
        let data = serde_json::to_vec(stored)?;
        redis::cmd("SETEX")
            .arg(&key)
            .arg(TOKEN_TTL_SECS)
            .arg(data)
            .exec_async(&mut self.conn.clone())
            .await?;
        redis::cmd("SADD")
            .arg(Self::user_index_key(stored.fichub_user_id))
            .arg(format!("{platform}:{ext_id}"))
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// Fetch a platform-scoped token.
    pub async fn get_for(&self, platform: &str, ext_id: &str) -> Result<Option<StoredToken>> {
        let key = Self::key_for(platform, ext_id);
        let raw: Option<Vec<u8>> = redis::cmd("GET")
            .arg(&key)
            .query_async(&mut self.conn.clone())
            .await?;
        match raw {
            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
            None => Ok(None),
        }
    }

    /// Delete a platform-scoped token (unlink). Removes the index entry too.
    pub async fn delete_for(&self, platform: &str, ext_id: &str) -> Result<()> {
        if let Some(stored) = self.get_for(platform, ext_id).await? {
            redis::cmd("SREM")
                .arg(Self::user_index_key(stored.fichub_user_id))
                .arg(format!("{platform}:{ext_id}"))
                .exec_async(&mut self.conn.clone())
                .await?;
        }
        redis::cmd("DEL")
            .arg(Self::key_for(platform, ext_id))
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// Refresh a platform-scoped token's TTL.
    pub async fn touch_for(&self, platform: &str, ext_id: &str) -> Result<()> {
        let key = Self::key_for(platform, ext_id);
        redis::cmd("EXPIRE")
            .arg(&key)
            .arg(TOKEN_TTL_SECS)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(())
    }

    /// List all linked platform identities for a FicHub user
    /// (`platform:ext_id` strings, for the web linked-accounts page).
    pub async fn list_for_user(&self, fichub_user_id: i64) -> Result<Vec<String>> {
        let members: Vec<String> = redis::cmd("SMEMBERS")
            .arg(Self::user_index_key(fichub_user_id))
            .query_async(&mut self.conn.clone())
            .await?;
        Ok(members)
    }
}

/// Pending `/link` code: Discord id + creation time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingLink {
    /// Platform user id that initiated the link.
    pub discord_id: u64,
    /// ISO timestamp when the code was created.
    pub created_at: String,
}

/// Store for pending one-time link codes (web-side confirmation flow).
#[derive(Clone)]
pub struct PendingLinkStore {
    conn: ConnectionManager,
    ttl_secs: u64,
}

impl PendingLinkStore {
    /// Connect to Redis (uses `config.redis_url` and `config.link_code_ttl_secs`).
    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,
            ttl_secs: config.link_code_ttl_secs,
        })
    }

    fn key(code: &str) -> String {
        format!("archivist:link:{code}")
    }

    /// Generate + store a pending link code.
    pub async fn create(&self, discord_id: u64) -> Result<String> {
        let code = Self::gen_code();
        let pending = PendingLink {
            discord_id,
            created_at: chrono::Utc::now().to_rfc3339(),
        };
        let data = serde_json::to_vec(&pending)?;
        redis::cmd("SETEX")
            .arg(Self::key(&code))
            .arg(self.ttl_secs)
            .arg(data)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(code)
    }

    /// Consume a pending link code (returns the Discord id it belongs to).
    pub async fn consume(&self, code: &str) -> Result<Option<PendingLink>> {
        let key = Self::key(code);
        let raw: Option<Vec<u8>> = redis::cmd("GET")
            .arg(&key)
            .query_async(&mut self.conn.clone())
            .await?;
        let pending: Option<PendingLink> = match raw {
            Some(bytes) => Some(serde_json::from_slice(&bytes)?),
            None => None,
        };
        // One-time use: delete regardless of whether it existed.
        redis::cmd("DEL")
            .arg(&key)
            .exec_async(&mut self.conn.clone())
            .await?;
        Ok(pending)
    }

    /// Generate a short alphanumeric code (8 chars, unambiguous alphabet).
    fn gen_code() -> String {
        const ALPHABET: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789"; // no I/O/0/1
        let mut out = String::with_capacity(8);
        for _ in 0..8 {
            let idx = (uuid::Uuid::new_v4().as_u128() % ALPHABET.len() as u128) as usize;
            out.push(ALPHABET[idx] as char);
        }
        out
    }
}

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

    #[test]
    fn link_code_format() {
        let code = PendingLinkStore::gen_code();
        assert_eq!(code.len(), 8);
        assert!(code.chars().all(|c| c.is_ascii_alphanumeric()));
        // Alphabet excludes ambiguous chars.
        assert!(!code.contains('I'));
        assert!(!code.contains('O'));
        assert!(!code.contains('0'));
        assert!(!code.contains('1'));
    }

    #[test]
    fn stored_token_roundtrip() {
        let st = StoredToken {
            token: "abc".into(),
            fichub_user_id: 42,
            username: "tester".into(),
            linked_at: "now".into(),
        };
        let bytes = serde_json::to_vec(&st).unwrap();
        let back: StoredToken = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(back.token, "abc");
        assert_eq!(back.fichub_user_id, 42);
        assert_eq!(back.username, "tester");
    }

    #[test]
    fn platform_keys_are_namespaced() {
        assert_eq!(
            TokenStore::key_for("matrix", "@u:server"),
            "archivist:token:matrix:@u:server"
        );
        assert_eq!(
            TokenStore::key_for("telegram", "123"),
            "archivist:token:telegram:123"
        );
        // Legacy Discord key keeps old shape (no platform segment).
        assert_eq!(TokenStore::key("123"), "archivist:token:123");
        assert_ne!(TokenStore::key_for("discord", "123"), TokenStore::key("123"));
    }
}