logo
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
117
118
119
120
121

use std::{borrow::Cow, time::Duration};

use redis::Client;

use crate::{errors::Error, Limiter};

/// Rate limit builder.
#[derive(Debug)]
pub struct Builder<'a> {
    pub(crate) redis_url: &'a str,
    pub(crate) limit: usize,
    pub(crate) period: Duration,
    pub(crate) cookie_name: Cow<'static, str>,
    pub(crate) session_key: Cow<'static, str>,
}

impl Builder<'_> {
    /// Set upper limit.
    pub fn limit(&mut self, limit: usize) -> &mut Self {
        self.limit = limit;
        self
    }

    /// Set limit window/period.
    pub fn period(&mut self, period: Duration) -> &mut Self {
        self.period = period;
        self
    }

    /// Set name of cookie to be sent.
    pub fn cookie_name(&mut self, cookie_name: impl Into<Cow<'static, str>>) -> &mut Self {
        self.cookie_name = cookie_name.into();
        self
    }

    /// Set session key to be used in backend.
    pub fn session_key(&mut self, session_key: impl Into<Cow<'static, str>>) -> &mut Self {
        self.session_key = session_key.into();
        self
    }

    /// Finalizes and returns a `Limiter`.
    ///
    /// Note that this method will connect to the Redis server to test its connection which is a
    /// **synchronous** operation.
    pub fn build(&self) -> Result<Limiter, Error> {
        Ok(Limiter {
            client: Client::open(self.redis_url)?,
            limit: self.limit,
            period: self.period,
            cookie_name: self.cookie_name.clone(),
            session_key: self.session_key.clone(),
        })
    }
}

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

    #[test]
    fn test_create_builder() {
        let redis_url = "redis://127.0.0.1";
        let period = Duration::from_secs(10);
        let builder = Builder {
            redis_url,
            limit: 100,
            period,
            cookie_name: Cow::Owned("session".to_string()),
            session_key: Cow::Owned("rate-api".to_string()),
        };

        assert_eq!(builder.redis_url, redis_url);
        assert_eq!(builder.limit, 100);
        assert_eq!(builder.period, period);
        assert_eq!(builder.session_key, "rate-api");
        assert_eq!(builder.cookie_name, "session");
    }

    #[test]
    fn test_create_limiter() {
        let redis_url = "redis://127.0.0.1";
        let period = Duration::from_secs(20);
        let mut builder = Builder {
            redis_url,
            limit: 100,
            period: Duration::from_secs(10),
            session_key: Cow::Borrowed("key"),
            cookie_name: Cow::Borrowed("sid"),
        };

        let limiter = builder
            .limit(200)
            .period(period)
            .cookie_name("session".to_string())
            .session_key("rate-api".to_string())
            .build()
            .unwrap();

        assert_eq!(limiter.limit, 200);
        assert_eq!(limiter.period, period);
        assert_eq!(limiter.session_key, "rate-api");
        assert_eq!(limiter.cookie_name, "session");
    }

    #[test]
    #[should_panic = "Redis URL did not parse"]
    fn test_create_limiter_error() {
        let redis_url = "127.0.0.1";
        let period = Duration::from_secs(20);
        let mut builder = Builder {
            redis_url,
            limit: 100,
            period: Duration::from_secs(10),
            session_key: Cow::Borrowed("key"),
            cookie_name: Cow::Borrowed("sid"),
        };

        builder.limit(200).period(period).build().unwrap();
    }
}