tower_rate_limiter/limiter/
builder.rs

1//! Builder and immutable configuration for the rate-limit layer.
2
3use std::{fmt, sync::Arc, time::Duration};
4
5use http::Request;
6
7use super::{
8    error::ConfigError,
9    layer::RateLimitLayer,
10    limit::LimitProvider,
11    response::{DefaultResponseFactory, RateLimitFields},
12    store::{Store, StoreFailureMode},
13};
14
15/// The minimum window duration allowed.
16const MINIMUM_WINDOW: Duration = Duration::from_millis(1);
17
18/// A callback to encode the scoped key before passing it to the [`crate::Store`].
19pub(crate) type KeyEncoder = Box<dyn Fn(&str) -> String + Send + Sync>;
20
21/// A callback that decides whether a request bypasses rate limiting.
22pub(crate) type SkipPredicate = Box<dyn Fn(&Request<()>) -> bool + Send + Sync>;
23
24/// Check if the request should be skipped based on the skip predicate.
25pub(crate) fn check_skip_predicate<B>(predicate: Option<&SkipPredicate>, request: Request<B>) -> (bool, Request<B>) {
26    let Some(predicate) = predicate else {
27        return (false, request);
28    };
29
30    // This is a hack to get the request head. maybe there's a better way to do this.
31    let (parts, body) = request.into_parts();
32    let request_head = Request::from_parts(parts, ());
33    let should_skip = predicate(&request_head);
34    let (parts, ()) = request_head.into_parts();
35
36    (should_skip, Request::from_parts(parts, body))
37}
38
39/// Builder for a rate-limit layer with compile-time store/resolver/factory types.
40#[derive(Debug)]
41#[must_use]
42pub struct RateLimitBuilder<K, S = (), P = u64, F = DefaultResponseFactory> {
43    key_extractor: K,
44    store: S,
45    limit_provider: P,
46    response_factory: F,
47    config: RateLimitConfig,
48}
49
50impl<K> RateLimitBuilder<K> {
51    pub(crate) fn new(key_extractor: K) -> Self {
52        Self {
53            key_extractor,
54            store: (),
55            limit_provider: 1,
56            response_factory: DefaultResponseFactory,
57            config: RateLimitConfig {
58                policy_name: String::from("default-policy"),
59                window: Duration::from_secs(60),
60                key_encoder: None,
61                skip_predicate: None,
62                store_failure_mode: StoreFailureMode::default(),
63                #[cfg(feature = "tracing")]
64                store_failure_tracing_level: tracing::Level::WARN,
65                rate_limit_fields: RateLimitFields::default(),
66            },
67        }
68    }
69}
70
71impl<K, S, P, F> RateLimitBuilder<K, S, P, F> {
72    /// Inject the rate-limit store and update the builder's store type state.
73    pub fn with_store<S2>(self, store: S2) -> RateLimitBuilder<K, S2, P, F> {
74        let Self {
75            key_extractor,
76            limit_provider,
77            response_factory,
78            config,
79            ..
80        } = self;
81        RateLimitBuilder {
82            key_extractor,
83            store,
84            limit_provider,
85            response_factory,
86            config,
87        }
88    }
89
90    /// Set a fixed quota limit and update the provider type state.
91    pub fn limit(self, limit: u64) -> RateLimitBuilder<K, S, u64, F> {
92        let Self {
93            key_extractor,
94            store,
95            response_factory,
96            config,
97            ..
98        } = self;
99        RateLimitBuilder {
100            key_extractor,
101            store,
102            limit_provider: limit,
103            response_factory,
104            config,
105        }
106    }
107
108    /// Replace the fixed provider with a custom asynchronous limit provider.
109    pub fn limit_provider<P2>(self, limit_provider: P2) -> RateLimitBuilder<K, S, P2, F> {
110        let Self {
111            key_extractor,
112            store,
113            response_factory,
114            config,
115            ..
116        } = self;
117        RateLimitBuilder {
118            key_extractor,
119            store,
120            limit_provider,
121            response_factory,
122            config,
123        }
124    }
125
126    /// Replace the response factory and update its type state.
127    pub fn response_factory<F2>(self, response_factory: F2) -> RateLimitBuilder<K, S, P, F2> {
128        let Self {
129            key_extractor,
130            store,
131            limit_provider,
132            config,
133            ..
134        } = self;
135        RateLimitBuilder {
136            key_extractor,
137            store,
138            limit_provider,
139            response_factory,
140            config,
141        }
142    }
143
144    /// Set the fixed-window duration.
145    pub fn window(mut self, window: Duration) -> Self {
146        self.config.window = window;
147        self
148    }
149
150    /// Set the stable policy identifier used in the scoped key and response metadata.
151    pub fn policy_name(mut self, policy_name: impl Into<String>) -> Self {
152        self.config.policy_name = policy_name.into();
153        self
154    }
155
156    /// Encode the scoped key before passing it to the [`crate::Store`].
157    ///
158    /// The callback runs in the middleware future's polling path. It must be deterministic,
159    /// non-blocking, free of I/O, collision-resistant for the caller's key space, and
160    /// non-panicking. Without this method, the complete scoped key is passed to the Store
161    /// unchanged.
162    pub fn with_key_encoder<E>(mut self, encoder: E) -> Self
163    where
164        E: Fn(&str) -> String + Send + Sync + 'static,
165    {
166        self.config.key_encoder = Some(Box::new(encoder));
167        self
168    }
169
170    /// Bypass rate limiting when `predicate` returns `true` for the request.
171    ///
172    /// The predicate receives the request head and extensions with a unit body. It runs
173    /// synchronously before client-key extraction and must be non-blocking, free of I/O, and
174    /// non-panicking. A bypassed request calls the inner service without resolving a limit,
175    /// charging the Store, or adding rate-limit context or response fields.
176    pub fn skip<Predicate>(mut self, predicate: Predicate) -> Self
177    where
178        Predicate: Fn(&Request<()>) -> bool + Send + Sync + 'static,
179    {
180        self.config.skip_predicate = Some(Box::new(predicate));
181        self
182    }
183
184    /// Select the mode to use when the Store fails.
185    pub fn store_failure_mode(mut self, mode: StoreFailureMode) -> Self {
186        self.config.store_failure_mode = mode;
187        self
188    }
189
190    /// Select the tracing level used for Store failure events.
191    ///
192    /// This configuration is available with the `tracing` Cargo feature. The default is
193    /// [`tracing::Level::WARN`]. It affects both [`StoreFailureMode::Allow`] and
194    /// [`StoreFailureMode::Reject`] events.
195    #[cfg(feature = "tracing")]
196    pub fn store_failure_tracing_level(mut self, level: tracing::Level) -> Self {
197        self.config.store_failure_tracing_level = level;
198        self
199    }
200
201    /// Select the Rate Limit Fields revision emitted in responses.
202    pub fn rate_limit_fields(mut self, fields: RateLimitFields) -> Self {
203        self.config.rate_limit_fields = fields;
204        self
205    }
206
207    fn validate(&self) -> Result<(), ConfigError> {
208        if self.config.window < MINIMUM_WINDOW {
209            return Err(ConfigError::WindowTooShort(self.config.window, MINIMUM_WINDOW));
210        }
211        if self.config.policy_name.is_empty() {
212            return Err(ConfigError::EmptyPolicyName);
213        }
214        Ok(())
215    }
216}
217
218impl<K> RateLimitLayer<K, (), u64, DefaultResponseFactory> {
219    /// Start a typed rate-limit layer builder.
220    pub fn builder(key_extractor: K) -> RateLimitBuilder<K> {
221        RateLimitBuilder::new(key_extractor)
222    }
223}
224
225impl<K, S, P, F> RateLimitBuilder<K, S, P, F>
226where
227    S: Store,
228    P: LimitProvider,
229{
230    /// Validate the builder and produce a configured layer.
231    pub fn build(self) -> Result<RateLimitLayer<K, S, P, F>, ConfigError> {
232        self.validate()?;
233        Ok(RateLimitLayer {
234            key_extractor: self.key_extractor,
235            store: self.store,
236            limit_provider: self.limit_provider,
237            response_factory: self.response_factory,
238            config: Arc::new(self.config),
239        })
240    }
241}
242
243/// Immutable configuration shared by every service produced from a layer.
244pub(crate) struct RateLimitConfig {
245    /// The stable policy identifier used in the scoped key and response metadata.
246    pub(crate) policy_name: String,
247    /// The fixed-window duration.
248    pub(crate) window: Duration,
249    /// Encode the scoped key before passing it to the [`crate::Store`].
250    pub(crate) key_encoder: Option<KeyEncoder>,
251    /// Decide whether a request bypasses rate limiting.
252    pub(crate) skip_predicate: Option<SkipPredicate>,
253    /// Select the mode to use when the Store fails.
254    pub(crate) store_failure_mode: StoreFailureMode,
255    /// Select the tracing level used when the Store fails.
256    #[cfg(feature = "tracing")]
257    pub(crate) store_failure_tracing_level: tracing::Level,
258    /// Select the [`RateLimitFields`] revision emitted in responses.
259    pub(crate) rate_limit_fields: RateLimitFields,
260}
261
262impl fmt::Debug for RateLimitConfig {
263    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
264        let mut debug = f.debug_struct("RateLimitConfig");
265        debug
266            .field("policy_name", &self.policy_name)
267            .field("window", &self.window)
268            .field("has_key_encoder", &self.key_encoder.is_some())
269            .field("has_skip_predicate", &self.skip_predicate.is_some())
270            .field("store_failure_mode", &self.store_failure_mode);
271        #[cfg(feature = "tracing")]
272        debug.field("store_failure_tracing_level", &self.store_failure_tracing_level);
273        debug.field("rate_limit_fields", &self.rate_limit_fields).finish()
274    }
275}
tower_rate_limiter/limiter/builder.rs

tower_rate_limiter/limiter/
builder.rs
