From 07f124fb82662f81c070ce77cd8a3317e47cace8 Mon Sep 17 00:00:00 2001 From: realaravinth Date: Wed, 9 Jun 2021 14:03:19 +0530 Subject: [PATCH] redis generic client --- examples/simple.rs | 2 +- src/lib.rs | 17 +++-- src/master/embedded/counter.rs | 2 +- src/master/embedded/master.rs | 4 +- src/master/mod.rs | 77 ++++++++++------------- src/master/redis/connection.rs | 107 +++++++++++++++---------------- src/master/redis/master.rs | 103 +++++++----------------------- src/master/redis/mod.rs | 21 +------ src/redis.rs | 112 +++++++++++++++++++++++++++++++++ src/system.rs | 7 +-- 10 files changed, 239 insertions(+), 213 deletions(-) create mode 100644 src/redis.rs diff --git a/examples/simple.rs b/examples/simple.rs index 9b9e2c1..09f35d9 100644 --- a/examples/simple.rs +++ b/examples/simple.rs @@ -1,7 +1,7 @@ use libmcaptcha::{ cache::{messages::VerifyCaptchaResult, HashCache}, master::embedded::master::Master, - master::AddSiteBuilder, + master::messages::AddSiteBuilder, pow::{ConfigBuilder, Work}, system::SystemBuilder, DefenseBuilder, LevelBuilder, MCaptchaBuilder, diff --git a/src/lib.rs b/src/lib.rs index 9bce453..e2f8e23 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -49,7 +49,7 @@ //! use libmcaptcha::{ //! cache::{messages::VerifyCaptchaResult, HashCache}, //! master::embedded::master:: Master, -//! master::AddSiteBuilder, +//! master::messages::AddSiteBuilder, //! pow::{ConfigBuilder, Work}, //! system::SystemBuilder, //! DefenseBuilder, LevelBuilder, MCaptchaBuilder, @@ -187,6 +187,9 @@ pub mod defense; pub mod errors; pub mod master; +#[cfg(feature = "full")] +mod redis; + /// message datatypes to interact with [MCaptcha] actor #[cfg(feature = "full")] pub mod cache; @@ -200,18 +203,20 @@ mod utils; #[cfg(feature = "full")] pub use crate::cache::hashcache::HashCache; - -pub use crate::defense::{Defense, DefenseBuilder, LevelBuilder}; -pub use crate::master::{AddVisitorResult, CreateMCaptcha}; -pub use crate::mcaptcha::{MCaptcha, MCaptchaBuilder}; #[cfg(feature = "full")] pub use master::embedded::counter::Counter; +pub use crate::defense::{Defense, DefenseBuilder, LevelBuilder}; +pub use crate::master::AddVisitorResult; +pub use crate::master::CreateMCaptcha; +pub use crate::mcaptcha::{MCaptcha, MCaptchaBuilder}; + #[cfg(feature = "minimal")] pub mod dev { + pub use super::AddVisitorResult; + pub use super::CreateMCaptcha; pub use crate::defense; pub use crate::defense::{Defense, DefenseBuilder, LevelBuilder}; - pub use crate::master::{AddVisitorResult, CreateMCaptcha}; pub use crate::mcaptcha; pub use crate::mcaptcha::{MCaptcha, MCaptchaBuilder}; } diff --git a/src/master/embedded/counter.rs b/src/master/embedded/counter.rs index dc67e7c..f6f83a0 100644 --- a/src/master/embedded/counter.rs +++ b/src/master/embedded/counter.rs @@ -212,7 +212,7 @@ pub mod tests { use super::*; use crate::defense::*; use crate::errors::*; - use crate::master::MCaptchaBuilder; + use crate::mcaptcha::MCaptchaBuilder; // constants for testing // (visitor count, level) diff --git a/src/master/embedded/master.rs b/src/master/embedded/master.rs index 9506e88..9da880d 100644 --- a/src/master/embedded/master.rs +++ b/src/master/embedded/master.rs @@ -27,8 +27,8 @@ use log::info; use super::counter::Counter; use crate::errors::*; +use crate::master::messages::{AddSite, AddVisitor}; use crate::master::Master as MasterTrait; -use crate::master::{AddSite, AddVisitor}; /// This Actor manages the [Counter] actors. /// A service can have several [Counter] actors with @@ -191,7 +191,7 @@ impl Handler for Master { mod tests { use super::*; use crate::master::embedded::counter::tests::*; - use crate::master::AddSiteBuilder; + use crate::master::messages::AddSiteBuilder; #[actix_rt::test] async fn master_actor_works() { diff --git a/src/master/mod.rs b/src/master/mod.rs index 1c10475..269ad5d 100644 --- a/src/master/mod.rs +++ b/src/master/mod.rs @@ -16,49 +16,28 @@ * along with this program. If not, see . */ //! [Master] actor module that manages [MCaptcha] actors -#[cfg(feature = "full")] -use std::sync::mpsc::Receiver; -#[cfg(feature = "full")] -use actix::dev::*; -#[cfg(feature = "full")] -use derive_builder::Builder; use serde::{Deserialize, Serialize}; -#[cfg(feature = "full")] -use crate::errors::CaptchaResult; +use crate::mcaptcha::*; #[cfg(feature = "full")] pub mod embedded; -#[allow( - unused_variables, - unused_imports, - unused_variables, - dead_code, - unused_macros -)] -use crate::mcaptcha::*; -#[allow( - unused_variables, - unused_imports, - unused_variables, - dead_code, - unused_macros -)] #[cfg(feature = "full")] pub mod redis; #[cfg(feature = "full")] /// Describes actor handler trait impls that are required by a cache implementation -pub trait Master: actix::Actor + actix::Handler + actix::Handler {} +pub trait Master: + actix::Actor + actix::Handler + actix::Handler +{ +} -//+ actix::Handler - -/// Message to add visitor to an [MCaptcha] actor -#[derive(Message)] -#[cfg(feature = "full")] -#[rtype(result = "Receiver>>")] -pub struct AddVisitor(pub String); +#[derive(Serialize, Deserialize)] +pub struct CreateMCaptcha { + pub levels: Vec, + pub duration: u64, +} /// Struct representing the return datatime of /// [AddVisitor] message. Contains MCaptcha lifetime @@ -69,15 +48,6 @@ pub struct AddVisitorResult { pub difficulty_factor: u32, } -/// Message to add an [Counter] actor to [Master] -#[derive(Message, Builder)] -#[rtype(result = "()")] -#[cfg(feature = "full")] -pub struct AddSite { - pub id: String, - pub mcaptcha: MCaptcha, -} - impl AddVisitorResult { pub fn new(m: &MCaptcha) -> Self { AddVisitorResult { @@ -87,9 +57,26 @@ impl AddVisitorResult { } } -#[cfg(feature = "minimal")] -#[derive(Serialize, Deserialize)] -pub struct CreateMCaptcha { - pub levels: Vec, - pub duration: u64, +#[cfg(feature = "full")] +pub mod messages { + use std::sync::mpsc::Receiver; + + use actix::dev::*; + use derive_builder::Builder; + + use crate::errors::CaptchaResult; + use crate::mcaptcha::MCaptcha; + + /// Message to add visitor to an [MCaptcha] actor + #[derive(Message)] + #[rtype(result = "Receiver>>")] + pub struct AddVisitor(pub String); + + /// Message to add an [Counter] actor to [Master] + #[derive(Message, Builder)] + #[rtype(result = "()")] + pub struct AddSite { + pub id: String, + pub mcaptcha: MCaptcha, + } } diff --git a/src/master/redis/connection.rs b/src/master/redis/connection.rs index 0678b34..cfacc5d 100644 --- a/src/master/redis/connection.rs +++ b/src/master/redis/connection.rs @@ -15,51 +15,50 @@ * 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 crate::errors::*; -use crate::master::{AddSite, AddVisitor, AddVisitorResult, CreateMCaptcha}; +use crate::master::messages::{AddSite, AddVisitor}; +use crate::master::AddVisitorResult; +use crate::master::CreateMCaptcha; +use crate::redis::Redis; +use crate::redis::RedisConfig; +use crate::redis::RedisConnection; -pub enum RedisConnection { - Single(Rc>), - Cluster(Rc>), -} +/// Redis instance with mCaptcha Redis module loaded +pub struct MCaptchaRedis(Redis); + +/// Redis instance with mCaptcha Redis module loaded +pub struct MCaptchaRedisConnection(RedisConnection); -#[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 MCaptchaRedis { + pub async fn new(redis: RedisConfig) -> CaptchaResult { + let redis = Redis::new(redis).await?; + let m = MCaptchaRedis(redis); + m.get_client().is_module_loaded().await?; + Ok(m) + } + + pub fn get_client(&self) -> MCaptchaRedisConnection { + MCaptchaRedisConnection(self.0.get_client()) + } } -impl RedisConnection { +impl MCaptchaRedisConnection { pub async fn is_module_loaded(&self) -> CaptchaResult<()> { - let modules: Vec> = exec!(redis::cmd("MODULE").arg(&["LIST"]), &self).unwrap(); + let modules: Vec> = self + .0 + .exec(redis::cmd("MODULE").arg(&["LIST"])) + .await + .unwrap(); for list in modules.iter() { match list.iter().find(|module| module.as_str() == MODULE_NAME) { @@ -71,7 +70,12 @@ impl RedisConnection { 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() { + match self + .0 + .exec(redis::cmd("COMMAND").arg(&["INFO", cmd])) + .await + .unwrap() + { Value::Bulk(mut val) => { match val.pop() { Some(Value::Nil) => { @@ -91,7 +95,7 @@ impl RedisConnection { } pub async fn add_visitor(&self, msg: AddVisitor) -> CaptchaResult> { - let res: String = exec!(redis::cmd(ADD_VISITOR).arg(&[msg.0]), &self)?; + let res: String = self.0.exec(redis::cmd(ADD_VISITOR).arg(&[msg.0])).await?; let res: AddVisitorResult = serde_json::from_str(&res).unwrap(); Ok(Some(res)) } @@ -100,12 +104,17 @@ impl RedisConnection { 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)?; + self.0 + .exec(redis::cmd(ADD_CAPTCHA).arg(&[name, payload])) + .await?; Ok(()) } pub async fn check_captcha_exists(&self, captcha: &str) -> CaptchaResult { - let exists: usize = exec!(redis::cmd(CAPTCHA_EXISTS).arg(&[captcha]), &self)?; + let exists: usize = self + .0 + .exec(redis::cmd(CAPTCHA_EXISTS).arg(&[captcha])) + .await?; if exists == 1 { Ok(false) } else if exists == 0 { @@ -121,12 +130,12 @@ impl RedisConnection { } pub async fn delete_captcha(&self, captcha: &str) -> CaptchaResult<()> { - exec!(redis::cmd(DEL).arg(&[captcha]), &self)?; + self.0.exec(redis::cmd(DEL).arg(&[captcha])).await?; Ok(()) } pub async fn get_visitors(&self, captcha: &str) -> CaptchaResult { - let visitors: usize = exec!(redis::cmd(GET).arg(&[captcha]), &self)?; + let visitors: usize = self.0.exec(redis::cmd(GET).arg(&[captcha])).await?; Ok(visitors) } } @@ -134,32 +143,20 @@ impl RedisConnection { #[cfg(test)] pub mod tests { use super::*; - use crate::defense::{Level, LevelBuilder}; use crate::master::embedded::counter::tests::get_mcaptcha; - use crate::master::redis::master::{Master, Redis, RedisClient}; - - pub async fn connect(redis: &Redis) -> RedisConnection { - let redis = redis.connect(); - match &redis { - RedisClient::Single(c) => { - let con = c.get_async_connection().await.unwrap(); - RedisConnection::Single(Rc::new(RefCell::new(con))) - } - RedisClient::Cluster(c) => { - let con = c.get_connection().unwrap(); - RedisConnection::Cluster(Rc::new(RefCell::new(con))) - } - } - } + use crate::redis::*; const CAPTCHA_NAME: &str = "REDIS_CAPTCHA_TEST"; - const DURATION: usize = 10; const REDIS_URL: &str = "redis://127.0.1.1/"; #[actix_rt::test] async fn redis_master_works() { - let client = redis::Client::open("redis://127.0.0.1/").unwrap(); - let r = connect(&Redis::Single(REDIS_URL.into())).await; + let redis = Redis::new(RedisConfig::Single(REDIS_URL.into())) + .await + .unwrap(); + + let r = MCaptchaRedis(redis); + let r = r.get_client(); { let _ = r.delete_captcha(CAPTCHA_NAME).await; } diff --git a/src/master/redis/master.rs b/src/master/redis/master.rs index 077dd1c..99a0060 100644 --- a/src/master/redis/master.rs +++ b/src/master/redis/master.rs @@ -15,84 +15,26 @@ * 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 super::connection::MCaptchaRedis; use crate::errors::*; -use crate::master::AddVisitorResult; -use crate::master::{AddSite, AddVisitor, CreateMCaptcha, Master as MasterTrait}; - -use super::connection::RedisConnection; - -#[derive(Clone)] -pub enum Redis { - Single(String), - Cluster(Vec), -} - -impl Redis { - pub fn connect(&self) -> RedisClient { - match self { - Self::Single(url) => { - let client = Client::open("redis://127.0.0.1/").unwrap(); - RedisClient::Single(client) - } - Self::Cluster(nodes) => { - let cluster_client = ClusterClient::open(nodes.to_owned()).unwrap(); - RedisClient::Cluster(cluster_client) - } - } - } -} - -#[derive(Clone)] -pub enum RedisClient { - Single(Client), - Cluster(ClusterClient), -} +use crate::master::messages::{AddSite, AddVisitor}; +use crate::master::Master as MasterTrait; +use crate::redis::RedisConfig; pub struct Master { - pub redis: RedisClient, - pub con: Rc, + pub redis: MCaptchaRedis, } impl Master { - pub async fn new(redis: Redis) -> CaptchaResult { - let (redis, con) = Self::connect(redis).await; - con.is_module_loaded().await?; - let con = Rc::new(con); - let master = Self { redis, con }; + pub async fn new(redis: RedisConfig) -> CaptchaResult { + let redis = MCaptchaRedis::new(redis).await?; + let master = Self { redis }; Ok(master) } - - async fn connect(redis: Redis) -> (RedisClient, RedisConnection) { - let redis = redis.connect(); - let client = match &redis { - RedisClient::Single(c) => { - let con = c.get_async_connection().await.unwrap(); - RedisConnection::Single(Rc::new(RefCell::new(con))) - } - RedisClient::Cluster(c) => { - let con = c.get_connection().unwrap(); - RedisConnection::Cluster(Rc::new(RefCell::new(con))) - } - }; - (redis, client) - } } impl MasterTrait for Master {} @@ -107,7 +49,7 @@ impl Handler for Master { fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result { let (tx, rx) = mpsc::channel(); - let con = Rc::clone(&self.con); + let con = self.redis.get_client(); let fut = async move { let res = con.add_visitor(m).await; tx.send(res).unwrap() @@ -123,10 +65,9 @@ impl Handler for Master { fn handle(&mut self, m: AddSite, ctx: &mut Self::Context) -> Self::Result { //let (tx, rx) = mpsc::channel(); - - let con = Rc::clone(&self.con); + let con = self.redis.get_client(); let fut = async move { - let res = con.add_mcaptcha(m).await; + let _res = con.add_mcaptcha(m).await; //tx.send(res).unwrap(); } .into_actor(self); @@ -137,33 +78,34 @@ impl Handler for Master { #[cfg(test)] mod tests { use super::*; - use crate::defense::{Level, LevelBuilder}; use crate::master::embedded::counter::tests::get_mcaptcha; - use crate::master::redis::connection::tests::connect; - use crate::master::redis::master::{Master, Redis}; + use crate::master::redis::master::Master; + use crate::redis::RedisConfig; const CAPTCHA_NAME: &str = "REDIS_MASTER_CAPTCHA_TEST"; - const DURATION: usize = 10; const REDIS_URL: &str = "redis://127.0.1.1/"; #[actix_rt::test] async fn redis_master_works() { - let master = Master::new(Redis::Single(REDIS_URL.into())).await; - let client = redis::Client::open("redis://127.0.0.1/").unwrap(); - let r = connect(&Redis::Single(REDIS_URL.into())).await; + let master = Master::new(RedisConfig::Single(REDIS_URL.into())).await; + let sec_master = Master::new(RedisConfig::Single(REDIS_URL.into())).await; + let r = sec_master.unwrap().redis.get_client(); assert!(master.is_ok()); let master = master.unwrap(); { - let _ = r.delete_captcha(CAPTCHA_NAME).await; + let _ = master.redis.get_client().delete_captcha(CAPTCHA_NAME).await; } let addr = master.start(); + let mcaptcha = get_mcaptcha(); + let duration = mcaptcha.get_duration(); + let add_mcaptcha_msg = AddSite { id: CAPTCHA_NAME.into(), - mcaptcha: get_mcaptcha(), + mcaptcha, }; addr.send(add_mcaptcha_msg).await.unwrap(); @@ -171,5 +113,8 @@ mod tests { addr.send(add_visitor_msg).await.unwrap(); let visitors = r.get_visitors(CAPTCHA_NAME).await.unwrap(); assert_eq!(visitors, 1); + + let timer_expire = std::time::Duration::new(duration, 0); + actix::clock::delay_for(timer_expire).await; } } diff --git a/src/master/redis/mod.rs b/src/master/redis/mod.rs index 06d0f62..fbed376 100644 --- a/src/master/redis/mod.rs +++ b/src/master/redis/mod.rs @@ -15,23 +15,6 @@ * 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::*; +pub mod connection; +pub mod master; diff --git a/src/redis.rs b/src/redis.rs new file mode 100644 index 0000000..4c3f148 --- /dev/null +++ b/src/redis.rs @@ -0,0 +1,112 @@ +/* + * 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::rc::Rc; + +use redis::cluster::ClusterClient; +use redis::FromRedisValue; +//use redis::cluster::ClusterConnection; +use redis::Client; +//use redis::Connection; +use redis::{aio::Connection, cluster::ClusterConnection}; + +use crate::errors::*; + +#[derive(Clone)] +pub enum RedisConfig { + Single(String), + Cluster(Vec), +} + +impl RedisConfig { + pub fn connect(&self) -> RedisClient { + match self { + Self::Single(url) => { + let client = Client::open(url.as_str()).unwrap(); + RedisClient::Single(client) + } + Self::Cluster(nodes) => { + let cluster_client = ClusterClient::open(nodes.to_owned()).unwrap(); + RedisClient::Cluster(cluster_client) + } + } + } +} + +pub enum RedisConnection { + Single(Rc>), + Cluster(Rc>), +} + +impl RedisConnection { + #[inline] + pub fn get_client(&self) -> Self { + match self { + Self::Single(con) => Self::Single(Rc::clone(&con)), + Self::Cluster(con) => Self::Cluster(Rc::clone(&con)), + } + } + #[inline] + pub async fn exec(&self, cmd: &mut redis::Cmd) -> redis::RedisResult { + match self { + RedisConnection::Single(con) => cmd.query_async(&mut *con.borrow_mut()).await, + RedisConnection::Cluster(con) => cmd.query(&mut *con.borrow_mut()), + } + } +} + +#[derive(Clone)] +pub enum RedisClient { + Single(Client), + Cluster(ClusterClient), +} + +pub struct Redis { + _client: RedisClient, + connection: RedisConnection, +} + +impl Redis { + pub async fn new(redis: RedisConfig) -> CaptchaResult { + let (_client, connection) = Self::connect(redis).await; + let master = Self { + _client, + connection, + }; + Ok(master) + } + + pub fn get_client(&self) -> RedisConnection { + self.connection.get_client() + } + + async fn connect(redis: RedisConfig) -> (RedisClient, RedisConnection) { + let redis = redis.connect(); + let client = match &redis { + RedisClient::Single(c) => { + let con = c.get_async_connection().await.unwrap(); + RedisConnection::Single(Rc::new(RefCell::new(con))) + } + RedisClient::Cluster(c) => { + let con = c.get_connection().unwrap(); + RedisConnection::Cluster(Rc::new(RefCell::new(con))) + } + }; + (redis, client) + } +} diff --git a/src/system.rs b/src/system.rs index c869f39..a785d50 100644 --- a/src/system.rs +++ b/src/system.rs @@ -23,6 +23,7 @@ use pow_sha256::Config; use crate::cache::messages::*; use crate::cache::Save; use crate::errors::*; +use crate::master::messages::*; use crate::master::Master; use crate::pow::*; @@ -42,13 +43,10 @@ where + ToEnvelope + ToEnvelope, X: Master, - ::Context: - ToEnvelope + ToEnvelope, + ::Context: ToEnvelope + ToEnvelope, { /// utility function to get difficulty factor of site `id` and cache it pub async fn get_pow(&self, id: String) -> Option { - use crate::master::AddVisitor; - match self .master .send(AddVisitor(id.clone())) @@ -130,7 +128,6 @@ mod tests { use crate::cache::HashCache; use crate::master::embedded::counter::tests::*; use crate::master::embedded::master::Master; - use crate::master::*; const MCAPTCHA_NAME: &str = "batsense.net";