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
93
94
95
96
97
98
99
100

//! Provides a similiar catcher for html/json like the Rocket framework

use rocket::{
    http::{ContentType, Status},
    Request, Response,
};
use std::{borrow::Cow, io::Cursor};

// copied from rocket/src/catcher/catcher.rs because it is also pub(crate)

macro_rules! html_error_template {
    ($code:expr, $reason:expr, $description:expr) => {
        concat!(
            r#"<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <title>"#,
            $code,
            " ",
            $reason,
            r#"</title>
</head>
<body align="center">
    <div role="main" align="center">
        <h1>"#,
            $code,
            ": ",
            $reason,
            r#"</h1>
        <p>"#,
            $description,
            r#"</p>
        <hr />
    </div>
    <div role="contentinfo" align="center">
        <small>Rocket</small>
    </div>
</body>
</html>"#
        )
    };
}

macro_rules! json_error_template {
    ($code:expr, $reason:expr, $description:expr) => {
        concat!(
            r#"{
  "error": {
    "code": "#,
            $code,
            r#",
    "reason": ""#,
            $reason,
            r#"",
    "description": ""#,
            $description,
            r#""
  }
}"#
        )
    };
}

// also copied and modified

/// Create the handler for `Status::TooManyRequests`.
pub(crate) fn too_many_requests_handler<'r, 'o: 'r>(req: &'r Request<'_>) -> Response<'o> {
    let preferred = req.accept().map(|a| a.preferred());
    let (mime, text) = if preferred.map_or(false, |a| a.is_json()) {
        let json: Cow<'_, str> = json_error_template!(
            429,
            "Too Many Requests",
            "Too many requests have been received recently."
        )
        .into();

        (ContentType::JSON, json)
    } else {
        let html: Cow<'_, str> = html_error_template!(
            429,
            "Too Many Requests",
            "Too many requests have been received recently."
        )
        .into();

        (ContentType::HTML, html)
    };

    let mut r = Response::build()
        .status(Status::TooManyRequests)
        .header(mime)
        .finalize();
    match text {
        Cow::Owned(v) => r.set_sized_body(v.len(), Cursor::new(v)),
        Cow::Borrowed(v) => r.set_sized_body(v.len(), Cursor::new(v)),
    };

    r
}