use rand::Rng; use scc::{HashMap, HashSet}; use std::borrow::Cow; use std::fmt::Write; use std::path::Path; use std::sync::Arc; use std::time::UNIX_EPOCH; use crate::error::AppError; fn get_current_time() -> u64 { UNIX_EPOCH.elapsed().unwrap().as_secs() } fn hash_string(data: &[u8]) -> Result { argon2::hash_encoded(data, "11111111".as_bytes(), &argon2::Config::default()) } fn generate_random_string(len: usize) -> String { rand::thread_rng() .sample_iter(rand::distributions::Alphanumeric) .take(len) .map(|c| c as char) .collect::() } #[derive(Debug, Clone)] pub(crate) struct Tokens { hashed: Arc>>, raw_contents: &'static str, } impl Tokens { pub async fn read(path: impl AsRef) -> Result { let tokens = tokio::fs::read_to_string(path).await?; let this = Self { hashed: Arc::new(HashSet::new()), // this is okay since we only call this once and it will be // used for all of it's lifetime raw_contents: Box::leak(tokens.into_boxed_str()), }; for token in this.raw_contents.lines() { let token = token.trim(); if !token.is_empty() { this.hashed .insert_async(Cow::Borrowed(token)) .await .expect("the set will be empty"); } } Ok(this) } pub async fn write(&self, path: impl AsRef) -> Result<(), AppError> { let mut contents = String::new(); self.hashed .scan_async(|hash| { writeln!(&mut contents, "{hash}").expect("if this fails then too bad") }) .await; tokio::fs::write(path, contents).await.map_err(Into::into) } pub async fn generate(&self) -> Result { let token = generate_random_string(30); let token_hash = hash_string(token.as_bytes())?; self.hashed.insert_async(token_hash.into()).await?; Ok(token) } pub async fn verify(&self, token: impl AsRef) -> Result { let token = token.as_ref(); let token_hash = hash_string(token.as_bytes())?; Ok(self.hashed.contains_async(&Cow::Owned(token_hash)).await) } pub async fn revoke_all(&self) { self.hashed.clear_async().await; } } #[derive(Clone)] pub(crate) struct MusicScopedTokens { map: Arc>, expiry_time: u64, } impl MusicScopedTokens { pub fn new(expiry_time: u64) -> Self { Self { map: Arc::new(HashMap::new()), expiry_time, } } pub async fn generate_for_id(&self, music_id: String) -> String { let data = MusicScopedToken { creation: get_current_time(), music_id, }; let token = generate_random_string(12); let _ = self.map.insert_async(token.clone(), data).await; return token; } pub async fn verify(&self, token: impl AsRef) -> Option { let token = token.as_ref(); let data = self.map.read_async(token, |_, v| v.clone()).await?; if get_current_time() - data.creation > self.expiry_time { self.map.remove_async(token).await; return None; } Some(data.music_id) } } #[derive(Clone)] struct MusicScopedToken { creation: u64, music_id: String, }