From f76569e0b70d1bf32c5baa5f2f30caa84cc91c7d Mon Sep 17 00:00:00 2001 From: realaravinth Date: Mon, 8 Mar 2021 17:02:36 +0530 Subject: [PATCH] utility method for pow gen --- src/cache/hashcache.rs | 2 +- src/cache/mod.rs | 50 +++++++----------- src/lib.rs | 1 + src/master.rs | 8 +-- src/mcaptcha.rs | 12 ++--- src/pow.rs | 116 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 147 insertions(+), 42 deletions(-) create mode 100644 src/pow.rs diff --git a/src/cache/hashcache.rs b/src/cache/hashcache.rs index 0f7e24d..a6ea4d7 100644 --- a/src/cache/hashcache.rs +++ b/src/cache/hashcache.rs @@ -1,11 +1,11 @@ use std::collections::HashMap; -use std::sync::Arc; use actix::prelude::*; use super::messages::*; use super::Save; use crate::errors::*; +use crate::pow::PoWConfig; #[derive(Clone, Default)] pub struct HashCache(HashMap); diff --git a/src/cache/mod.rs b/src/cache/mod.rs index 52728e2..fe6a828 100644 --- a/src/cache/mod.rs +++ b/src/cache/mod.rs @@ -1,3 +1,21 @@ +/* + * 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 . + */ + pub use hashcache::HashCache; use messages::*; @@ -6,10 +24,8 @@ pub mod hashcache; pub trait Save: actix::Actor + actix::Handler + actix::Handler {} pub mod messages { - use std::sync::Arc; - + use crate::pow::PoWConfig; use actix::dev::*; - use serde::Serialize; use crate::errors::*; @@ -22,32 +38,4 @@ pub mod messages { #[derive(Message)] #[rtype(result = "CaptchaResult>")] pub struct Retrive(pub String); - - /// PoW Config that will be sent to clients for generating PoW - #[derive(Clone, Serialize, Debug)] - pub struct PoWConfig { - pub string: String, - pub difficulty_factor: u32, - } - - impl PoWConfig { - pub fn new(m: u32) -> Self { - use std::iter; - - use rand::{distributions::Alphanumeric, rngs::ThreadRng, thread_rng, Rng}; - - let mut rng: ThreadRng = thread_rng(); - - let string = iter::repeat(()) - .map(|()| rng.sample(Alphanumeric)) - .map(char::from) - .take(32) - .collect::(); - - PoWConfig { - string, - difficulty_factor: m, - } - } - } } diff --git a/src/lib.rs b/src/lib.rs index 4bd5380..e5f0b7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -105,6 +105,7 @@ pub mod message { /// message datatypes to interact with [MCaptcha] actor pub mod cache; +pub mod pow; pub use crate::cache::hashcache::HashCache; diff --git a/src/master.rs b/src/master.rs index 05ad0f3..b7a04a1 100644 --- a/src/master.rs +++ b/src/master.rs @@ -52,13 +52,13 @@ impl Actor for Master<'static> { /// Message to increment the visitor count #[derive(Message)] #[rtype(result = "Option>")] -pub struct GetSite<'a>(pub &'a str); +pub struct GetSite(pub String); -impl<'b> Handler> for Master<'static> { - type Result = MessageResult>; +impl Handler for Master<'static> { + type Result = MessageResult; fn handle(&mut self, m: GetSite, _ctx: &mut Self::Context) -> Self::Result { - let addr = self.get_site(m.0); + let addr = self.get_site(&m.0); if addr.is_none() { return MessageResult(None); } else { diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs index 55d12a5..c483bd8 100644 --- a/src/mcaptcha.rs +++ b/src/mcaptcha.rs @@ -172,19 +172,19 @@ impl Handler for MCaptcha { } #[cfg(test)] -mod tests { +pub mod tests { use super::*; use crate::defense::*; // constants for testing // (visitor count, level) - const LEVEL_1: (u32, u32) = (50, 50); - const LEVEL_2: (u32, u32) = (500, 500); - const DURATION: u64 = 10; + pub const LEVEL_1: (u32, u32) = (50, 50); + pub const LEVEL_2: (u32, u32) = (500, 500); + pub const DURATION: u64 = 10; type MyActor = Addr; - fn get_defense() -> Defense { + pub fn get_defense() -> Defense { DefenseBuilder::default() .add_level( LevelBuilder::default() @@ -214,7 +214,7 @@ mod tests { } } - fn get_counter() -> MCaptcha { + pub fn get_counter() -> MCaptcha { MCaptchaBuilder::default() .defense(get_defense()) .duration(DURATION) diff --git a/src/pow.rs b/src/pow.rs new file mode 100644 index 0000000..c00c3bc --- /dev/null +++ b/src/pow.rs @@ -0,0 +1,116 @@ +/* + * 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 . + */ +use actix::dev::*; +use derive_builder::Builder; +use serde::Serialize; + +use crate::cache::messages; +use crate::cache::Save; +use crate::master::Master; + +/// PoW Config that will be sent to clients for generating PoW +#[derive(Clone, Serialize, Debug)] +pub struct PoWConfig { + pub string: String, + pub difficulty_factor: u32, +} +impl PoWConfig { + pub fn new(m: u32) -> Self { + use std::iter; + + use rand::{distributions::Alphanumeric, rngs::ThreadRng, thread_rng, Rng}; + + let mut rng: ThreadRng = thread_rng(); + + let string = iter::repeat(()) + .map(|()| rng.sample(Alphanumeric)) + .map(char::from) + .take(32) + .collect::(); + + PoWConfig { + string, + difficulty_factor: m, + } + } +} + +#[derive(Clone, Builder)] +pub struct Actors { + master: Addr>, + cache: Addr, +} + +impl Actors +where + T: Save, + ::Context: ToEnvelope, +{ + pub async fn get_pow(&self, id: String) -> Option { + use crate::cache::messages::Cache; + use crate::master::GetSite; + use crate::mcaptcha::Visitor; + + let site_addr = self.master.send(GetSite(id)).await.unwrap(); + if site_addr.is_none() { + return None; + } + let difficulty_factor = site_addr.unwrap().send(Visitor).await.unwrap(); + let pow_config = PoWConfig::new(difficulty_factor); + self.cache + .send(Cache(pow_config.clone())) + .await + .unwrap() + .unwrap(); + Some(pow_config) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::HashCache; + use crate::master::*; + use crate::mcaptcha::tests::*; + + #[actix_rt::test] + async fn get_pow_works() { + let master = Master::new().start(); + let mcaptcha = get_counter().start(); + let mcaptcha_name = "batsense.net"; + + let cache = HashCache::default().start(); + let msg = AddSiteBuilder::default() + .id(mcaptcha_name.into()) + .addr(mcaptcha.clone()) + .build() + .unwrap(); + + master.send(msg).await.unwrap(); + + let actors = ActorsBuilder::default() + .master(master) + .cache(cache) + .build() + .unwrap(); + + let pow = actors.get_pow(mcaptcha_name.into()).await.unwrap(); + + assert_eq!(pow.difficulty_factor, LEVEL_1.0); + } +}