actix_extensible_rate_limit/backend/
mod.rs

1mod input_builder;
2
3#[cfg(feature = "dashmap")]
4#[cfg_attr(docsrs, doc(cfg(feature = "dashmap")))]
5pub mod memory;
6
7#[cfg(feature = "redis")]
8#[cfg_attr(docsrs, doc(cfg(feature = "redis")))]
9pub mod redis;
10
11pub use input_builder::{SimpleInputFunctionBuilder, SimpleInputFuture};
12use std::future::Future;
13
14use crate::HeaderCompatibleOutput;
15use actix_web::rt::time::Instant;
16use std::time::Duration;
17
18#[derive(Copy, Clone, Debug, Eq, PartialEq)]
19pub enum Decision {
20    Allowed,
21    Denied,
22}
23
24impl Decision {
25    pub fn from_allowed(allowed: bool) -> Self {
26        if allowed {
27            Self::Allowed
28        } else {
29            Self::Denied
30        }
31    }
32
33    pub fn is_allowed(self) -> bool {
34        matches!(self, Self::Allowed)
35    }
36
37    pub fn is_denied(self) -> bool {
38        matches!(self, Self::Denied)
39    }
40}
41
42/// Describes an implementation of a rate limiting store and algorithm.
43///
44/// A Backend is required to implement [Clone], usually this means wrapping your data store within
45/// an [Arc](std::sync::Arc), although many connection pools already do so internally; there is no
46/// need to wrap it twice.
47pub trait Backend<I: 'static = SimpleInput>: Clone {
48    type Output;
49    type RollbackToken;
50    type Error;
51
52    /// Process an incoming request.
53    ///
54    /// The input could include such things as a rate limit key, and the rate limit policy to be
55    /// applied.
56    ///
57    /// Returns a boolean of whether to allow or deny the request, arbitrary output that can be used
58    /// to transform the allowed and denied responses, and a token to allow the rate limit counter
59    /// to be rolled back in certain conditions.
60    fn request(
61        &self,
62        input: I,
63    ) -> impl Future<Output = Result<(Decision, Self::Output, Self::RollbackToken), Self::Error>>;
64
65    /// Under certain conditions we may not want to rollback the request operation.
66    ///
67    /// E.g. We may want to exclude 5xx errors from counting against a user's rate limit,
68    /// we can only exclude them after having already allowed the request through the rate limiter
69    /// in the first place, so we must therefore deduct from the rate limit counter afterwards.
70    ///
71    /// Note that if this function fails there is not much the [RateLimiter](crate::RateLimiter)
72    /// can do about it, given that the request has already been allowed.
73    ///
74    /// # Arguments
75    ///
76    /// * `token`: The token returned from the initial call to [Backend::request()].
77    fn rollback(&self, token: Self::RollbackToken)
78        -> impl Future<Output = Result<(), Self::Error>>;
79}
80
81/// A default [Backend] Input structure.
82///
83/// This may not be suitable for all use-cases.
84#[derive(Debug, Clone)]
85pub struct SimpleInput {
86    /// The rate limiting interval.
87    pub interval: Duration,
88    /// The total requests to be allowed within the interval.
89    pub max_requests: u64,
90    /// The rate limit key to be used for this request.
91    pub key: String,
92}
93
94/// A default [Backend::Output] structure.
95///
96/// This may not be suitable for all use-cases.
97#[derive(Debug, Clone)]
98pub struct SimpleOutput {
99    /// Total number of requests that are permitted within the rate limit interval.
100    pub limit: u64,
101    /// Number of requests that will be permitted until the limit resets.
102    pub remaining: u64,
103    /// Time at which the rate limit resets.
104    pub reset: Instant,
105}
106
107/// Additional functions for a [Backend] that uses [SimpleInput] and [SimpleOutput].
108pub trait SimpleBackend: Backend<SimpleInput, Output = SimpleOutput> {
109    /// Removes the bucket for a given rate limit key.
110    ///
111    /// Intended to be used to reset a key before changing the interval.
112    fn remove_key(&self, key: &str) -> impl Future<Output = Result<(), Self::Error>>;
113}
114
115impl HeaderCompatibleOutput for SimpleOutput {
116    fn limit(&self) -> u64 {
117        self.limit
118    }
119
120    fn remaining(&self) -> u64 {
121        self.remaining
122    }
123
124    /// Seconds until the rate limit resets (rounded upwards, so that it is guaranteed to be reset
125    /// after waiting for the duration).
126    fn seconds_until_reset(&self) -> u64 {
127        let millis = self
128            .reset
129            .saturating_duration_since(Instant::now())
130            .as_millis() as f64;
131        (millis / 1000f64).ceil() as u64
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[actix_web::test]
140    async fn test_seconds_until_reset() {
141        tokio::time::pause();
142        let output = SimpleOutput {
143            limit: 0,
144            remaining: 0,
145            reset: Instant::now() + Duration::from_secs(60),
146        };
147        tokio::time::advance(Duration::from_secs_f64(29.9)).await;
148        // Verify rounded upwards from 30.1
149        assert_eq!(output.seconds_until_reset(), 31);
150    }
151}
actix_extensible_rate_limit/backend/mod.rs

actix_extensible_rate_limit/backend/
mod.rs
