dcache/src/mcaptcha/cache.rs

/*
 * mCaptcha - A proof of work based DoS protection system
 * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
 *
 * 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 <http://www.gnu.org/licenses/>.
 */
//! 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<DashMap<String, CachedPoWConfig>>,
    result_map: Arc<DashMap<String, (String, u64)>>,
}

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<CachedPoWConfig> {
        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<CachedPoWConfig> {
        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<String> {
        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));
    }
}