tmp22/src/master/embedded/master.rs

267 lines
8.1 KiB
Rust
Raw Normal View History

2021-03-07 19:54:41 +05:30
/*
* 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/>.
*/
2021-06-09 14:18:31 +05:30
//! Embedded [Master] actor module that manages [Counter] actors
2021-03-07 19:54:41 +05:30
use std::collections::BTreeMap;
use std::time::Duration;
2021-03-07 19:54:41 +05:30
use actix::clock::sleep;
//use actix::clock::delay_for;
2021-03-07 19:54:41 +05:30
use actix::dev::*;
use log::info;
use tokio::sync::oneshot::channel;
2021-03-07 19:54:41 +05:30
2021-06-03 19:40:45 +05:30
use super::counter::Counter;
2021-06-08 18:13:01 +05:30
use crate::errors::*;
2021-07-20 12:36:42 +05:30
use crate::master::messages::{AddSite, AddVisitor, RemoveCaptcha, Rename};
2021-06-03 18:19:57 +05:30
use crate::master::Master as MasterTrait;
2021-03-07 19:54:41 +05:30
2021-06-03 19:40:45 +05:30
/// This Actor manages the [Counter] actors.
/// A service can have several [Counter] actors with
2021-03-08 19:43:26 +05:30
/// varying [Defense][crate::defense::Defense] configurations
/// so a "master" actor is needed to manage them all
2021-06-03 18:08:43 +05:30
#[derive(Clone, Default)]
pub struct Master {
2021-06-03 19:40:45 +05:30
sites: BTreeMap<String, (Option<()>, Addr<Counter>)>,
gc: u64,
2021-03-07 19:54:41 +05:30
}
2021-06-03 18:19:57 +05:30
impl MasterTrait for Master {}
2021-06-03 18:08:43 +05:30
impl Master {
2021-06-03 19:40:45 +05:30
/// add [Counter] actor to [Master]
pub fn add_site(&mut self, addr: Addr<Counter>, id: String) {
2021-06-09 20:18:46 +05:30
self.sites.insert(id, (None, addr));
2021-03-07 19:54:41 +05:30
}
2021-03-08 19:43:26 +05:30
/// create new master
/// accepts a `u64` to configure garbage collection period
pub fn new(gc: u64) -> Self {
2021-03-07 19:54:41 +05:30
Master {
sites: BTreeMap::new(),
gc,
2021-03-07 19:54:41 +05:30
}
}
2021-06-03 19:40:45 +05:30
/// get [Counter] actor from [Master]
2021-06-09 20:18:46 +05:30
pub fn get_site(&mut self, id: &str) -> Option<Addr<Counter>> {
let mut r = None;
if let Some((read_val, addr)) = self.sites.get_mut(id) {
r = Some(addr.clone());
*read_val = Some(());
};
r
}
2021-06-03 19:40:45 +05:30
/// remvoes [Counter] actor from [Master]
pub fn rm_site(&mut self, id: &str) {
self.sites.remove(id);
2021-03-07 19:54:41 +05:30
}
2021-07-19 13:07:46 +05:30
/// renames [Counter] actor
pub fn rename(&mut self, msg: Rename) {
// If actor isn't present, it's okay to not throw an error
// since actors are lazyily initialized and are cleaned up when inactive
if let Some((_, counter)) = self.sites.remove(&msg.name) {
self.add_site(counter, msg.rename_to);
}
}
2021-03-07 19:54:41 +05:30
}
impl Actor for Master {
2021-03-07 19:54:41 +05:30
type Context = Context<Self>;
fn started(&mut self, ctx: &mut Self::Context) {
let addr = ctx.address();
let task = async move {
addr.send(CleanUp).await.unwrap();
}
.into_actor(self);
ctx.spawn(task);
}
2021-03-07 19:54:41 +05:30
}
2021-06-03 18:08:43 +05:30
impl Handler<AddVisitor> for Master {
type Result = MessageResult<AddVisitor>;
fn handle(&mut self, m: AddVisitor, ctx: &mut Self::Context) -> Self::Result {
2021-06-08 18:13:01 +05:30
let (tx, rx) = channel();
match self.get_site(&m.0) {
2021-06-09 14:30:41 +05:30
None => {
let _ = tx.send(Ok(None));
}
2021-06-08 18:13:01 +05:30
Some(addr) => {
let fut = async move {
match addr.send(super::counter::AddVisitor).await {
2021-06-09 14:30:41 +05:30
Ok(val) => {
let _ = tx.send(Ok(Some(val)));
}
2021-06-08 18:13:01 +05:30
Err(e) => {
let err: CaptchaError = e.into();
2021-06-09 14:30:41 +05:30
let _ = tx.send(Err(err));
2021-06-08 18:13:01 +05:30
}
}
}
.into_actor(self);
ctx.spawn(fut);
2021-06-03 18:08:43 +05:30
}
}
2021-06-09 20:18:46 +05:30
MessageResult(rx)
2021-06-03 18:08:43 +05:30
}
}
2021-07-19 13:07:46 +05:30
impl Handler<Rename> for Master {
type Result = MessageResult<Rename>;
2021-07-20 12:36:42 +05:30
fn handle(&mut self, m: Rename, _ctx: &mut Self::Context) -> Self::Result {
2021-07-19 13:07:46 +05:30
self.rename(m);
let (tx, rx) = channel();
2021-07-20 12:36:42 +05:30
let _ = tx.send(Ok(()));
2021-07-19 13:07:46 +05:30
MessageResult(rx)
}
}
2021-06-03 19:40:45 +05:30
/// Message to get an [Counter] actor from master
2021-03-07 19:54:41 +05:30
#[derive(Message)]
2021-06-03 19:40:45 +05:30
#[rtype(result = "Option<Addr<Counter>>")]
2021-03-08 17:02:36 +05:30
pub struct GetSite(pub String);
2021-03-07 19:54:41 +05:30
impl Handler<GetSite> for Master {
2021-03-08 17:02:36 +05:30
type Result = MessageResult<GetSite>;
2021-03-07 19:54:41 +05:30
fn handle(&mut self, m: GetSite, _ctx: &mut Self::Context) -> Self::Result {
2021-03-08 17:02:36 +05:30
let addr = self.get_site(&m.0);
2021-06-09 20:18:46 +05:30
match addr {
None => MessageResult(None),
Some(addr) => MessageResult(Some(addr)),
2021-03-07 19:54:41 +05:30
}
}
}
2021-06-03 19:40:45 +05:30
/// Message to clean up master of [Counter] actors with zero visitor count
#[derive(Message)]
#[rtype(result = "()")]
pub struct CleanUp;
impl Handler<CleanUp> for Master {
type Result = ();
fn handle(&mut self, _: CleanUp, ctx: &mut Self::Context) -> Self::Result {
let sites = self.sites.clone();
let gc = self.gc;
let master = ctx.address();
info!("init master actor cleanup up");
let task = async move {
for (id, (new, addr)) in sites.iter() {
2021-06-03 19:40:45 +05:30
use super::counter::{GetCurrentVisitorCount, Stop};
let visitor_count = addr.send(GetCurrentVisitorCount).await.unwrap();
println!("{}", visitor_count);
if visitor_count == 0 && new.is_some() {
addr.send(Stop).await.unwrap();
2021-07-20 12:36:42 +05:30
master.send(RemoveCaptcha(id.to_owned())).await.unwrap();
println!("cleaned up");
}
}
let duration = Duration::new(gc, 0);
sleep(duration).await;
//delay_for(duration).await;
master.send(CleanUp).await.unwrap();
}
.into_actor(self);
ctx.spawn(task);
}
}
2021-07-20 12:36:42 +05:30
impl Handler<RemoveCaptcha> for Master {
type Result = MessageResult<RemoveCaptcha>;
2021-07-20 12:36:42 +05:30
fn handle(&mut self, m: RemoveCaptcha, _ctx: &mut Self::Context) -> Self::Result {
let (tx, rx) = channel();
self.rm_site(&m.0);
2021-07-20 12:36:42 +05:30
tx.send(Ok(())).unwrap();
MessageResult(rx)
}
}
impl Handler<AddSite> for Master {
type Result = MessageResult<AddSite>;
2021-03-07 19:54:41 +05:30
fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result {
let (tx, rx) = channel();
2021-06-03 19:40:45 +05:30
let counter: Counter = m.mcaptcha.into();
let addr = counter.start();
self.add_site(addr, m.id);
tx.send(Ok(())).unwrap();
MessageResult(rx)
2021-03-07 19:54:41 +05:30
}
}
#[cfg(test)]
mod tests {
use super::*;
2021-06-03 19:40:45 +05:30
use crate::master::embedded::counter::tests::*;
2021-06-09 14:03:19 +05:30
use crate::master::messages::AddSiteBuilder;
2021-07-19 13:07:46 +05:30
use crate::master::messages::RenameBuilder;
use crate::MCaptcha;
2021-03-07 19:54:41 +05:30
#[actix_rt::test]
2021-03-08 19:43:26 +05:30
async fn master_actor_works() {
let addr = Master::new(1).start();
2021-03-07 19:54:41 +05:30
2021-07-19 13:07:46 +05:30
let get_add_site_msg = |id: String, mcaptcha: MCaptcha| {
AddSiteBuilder::default()
.id(id)
.mcaptcha(mcaptcha)
.build()
.unwrap()
};
2021-03-07 19:54:41 +05:30
let id = "yo";
2021-07-19 13:07:46 +05:30
let msg = get_add_site_msg(id.into(), get_mcaptcha());
2021-03-07 19:54:41 +05:30
addr.send(msg).await.unwrap();
2021-03-08 19:43:26 +05:30
let mcaptcha_addr = addr.send(GetSite(id.into())).await.unwrap();
2021-06-03 19:40:45 +05:30
assert!(mcaptcha_addr.is_some());
2021-03-08 19:43:26 +05:30
2021-07-19 13:07:46 +05:30
let new_id = "yoyo";
let rename = RenameBuilder::default()
.name(id.into())
.rename_to(new_id.into())
.build()
.unwrap();
addr.send(rename).await.unwrap();
let mcaptcha_addr = addr.send(GetSite(new_id.into())).await.unwrap();
assert!(mcaptcha_addr.is_some());
2021-03-08 19:43:26 +05:30
let addr_doesnt_exist = addr.send(GetSite("a".into())).await.unwrap();
assert!(addr_doesnt_exist.is_none());
let timer_expire = Duration::new(DURATION, 0);
sleep(timer_expire).await;
sleep(timer_expire).await;
// delay_for(timer_expire).await;
// delay_for(timer_expire).await;
2021-07-19 13:07:46 +05:30
let mcaptcha_addr = addr.send(GetSite(new_id.into())).await.unwrap();
assert_eq!(mcaptcha_addr, None);
2021-07-20 12:36:42 +05:30
assert!(addr.send(RemoveCaptcha(new_id.into())).await.is_ok());
2021-03-07 19:54:41 +05:30
}
}