use redis::aio::ConnectionManager;
use serde::{Deserialize, Serialize};
use crate::config::BotConfig;
use crate::error::Result;
pub const TOKEN_TTL_SECS: u64 = 60 * 60 * 24 * 30;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredToken {
pub token: String,
pub fichub_user_id: i64,
pub username: String,
pub linked_at: String,
}
#[derive(Clone)]
pub struct TokenStore {
conn: ConnectionManager,
}
impl TokenStore {
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 })
}
fn key(discord_id: &str) -> String {
format!("archivist:token:{discord_id}")
}
fn key_for(platform: &str, ext_id: &str) -> String {
format!("archivist:token:{platform}:{ext_id}")
}
fn user_index_key(fichub_user_id: i64) -> String {
format!("archivist:user_tokens:{fichub_user_id}")
}
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(())
}
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),
}
}
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(())
}
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(())
}
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(())
}
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),
}
}
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(())
}
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(())
}
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)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingLink {
pub discord_id: u64,
pub created_at: String,
}
#[derive(Clone)]
pub struct PendingLinkStore {
conn: ConnectionManager,
ttl_secs: u64,
}
impl PendingLinkStore {
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}")
}
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)
}
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,
};
redis::cmd("DEL")
.arg(&key)
.exec_async(&mut self.conn.clone())
.await?;
Ok(pending)
}
fn gen_code() -> String {
const ALPHABET: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789"; 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()));
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"
);
assert_eq!(TokenStore::key("123"), "archivist:token:123");
assert_ne!(TokenStore::key_for("discord", "123"), TokenStore::key("123"));
}
}