From 470f1fcba4cd827198069b53606054f3b8db6e0c Mon Sep 17 00:00:00 2001 From: realaravinth Date: Sun, 7 Mar 2021 19:54:41 +0530 Subject: [PATCH] site manager --- Cargo.lock | 79 +++++++++++++++--- Cargo.toml | 8 +- src/cache/hashcache.rs | 25 +++--- src/cache/mod.rs | 29 ++++--- src/lib.rs | 5 +- src/master.rs | 180 +++++++++++++++++++++++++++++++++++++++++ src/mcaptcha.rs | 92 ++++++--------------- 7 files changed, 304 insertions(+), 114 deletions(-) create mode 100644 src/master.rs diff --git a/Cargo.lock b/Cargo.lock index e524e10..bd10010 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -541,7 +541,18 @@ checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" dependencies = [ "cfg-if 1.0.0", "libc", - "wasi", + "wasi 0.9.0+wasi-snapshot-preview1", +] + +[[package]] +name = "getrandom" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9495705279e7140bf035dde1f6e750c162df8b625267cd52cc44e0b156732c8" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "wasi 0.10.2+wasi-snapshot-preview1", ] [[package]] @@ -716,12 +727,10 @@ version = "0.1.0" dependencies = [ "actix", "actix-rt", - "async-trait", "derive_builder", "derive_more", - "lazy_static", "pow_sha256 0.1.0", - "rand", + "rand 0.8.3", "serde", "serde_json", ] @@ -1030,11 +1039,23 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" dependencies = [ - "getrandom", + "getrandom 0.1.16", "libc", - "rand_chacha", - "rand_core", - "rand_hc", + "rand_chacha 0.2.2", + "rand_core 0.5.1", + "rand_hc 0.2.0", +] + +[[package]] +name = "rand" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ef9e7e66b4468674bfcb0c81af8b7fa0bb154fa9f28eb840da5c447baeb8d7e" +dependencies = [ + "libc", + "rand_chacha 0.3.0", + "rand_core 0.6.2", + "rand_hc 0.3.0", ] [[package]] @@ -1044,7 +1065,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.5.1", +] + +[[package]] +name = "rand_chacha" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e12735cf05c9e10bf21534da50a147b924d555dc7a547c42e6bb2d5b6017ae0d" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.2", ] [[package]] @@ -1053,7 +1084,16 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" dependencies = [ - "getrandom", + "getrandom 0.1.16", +] + +[[package]] +name = "rand_core" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34cf66eb183df1c5876e2dcf6b13d57340741e8dc255b48e40a26de954d06ae7" +dependencies = [ + "getrandom 0.2.2", ] [[package]] @@ -1062,7 +1102,16 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" dependencies = [ - "rand_core", + "rand_core 0.5.1", +] + +[[package]] +name = "rand_hc" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3190ef7066a446f2e7f42e239d161e905420ccab01eb967c9eb27d21b2322a73" +dependencies = [ + "rand_core 0.6.2", ] [[package]] @@ -1321,7 +1370,7 @@ dependencies = [ "idna", "lazy_static", "log", - "rand", + "rand 0.7.3", "smallvec", "thiserror", "tokio", @@ -1420,6 +1469,12 @@ version = "0.9.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" +[[package]] +name = "wasi" +version = "0.10.2+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6" + [[package]] name = "wasm-bindgen" version = "0.2.71" diff --git a/Cargo.toml b/Cargo.toml index b9e59b5..6b281e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,6 @@ members = [ ".", "browser", "cli" ] [dependencies] actix = "0.10" -actix-rt = "1" serde = "1.0.114" serde_json = "1" @@ -24,9 +23,8 @@ serde_json = "1" derive_builder = "0.9" derive_more = "0.99" -lazy_static = "1.4" - -rand = "0.7" +rand = "0.8" pow_sha256 = { version = "0.1", git = "https://github.com/mcaptcha/pow_sha256" } -async-trait = "0.1.44" +[dev-dependencies] +actix-rt = "1" diff --git a/src/cache/hashcache.rs b/src/cache/hashcache.rs index 4e54eff..0f7e24d 100644 --- a/src/cache/hashcache.rs +++ b/src/cache/hashcache.rs @@ -11,14 +11,13 @@ use crate::errors::*; pub struct HashCache(HashMap); impl HashCache { - fn save(&mut self, config: Arc) -> CaptchaResult<()> { - self.0 - .insert(config.string.clone(), config.difficulty_factor); + fn save(&mut self, config: PoWConfig) -> CaptchaResult<()> { + self.0.insert(config.string, config.difficulty_factor); Ok(()) } - fn retrive(&mut self, string: Arc) -> CaptchaResult> { - if let Some(difficulty_factor) = self.0.get(&*string) { + fn retrive(&mut self, string: String) -> CaptchaResult> { + if let Some(difficulty_factor) = self.0.get(&string) { Ok(Some(difficulty_factor.to_owned())) } else { Ok(None) @@ -35,11 +34,6 @@ impl Actor for HashCache { impl Handler for HashCache { type Result = MessageResult; fn handle(&mut self, msg: Cache, _ctx: &mut Self::Context) -> Self::Result { - // if let Err(e) = self.save(msg.0.clone()) { - // MessageResult(Err(e)) - // } else { - // MessageResult(Ok(msg.0)) - // } MessageResult(self.save(msg.0)) } } @@ -47,7 +41,7 @@ impl Handler for HashCache { impl Handler for HashCache { type Result = MessageResult; fn handle(&mut self, msg: Retrive, _ctx: &mut Self::Context) -> Self::Result { - MessageResult(self.retrive(msg.0.clone())) + MessageResult(self.retrive(msg.0)) } } @@ -56,10 +50,13 @@ mod tests { use super::*; #[actix_rt::test] - async fn counter_defense_tightenup_works() { + async fn hashcache_works() { let addr = HashCache::default().start(); - let p = Arc::new("ewerw".to_string()); - addr.send(Retrive(p)).await.unwrap().unwrap(); + let cache: PoWConfig = PoWConfig::new(54); + let string = cache.string.clone(); + addr.send(Cache(cache)).await.unwrap().unwrap(); + let difficulty_factor = addr.send(Retrive(string)).await.unwrap().unwrap(); + assert_eq!(difficulty_factor.unwrap(), 54); } // // #[actix_rt::test] diff --git a/src/cache/mod.rs b/src/cache/mod.rs index 2db9cea..52728e2 100644 --- a/src/cache/mod.rs +++ b/src/cache/mod.rs @@ -9,22 +9,19 @@ pub mod messages { use std::sync::Arc; use actix::dev::*; - use rand::{distributions::Alphanumeric, thread_rng, Rng}; use serde::Serialize; - use super::Save; use crate::errors::*; - use crate::mcaptcha::MCaptcha; /// Message to decrement the visitor count #[derive(Message)] #[rtype(result = "CaptchaResult<()>")] - pub struct Cache(pub Arc); + pub struct Cache(pub PoWConfig); /// Message to decrement the visitor count #[derive(Message)] #[rtype(result = "CaptchaResult>")] - pub struct Retrive(pub Arc); + pub struct Retrive(pub String); /// PoW Config that will be sent to clients for generating PoW #[derive(Clone, Serialize, Debug)] @@ -34,14 +31,22 @@ pub mod messages { } impl PoWConfig { - pub fn new(m: &MCaptcha) -> Self - where - T: Save, - ::Context: ToEnvelope + ToEnvelope, - { + 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: thread_rng().sample_iter(&Alphanumeric).take(32).collect(), - difficulty_factor: m.get_difficulty(), + string, + difficulty_factor: m, } } } diff --git a/src/lib.rs b/src/lib.rs index 7350cac..4bd5380 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -74,14 +74,14 @@ //! .build() //! .unwrap(); //! -//! let cache = HashCache::default().start(); +//! //let cache = HashCache::default().start(); //! //! // create and start MCaptcha actor //! let mcaptcha = MCaptchaBuilder::default() //! .defense(defense) //! // leaky bucket algorithm's emission interval //! .duration(30) -//! .cache(cache) +//! // .cache(cache) //! .build() //! .unwrap() //! .start(); @@ -95,6 +95,7 @@ pub mod defense; pub mod errors; +pub mod master; pub mod mcaptcha; /// message datatypes to interact with [MCaptcha] actor diff --git a/src/master.rs b/src/master.rs new file mode 100644 index 0000000..05ad0f3 --- /dev/null +++ b/src/master.rs @@ -0,0 +1,180 @@ +/* + * 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 std::collections::BTreeMap; + +use actix::dev::*; +use derive_builder::Builder; + +use crate::mcaptcha::MCaptcha; + +/// This struct represents the mCaptcha state and is used +/// to configure leaky-bucket lifetime and manage defense +#[derive(Clone)] +pub struct Master<'a> { + sites: BTreeMap<&'a str, Addr>, +} + +impl Master<'static> { + pub fn add_site(&mut self, details: AddSite) { + self.sites.insert(details.id, details.addr.to_owned()); + } + + pub fn new() -> Self { + Master { + sites: BTreeMap::new(), + } + } + + pub fn get_site<'a, 'b>(&'a self, id: &'b str) -> Option<&'a Addr> { + self.sites.get(id) + } +} + +impl Actor for Master<'static> { + type Context = Context; +} + +/// Message to increment the visitor count +#[derive(Message)] +#[rtype(result = "Option>")] +pub struct GetSite<'a>(pub &'a str); + +impl<'b> 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); + if addr.is_none() { + return MessageResult(None); + } else { + return MessageResult(Some(addr.unwrap().clone())); + } + } +} + +/// Message to increment the visitor count +#[derive(Message, Builder)] +#[rtype(result = "()")] +pub struct AddSite { + pub id: &'static str, + pub addr: Addr, +} + +impl Handler for Master<'static> { + type Result = (); + + fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result { + self.add_site(m); + } +} + +///// Message to decrement the visitor count +//#[derive(Message, Deserialize)] +//#[rtype(result = "()")] +//pub struct VerifyPoW { +// pow: ShaPoW>, +// id: String, +//} +// +//impl Handler for MCaptcha { +// type Result = (); +// fn handle(&mut self, msg: VerifyPoW, _ctx: &mut Self::Context) -> Self::Result { +// self.decrement_visiotr(); +// } +//} + +#[cfg(test)] +mod tests { + use super::*; + use crate::defense::*; + + // use crate::cache::HashCache; + // + // // constants for testing + // // (visitor count, level) + const LEVEL_1: (u32, u32) = (50, 50); + const LEVEL_2: (u32, u32) = (500, 500); + const DURATION: u64 = 10; + + type MyActor = Addr; + + 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() + } + + fn get_counter() -> MCaptcha { + use crate::MCaptchaBuilder; + + MCaptchaBuilder::default() + .defense(get_defense()) + .duration(DURATION) + .build() + .unwrap() + } + + #[actix_rt::test] + async fn master() { + let addr = Master::new().start(); + + let id = "yo"; + let mcaptcha = get_counter().start(); + let msg = AddSiteBuilder::default() + .id(id) + .addr(mcaptcha) + .build() + .unwrap(); + addr.send(msg).await.unwrap(); + } + // + // #[actix_rt::test] + // async fn counter_defense_loosenup_works() { + // 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 difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); + // assert_eq!(difficulty_factor, LEVEL_2.1); + // + // let duration = Duration::new(DURATION, 0); + // delay_for(duration).await; + // + // difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); + // assert_eq!(difficulty_factor, LEVEL_1.1); + // } +} diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs index 03867a1..55d12a5 100644 --- a/src/mcaptcha.rs +++ b/src/mcaptcha.rs @@ -56,12 +56,11 @@ //! .unwrap(); //! //! // create and start MCaptcha actor -//! let cache = HashCache::default().start(); +//! //let cache = HashCache::default().start(); //! let mcaptcha = MCaptchaBuilder::default() //! .defense(defense) //! // leaky bucket algorithm's emission interval //! .duration(30) -//! .cache(cache) //! .build() //! .unwrap() //! .start(); @@ -73,37 +72,26 @@ //! } //! ``` -use std::sync::Arc; use std::time::Duration; use actix::dev::*; use derive_builder::Builder; use pow_sha256::PoW as ShaPoW; -use rand::{distributions::Alphanumeric, thread_rng, Rng}; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; -use crate::cache::{messages::*, HashCache, Save}; use crate::defense::Defense; -use crate::errors::*; /// This struct represents the mCaptcha state and is used /// to configure leaky-bucket lifetime and manage defense #[derive(Clone, Builder)] -pub struct MCaptcha -where - T: Save, -{ +pub struct MCaptcha { #[builder(default = "0", setter(skip))] visitor_threshold: u32, defense: Defense, duration: u64, - cache: Addr, } -impl MCaptcha -where - T: Save, -{ +impl MCaptcha { /// incerment visiotr count by one pub fn add_visitor(&mut self) { self.visitor_threshold += 1; @@ -126,10 +114,7 @@ where self.defense.get_difficulty() } } -impl Actor for MCaptcha -where - T: Save, -{ +impl Actor for MCaptcha { type Context = Context; } @@ -138,11 +123,7 @@ where #[rtype(result = "()")] struct DeleteVisitor; -impl Handler for MCaptcha -where - T: Save, - // ::Context: ToEnvelope + ToEnvelope, -{ +impl Handler for MCaptcha { type Result = (); fn handle(&mut self, _msg: DeleteVisitor, _ctx: &mut Self::Context) -> Self::Result { self.decrement_visiotr(); @@ -151,19 +132,14 @@ where /// Message to increment the visitor count #[derive(Message)] -#[rtype(result = "CaptchaResult")] +#[rtype(result = "u32")] pub struct Visitor; -impl Handler for MCaptcha -where - T: Save, - ::Context: ToEnvelope + ToEnvelope, -{ - type Result = ResponseActFuture>; +impl Handler for MCaptcha { + type Result = u32; fn handle(&mut self, _: Visitor, ctx: &mut Self::Context) -> Self::Result { use actix::clock::delay_for; - use actix::fut::wrap_future; let addr = ctx.address(); @@ -176,16 +152,7 @@ where ctx.spawn(wait_for); self.add_visitor(); - let res = Arc::new(PoWConfig::new(&self)); - - let act_fut = wrap_future::<_, Self>(self.cache.send(Cache(res.clone()))).map( - |result, _actor, _ctx| match result { - Ok(Ok(())) => Ok(Arc::try_unwrap(res).unwrap()), - Ok(Err(e)) => Err(e), - Err(_) => Err(CaptchaError::MailboxError), //TODO do typecasting from mailbox error to captcha error - }, - ); - Box::pin(act_fut) + self.get_difficulty() } } @@ -197,13 +164,9 @@ pub struct VerifyPoW { id: String, } -impl Handler for MCaptcha -where - T: Save, - ::Context: ToEnvelope + ToEnvelope, -{ +impl Handler for MCaptcha { type Result = (); - fn handle(&mut self, msg: VerifyPoW, _ctx: &mut Self::Context) -> Self::Result { + fn handle(&mut self, _: VerifyPoW, _ctx: &mut Self::Context) -> Self::Result { self.decrement_visiotr(); } } @@ -219,8 +182,7 @@ mod tests { const LEVEL_2: (u32, u32) = (500, 500); const DURATION: u64 = 10; - type MyActor = Addr>; - type CacheAddr = Addr; + type MyActor = Addr; fn get_defense() -> Defense { DefenseBuilder::default() @@ -246,23 +208,15 @@ mod tests { .unwrap() } - async fn race(addr: Addr>, count: (u32, u32)) - where - // Actor + Handler, - T: Save, - ::Context: ToEnvelope + ToEnvelope, - { + async fn race(addr: Addr, count: (u32, u32)) { for _ in 0..count.0 as usize - 1 { let _ = addr.send(Visitor).await.unwrap(); } } - fn get_counter() -> MCaptcha { - use actix::prelude::*; - let cache: CacheAddr = HashCache::default().start(); + fn get_counter() -> MCaptcha { MCaptchaBuilder::default() .defense(get_defense()) - .cache(cache) .duration(DURATION) .build() .unwrap() @@ -272,12 +226,12 @@ mod tests { async fn counter_defense_tightenup_works() { let addr: MyActor = get_counter().start(); - let mut difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); - assert_eq!(difficulty_factor.difficulty_factor, LEVEL_1.0); + let mut difficulty_factor = addr.send(Visitor).await.unwrap(); + assert_eq!(difficulty_factor, LEVEL_1.0); race(addr.clone(), LEVEL_2).await; - difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); - assert_eq!(difficulty_factor.difficulty_factor, LEVEL_2.1); + difficulty_factor = addr.send(Visitor).await.unwrap(); + assert_eq!(difficulty_factor, LEVEL_2.1); } #[actix_rt::test] @@ -287,13 +241,13 @@ mod tests { race(addr.clone(), LEVEL_2).await; race(addr.clone(), LEVEL_2).await; - let mut difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); - assert_eq!(difficulty_factor.difficulty_factor, LEVEL_2.1); + let mut difficulty_factor = addr.send(Visitor).await.unwrap(); + assert_eq!(difficulty_factor, LEVEL_2.1); let duration = Duration::new(DURATION, 0); delay_for(duration).await; - difficulty_factor = addr.send(Visitor).await.unwrap().unwrap(); - assert_eq!(difficulty_factor.difficulty_factor, LEVEL_1.1); + difficulty_factor = addr.send(Visitor).await.unwrap(); + assert_eq!(difficulty_factor, LEVEL_1.1); } }