krustie 0.3.0

Krustie is a simple and easy-to-use backend framework.
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

use std::{
    collections::HashMap,
    net::IpAddr,
    time::{Duration, Instant},
};

use crate::{server::route_handler::HandlerResult, Middleware, StatusCode};

/// A rate limiter middleware
///
/// Limits the number of requests from an IP address based on the token number and token refill time.
#[derive(Clone, Debug)]
pub struct RateLimiter {
    token_number: u16,
    token_refill_duration: Duration,
    requests: HashMap<IpAddr, (Instant, u16)>,
}

impl RateLimiter {
    /// Creates a new instance of RateLimiter
    ///
    /// # Example
    ///
    /// ```rust
    /// use krustie::{ Server, middleware::RateLimiter };
    ///
    /// let mut server = Server::create();
    ///
    /// let rate_limiter = RateLimiter::new(10, 1000);
    ///
    /// server.use_handler(rate_limiter);
    /// ```
    pub fn new(token_number: u16, token_refill_ms: u64) -> Self {
        Self {
            token_number,
            token_refill_duration: Duration::from_millis(token_refill_ms),
            requests: HashMap::new(),
        }
    }

    fn check(&mut self, ip: IpAddr) -> bool {
        let now: Instant = Instant::now();

        let entry = self
            .requests
            .entry(ip)
            .or_insert_with(|| (now, self.token_number));

        if now.duration_since(entry.0) >= self.token_refill_duration {
            // Duration passed, token refreshed
            entry.0 = now;
            entry.1 = self.token_number - 1;
            true
        } else if entry.1 > 0 {
            // Duration hasn't passed, token reduced
            entry.1 -= 1;
            true
        } else {
            // No tokens left
            false
        }
    }
}

impl Middleware for RateLimiter {
    fn middleware(
        &mut self,
        request: &crate::Request,
        response: &mut crate::Response,
    ) -> HandlerResult {
        match Self::check(self, request.get_peer_addr().ip()) {
            true => HandlerResult::Next,
            false => {
                response.status(StatusCode::TooManyRequests);
                return HandlerResult::End;
            }
        }
    }
}