diff --git a/src/defense.rs b/src/defense.rs index 8294292..60414e0 100644 --- a/src/defense.rs +++ b/src/defense.rs @@ -114,6 +114,12 @@ pub struct Defense { current_visitor_threshold: usize, } +impl From for Vec { + fn from(d: Defense) -> Self { + d.levels + } +} + /// Builder struct for [Defense] #[derive(Debug, Clone, PartialEq)] pub struct DefenseBuilder { diff --git a/src/errors.rs b/src/errors.rs index 130a87b..b9232e3 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -18,9 +18,10 @@ //! Errors and Result module use derive_more::{Display, Error}; +use redis::RedisError; /// Error datatype -#[derive(Debug, PartialEq, Display, Clone, Error)] +#[derive(Debug, PartialEq, Display, Error)] #[cfg(not(tarpaulin_include))] pub enum CaptchaError { /// When configuring libmcaptcha, [DefenseBuilder][crate::defense::DefenseBuilder] @@ -80,6 +81,28 @@ pub enum CaptchaError { /// Used in builder structs when a value is not set #[display(fmt = "Please set value: {}", _0)] PleaseSetValue(#[error(not(source))] String), + + /// RedisError + #[display(fmt = "{}", _0)] + RedisError(RedisError), + + /// Weird behaviour from mcaptcha redis module + #[display( + fmt = "Something weird happening with mCaptcha redis module. Please file bug report" + )] + MCaptchaRedisModuleError, +} + +impl From for CaptchaError { + fn from(e: RedisError) -> Self { + Self::RedisError(e) + } +} + +impl From for CaptchaError { + fn from(_: actix::MailboxError) -> Self { + Self::MailboxError + } } /// [Result] datatype for libmcaptcha diff --git a/src/master/embedded/master.rs b/src/master/embedded/master.rs index 438c18a..9506e88 100644 --- a/src/master/embedded/master.rs +++ b/src/master/embedded/master.rs @@ -26,6 +26,7 @@ use actix::dev::*; use log::info; use super::counter::Counter; +use crate::errors::*; use crate::master::Master as MasterTrait; use crate::master::{AddSite, AddVisitor}; @@ -89,24 +90,24 @@ impl Handler for Master { type Result = MessageResult; fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { - let addr = self.get_site(&m.0); - if addr.is_none() { - return MessageResult(None); - } else { - let (tx, rx) = channel(); - let fut = async move { - let config = addr - .unwrap() - .send(super::counter::AddVisitor) - .await - .unwrap(); - - tx.send(config).unwrap(); + let (tx, rx) = channel(); + match self.get_site(&m.0) { + None => tx.send(Ok(None)).unwrap(), + Some(addr) => { + let fut = async move { + match addr.send(super::counter::AddVisitor).await { + Ok(val) => tx.send(Ok(Some(val))).unwrap(), + Err(e) => { + let err: CaptchaError = e.into(); + tx.send(Err(err)).unwrap(); + } + } + } + .into_actor(self); + ctx.spawn(fut); } - .into_actor(self); - ctx.spawn(fut); - return MessageResult(Some(rx)); } + return MessageResult(rx); } } diff --git a/src/master/mod.rs b/src/master/mod.rs index ec524f2..b5ad770 100644 --- a/src/master/mod.rs +++ b/src/master/mod.rs @@ -22,6 +22,8 @@ use actix::dev::*; use derive_builder::Builder; use serde::{Deserialize, Serialize}; +use crate::errors::CaptchaResult; + pub mod embedded; #[allow( unused_variables, @@ -31,6 +33,14 @@ pub mod embedded; unused_macros )] use crate::mcaptcha::*; +#[allow( + unused_variables, + unused_imports, + unused_variables, + dead_code, + unused_macros +)] +pub mod redis; /// Describes actor handler trait impls that are required by a cache implementation pub trait Master: actix::Actor + actix::Handler + actix::Handler {} @@ -39,7 +49,7 @@ pub trait Master: actix::Actor + actix::Handler + actix::Handler + * + * 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::cell::RefCell; +use std::cell::RefMut; +use std::rc::Rc; + +use redis::cluster::ClusterClient; +use redis::RedisError; +//use redis::cluster::ClusterConnection; +use redis::Client; +//use redis::Connection; +use redis::RedisResult; +use redis::Value; +use redis::{aio::Connection, cluster::ClusterConnection}; + +use super::CreateMCaptcha; +use crate::errors::*; +use crate::master::{AddSite, AddVisitor, AddVisitorResult}; + +pub enum RedisConnection { + Single(Rc>), + Cluster(Rc>), +} + +#[allow(dead_code)] +const GET: &str = "MCAPTCHA_CACHE.GET"; +#[allow(dead_code)] +const ADD_VISITOR: &str = "MCAPTCHA_CACHE.ADD_VISITOR"; +#[allow(dead_code)] +const DEL: &str = "MCAPTCHA_CACHE.DELETE_CAPTCHA"; +#[allow(dead_code)] +const ADD_CAPTCHA: &str = "MCAPTCHA_CACHE.ADD_CAPTCHA"; +#[allow(dead_code)] +const CAPTCHA_EXISTS: &str = "MCAPTCHA_CACHE.CAPTCHA_EXISTS"; + +const MODULE_NAME: &str = "mcaptcha_cahce"; +macro_rules! exec { + ($cmd:expr, $con:expr) => { + match *$con { + RedisConnection::Single(con) => $cmd.query_async(&mut *con.borrow_mut()).await, + RedisConnection::Cluster(con) => $cmd.query(&mut *con.borrow_mut()), + } + }; +} + +impl RedisConnection { + pub async fn is_module_loaded(&self) { + let modules: Vec> = exec!(redis::cmd("MODULE").arg(&["LIST"]), &self).unwrap(); + + for list in modules.iter() { + match list.iter().find(|module| module.as_str() == MODULE_NAME) { + Some(_) => println!("module exists"), + None => println!("Module doesn't exist"), + } + } + + let commands = vec![ADD_VISITOR, ADD_CAPTCHA, DEL, CAPTCHA_EXISTS, GET]; + + for cmd in commands.iter() { + match exec!(redis::cmd("COMMAND").arg(&["INFO", cmd]), &self).unwrap() { + Value::Bulk(mut val) => { + let x = val.pop(); + match x { + Some(Value::Nil) => println!("Command: {} doesn't exist", &cmd), + _ => println!("commands {} exists", &cmd), + }; + } + + _ => println!("commands exists"), + }; + } + } + + pub async fn add_visitor(&self, msg: AddVisitor) -> CaptchaResult> { + let res: String = exec!(redis::cmd(ADD_VISITOR).arg(&[msg.0]), &self)?; + let res: AddVisitorResult = serde_json::from_str(&res).unwrap(); + Ok(Some(res)) + } + + pub async fn add_mcaptcha(&self, msg: AddSite) -> CaptchaResult<()> { + let name = msg.id; + let captcha: CreateMCaptcha = msg.mcaptcha.into(); + let payload = serde_json::to_string(&captcha).unwrap(); + exec!(redis::cmd(ADD_CAPTCHA).arg(&[name, payload]), &self)?; + Ok(()) + } + + pub async fn check_captcha_exists(&self, captcha: &str) -> CaptchaResult { + let exists: usize = exec!(redis::cmd(CAPTCHA_EXISTS).arg(&[captcha]), &self)?; + if exists == 1 { + Ok(false) + } else if exists == 0 { + Ok(true) + } else { + log::error!( + "mCaptcha redis module responded with {} when for {}", + exists, + CAPTCHA_EXISTS + ); + Err(CaptchaError::MCaptchaRedisModuleError) + } + } + + pub async fn delete_captcha(&self, captcha: &str) -> CaptchaResult<()> { + exec!(redis::cmd(DEL).arg(&[captcha]), &self)?; + Ok(()) + } + + pub async fn get_visitors(&self, captcha: &str) -> CaptchaResult { + let visitors: usize = exec!(redis::cmd(GET).arg(&[captcha]), &self)?; + Ok(visitors) + } +} diff --git a/src/master/redis/master.rs b/src/master/redis/master.rs new file mode 100644 index 0000000..82e3c67 --- /dev/null +++ b/src/master/redis/master.rs @@ -0,0 +1,118 @@ +/* + * 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::cell::RefCell; +use std::cell::RefMut; +use std::rc::Rc; +use std::sync::mpsc; + +use actix::dev::*; +use redis::cluster::ClusterClient; +use redis::RedisError; +//use redis::cluster::ClusterConnection; +use redis::Client; +//use redis::Connection; +use redis::RedisResult; +use redis::Value; +use redis::{aio::Connection, cluster::ClusterConnection}; +use serde::{Deserialize, Serialize}; + +use crate::defense::Level; +use crate::errors::*; +use crate::master::AddVisitorResult; +use crate::master::{AddSite, AddVisitor, Master as MasterTrait}; + +use super::connection::RedisConnection; + + +#[derive(Clone)] +pub enum Redis { + Single(Client), + Cluster(ClusterClient), +} + +#[derive(Serialize, Deserialize)] +pub struct CreateMCaptcha { + pub levels: Vec, + pub duration: u64, +} +pub struct Master { + pub redis: Redis, + pub con: Rc, +} + +impl Master { + async fn new(redis: Redis) -> Self { + let con = Self::connect(&redis).await; + con.is_module_loaded().await; + let con = Rc::new(con); + let master = Self { redis, con }; + master + } + + async fn connect(redis: &Redis) -> RedisConnection { + match &redis { + Redis::Single(c) => { + let con = c.get_async_connection().await.unwrap(); + RedisConnection::Single(Rc::new(RefCell::new(con))) + } + Redis::Cluster(c) => { + let con = c.get_connection().unwrap(); + RedisConnection::Cluster(Rc::new(RefCell::new(con))) + } + } + } +} + +impl MasterTrait for Master {} + +impl Actor for Master { + type Context = Context; +} + +impl Handler for Master { + type Result = MessageResult; + + fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { + let (tx, rx) = mpsc::channel(); + + let con = Rc::clone(&self.con); + let fut = async move { + let res = con.add_visitor(m).await; + tx.send(res).unwrap() + } + .into_actor(self); + ctx.wait(fut); + MessageResult(rx) + } +} + +impl Handler for Master { + type Result = (); + + fn handle(&mut self, m: AddSite, ctx: &mut Self::Context) -> Self::Result { + let (tx, rx) = mpsc::channel(); + + let con = Rc::clone(&self.con); + let fut = async move { + let res = con.add_mcaptcha(m).await; + tx.send(res).unwrap(); + } + .into_actor(self); + ctx.wait(fut); + } +} diff --git a/src/master/redis/mod.rs b/src/master/redis/mod.rs new file mode 100644 index 0000000..06d0f62 --- /dev/null +++ b/src/master/redis/mod.rs @@ -0,0 +1,37 @@ +/* + * 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::cell::RefCell; +use std::cell::RefMut; +use std::rc::Rc; +use std::sync::mpsc; + +use actix::dev::*; +use redis::cluster::ClusterClient; +use redis::RedisError; +//use redis::cluster::ClusterConnection; +use redis::Client; +//use redis::Connection; +use redis::RedisResult; +use redis::Value; +use redis::{aio::Connection, cluster::ClusterConnection}; +use serde::{Deserialize, Serialize}; + +mod connection; +mod master; + +pub use master::*; diff --git a/src/master/redis_master/mod.rs b/src/master/redis_master/mod.rs deleted file mode 100644 index 43e1119..0000000 --- a/src/master/redis_master/mod.rs +++ /dev/null @@ -1,206 +0,0 @@ -/* - * 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 redis::cluster::ClusterClient; -//use redis::cluster::ClusterConnection; -use redis::Client; -//use redis::Connection; -use redis::RedisResult; -use redis::{aio::Connection, cluster::ClusterConnection}; - -//use crate::errors::*; -use crate::master::{AddSite, AddVisitor, Master as MasterTrait}; - -#[derive(Clone)] -pub enum Redis { - Single(Client), - Cluster(ClusterClient), -} - -pub enum RedisConnection { - Single(Connection), - Cluster(ClusterConnection), -} - -const INCR: &str = "MCAPTCHA_CACHE.COUNT"; -const GET: &str = "MCAPTCHA_CACHE.GET"; - -#[derive(Clone)] -pub struct Master { - pub redis: Redis, -} - -macro_rules! exec { - ($cmd:expr, $con:expr) => { - match $con { - RedisConnection::Single(mut con) => $cmd.query_async(&mut con).await, - RedisConnection::Cluster(mut con) => $cmd.query(&mut con), - } - }; -} - -impl Master { - async fn add_visitor(&mut self, key: &str) { - let mut cmd = redis::cmd(INCR); - cmd.arg(&[key]); - let a: RedisResult = exec!(cmd, self.get_connection().await); - - unimplemented!("Have to check return types of INCR command") - } - - async fn get_visitors(&mut self, key: &str) { - let mut cmd = redis::cmd(GET); - cmd.arg(&[key]); - let a: RedisResult = exec!(cmd, self.get_connection().await); - - unimplemented!("Have to check return types of GET command") - } - - async fn get_connection(&mut self) -> RedisConnection { - match &self.redis { - Redis::Single(c) => { - let con = c.get_async_connection().await.unwrap(); - RedisConnection::Single(con) - } - Redis::Cluster(c) => { - let con = c.get_connection().unwrap(); - RedisConnection::Cluster(con) - } - } - } - - async fn module_is_loaded(&mut self) -> () { - let mut cmd = redis::cmd("COMMAND"); - cmd.arg(&["INFO", INCR]); - let a: RedisResult = exec!(cmd, self.get_connection().await); - - let mut cmd = redis::cmd("COMMAND"); - cmd.arg(&["INFO", GET]); - let a: RedisResult = exec!(cmd, self.get_connection().await); - - unimplemented!("Have to check return types of INFO command") - } -} - -impl MasterTrait for Master {} - -impl Actor for Master { - type Context = Context; -} - -impl Handler for Master { - type Result = MessageResult; - - fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { - let fut = async { - self.add_visitor(&m.0).await; - }; - unimplemented!(); - } -} - -impl Handler for Master { - type Result = (); - - fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result { - unimplemented!(); - } -} - -//#[derive(Clone)] -//pub struct Cluster(pub ClusterClient); -//#[derive(Clone)] -//pub struct Single(pub Client); -// -//pub trait Redis: 'static + Unpin + Clone { -// type Result; -// fn get_connection(&'static self) -> Self::Result; -//} -//impl Redis for Cluster { -// type Result = RedisResult; -// fn get_connection(&self) -> Self::Result { -// self.0.get_connection() -// } -//} -// -//impl Redis for Single { -// type Result = impl Future; -// fn get_connection(&'static self) -> Self::Result { -// self.0.get_async_connection() -// } -//} - -//#[derive(Clone)] -//pub struct Master { -// pub redis: usize, -//} - -//impl MasterTrait for Master {} -// -//impl Actor for Master { -// type Context = Context; -//} -// -//impl Handler for Master { -// type Result = MessageResult; -// -// fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { -// let fut = async { -// let test = "1"; -// }; -// unimplemented!(); -// } -//} -// -//impl Handler for Master { -// type Result = (); -// -// fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result { -// unimplemented!(); -// } -//} - -//pub struct Master { -// pub redis: T, -//} -// -//impl MasterTrait for Master {} -// -//impl Actor for Master { -// type Context = Context; -//} -// -//impl Handler for Master { -// type Result = MessageResult; -// -// fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { -// let fut = async { -// self.redis.get_connection(); -// let test = "1"; -// }; -// unimplemented!(); -// } -//} -// -//impl Handler for Master { -// type Result = (); -// -// fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result { -// unimplemented!(); -// } -//} diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs index a5919c2..fdb9c5a 100644 --- a/src/mcaptcha.rs +++ b/src/mcaptcha.rs @@ -78,6 +78,15 @@ pub struct MCaptcha { duration: u64, } +impl From for crate::master::redis::CreateMCaptcha { + fn from(m: MCaptcha) -> Self { + Self { + levels: m.defense.into(), + duration: m.duration, + } + } +} + impl MCaptcha { /// increments the visitor count by one #[inline] diff --git a/src/system.rs b/src/system.rs index 72ced7c..c869f39 100644 --- a/src/system.rs +++ b/src/system.rs @@ -49,26 +49,30 @@ where pub async fn get_pow(&self, id: String) -> Option { use crate::master::AddVisitor; - let rx = self.master.send(AddVisitor(id.clone())).await.unwrap(); + match self + .master + .send(AddVisitor(id.clone())) + .await + .unwrap() + .recv() + .unwrap() + { + Ok(Some(mcaptcha)) => { + let pow_config = PoWConfig::new(mcaptcha.difficulty_factor, self.pow.salt.clone()); - if rx.is_none() { - return None; + let cache_msg = CachePoWBuilder::default() + .string(pow_config.string.clone()) + .difficulty_factor(mcaptcha.difficulty_factor) + .duration(mcaptcha.duration) + .key(id) + .build() + .unwrap(); + + self.cache.send(cache_msg).await.unwrap().unwrap(); + Some(pow_config) + } + _ => None, } - - let mcaptcha = rx.unwrap().recv().unwrap(); - - let pow_config = PoWConfig::new(mcaptcha.difficulty_factor, self.pow.salt.clone()); - - let cache_msg = CachePoWBuilder::default() - .string(pow_config.string.clone()) - .difficulty_factor(mcaptcha.difficulty_factor) - .duration(mcaptcha.duration) - .key(id) - .build() - .unwrap(); - - self.cache.send(cache_msg).await.unwrap().unwrap(); - Some(pow_config) } /// utility function to verify [Work]