From b8a2a026d638e5f0b6d87c8bda900a0d0eb90540 Mon Sep 17 00:00:00 2001 From: Aravinth Manivannan Date: Sat, 30 Dec 2023 20:28:05 +0530 Subject: [PATCH] feat: test variable difficulty driver --- src/mcaptcha/cache.rs | 83 ++++++++- src/mcaptcha/defense.rs | 3 + src/mcaptcha/mcaptcha.rs | 376 +++++++++++++++++++++++++++++++++++++-- src/store/mod.rs | 4 +- 4 files changed, 449 insertions(+), 17 deletions(-) diff --git a/src/mcaptcha/cache.rs b/src/mcaptcha/cache.rs index 00dd063..cc11cbf 100644 --- a/src/mcaptcha/cache.rs +++ b/src/mcaptcha/cache.rs @@ -47,13 +47,13 @@ impl HashCache { } pub async fn clean_all_after_cold_start(&self, updated: HashCache) { - updated.difficulty_map.iter().map(|x| { + updated.difficulty_map.iter().for_each(|x| { self.difficulty_map - .insert(x.key().to_owned(), x.value().to_owned()) + .insert(x.key().to_owned(), x.value().to_owned()); }); - updated.result_map.iter().map(|x| { + updated.result_map.iter().for_each(|x| { self.result_map - .insert(x.key().to_owned(), x.value().to_owned()) + .insert(x.key().to_owned(), x.value().to_owned()); }); let cache = self.clone(); let fut = async move { @@ -169,6 +169,81 @@ mod tests { 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.retrive_pow_config(msg.clone()).unwrap(); + let new_cache_difficulty_factor = new_cache.retrive_pow_config(msg.clone()).unwrap(); + + assert_eq!( + DIFFICULTY_FACTOR, + cache_difficulty_factor.unwrap().difficulty_factor + ); + assert_eq!( + DIFFICULTY_FACTOR, + new_cache_difficulty_factor.unwrap().difficulty_factor + ); + + let verify_msg = VerifyCaptchaResult { + key: KEY.into(), + token: RES.into(), + }; + + assert!(new_cache.verify_captcha_result(verify_msg.clone()).unwrap()); + assert!(!new_cache.verify_captcha_result(verify_msg.clone()).unwrap()); + + let duration: Duration = Duration::new(5, 0); + //sleep(DURATION + DURATION).await; + tokio::time::sleep(duration + duration).await; + + let expired_string = cache.retrive_pow_config(msg.clone()).unwrap(); + assert_eq!(None, expired_string); + let expired_string = new_cache.retrive_pow_config(msg).unwrap(); + 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()).unwrap()); + assert!(!cache.verify_captcha_result(verify_msg).unwrap()); + } #[actix_rt::test] async fn hashcache_pow_cache_works() { diff --git a/src/mcaptcha/defense.rs b/src/mcaptcha/defense.rs index 2445ca9..6e7bb40 100644 --- a/src/mcaptcha/defense.rs +++ b/src/mcaptcha/defense.rs @@ -184,6 +184,9 @@ impl Defense { // self.current_visitor_threshold = 0; // } // + pub fn get_levels(&self) -> Vec { + self.levels.clone() + } /// Get current level's visitor threshold pub fn current_level(&self, current_visitor_level: u32) -> &Level { for level in self.levels.iter() { diff --git a/src/mcaptcha/mcaptcha.rs b/src/mcaptcha/mcaptcha.rs index 67e7263..816a481 100644 --- a/src/mcaptcha/mcaptcha.rs +++ b/src/mcaptcha/mcaptcha.rs @@ -24,7 +24,7 @@ use serde::{Deserialize, Serialize}; use super::defense::Defense; use libmcaptcha::errors::*; -use libmcaptcha::master::messages as MasterMessages; +use libmcaptcha::master::messages as ManagerMessages; /// Builder for [MCaptcha] #[derive(Clone, Serialize, Deserialize, Debug)] @@ -123,6 +123,11 @@ impl MCaptcha { } }); } + + /// get [Counter]'s current visitor_threshold + pub fn get_visitors(&self) -> u32 { + self.visitor_threshold.load(Ordering::Relaxed) + } } #[derive(Clone, Serialize, Deserialize)] @@ -132,7 +137,7 @@ pub struct Manager { } impl Manager { - /// add [Counter] actor to [Master] + /// add [Counter] actor to [Manager] pub fn add_captcha(&self, m: Arc, id: String) { self.captchas.insert(id, m); } @@ -146,7 +151,18 @@ impl Manager { } } - /// get [Counter] actor from [Master] + fn gc(captchas: Arc>>) { + println!("running gc"); + for captcha in captchas.iter() { + let visitor = { captcha.value().get_visitors() }; + if visitor == 0 { + println!("Removing captcha"); + captchas.remove(captcha.key()); + } + } + } + + /// get [Counter] actor from [Manager] pub fn get_captcha(&self, id: &str) -> Option> { if let Some(captcha) = self.captchas.get(id) { Some(captcha.clone()) @@ -155,7 +171,7 @@ impl Manager { } } - /// removes [Counter] actor from [Master] + /// removes [Counter] actor from [Manager] pub fn rm_captcha(&self, id: &str) -> Option<(String, Arc)> { self.captchas.remove(id) } @@ -170,9 +186,9 @@ impl Manager { } pub async fn clean_all_after_cold_start(&self, updated: Manager) { - updated.captchas.iter().map(|x| { + updated.captchas.iter().for_each(|x| { self.captchas - .insert(x.key().to_owned(), x.value().to_owned()) + .insert(x.key().to_owned(), x.value().to_owned()); }); let captchas = self.clone(); let keys: Vec = captchas @@ -193,15 +209,22 @@ impl Manager { pub fn add_visitor( &self, - msg: &MasterMessages::AddVisitor, + msg: &ManagerMessages::AddVisitor, ) -> Option { if let Some(captcha) = self.captchas.get(&msg.0) { let difficulty_factor = captcha.add_visitor(); + // let id = msg.0.clone(); let c = captcha.clone(); + let captchas = self.captchas.clone(); let fut = async move { tokio::time::sleep(Duration::new(c.duration, 0)).await; c.decrement_visitor_by(1); + // Self::gc(captchas); + // if c.get_visitors() == 0 { + // println!("Removing captcha addvivi"); + // captchas.remove(&id); + // } }; tokio::spawn(fut); @@ -215,16 +238,25 @@ impl Manager { } } - pub fn get_internal_data(&self) -> Arc>> { - self.captchas.clone() + pub fn get_internal_data(&self) -> HashMap { + let mut res = HashMap::with_capacity(self.captchas.len()); + for value in self.captchas.iter() { + res.insert(value.key().to_owned(), value.value().as_ref().into()); + } + res } pub fn set_internal_data(&self, mut map: HashMap) { for (id, captcha) in map.drain() { let visitors = captcha.get_visitors(); let new_captcha: MCaptcha = (&captcha).into(); - new_captcha.set_visitor_count(visitors); - self.captchas.insert(id, Arc::new(new_captcha)); + let new_captcha = Arc::new(new_captcha); + self.captchas.insert(id.clone(), new_captcha.clone()); + let msg = ManagerMessages::AddVisitor(id); + for _ in 0..visitors { + self.add_visitor(&msg); + println!("Adding visitor"); + } } } } @@ -241,6 +273,328 @@ impl From<&libmcaptcha::mcaptcha::MCaptcha> for MCaptcha { .duration(value.get_duration()) .build() .unwrap(); + // for _ in 0..value.get_visitors() { + // new_captcha.add_visitor(); + // } + new_captcha } } + +impl From<&MCaptcha> for libmcaptcha::mcaptcha::MCaptcha { + fn from(value: &MCaptcha) -> Self { + let mut defense = libmcaptcha::defense::DefenseBuilder::default(); + for level in value.defense.get_levels().drain(0..) { + let _ = defense.add_level(level); + } + let defense = defense.build().unwrap(); + let mut new_captcha = libmcaptcha::mcaptcha::MCaptchaBuilder::default() + .defense(defense) + .duration(value.duration) + .build() + .unwrap(); + for _ in 0..value.get_visitors() { + new_captcha.add_visitor(); + } + new_captcha + } +} + +#[cfg(test)] +mod tests { + use super::*; + use libmcaptcha::defense::LevelBuilder; + use libmcaptcha::master::messages::*; + + pub const LEVEL_1: (u32, u32) = (50, 50); + pub const LEVEL_2: (u32, u32) = (500, 500); + pub const DURATION: u64 = 5; + + use crate::mcaptcha::defense::*; + + pub fn get_defense() -> Defense { + DefenseBuilder::default() + .add_level( + LevelBuilder::default() + .visitor_threshold(LEVEL_1.0) + .difficulty_factor(LEVEL_1.1) + .unwrap() + .build() + .unwrap(), + ) + .unwrap() + .add_level( + LevelBuilder::default() + .visitor_threshold(LEVEL_2.0) + .difficulty_factor(LEVEL_2.1) + .unwrap() + .build() + .unwrap(), + ) + .unwrap() + .build() + .unwrap() + } + + async fn race(manager: &Manager, id: String, count: (u32, u32)) { + let msg = ManagerMessages::AddVisitor(id); + for _ in 0..count.0 as usize - 1 { + manager.add_visitor(&msg); + } + } + + // pub fn get_counter() -> Counter { + // get_mcaptcha().into() + // } + + pub fn get_mcaptcha() -> MCaptcha { + MCaptchaBuilder::default() + .defense(get_defense()) + .duration(DURATION) + .build() + .unwrap() + } + + #[actix_rt::test] + async fn manager_works() { + let manager = Manager::new(1); + + // let get_add_site_msg = |id: String, mcaptcha: MCaptcha| { + // AddSiteBuilder::default() + // .id(id) + // .mcaptcha(mcaptcha) + // .build() + // .unwrap() + // }; + + let id = "yo"; + manager.add_captcha(Arc::new(get_mcaptcha()), id.into()); + + let mcaptcha_addr = manager.get_captcha(id); + assert!(mcaptcha_addr.is_some()); + + let mut mcaptcha_data = manager.get_internal_data(); + mcaptcha_data.get_mut(id).unwrap().add_visitor(); + mcaptcha_data.get_mut(id).unwrap().add_visitor(); + mcaptcha_data.get_mut(id).unwrap().add_visitor(); + println!("{:?}", mcaptcha_data); + // let mcaptcha_data: HashMap = { + // let serialized = serde_json::to_string(&mcaptcha_data).unwrap(); + // serde_json::from_str(&serialized).unwrap() + // }; + // println!("{:?}", mcaptcha_data); + manager.set_internal_data(mcaptcha_data); + + let mcaptcha_data = manager.get_internal_data(); + println!("{:?}", mcaptcha_data); + assert_eq!( + manager.get_captcha(id).unwrap().get_visitors(), + mcaptcha_data.get(id).unwrap().get_visitors() + ); + + let new_id = "yoyo"; + manager.rename(id, new_id.into()); + + { + let mcaptcha_addr = manager.get_captcha(new_id); + assert!(mcaptcha_addr.is_some()); + + let addr_doesnt_exist = manager.get_captcha(id); + assert!(addr_doesnt_exist.is_none()); + + let timer_expire = Duration::new(DURATION, 0); + tokio::time::sleep(timer_expire).await; + tokio::time::sleep(timer_expire).await; + } + + // Manager::gc(manager.captchas.clone()); + // let mcaptcha_addr = manager.get_captcha(new_id); + // assert_eq!(mcaptcha_addr.as_ref().unwrap().get_visitors(), 0); + // assert!(mcaptcha_addr.is_none()); + // + // assert!( + // manager.rm_captcha(new_id.into()).is_some()); + } + + #[actix_rt::test] + async fn counter_defense_works() { + let manager = Manager::new(1); + let id = "yo"; + manager.add_captcha(Arc::new(get_mcaptcha()), id.into()); + + let mut mcaptcha = manager + .add_visitor(&ManagerMessages::AddVisitor(id.to_string())) + .unwrap(); + assert_eq!(mcaptcha.difficulty_factor, LEVEL_1.0); + + race(&manager, id.to_string(), LEVEL_2).await; + mcaptcha = manager + .add_visitor(&ManagerMessages::AddVisitor(id.to_string())) + .unwrap(); + assert_eq!(mcaptcha.difficulty_factor, LEVEL_2.1); + tokio::time::sleep(Duration::new(DURATION * 2, 0)).await; + assert_eq!(manager.get_captcha(id).unwrap().get_visitors(), 0); + } +} +// +//#[cfg(test)] +//pub mod tests { +// use super::*; +// use crate::defense::*; +// use crate::errors::*; +// use crate::mcaptcha; +// use crate::mcaptcha::MCaptchaBuilder; +// +// // constants for testing +// // (visitor count, level) +// pub const LEVEL_1: (u32, u32) = (50, 50); +// pub const LEVEL_2: (u32, u32) = (500, 500); +// pub const DURATION: u64 = 5; +// +// type MyActor = Addr; +// +// pub fn get_defense() -> Defense { +// DefenseBuilder::default() +// .add_level( +// LevelBuilder::default() +// .visitor_threshold(LEVEL_1.0) +// .difficulty_factor(LEVEL_1.1) +// .unwrap() +// .build() +// .unwrap(), +// ) +// .unwrap() +// .add_level( +// LevelBuilder::default() +// .visitor_threshold(LEVEL_2.0) +// .difficulty_factor(LEVEL_2.1) +// .unwrap() +// .build() +// .unwrap(), +// ) +// .unwrap() +// .build() +// .unwrap() +// } +// +// async fn race(addr: Addr, count: (u32, u32)) { +// for _ in 0..count.0 as usize - 1 { +// let _ = addr.send(AddVisitor).await.unwrap(); +// } +// } +// +// pub fn get_counter() -> Counter { +// get_mcaptcha().into() +// } +// +// pub fn get_mcaptcha() -> MCaptcha { +// MCaptchaBuilder::default() +// .defense(get_defense()) +// .duration(DURATION) +// .build() +// .unwrap() +// } +// +// #[test] +// fn mcaptcha_decrement_by_works() { +// let mut m = get_mcaptcha(); +// for _ in 0..100 { +// m.add_visitor(); +// } +// m.decrement_visitor_by(50); +// assert_eq!(m.get_visitors(), 50); +// m.decrement_visitor_by(500); +// assert_eq!(m.get_visitors(), 0); +// } +// + +// +// #[actix_rt::test] +// async fn counter_defense_loosenup_works() { +// //use actix::clock::sleep; +// //use actix::clock::delay_for; +// let addr: MyActor = get_counter().start(); +// +// race(addr.clone(), LEVEL_2).await; +// race(addr.clone(), LEVEL_2).await; +// let mut mcaptcha = addr.send(AddVisitor).await.unwrap(); +// assert_eq!(mcaptcha.difficulty_factor, LEVEL_2.1); +// +// let duration = Duration::new(DURATION, 0); +// sleep(duration).await; +// //delay_for(duration).await; +// +// mcaptcha = addr.send(AddVisitor).await.unwrap(); +// assert_eq!(mcaptcha.difficulty_factor, LEVEL_1.1); +// } +// +// #[test] +// fn test_mcatcptha_builder() { +// let defense = get_defense(); +// let m = MCaptchaBuilder::default() +// .duration(0) +// .defense(defense.clone()) +// .build(); +// +// assert_eq!(m.err(), Some(CaptchaError::CaptchaDurationZero)); +// +// let m = MCaptchaBuilder::default().duration(30).build(); +// assert_eq!( +// m.err(), +// Some(CaptchaError::PleaseSetValue("defense".into())) +// ); +// +// let m = MCaptchaBuilder::default().defense(defense).build(); +// assert_eq!( +// m.err(), +// Some(CaptchaError::PleaseSetValue("duration".into())) +// ); +// } +// +// #[actix_rt::test] +// async fn get_current_visitor_count_works() { +// let addr: MyActor = get_counter().start(); +// +// addr.send(AddVisitor).await.unwrap(); +// addr.send(AddVisitor).await.unwrap(); +// addr.send(AddVisitor).await.unwrap(); +// addr.send(AddVisitor).await.unwrap(); +// let count = addr.send(GetCurrentVisitorCount).await.unwrap(); +// +// assert_eq!(count, 4); +// } +// +// #[actix_rt::test] +// #[should_panic] +// async fn stop_works() { +// let addr: MyActor = get_counter().start(); +// addr.send(Stop).await.unwrap(); +// addr.send(AddVisitor).await.unwrap(); +// } +// +// #[actix_rt::test] +// async fn get_set_internal_data_works() { +// let addr: MyActor = get_counter().start(); +// let mut mcaptcha = addr.send(GetInternalData).await.unwrap(); +// mcaptcha.add_visitor(); +// addr.send(SetInternalData(mcaptcha.clone())).await.unwrap(); +// assert_eq!( +// addr.send(GetInternalData).await.unwrap().get_visitors(), +// mcaptcha.get_visitors() +// ); +// +// let duration = Duration::new(mcaptcha.get_duration() + 3, 0); +// sleep(duration).await; +// assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 0); +// } +// +// #[actix_rt::test] +// async fn bulk_delete_works() { +// let addr: MyActor = get_counter().start(); +// addr.send(AddVisitor).await.unwrap(); +// addr.send(AddVisitor).await.unwrap(); +// assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 2); +// addr.send(BulkDecrement(3)).await.unwrap(); +// assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 0); +// } +//} diff --git a/src/store/mod.rs b/src/store/mod.rs index 67d974d..f66fbe8 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -515,7 +515,7 @@ mod tests { } #[test] - pub fn test_dcche_store() { - openraft::testing::Suite::test_all(provision_dcache_store).unwrap() + fn test_dcache_store() { + openraft::testing::Suite::test_all(provision_dcache_store).unwrap(); } }