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

use core::fmt;

use std::{error, time::Instant};

use http::{HeaderName, HeaderValue, Response, StatusCode};

use crate::{
    gcra::NotUntil,
    timer::{DefaultTimer, Timer},
};

/// Error happen when client exceeds rate limit.
#[derive(Debug)]
pub struct TooManyRequests {
    after_seconds: u64,
}

impl fmt::Display for TooManyRequests {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "too many requests. wait for {}", self.after_seconds)
    }
}

impl error::Error for TooManyRequests {}

impl From<NotUntil<Instant>> for TooManyRequests {
    fn from(e: NotUntil<Instant>) -> Self {
        let after_seconds = e.wait_time_from(DefaultTimer.now()).as_secs();

        Self { after_seconds }
    }
}

const X_RT_AFTER: HeaderName = HeaderName::from_static("x-ratelimit-after");

impl TooManyRequests {
    /// extend response headers with status code and headers
    /// StatusCode: 429
    /// Header: `x-ratelimit-after: <num in second>`
    pub fn extend_response<Ext>(&self, res: &mut Response<Ext>) {
        *res.status_mut() = StatusCode::TOO_MANY_REQUESTS;
        res.headers_mut()
            .insert(X_RT_AFTER, HeaderValue::from(self.after_seconds));
    }
}

/// Error indicating that the number of cells tested (the first
/// argument) is larger than the bucket's capacity.
///
/// This means the decision can never have a conforming result. The
/// argument gives the maximum number of cells that could ever have a
/// conforming result.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InsufficientCapacity(pub u32);

impl fmt::Display for InsufficientCapacity {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "required number of cells {} exceeds bucket's capacity", self.0)
    }
}

impl std::error::Error for InsufficientCapacity {}

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

    #[test]
    fn coverage() {
        let display_output = format!("{}", InsufficientCapacity(3));
        assert!(display_output.contains('3'));
        let debug_output = format!("{:?}", InsufficientCapacity(3));
        assert!(debug_output.contains('3'));
        assert_eq!(InsufficientCapacity(3), InsufficientCapacity(3));
    }
}