actix_governor/
lib.rs

1//! A middleware for [actionable](https://github.com/actix/actix-web) that provides
2//! rate-limiting backed by [governor](https://github.com/antifuchs/governor).
3//!
4//! # Features:
5//!
6//! + Simple to use
7//! + High customizability
8//! + High performance
9//! + Robust yet flexible API
10//!
11//! # How does it work?
12//!
13//! Each governor middleware has a configuration that stores a quota.
14//! The quota specifies how many requests can be sent from an IP address
15//! before the middleware starts blocking further requests.
16//!
17//! For example if the quota allowed ten requests a client could send a burst of
18//! ten requests in short time before the middleware starts blocking.
19//!
20//! Once at least one element of the quota was used the elements of the quota
21//! will be replenished after a specified period.
22//!
23//! For example if this period was 2 seconds and the quota was empty
24//! it would take 2 seconds to replenish one element of the quota.
25//! This means you could send one request every two seconds on average.
26//!
27//! If there was a quota that allowed ten requests with the same period
28//! a client could again send a burst of ten requests and then had to wait
29//! two seconds before sending further requests or 20 seconds before the full
30//! quota would be replenished and he could send another burst.
31//!
32//! # Example
33//! ```rust,no_run
34//! use actix_governor::{Governor, GovernorConfigBuilder};
35//! use actix_web::{web, App, HttpServer, Responder};
36//!
37//! async fn index() -> impl Responder {
38//!     "Hello world!"
39//! }
40//!
41//! #[actix_web::main]
42//! async fn main() -> std::io::Result<()> {
43//!     // Allow bursts with up to five requests per IP address
44//!     // and replenishes one element every two seconds
45//!     let governor_conf = GovernorConfigBuilder::default()
46//!         .seconds_per_request(2)
47//!         .burst_size(5)
48//!         .finish()
49//!         .unwrap();
50//!
51//!     HttpServer::new(move || {
52//!         App::new()
53//!             // Enable Governor middleware
54//!             .wrap(Governor::new(&governor_conf))
55//!             // Route hello world service
56//!             .route("/", web::get().to(index))
57//!    })
58//!     .bind("127.0.0.1:8080")?
59//!     .run()
60//!     .await
61//! }
62//! ```
63//!
64//! # Configuration presets
65//!
66//! Instead of using the configuration builder you can use predefined presets.
67//!
68//! + [`GovernorConfig::default()`]: The default configuration which is suitable for most services.
69//!   Allows bursts with up to eight requests and replenishes one element after 500ms, based on peer IP.
70//!
71//! + [`GovernorConfig::secure()`]: A default configuration for security related services.
72//!   Allows bursts with up to two requests and replenishes one element after four seconds, based on peer IP.
73//!
74//! For example the secure configuration can be used as a short version of this code:
75//!
76//! ```rust
77//! use actix_governor::GovernorConfigBuilder;
78//!
79//! let config = GovernorConfigBuilder::default()
80//!     .seconds_per_request(4)
81//!     .burst_size(2)
82//!     .finish()
83//!     .unwrap();
84//! ```
85//!
86//! # Customize rate limiting key
87//!
88//! By default, rate limiting is done using the peer IP address (i.e. the IP address of the HTTP client that requested your app: either your user or a reverse proxy, depending on your deployment setup).
89//! You can configure a different behavior which:
90//! 1. can be useful in itself
91//! 2. allows you to setup multiple instances of this middleware based on different keys (for example, if you want to apply rate limiting with different rates on IP and API keys at the same time)
92//!
93//! This is achieved by defining a [KeyExtractor] and giving it to a [Governor] instance.
94//! Two ready-to-use key extractors are provided:
95//! - [PeerIpKeyExtractor]: this is the default
96//! - [GlobalKeyExtractor]: uses the same key for all incoming requests
97//!
98//! Check out the [custom_key](https://github.com/AaronErhardt/actix-governor/blob/main/examples/custom_key.rs) example to see how a custom key extractor can be implemented.
99//!
100//! # Customizing error responses
101//!
102//! There are two errors that might occur during rate-limiting.
103//! The first error occurs when the key can't be extracted from the request, for example when a session id is missing.
104//! The second error occurs when the rate limit is exceeded.
105//!
106//! ## The response when key-extractions fails
107//! The response for this error is generated by the implementation of [`ResponseError::error_response`] for [KeyExtractor::KeyExtractionError].
108//! With this method you can generate any [`HttpResponse`] you want, [for example to return json].
109//!
110//! But it has a simplistic problem that you can't access the request directly, but you can solve it by creating a
111//! `new` method in the struct and passing the request that is given to you is in a method [`KeyExtractor::extract`] to it and do whatever you want with it
112//!
113//! For most cases can simply use [`SimpleKeyExtractionError`] to return an error.
114//! It returns a response with the body you set in `new`, type `text/plain` and `500 Internal Server Error` status by default,
115//! but you can customize the content type with [`set_content_type`] and the status with [`set_status_code`].
116//!
117//! [`ResponseError::error_response`]: actix_web::error::ResponseError::error_response
118//! [`set_status_code`]: SimpleKeyExtractionError::set_status_code
119//! [`set_content_type`]: SimpleKeyExtractionError::set_content_type
120//! [for example to return json]: https://github.com/AaronErhardt/actix-governor/blob/main/examples/custom_key_bearer.rs
121//! [`HttpResponse`]: actix_web::HttpResponse
122//!
123//! ## The response of exceeding the rate limit
124//! The response of this error is generated by [`KeyExtractor::exceed_rate_limit_response`].
125//! This method will give you a [`HttpResponseBuilder`] and return a [`HttpResponse`].
126//! This allows you to fully customize the response.
127//!
128//! Check out the [custom_key_bearer] example for more information.
129//!
130//! [`HttpResponseBuilder`]: actix_web::HttpResponseBuilder
131//! [`HttpResponse`]: actix_web::HttpResponse
132//! [custom_key_bearer]: https://github.com/AaronErhardt/actix-governor/blob/main/examples/custom_key_bearer.rs
133//!
134//! # Add x-ratelimit headers
135//!
136//! By default, `retry-after` and `x-ratelimit-after` are enabled but if you want to enable `x-ratelimit-limit`, `x-ratelimit-whitelisted` and `x-ratelimit-remaining` use [`use_headers`] method
137//!
138//! [`use_headers`]: crate::GovernorConfigBuilder::use_headers()
139//!
140//! # Common pitfalls
141//!
142//! ## Creating independant rate limiters
143//!
144//! Do not construct the same configuration multiple times, unless explicitly wanted!
145//! This will create an independent rate limiter for each configuration!
146//!
147//! Instead pass the same configuration reference into [`Governor::new()`],
148//! like it is described in the example.
149//!
150//! ## Memory leak (because of keys accumulation)
151//!
152//! If your application gets a lot of traffic and your rate limiter ends up with a lot of keys (for example, user tokens), you may observe some kind of memory leak: your application consumes more and more memory over time.
153//!
154//! In this case, you may want to regularly call [`governor::RateLimiter::retain_recent`] followed by [`governor::RateLimiter::shrink_to_fit`].
155//! These methods need to be called on the underlying [`governor::RateLimiter`] which can be accessed using [`GovernorConfig::limiter`].
156//!
157//! As you probably already have a Tokio context, you can do this from inside a task that is spawned after creating your rate limiter (`GovernorConfig` can be cloned while keeping the same underlying rate limiter).
158
159#![warn(
160    rust_2018_idioms,
161    unreachable_pub,
162    missing_docs,
163    clippy::must_use_candidate,
164    clippy::cargo
165)]
166
167#[cfg(test)]
168mod tests;
169
170use governor::{
171    clock::{DefaultClock, QuantaInstant},
172    middleware::{NoOpMiddleware, RateLimitingMiddleware, StateInformationMiddleware},
173    state::keyed::DefaultKeyedStateStore,
174    Quota, RateLimiter,
175};
176
177use actix_http::body::EitherBody;
178use std::{cell::RefCell, marker::PhantomData, num::NonZeroU32, rc::Rc, sync::Arc, time::Duration};
179
180use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
181use actix_web::http::Method;
182use actix_web::{body::MessageBody, Error};
183use futures::future;
184
185mod extractor;
186mod key_extractor;
187mod service;
188
189type SharedRateLimiter<Key, M> =
190    Arc<RateLimiter<Key, DefaultKeyedStateStore<Key>, DefaultClock, M>>;
191
192/// Re-export governor
193pub use governor;
194
195pub use extractor::GovernorExtractor;
196pub use key_extractor::{
197    GlobalKeyExtractor, KeyExtractor, PeerIpKeyExtractor, SimpleKeyExtractionError,
198};
199
200const DEFAULT_PERIOD: Duration = Duration::from_millis(500);
201const DEFAULT_BURST_SIZE: u32 = 8;
202
203/// Helper struct for building a configuration for the governor middleware.
204///
205/// # Example
206///
207/// Create a configuration with a quota of ten requests per IP address
208/// that replenishes one element every minute.
209///
210/// ```rust
211/// use actix_governor::GovernorConfigBuilder;
212///
213/// let config = GovernorConfigBuilder::default()
214///     .seconds_per_request(60)
215///     .burst_size(10)
216///     .finish()
217///     .unwrap();
218/// ```
219///
220/// with x-ratelimit headers
221///
222/// ```rust
223/// use actix_governor::GovernorConfigBuilder;
224///
225/// let config = GovernorConfigBuilder::default()
226///     .seconds_per_request(60)
227///     .burst_size(10)
228///     .use_headers() // Add this
229///     .finish()
230///     .unwrap();
231/// ```
232#[must_use]
233#[derive(Debug, Eq)]
234pub struct GovernorConfigBuilder<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> {
235    period: Duration,
236    burst_size: u32,
237    methods: Option<Vec<Method>>,
238    key_extractor: K,
239    middleware: PhantomData<M>,
240    permissive: bool,
241}
242
243impl<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> Clone
244    for GovernorConfigBuilder<K, M>
245{
246    fn clone(&self) -> Self {
247        Self {
248            period: self.period,
249            burst_size: self.burst_size,
250            methods: self.methods.clone(),
251            key_extractor: self.key_extractor.clone(),
252            middleware: self.middleware,
253            permissive: self.permissive,
254        }
255    }
256}
257
258impl<K: KeyExtractor + PartialEq, M: RateLimitingMiddleware<QuantaInstant>> PartialEq
259    for GovernorConfigBuilder<K, M>
260{
261    fn eq(&self, other: &Self) -> bool {
262        self.period == other.period
263            && self.burst_size == other.burst_size
264            && self.methods == other.methods
265            && self.key_extractor == other.key_extractor
266            && self.permissive == other.permissive
267    }
268}
269
270impl Default for GovernorConfigBuilder<PeerIpKeyExtractor, NoOpMiddleware> {
271    /// The default configuration which is suitable for most services.
272    /// Allows burst with up to eight requests and replenishes one element after 500ms, based on peer IP.
273    /// The values can be modified by calling other methods on this struct.
274    fn default() -> Self {
275        Self::const_default()
276    }
277}
278
279impl<M: RateLimitingMiddleware<QuantaInstant>> GovernorConfigBuilder<PeerIpKeyExtractor, M> {
280    /// Returns the default configuration.
281    pub const fn const_default() -> Self {
282        GovernorConfigBuilder {
283            period: DEFAULT_PERIOD,
284            burst_size: DEFAULT_BURST_SIZE,
285            methods: None,
286            key_extractor: PeerIpKeyExtractor,
287            middleware: PhantomData,
288            permissive: false,
289        }
290    }
291    /// Set the interval after which one element of the quota is replenished.
292    ///
293    /// **The interval must not be zero.**
294    pub const fn const_period(mut self, duration: Duration) -> Self {
295        self.period = duration;
296        self
297    }
298    /// Set the number of quota elements to replenish per second.
299    pub fn const_requests_per_second(self, count: u64) -> Self {
300        let replenish_interval_ns = Duration::from_secs(1).as_nanos() / count as u128;
301        self.const_nanoseconds_per_request(replenish_interval_ns as u64)
302    }
303    /// Set the number of quota elements to replenish per minute.
304    pub fn const_requests_per_minute(self, count: u64) -> Self {
305        let replenish_interval_ns = Duration::from_secs(60).as_nanos() / count as u128;
306        self.const_nanoseconds_per_request(replenish_interval_ns as u64)
307    }
308    /// Set the number of quota elements to replenish per hour.
309    pub fn const_requests_per_hour(self, count: u64) -> Self {
310        let replenish_interval_ns = Duration::from_secs(60 * 60).as_nanos() / count as u128;
311        self.const_nanoseconds_per_request(replenish_interval_ns as u64)
312    }
313    /// Renamed to `const_seconds_per_request`.
314    ///
315    /// **The interval must not be zero.**
316    #[deprecated(
317        since = "0.6.0",
318        note = "Might be the inverse of what's expected. Use `const_seconds_per_request` as an exact replacement."
319    )]
320    pub const fn const_per_second(mut self, seconds: u64) -> Self {
321        self.period = Duration::from_secs(seconds);
322        self
323    }
324    /// Set the interval after which one element of the quota is replenished in seconds.
325    ///
326    /// **The interval must not be zero.**
327    pub const fn const_second_per_request(mut self, seconds: u64) -> Self {
328        self.period = Duration::from_secs(seconds);
329        self
330    }
331    /// Renamed to `const_milliseconds_per_request`.
332    ///
333    /// **The interval must not be zero.**
334    #[deprecated(
335        since = "0.6.0",
336        note = "Might be the inverse of what's expected. Use `const_milliseconds_per_request` as an exact replacement."
337    )]
338    pub const fn const_per_millisecond(mut self, milliseconds: u64) -> Self {
339        self.period = Duration::from_millis(milliseconds);
340        self
341    }
342    /// Set the interval after which one element of the quota is replenished in milliseconds.
343    ///
344    /// **The interval must not be zero.**
345    pub const fn const_milliseconds_per_request(mut self, milliseconds: u64) -> Self {
346        self.period = Duration::from_millis(milliseconds);
347        self
348    }
349    /// Renamed to `const_nanoseconds_per_request`.
350    ///
351    /// **The interval must not be zero.**
352    #[deprecated(
353        since = "0.6.0",
354        note = "Might be the inverse of what's expected. Use `const_nanoseconds_per_request` as an exact replacement."
355    )]
356    pub const fn const_per_nanosecond(mut self, nanoseconds: u64) -> Self {
357        self.period = Duration::from_nanos(nanoseconds);
358        self
359    }
360    /// Set the interval after which one element of the quota is replenished in nanoseconds.
361    ///
362    /// **The interval must not be zero.**
363    pub const fn const_nanoseconds_per_request(mut self, nanoseconds: u64) -> Self {
364        self.period = Duration::from_nanos(nanoseconds);
365        self
366    }
367    /// Set quota size that defines how many requests can occur
368    /// before the governor middleware starts blocking requests from an IP address and
369    /// clients have to wait until the elements of the quota are replenished.
370    ///
371    /// **The burst_size must not be zero.**
372    pub const fn const_burst_size(mut self, burst_size: u32) -> Self {
373        self.burst_size = burst_size;
374        self
375    }
376    /// Set the mode of the governor middleware.
377    ///
378    /// If permissive is set to true, the middleware will not block requests.
379    /// See also [`GovernorExtractor`].
380    pub const fn const_permissive(mut self, permissive: bool) -> Self {
381        self.permissive = permissive;
382        self
383    }
384}
385
386impl<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> GovernorConfigBuilder<K, M> {
387    /// Set the interval after which one element of the quota is replenished.
388    ///
389    /// **The interval must not be zero.**
390    pub fn period(&mut self, duration: Duration) -> &mut Self {
391        self.period = duration;
392        self
393    }
394    /// Set the number of quota elements to replenish per second.
395    pub fn requests_per_second(&mut self, count: u64) -> &mut Self {
396        let replenish_interval_ns = Duration::from_secs(1).as_nanos() / count as u128;
397        self.nanoseconds_per_request(replenish_interval_ns as u64)
398    }
399    /// Set the number of quota elements to replenish per minute.
400    pub fn requests_per_minute(&mut self, count: u64) -> &mut Self {
401        let replenish_interval_ns = Duration::from_secs(60).as_nanos() / count as u128;
402        self.nanoseconds_per_request(replenish_interval_ns as u64)
403    }
404    /// Set the number of quota elements to replenish per hour.
405    pub fn requests_per_hour(&mut self, count: u64) -> &mut Self {
406        let replenish_interval_ns = Duration::from_secs(60 * 60).as_nanos() / count as u128;
407        self.nanoseconds_per_request(replenish_interval_ns as u64)
408    }
409    /// Renamed to `seconds_per_request`.
410    ///
411    /// **The interval must not be zero.**
412    #[deprecated(
413        since = "0.6.0",
414        note = "Might be the inverse of what's expected. Use `seconds_per_request` as an exact replacement."
415    )]
416    pub fn per_second(&mut self, seconds: u64) -> &mut Self {
417        self.period = Duration::from_secs(seconds);
418        self
419    }
420    /// Set the interval after which one element of the quota is replenished in seconds.
421    ///
422    /// **The interval must not be zero.**
423    pub fn seconds_per_request(&mut self, seconds: u64) -> &mut Self {
424        self.period = Duration::from_secs(seconds);
425        self
426    }
427    /// Renamed to `milliseconds_per_request`.
428    ///
429    /// **The interval must not be zero.**
430    #[deprecated(
431        since = "0.6.0",
432        note = "Might be the inverse of what's expected. Use `milliseconds_per_request` as an exact replacement."
433    )]
434    pub fn per_millisecond(&mut self, milliseconds: u64) -> &mut Self {
435        self.period = Duration::from_millis(milliseconds);
436        self
437    }
438    /// Set the interval after which one element of the quota is replenished in milliseconds.
439    ///
440    /// **The interval must not be zero.**
441    pub fn milliseconds_per_request(&mut self, milliseconds: u64) -> &mut Self {
442        self.period = Duration::from_millis(milliseconds);
443        self
444    }
445    /// Renamed to `nanoseconds_per_request`.
446    ///
447    /// **The interval must not be zero.**
448    #[deprecated(
449        since = "0.6.0",
450        note = "Might be the inverse of what's expected. Use `nanoseconds_per_request` as an exact replacement."
451    )]
452    pub fn per_nanosecond(&mut self, nanoseconds: u64) -> &mut Self {
453        self.period = Duration::from_nanos(nanoseconds);
454        self
455    }
456    /// Set the interval after which one element of the quota is replenished in nanoseconds.
457    ///
458    /// **The interval must not be zero.**
459    pub fn nanoseconds_per_request(&mut self, nanoseconds: u64) -> &mut Self {
460        self.period = Duration::from_nanos(nanoseconds);
461        self
462    }
463    /// Set quota size that defines how many requests can occur
464    /// before the governor middleware starts blocking requests from an IP address and
465    /// clients have to wait until the elements of the quota are replenished.
466    ///
467    /// **The burst_size must not be zero.**
468    pub fn burst_size(&mut self, burst_size: u32) -> &mut Self {
469        self.burst_size = burst_size;
470        self
471    }
472    /// Set the mode of the governor middleware.
473    ///
474    /// If permissive is set to true, the middleware will not block requests.
475    /// See also [`GovernorExtractor`].
476    pub fn permissive(&mut self, permissive: bool) -> &mut Self {
477        self.permissive = permissive;
478        self
479    }
480
481    /// Set the HTTP methods this configuration should apply to.
482    /// By default this is all methods.
483    pub fn methods(&mut self, methods: Vec<Method>) -> &mut Self {
484        self.methods = Some(methods);
485        self
486    }
487
488    /// Set the key extractor this configuration should use.
489    /// By default this is using the [PeerIpKeyExtractor].
490    pub fn key_extractor<K2: KeyExtractor>(
491        &mut self,
492        key_extractor: K2,
493    ) -> GovernorConfigBuilder<K2, M> {
494        GovernorConfigBuilder {
495            period: self.period,
496            burst_size: self.burst_size,
497            methods: self.methods.to_owned(),
498            key_extractor,
499            middleware: PhantomData,
500            permissive: self.permissive,
501        }
502    }
503
504    /// Set x-ratelimit headers to response, the headers is
505    /// - `retry-after`             - Number of seconds in which the API will become available after its rate limit has been exceeded
506    /// - `x-ratelimit-after`       - Number of seconds in which the API will become available after its rate limit has been exceeded
507    /// - `x-ratelimit-limit`       - Request limit
508    /// - `x-ratelimit-remaining`   - The number of requests left for the time window
509    /// - `x-ratelimit-whitelisted` - If the request method not in methods, this header will be add it, use [`methods`] to add methods
510    ///
511    /// By default `retry-after` and `x-ratelimit-after` are enabled, with [`use_headers`] will enable `x-ratelimit-limit`, `x-ratelimit-whitelisted` and `x-ratelimit-remaining`
512    ///
513    /// [`methods`]: crate::GovernorConfigBuilder::methods()
514    /// [`use_headers`]: Self::use_headers
515    pub fn use_headers(&mut self) -> GovernorConfigBuilder<K, StateInformationMiddleware> {
516        GovernorConfigBuilder {
517            period: self.period,
518            burst_size: self.burst_size,
519            methods: self.methods.to_owned(),
520            key_extractor: self.key_extractor.clone(),
521            middleware: PhantomData,
522            permissive: self.permissive,
523        }
524    }
525
526    /// Finish building the configuration and return the configuration for the middleware.
527    /// Returns `None` if either burst size or period interval are zero.
528    pub fn finish(&mut self) -> Option<GovernorConfig<K, M>> {
529        if self.burst_size != 0 && self.period.as_nanos() != 0 {
530            Some(GovernorConfig {
531                key_extractor: self.key_extractor.clone(),
532                limiter: Arc::new(
533                    RateLimiter::keyed(
534                        Quota::with_period(self.period)
535                            .unwrap()
536                            .allow_burst(NonZeroU32::new(self.burst_size).unwrap()),
537                    )
538                    .with_middleware::<M>(),
539                ),
540                methods: self.methods.clone(),
541                permissive: self.permissive,
542            })
543        } else {
544            None
545        }
546    }
547}
548
549#[derive(Debug)]
550#[must_use]
551/// Configuration for the Governor middleware.
552pub struct GovernorConfig<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> {
553    key_extractor: K,
554    limiter: SharedRateLimiter<K::Key, M>,
555    methods: Option<Vec<Method>>,
556    permissive: bool,
557}
558
559impl<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> Clone for GovernorConfig<K, M> {
560    fn clone(&self) -> Self {
561        GovernorConfig {
562            key_extractor: self.key_extractor.clone(),
563            limiter: self.limiter.clone(),
564            methods: self.methods.clone(),
565            permissive: self.permissive,
566        }
567    }
568}
569
570impl Default for GovernorConfig<PeerIpKeyExtractor, NoOpMiddleware> {
571    /// The default configuration which is suitable for most services.
572    /// Allows bursts with up to eight requests and replenishes one element after 500ms, based on peer IP.
573    fn default() -> Self {
574        GovernorConfigBuilder::default().finish().unwrap()
575    }
576}
577
578impl<M: RateLimitingMiddleware<QuantaInstant>> GovernorConfig<PeerIpKeyExtractor, M> {
579    /// A default configuration for security related services.
580    /// Allows bursts with up to two requests and replenishes one element after four seconds, based on peer IP.
581    ///
582    /// This prevents brute-forcing passwords or security tokens
583    /// yet allows to quickly retype a wrong password once before the quota is exceeded.
584    pub fn secure() -> Self {
585        GovernorConfigBuilder {
586            period: Duration::from_secs(4),
587            burst_size: 2,
588            methods: None,
589            key_extractor: PeerIpKeyExtractor,
590            middleware: PhantomData,
591            permissive: false,
592        }
593        .finish()
594        .unwrap()
595    }
596}
597
598impl<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> GovernorConfig<K, M> {
599    /// Access internal [`governor::RateLimiter`]
600    pub fn limiter(&self) -> SharedRateLimiter<K::Key, M> {
601        self.limiter.clone()
602    }
603}
604
605/// Governor middleware factory.
606pub struct Governor<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> {
607    key_extractor: K,
608    limiter: SharedRateLimiter<K::Key, M>,
609    methods: Option<Vec<Method>>,
610    permissive: bool,
611}
612
613impl<K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> Governor<K, M> {
614    /// Create new governor middleware factory from configuration.
615    pub fn new(config: &GovernorConfig<K, M>) -> Self {
616        Governor {
617            key_extractor: config.key_extractor.clone(),
618            limiter: config.limiter.clone(),
619            methods: config.methods.clone(),
620            permissive: config.permissive,
621        }
622    }
623}
624
625impl<S, B, K> Transform<S, ServiceRequest> for Governor<K, NoOpMiddleware>
626where
627    K: KeyExtractor,
628    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
629    B: MessageBody,
630{
631    type Response = ServiceResponse<EitherBody<B>>;
632    type Error = Error;
633    type Transform = GovernorMiddleware<S, K, NoOpMiddleware>;
634    type InitError = ();
635    type Future = future::Ready<Result<Self::Transform, Self::InitError>>;
636
637    fn new_transform(&self, service: S) -> Self::Future {
638        future::ok(GovernorMiddleware::<S, K, NoOpMiddleware> {
639            service: Rc::new(RefCell::new(service)),
640            key_extractor: self.key_extractor.clone(),
641            limiter: self.limiter.clone(),
642            methods: self.methods.clone(),
643            permissive: self.permissive,
644        })
645    }
646}
647
648impl<S, B, K> Transform<S, ServiceRequest> for Governor<K, StateInformationMiddleware>
649where
650    K: KeyExtractor,
651    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
652    B: MessageBody,
653    <S as Service<ServiceRequest>>::Future: Unpin,
654{
655    type Response = ServiceResponse<EitherBody<B>>;
656    type Error = Error;
657    type Transform = GovernorMiddleware<S, K, StateInformationMiddleware>;
658    type InitError = ();
659    type Future = future::Ready<Result<Self::Transform, Self::InitError>>;
660
661    fn new_transform(&self, service: S) -> Self::Future {
662        future::ok(GovernorMiddleware::<S, K, StateInformationMiddleware> {
663            service: Rc::new(RefCell::new(service)),
664            key_extractor: self.key_extractor.clone(),
665            limiter: self.limiter.clone(),
666            methods: self.methods.clone(),
667            permissive: self.permissive,
668        })
669    }
670}
671
672#[derive(Debug, Copy, Clone, PartialEq, Eq)]
673/// The result of a [`GovernorExtractor`].
674pub enum GovernorResult<E> {
675    /// The request does not exceed the rate limit.
676    Ok {
677        /// The maximum burst size.
678        burst_size: Option<u32>,
679        /// The remaining burst capacity.
680        remaining: Option<u32>,
681    },
682    /// The request method was whitelisted.
683    Whitelisted,
684    /// The request exceeds the rate limit.
685    Wait {
686        /// The maximum burst size.
687        burst_size: Option<u32>,
688        /// The time to wait for new requests in seconds.
689        wait: u64,
690    },
691    /// Internal error.
692    Err(E),
693}
694
695impl<E> GovernorResult<E> {
696    const fn whitelist() -> Self {
697        Self::Whitelisted
698    }
699
700    const fn ok() -> Self {
701        Self::Ok {
702            burst_size: None,
703            remaining: None,
704        }
705    }
706
707    const fn ok_with_info(burst_size: u32, remaining: u32) -> Self {
708        Self::Ok {
709            burst_size: Some(burst_size),
710            remaining: Some(remaining),
711        }
712    }
713
714    const fn wait(wait: u64) -> Self {
715        Self::Wait {
716            wait,
717            burst_size: None,
718        }
719    }
720
721    const fn wait_with_info(wait: u64, burst_size: u32) -> Self {
722        Self::Wait {
723            wait,
724            burst_size: Some(burst_size),
725        }
726    }
727
728    const fn err(e: E) -> Self {
729        Self::Err(e)
730    }
731}
732
733impl<E> GovernorResult<E> {
734    /// Check if this request is rate limited.
735    ///
736    /// Returns `Ok(Some(wait))` if the request is rate limited and should be retried after `wait` milliseconds.
737    ///
738    /// Returns `Ok(None)` if the request is not rate limited.
739    ///
740    /// # Errors
741    /// Returns error if key extraction failed.
742    pub const fn check(&self) -> Result<Option<u64>, &E> {
743        match self {
744            Self::Wait { wait, .. } => Ok(Some(*wait)),
745            Self::Err(e) => Err(e),
746            _ => Ok(None),
747        }
748    }
749}
750
751/// A middleware that implements rate limiting based on the governor crate.
752pub struct GovernorMiddleware<S, K: KeyExtractor, M: RateLimitingMiddleware<QuantaInstant>> {
753    service: std::rc::Rc<std::cell::RefCell<S>>,
754    key_extractor: K,
755    limiter: SharedRateLimiter<K::Key, M>,
756    methods: Option<Vec<Method>>,
757    permissive: bool,
758}