redis master

src/defense.rs

@@ -114,6 +114,12 @@ pub struct Defense {
     current_visitor_threshold: usize,
 }
 
+impl From<Defense> for Vec<Level> {
+    fn from(d: Defense) -> Self {
+        d.levels
+    }
+}
+
 /// Builder struct for [Defense]
 #[derive(Debug, Clone, PartialEq)]
 pub struct DefenseBuilder {

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<RedisError> for CaptchaError {
+    fn from(e: RedisError) -> Self {
+        Self::RedisError(e)
+    }
+}
+
+impl From<actix::MailboxError> for CaptchaError {
+    fn from(_: actix::MailboxError) -> Self {
+        Self::MailboxError
+    }
 }
 
 /// [Result] datatype for libmcaptcha

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,25 +90,25 @@ impl Handler<AddVisitor> for Master {
     type Result = MessageResult<AddVisitor>;
 
     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();
+        match self.get_site(&m.0) {
+            None => tx.send(Ok(None)).unwrap(),
+            Some(addr) => {
                 let fut = async move {
-                let config = addr
-                    .unwrap()
-                    .send(super::counter::AddVisitor)
-                    .await
-                    .unwrap();
-
-                tx.send(config).unwrap();
+                    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);
-            return MessageResult(Some(rx));
             }
         }
+        return MessageResult(rx);
+    }
 }
 
 /// Message to get an [Counter] actor from master

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<AddVisitor> + actix::Handler<AddSite> {}
@@ -39,7 +49,7 @@ pub trait Master: actix::Actor + actix::Handler<AddVisitor> + actix::Handler<Add
 
 /// Message to add visitor to an [MCaptcha] actor
 #[derive(Message)]
-#[rtype(result = "Option<Receiver<AddVisitorResult>>")]
+#[rtype(result = "Receiver<CaptchaResult<Option<AddVisitorResult>>>")]
 pub struct AddVisitor(pub String);
 
 /// Struct representing the return datatime of

src/master/redis/connection.rs

@@ -0,0 +1,128 @@
+/*
+ * mCaptcha - A proof of work based DoS protection system
+ * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
+ *
+ * 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 <http://www.gnu.org/licenses/>.
+ */
+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<RefCell<Connection>>),
+    Cluster(Rc<RefCell<ClusterConnection>>),
+}
+
+#[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<Vec<String>> = 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<Option<AddVisitorResult>> {
+        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<bool> {
+        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<usize> {
+        let visitors: usize = exec!(redis::cmd(GET).arg(&[captcha]), &self)?;
+        Ok(visitors)
+    }
+}

src/master/redis/master.rs

@@ -0,0 +1,118 @@
+/*
+ * mCaptcha - A proof of work based DoS protection system
+ * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
+ *
+ * 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 <http://www.gnu.org/licenses/>.
+ */
+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<Level>,
+    pub duration: u64,
+}
+pub struct Master {
+    pub redis: Redis,
+    pub con: Rc<RedisConnection>,
+}
+
+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<Self>;
+}
+
+impl Handler<AddVisitor> for Master {
+    type Result = MessageResult<AddVisitor>;
+
+    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<AddSite> 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);
+    }
+}

src/master/redis/mod.rs

@@ -0,0 +1,37 @@
+/*
+ * mCaptcha - A proof of work based DoS protection system
+ * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
+ *
+ * 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 <http://www.gnu.org/licenses/>.
+ */
+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::*;

src/master/redis_master/mod.rs

@@ -1,206 +0,0 @@
-/*
- * mCaptcha - A proof of work based DoS protection system
- * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
- *
- * 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 <http://www.gnu.org/licenses/>.
- */
-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<usize> = 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<usize> = 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<usize> = exec!(cmd, self.get_connection().await);
-
-        let mut cmd = redis::cmd("COMMAND");
-        cmd.arg(&["INFO", GET]);
-        let a: RedisResult<usize> = 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<Self>;
-}
-
-impl Handler<AddVisitor> for Master {
-    type Result = MessageResult<AddVisitor>;
-
-    fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result {
-        let fut = async {
-            self.add_visitor(&m.0).await;
-        };
-        unimplemented!();
-    }
-}
-
-impl Handler<AddSite> 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<ClusterConnection>;
-//    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<Self>;
-//}
-//
-//impl Handler<AddVisitor> for Master {
-//    type Result = MessageResult<AddVisitor>;
-//
-//    fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result {
-//        let fut = async {
-//            let test = "1";
-//        };
-//        unimplemented!();
-//    }
-//}
-//
-//impl Handler<AddSite> for Master {
-//    type Result = ();
-//
-//    fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result {
-//        unimplemented!();
-//    }
-//}
-
-//pub struct Master<T: Redis> {
-//    pub redis: T,
-//}
-//
-//impl<T: Redis> MasterTrait for Master<T> {}
-//
-//impl<T: Redis> Actor for Master<T> {
-//    type Context = Context<Self>;
-//}
-//
-//impl<T: Redis> Handler<AddVisitor> for Master<T> {
-//    type Result = MessageResult<AddVisitor>;
-//
-//    fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result {
-//        let fut = async {
-//            self.redis.get_connection();
-//            let test = "1";
-//        };
-//        unimplemented!();
-//    }
-//}
-//
-//impl<T: Redis> Handler<AddSite> for Master<T> {
-//    type Result = ();
-//
-//    fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result {
-//        unimplemented!();
-//    }
-//}

src/mcaptcha.rs

@@ -78,6 +78,15 @@ pub struct MCaptcha {
     duration: u64,
 }
 
+impl From<MCaptcha> 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]

src/system.rs

@@ -49,14 +49,15 @@ where
     pub async fn get_pow(&self, id: String) -> Option<PoWConfig> {
         use crate::master::AddVisitor;
 
-        let rx = self.master.send(AddVisitor(id.clone())).await.unwrap();
-
-        if rx.is_none() {
-            return None;
-        }
-
-        let mcaptcha = rx.unwrap().recv().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());
 
                 let cache_msg = CachePoWBuilder::default()
@@ -70,6 +71,9 @@ where
                 self.cache.send(cache_msg).await.unwrap().unwrap();
                 Some(pow_config)
             }
+            _ => None,
+        }
+    }
 
     /// utility function to verify [Work]
     pub async fn verify_pow(&self, work: Work) -> CaptchaResult<String> {