redis master

This commit is contained in:
Aravinth Manivannan 2021-06-08 18:13:01 +05:30
parent 285be4d936
commit edcf55969c
Signed by untrusted user: realaravinth
GPG Key ID: AD9F0F08E855ED88
10 changed files with 372 additions and 242 deletions

View File

@ -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 {

View File

@ -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

View File

@ -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<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();
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);
}
}

View File

@ -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

View File

@ -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)
}
}

118
src/master/redis/master.rs Normal file
View File

@ -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);
}
}

37
src/master/redis/mod.rs Normal file
View File

@ -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::*;

View File

@ -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!();
// }
//}

View File

@ -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]

View File

@ -49,26 +49,30 @@ 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();
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]