/* * mCaptcha - A proof of work based DoS protection system * Copyright © 2021 Aravinth Manivannan * * This program is free software: you can redistribute it and/or modify * it under the terms of the GNU Affero General Public License as * published by the Free Software Foundation, either version 3 of the * License, or (at your option) any later version. * * This program is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU Affero General Public License for more details. * * You should have received a copy of the GNU Affero General Public License * along with this program. If not, see . */ //! In-memory cache implementation that uses [HashMap] use std::sync::Arc; use std::time::Duration; use dashmap::DashMap; use serde::{Deserialize, Serialize}; use libmcaptcha::cache::messages::*; use libmcaptcha::errors::*; #[derive(Clone, Default, Serialize, Deserialize)] /// cache datastructure implementing [Save] pub struct HashCache { difficulty_map: Arc>, result_map: Arc>, } impl HashCache { // save [PoWConfig] to cache fn save_pow_config(&self, config: CachePoW) -> CaptchaResult<()> { let challenge = config.string; let config: CachedPoWConfig = CachedPoWConfig { key: config.key, difficulty_factor: config.difficulty_factor, duration: config.duration, }; if self.difficulty_map.get(&challenge).is_none() { self.difficulty_map.insert(challenge, config); Ok(()) } else { Err(CaptchaError::InvalidPoW) } } pub async fn clean_all_after_cold_start(&self, updated: HashCache) { updated.difficulty_map.iter().for_each(|x| { self.difficulty_map .insert(x.key().to_owned(), x.value().to_owned()); }); updated.result_map.iter().for_each(|x| { self.result_map .insert(x.key().to_owned(), x.value().to_owned()); }); let cache = self.clone(); let fut = async move { for values in cache.result_map.iter() { let inner_cache = cache.clone(); let duration = values.value().1; let key = values.key().to_owned(); let inner_fut = async move { tokio::time::sleep(Duration::new(duration, 0)).await; inner_cache.remove_cache_result(&key); }; tokio::spawn(inner_fut); } for values in cache.difficulty_map.iter() { let inner_cache = cache.clone(); let duration = values.value().duration; let key = values.key().to_owned(); let inner_fut = async move { tokio::time::sleep(Duration::new(duration, 0)).await; inner_cache.remove_pow_config(&key); }; tokio::spawn(inner_fut); } }; tokio::spawn(fut); } // retrieve [PoWConfig] from cache. Deletes config post retrival pub fn retrieve_pow_config(&self, msg: VerifyCaptchaResult) -> Option { if let Some(difficulty_factor) = self.remove_pow_config(&msg.token) { Some(difficulty_factor) } else { None } } // delete [PoWConfig] from cache pub fn remove_pow_config(&self, string: &str) -> Option { self.difficulty_map.remove(string).map(|x| x.1) } // save captcha result fn save_captcha_result(&self, res: CacheResult) { self.result_map.insert(res.token, (res.key, res.duration)); } // verify captcha result pub fn verify_captcha_result(&self, challenge: VerifyCaptchaResult) -> bool { if let Some(captcha_id) = self.remove_cache_result(&challenge.token) { if captcha_id == challenge.key { true } else { false } } else { false } } // delete cache result pub fn remove_cache_result(&self, string: &str) -> Option { self.result_map.remove(string).map(|x| x.1 .0) } pub fn cache_pow(&self, msg: CachePoW) { use std::time::Duration; use tokio::time::sleep; let duration: Duration = Duration::new(msg.duration, 0); let string = msg.string.clone(); let cache = self.clone(); let wait_for = async move { sleep(duration).await; //delay_for(duration).await; cache.remove_pow_config(&string); }; let _ = self.save_pow_config(msg); tokio::spawn(wait_for); } /// cache PoW result pub fn cache_result(&self, msg: CacheResult) { use std::time::Duration; use tokio::time::sleep; let token = msg.token.clone(); msg.token.clone(); msg.token.clone(); msg.token.clone(); let duration: Duration = Duration::new(msg.duration, 0); let cache = self.clone(); let wait_for = async move { sleep(duration).await; //delay_for(duration).await; cache.remove_cache_result(&token); }; tokio::spawn(wait_for); let _ = self.save_captcha_result(msg); } } #[cfg(test)] mod tests { use super::*; use libmcaptcha::master::AddVisitorResult; use libmcaptcha::pow::PoWConfig; use std::time::Duration; #[actix_rt::test] async fn merge_works() { const DIFFICULTY_FACTOR: u32 = 54; const RES: &str = "b"; const DURATION: u64 = 5; const KEY: &str = "mcaptchakey"; let pow: PoWConfig = PoWConfig::new(DIFFICULTY_FACTOR, KEY.into()); //salt is dummy here let cache = HashCache::default(); let new_cache = HashCache::default(); let visitor_result = AddVisitorResult { difficulty_factor: DIFFICULTY_FACTOR, duration: DURATION, }; let string = pow.string.clone(); let msg = CachePoWBuilder::default() .string(pow.string.clone()) .difficulty_factor(DIFFICULTY_FACTOR) .duration(visitor_result.duration) .key(KEY.into()) .build() .unwrap(); cache.cache_pow(msg); let add_cache = CacheResult { key: KEY.into(), token: RES.into(), duration: DURATION, }; cache.cache_result(add_cache.clone()); new_cache.clean_all_after_cold_start(cache.clone()).await; let msg = VerifyCaptchaResult { token: string.clone(), key: KEY.into(), }; let cache_difficulty_factor = cache.retrieve_pow_config(msg.clone()).unwrap(); let new_cache_difficulty_factor = new_cache.retrieve_pow_config(msg.clone()).unwrap(); assert_eq!(DIFFICULTY_FACTOR, cache_difficulty_factor.difficulty_factor); assert_eq!( DIFFICULTY_FACTOR, new_cache_difficulty_factor.difficulty_factor ); let verify_msg = VerifyCaptchaResult { key: KEY.into(), token: RES.into(), }; assert!(new_cache.verify_captcha_result(verify_msg.clone())); assert!(!new_cache.verify_captcha_result(verify_msg.clone())); let duration: Duration = Duration::new(5, 0); //sleep(DURATION + DURATION).await; tokio::time::sleep(duration + duration).await; let expired_string = cache.retrieve_pow_config(msg.clone()); assert_eq!(None, expired_string); let expired_string = new_cache.retrieve_pow_config(msg); assert_eq!(None, expired_string); cache.cache_result(add_cache); new_cache.clean_all_after_cold_start(cache.clone()).await; tokio::time::sleep(duration + duration).await; assert!(!new_cache.verify_captcha_result(verify_msg.clone())); assert!(!cache.verify_captcha_result(verify_msg)); } #[actix_rt::test] async fn hashcache_pow_cache_works() { const DIFFICULTY_FACTOR: u32 = 54; const DURATION: u64 = 5; const KEY: &str = "mcaptchakey"; let cache = HashCache::default(); let pow: PoWConfig = PoWConfig::new(DIFFICULTY_FACTOR, KEY.into()); //salt is dummy here let visitor_result = AddVisitorResult { difficulty_factor: DIFFICULTY_FACTOR, duration: DURATION, }; let string = pow.string.clone(); let msg = CachePoWBuilder::default() .string(pow.string.clone()) .difficulty_factor(DIFFICULTY_FACTOR) .duration(visitor_result.duration) .key(KEY.into()) .build() .unwrap(); cache.cache_pow(msg); let msg = VerifyCaptchaResult { token: string.clone(), key: KEY.into(), }; let cache_difficulty_factor = cache.retrieve_pow_config(msg.clone()).unwrap(); assert_eq!(DIFFICULTY_FACTOR, cache_difficulty_factor.difficulty_factor); let duration: Duration = Duration::new(5, 0); //sleep(DURATION + DURATION).await; tokio::time::sleep(duration + duration).await; let expired_string = cache.retrieve_pow_config(msg); assert_eq!(None, expired_string); } #[actix_rt::test] async fn hashcache_result_cache_works() { const DURATION: u64 = 5; const KEY: &str = "a"; const RES: &str = "b"; let cache = HashCache::default(); // send value to cache // send another value to cache for auto delete // verify_captcha_result // delete // wait for timeout and verify_captcha_result against second value let add_cache = CacheResult { key: KEY.into(), token: RES.into(), duration: DURATION, }; cache.cache_result(add_cache); let verify_msg = VerifyCaptchaResult { key: KEY.into(), token: RES.into(), }; assert!(cache.verify_captcha_result(verify_msg.clone())); // duplicate assert!(!cache.verify_captcha_result(verify_msg)); let verify_msg = VerifyCaptchaResult { key: "cz".into(), token: RES.into(), }; assert!(!cache.verify_captcha_result(verify_msg)); let duration: Duration = Duration::new(5, 0); tokio::time::sleep(duration + duration).await; let verify_msg = VerifyCaptchaResult { key: KEY.into(), token: RES.into(), }; assert!(!cache.verify_captcha_result(verify_msg)); } }