actix_extensible_rate_limit/middleware/
builder.rs

1use crate::backend::Backend;
2use crate::middleware::{AllowedTransformation, DeniedResponse, RateLimiter, RollbackCondition};
3use actix_web::dev::ServiceRequest;
4use actix_web::http::header::{HeaderMap, HeaderName, HeaderValue, RETRY_AFTER};
5use actix_web::http::StatusCode;
6use actix_web::HttpResponse;
7use std::future::Future;
8use std::rc::Rc;
9
10#[allow(clippy::declare_interior_mutable_const)]
11pub const X_RATELIMIT_LIMIT: HeaderName = HeaderName::from_static("x-ratelimit-limit");
12#[allow(clippy::declare_interior_mutable_const)]
13pub const X_RATELIMIT_REMAINING: HeaderName = HeaderName::from_static("x-ratelimit-remaining");
14#[allow(clippy::declare_interior_mutable_const)]
15pub const X_RATELIMIT_RESET: HeaderName = HeaderName::from_static("x-ratelimit-reset");
16
17pub struct RateLimiterBuilder<BE, BO, F> {
18    backend: BE,
19    input_fn: F,
20    fail_open: bool,
21    allowed_transformation: Option<Rc<AllowedTransformation<BO>>>,
22    denied_response: Rc<DeniedResponse<BO>>,
23    rollback_condition: Option<Rc<RollbackCondition>>,
24}
25
26impl<BE, BI, BO, F, O> RateLimiterBuilder<BE, BO, F>
27where
28    BE: Backend<BI, Output = BO> + 'static,
29    BI: 'static,
30    F: Fn(&ServiceRequest) -> O,
31    O: Future<Output = Result<BI, actix_web::Error>>,
32{
33    pub(super) fn new(backend: BE, input_fn: F) -> Self {
34        Self {
35            backend,
36            input_fn,
37            fail_open: false,
38            allowed_transformation: None,
39            denied_response: Rc::new(|_| HttpResponse::TooManyRequests().finish()),
40            rollback_condition: None,
41        }
42    }
43
44    /// Choose whether to allow a request if the backend returns a failure.
45    ///
46    /// Default is false.
47    pub fn fail_open(mut self, fail_open: bool) -> Self {
48        self.fail_open = fail_open;
49        self
50    }
51
52    /// Sets the [RateLimiterBuilder::request_allowed_transformation] and
53    /// [RateLimiterBuilder::request_denied_response] functions, such that the following headers
54    /// are set in both the allowed and denied responses:
55    ///
56    /// - `x-ratelimit-limit`\
57    /// - `x-ratelimit-remaining`\
58    /// - `x-ratelimit-reset` (seconds until the reset)
59    /// - `retry-after` (denied only, seconds until the reset)
60    ///
61    /// This function requires the Backend Output to implement [HeaderCompatibleOutput]
62    pub fn add_headers(mut self) -> Self
63    where
64        BO: HeaderCompatibleOutput,
65    {
66        self.allowed_transformation = Some(Rc::new(|map, output, rolled_back| {
67            if let Some(status) = output {
68                map.insert(X_RATELIMIT_LIMIT, HeaderValue::from(status.limit()));
69                let remaining = if rolled_back {
70                    status.remaining() + 1
71                } else {
72                    status.remaining()
73                };
74                map.insert(X_RATELIMIT_REMAINING, HeaderValue::from(remaining));
75                map.insert(
76                    X_RATELIMIT_RESET,
77                    HeaderValue::from(status.seconds_until_reset()),
78                );
79            }
80        }));
81        self.denied_response = Rc::new(|status| {
82            let mut response = HttpResponse::TooManyRequests().finish();
83            let map = response.headers_mut();
84            map.insert(X_RATELIMIT_LIMIT, HeaderValue::from(status.limit()));
85            map.insert(X_RATELIMIT_REMAINING, HeaderValue::from(status.remaining()));
86            let seconds = status.seconds_until_reset();
87            map.insert(X_RATELIMIT_RESET, HeaderValue::from(seconds));
88            map.insert(RETRY_AFTER, HeaderValue::from(seconds));
89            response
90        });
91        self
92    }
93
94    /// In the event that the request is allowed:
95    ///
96    /// You can optionally mutate the response headers to include the rate limit status.
97    ///
98    /// By default no changes are made to the response.
99    ///
100    /// Note the [Backend::Output] will be [None] if the backend failed and
101    /// [RateLimiterBuilder::fail_open] is enabled.
102    ///
103    /// The boolean parameter indicates if the rate limit was rolled back (so the remaining
104    /// request count can be adjusted).
105    pub fn request_allowed_transformation<M>(mut self, mutation: Option<M>) -> Self
106    where
107        M: Fn(&mut HeaderMap, Option<&BO>, bool) + 'static,
108    {
109        self.allowed_transformation = mutation.map(|m| Rc::new(m) as Rc<AllowedTransformation<BO>>);
110        self
111    }
112
113    /// In the event that the request is denied, configure the [HttpResponse] returned.
114    ///
115    /// Defaults to an empty body with status 429.
116    pub fn request_denied_response<R>(mut self, denied_response: R) -> Self
117    where
118        R: Fn(&BO) -> HttpResponse + 'static,
119    {
120        self.denied_response = Rc::new(denied_response);
121        self
122    }
123
124    /// After processing a request, attempt to rollback the request count based on the status
125    /// of the service response.
126    ///
127    /// By default the rate limit is never rolled back.
128    pub fn rollback_condition<C>(mut self, condition: Option<C>) -> Self
129    where
130        C: Fn(StatusCode) -> bool + 'static,
131    {
132        self.rollback_condition = condition.map(|m| Rc::new(m) as Rc<RollbackCondition>);
133        self
134    }
135
136    /// Configures the [RateLimiterBuilder::rollback_condition] to rollback if the status code
137    /// is a server error (5xx).
138    pub fn rollback_server_errors(self) -> Self {
139        self.rollback_condition(Some(|status: StatusCode| status.is_server_error()))
140    }
141
142    pub fn build(self) -> RateLimiter<BE, BO, F> {
143        RateLimiter {
144            backend: self.backend,
145            input_fn: Rc::new(self.input_fn),
146            fail_open: self.fail_open,
147            allowed_mutation: self.allowed_transformation,
148            denied_response: self.denied_response,
149            rollback_condition: self.rollback_condition,
150        }
151    }
152}
153
154/// A trait that a [Backend::Output] should implement in order to use the
155/// [RateLimiterBuilder::add_headers] function.
156pub trait HeaderCompatibleOutput {
157    /// Value for the `x-ratelimit-limit` header.
158    fn limit(&self) -> u64;
159
160    /// Value for the `x-ratelimit-remaining` header.
161    fn remaining(&self) -> u64;
162
163    /// Value for the `x-ratelimit-reset` and `retry-at` headers.
164    ///
165    /// This should be the number of seconds from now until the limit resets.\
166    /// If the limit has already reset this should return 0.
167    fn seconds_until_reset(&self) -> u64;
168}
actix_extensible_rate_limit/middleware/builder.rs

actix_extensible_rate_limit/middleware/
builder.rs
