rocket-governor 0.2.0-rc.4

Implementation of Governor rate limiter for Rocket
Documentation 
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

use crate::logger::debug;
use governor::{
    clock::DefaultClock, middleware::StateInformationMiddleware,
    state::keyed::DefaultKeyedStateStore, Quota, RateLimiter,
};
use lazy_static::lazy_static;
use rocket::http::Method;
use std::{
    any::type_name,
    collections::HashMap,
    net::IpAddr,
    sync::{Arc, RwLock},
};

pub(crate) type RegisteredRateLimiter = Arc<
    RateLimiter<IpAddr, DefaultKeyedStateStore<IpAddr>, DefaultClock, StateInformationMiddleware>,
>;

#[derive(Debug)]
pub(crate) struct Registry {
    limiter: RwLock<HashMap<Method, HashMap<String, RegisteredRateLimiter>>>,
}

impl Registry {
    pub(crate) fn get_or_insert<T>(
        method: Method,
        route_name: &str,
        quota: Quota,
    ) -> RegisteredRateLimiter {
        let route_name = type_name::<T>().to_string() + "::" + route_name;

        // check if exist with readlock
        let limiter = if let Ok(rlock) = REG.limiter.read() {
            if let Some(meth_found) = rlock.get(&method) {
                if let Some(limiter) = meth_found.get(&route_name) {
                    debug!("limiter found method {} route {}", &method, &route_name);
                    return Arc::clone(limiter);
                }
            }
            None
        } else {
            None
        };

        // get the existing limiter or create the new one with writelock and return the created
        let limiter = if let Some(limiter) = limiter {
            limiter
        } else {
            let mut wlock_meth_map = REG.limiter.write().unwrap();
            if let Some(meth_found) = wlock_meth_map.get_mut(&method) {
                if let Some(limiter) = meth_found.get(&route_name) {
                    debug!("limiter found method {} route {}", &method, &route_name);
                    Arc::clone(limiter)
                } else {
                    debug!("new limiter method {} route {}", &method, &route_name);
                    let limiter = Arc::new(
                        RateLimiter::keyed(quota).with_middleware::<StateInformationMiddleware>(),
                    );
                    meth_found.insert(route_name, Arc::clone(&limiter));
                    limiter
                }
            } else {
                debug!("new limiter method {} route {}", &method, &route_name);
                let mut lim_map = HashMap::new();
                let limiter = Arc::new(
                    RateLimiter::keyed(quota).with_middleware::<StateInformationMiddleware>(),
                );
                lim_map.insert(route_name, Arc::clone(&limiter));
                wlock_meth_map.insert(method, lim_map);
                limiter
            }
        };

        // simple cleanup
        limiter.retain_recent();
        limiter.shrink_to_fit();

        limiter
    }
}

impl Default for Registry {
    fn default() -> Self {
        Self {
            limiter: RwLock::new(HashMap::new()),
        }
    }
}

lazy_static! {
    static ref REG: Registry = Registry::default();
}