rocket_governor/
limit_header_gen.rs

1//! Provides [`Fairing`](rocket::fairing::Fairing) in the implementation
2//! [`LimitHeaderGen`] which is [attachable](rocket::Rocket::attach()) to
3//! [`Rocket`](rocket::Rocket)-instance.
4
5use crate::{header::Header, ReqState};
6use rocket::{
7    fairing::{Fairing, Info, Kind},
8    Request, Response,
9};
10
11/// Provides [`Fairing`](rocket::fairing::Fairing) implementation
12/// which is [attachable](rocket::Rocket::attach()) to
13/// [`Rocket`](rocket::Rocket)-instance.
14///
15/// `LimitHeaderGen` generates [HTTP headers](crate::header) in
16/// [Reponses](rocket::Response) to [Requests](rocket::Request) at specific
17/// [mounted routes](rocket::Rocket::mount()).
18///
19/// The [`Route`](rocket::Route) needs a [RocketGovernor guard](crate).
20/// The guard is responsible for managing the governor rate limit.
21///
22/// Depending on the implementation of
23/// [`RocketGovernable::limit_info_allow()`](crate::RocketGovernable::limit_info_allow()),
24/// `LimitHeaderGen` fairing sets [HTTP headers](crate::header) like
25///
26/// * [X_RATELIMIT_LIMIT](crate::header::X_RATELIMIT_LIMIT)
27/// * [X_RATELIMIT_REMAINING](crate::header::X_RATELIMIT_REMAINING)
28///
29/// which can be used by HTTP clients to adjust service requests.
30///
31/// ## Example usage
32///
33/// ```rust
34/// use rocket;
35/// use rocket_governor;
36///
37/// #[rocket::launch]
38/// fn launch_rocket() -> _ {
39///     rocket::build().attach(rocket_governor::LimitHeaderGen::default())
40/// }
41/// ```
42///
43pub struct LimitHeaderGen;
44
45impl Default for LimitHeaderGen {
46    fn default() -> Self {
47        Self
48    }
49}
50
51#[rocket::async_trait]
52impl Fairing for LimitHeaderGen {
53    fn info(&self) -> Info {
54        Info {
55            name: "RateLimit Header Generator",
56            kind: Kind::Response,
57        }
58    }
59
60    /// Set rate limit headers if
61    /// [`RocketGovernable::limit_info_allow()`](crate::RocketGovernable::limit_info_allow())
62    /// returns true and [`ReqState`] has been cached in this case in
63    /// [`Request::local_cache`].
64    async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
65        let state_opt = ReqState::get_or_default(request);
66
67        if let Some(state) = state_opt {
68            response.set_header(Header::XRateLimitLimit(
69                state.quota.burst_size().get().into(),
70            ));
71            response.set_header(Header::XRateLimitRemaining(state.request_capacity.into()));
72        }
73    }
74}
rocket_governor/limit_header_gen.rs

rocket_governor/
limit_header_gen.rs
