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}