      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116

    /*
     * 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/>.
     */
    //! [Master] actor module that manages [MCaptcha] actors
    use std::collections::BTreeMap;

    use actix::dev::*;
    use derive_builder::Builder;

    use crate::mcaptcha::MCaptcha;

    /// This Actor manages the [MCaptcha] actors.
    /// A service can have several [MCaptcha] actors with
    /// varying [Defense][crate::defense::Defense] configurations
    /// so a "master" actor is needed to manage them all
    #[derive(Clone)]
    pub struct Master {
        sites: BTreeMap<String, Addr<MCaptcha>>,
    }

    impl Master {
        /// add [MCaptcha] actor to [Master]
        pub fn add_site(&mut self, details: AddSite) {
            self.sites.insert(details.id, details.addr.to_owned());
        }

        /// create new master
        pub fn new() -> Self {
            Master {
                sites: BTreeMap::new(),
            }
        }

        /// get [MCaptcha] actor from [Master]
        pub fn get_site<'a, 'b>(&'a self, id: &'b str) -> Option<&'a Addr<MCaptcha>> {
            self.sites.get(id)
        }
    }

    impl Actor for Master {
        type Context = Context<Self>;
    }

    /// Message to get an [MCaptcha] actor from master
    #[derive(Message)]
    #[rtype(result = "Option<Addr<MCaptcha>>")]
    pub struct GetSite(pub String);

    impl Handler<GetSite> for Master {
        type Result = MessageResult<GetSite>;

        fn handle(&mut self, m: GetSite, _ctx: &mut Self::Context) -> Self::Result {
            let addr = self.get_site(&m.0);
            if addr.is_none() {
                return MessageResult(None);
            } else {
                return MessageResult(Some(addr.unwrap().clone()));
            }
        }
    }

    /// Message to add an [MCaptcha] actor to [Master]
    #[derive(Message, Builder)]
    #[rtype(result = "()")]
    pub struct AddSite {
        pub id: String,
        pub addr: Addr<MCaptcha>,
    }

    impl Handler<AddSite> for Master {
        type Result = ();

        fn handle(&mut self, m: AddSite, _ctx: &mut Self::Context) -> Self::Result {
            self.add_site(m);
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use crate::mcaptcha::tests::*;

        #[actix_rt::test]
        async fn master_actor_works() {
            let addr = Master::new().start();

            let id = "yo";
            let mcaptcha = get_counter().start();
            let msg = AddSiteBuilder::default()
                .id(id.into())
                .addr(mcaptcha.clone())
                .build()
                .unwrap();
            addr.send(msg).await.unwrap();

            let mcaptcha_addr = addr.send(GetSite(id.into())).await.unwrap();
            assert_eq!(mcaptcha_addr, Some(mcaptcha));

            let addr_doesnt_exist = addr.send(GetSite("a".into())).await.unwrap();
            assert!(addr_doesnt_exist.is_none());
        }
    }