tower-rate-tier 0.2.0

Tier-based 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
108
109
110
111
112
113
114
115

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

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

use crate::check::{self, CheckOutcome};
use crate::cost::TierCost;
use crate::identifier::TierIdentifier;
use crate::layer::{OnLimitedFn, RateLimitedResponseFn};
use crate::on_storage_error::OnStorageError;
use crate::response;
use crate::tier::RateTier;

/// Tower service that enforces tier-based rate limiting.
///
/// Created by [`TierLimitLayer`](crate::layer::TierLimitLayer).
/// This service intercepts requests, identifies the user/tier,
/// checks the rate limit, and either forwards the request or returns 429.
pub struct TierLimitService<S> {
    pub(crate) inner: S,
    pub(crate) rate_tier: Arc<RateTier>,
    pub(crate) identifier: Arc<dyn TierIdentifier>,
    pub(crate) on_storage_error: OnStorageError,
    pub(crate) on_limited: Option<Arc<OnLimitedFn>>,
    pub(crate) rate_limited_response: Option<Arc<RateLimitedResponseFn>>,
}

impl<S: Clone> Clone for TierLimitService<S> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            rate_tier: self.rate_tier.clone(),
            identifier: self.identifier.clone(),
            on_storage_error: self.on_storage_error,
            on_limited: self.on_limited.clone(),
            rate_limited_response: self.rate_limited_response.clone(),
        }
    }
}

impl<S, B, ResBody> Service<Request<B>> for TierLimitService<S>
where
    S: Service<Request<B>, Response = Response<ResBody>> + Clone + Send + 'static,
    S::Future: Send,
    S::Error: Send,
    B: Send + 'static,
    ResBody: From<String> + Send,
{
    type Response = Response<ResBody>;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

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

    fn call(&mut self, req: Request<B>) -> Self::Future {
        let rate_tier = self.rate_tier.clone();
        let identifier = self.identifier.clone();
        let on_storage_error = self.on_storage_error;
        let on_limited = self.on_limited.clone();
        let rate_limited_response_fn = self.rate_limited_response.clone();
        let mut inner = self.inner.clone();
        // Swap to preserve readiness: the clone gets future calls, self keeps the ready one.
        std::mem::swap(&mut self.inner, &mut inner);

        Box::pin(async move {
            let identity = identifier.identify(req.headers()).await;

            let (user_id, tier_name) = match check::resolve_identity(identity, &rate_tier) {
                Ok(pair) => pair,
                Err(CheckOutcome::PassThrough) => return inner.call(req).await,
                Err(CheckOutcome::Deny(resp)) => return Ok(resp.map(Into::into)),
                Err(CheckOutcome::Allow(_)) => unreachable!(),
            };

            let quota = match check::resolve_quota(&rate_tier, &tier_name) {
                Ok(q) => q,
                Err(CheckOutcome::PassThrough) => return inner.call(req).await,
                Err(_) => unreachable!(),
            };

            let cost = req.extensions().get::<TierCost>().map(|c| c.0).unwrap_or(1);
            let now = rate_tier.clock().now();
            let storage_key = format!("{}:{}", user_id, tier_name);
            let result = rate_tier
                .storage()
                .check_and_update(&storage_key, quota, cost, now)
                .await;

            let unix_offset = rate_tier.clock().unix_offset_nanos();

            match check::process_result(
                result,
                &user_id,
                &tier_name,
                on_storage_error,
                &on_limited,
                &rate_limited_response_fn,
                unix_offset,
            ) {
                CheckOutcome::Allow(info) => {
                    let mut resp = inner.call(req).await?;
                    response::inject_headers(&mut resp, &info, unix_offset);
                    Ok(resp)
                }
                CheckOutcome::Deny(resp) => Ok(resp.map(Into::into)),
                CheckOutcome::PassThrough => inner.call(req).await,
            }
        })
    }
}