/* * 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, 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), } pub struct Master { pub redis: RedisClient, pub con: Rc, } 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 }; 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 {} 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); } } #[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}; 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; assert!(master.is_ok()); let master = master.unwrap(); { let _ = r.delete_captcha(CAPTCHA_NAME).await; } let addr = master.start(); let add_mcaptcha_msg = AddSite { id: CAPTCHA_NAME.into(), mcaptcha: get_mcaptcha(), }; addr.send(add_mcaptcha_msg).await.unwrap(); let add_visitor_msg = AddVisitor(CAPTCHA_NAME.into()); addr.send(add_visitor_msg).await.unwrap(); let visitors = r.get_visitors(CAPTCHA_NAME).await.unwrap(); assert_eq!(visitors, 1); } }