tower-rate-limiter 0.1.6

Keyed fixed-window rate limiting middleware for Tower
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
101
102
103
104
105
106
107

//! Tower service implementation.

use std::sync::Arc;
use std::task::{Context, Poll};

use http::{Request, Response};
use tower_service::Service;

use super::{
    builder::{RateLimitConfig, check_skip_predicate},
    future::ResponseFuture,
    key_extractor::KeyExtractor,
    limit::LimitProvider,
    response::ResponseFactory,
    store::Store,
};

/// Tower service produced by [`super::RateLimitLayer`].
#[must_use]
#[derive(Clone, Debug)]
pub struct RateLimit<Inner, K, S, P, F> {
    pub(crate) inner: Inner,
    pub(crate) key_extractor: K,
    pub(crate) store: S,
    pub(crate) limit_provider: P,
    pub(crate) response_factory: F,
    pub(crate) config: Arc<RateLimitConfig>,
}

impl<Inner, K, S, P, F> RateLimit<Inner, K, S, P, F> {
    /// Borrow the wrapped service.
    pub const fn get_ref(&self) -> &Inner {
        &self.inner
    }

    /// Mutably borrow the wrapped service.
    pub fn get_mut(&mut self) -> &mut Inner {
        &mut self.inner
    }

    /// Consume this middleware and return the wrapped service.
    pub fn into_inner(self) -> Inner {
        self.inner
    }
}

impl<Inner, K, S, P, F, ReqBody, ResBody> tower_service::Service<Request<ReqBody>> for RateLimit<Inner, K, S, P, F>
where
    Inner: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone,
    K: KeyExtractor,
    S: Store,
    P: LimitProvider,
    F: ResponseFactory<ReqBody, ResBody>,
{
    type Response = Inner::Response;
    type Error = Inner::Error;
    type Future = ResponseFuture<ReqBody, Inner, S, P, F>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
        let replacement = self.inner.clone();
        let inner = std::mem::replace(&mut self.inner, replacement);

        // Check if the request should be skipped.
        let request = match check_skip_predicate(self.config.skip_predicate.as_ref(), request) {
            (true, request) => {
                return ResponseFuture::skipped(
                    request,
                    inner,
                    self.store.clone(),
                    Arc::clone(&self.config),
                    self.response_factory.clone(),
                );
            },
            (false, request) => request,
        };

        // Extract the key.
        let key = match self.key_extractor.extract(&request) {
            Ok(key) => key,
            Err(error) => {
                return ResponseFuture::error(
                    request,
                    error,
                    inner,
                    self.store.clone(),
                    Arc::clone(&self.config),
                    self.response_factory.clone(),
                );
            },
        };

        let limit_future = self.limit_provider.limit(&request);
        ResponseFuture::new(
            request,
            inner,
            self.store.clone(),
            key.to_string(),
            limit_future,
            Arc::clone(&self.config),
            self.response_factory.clone(),
        )
    }
}