Skip to main content

tork_core/throttle/
key.rs

1//! Rate-limit keys: how a request is identified for counting.
2
3use std::future::Future;
4
5use crate::error::Result;
6use crate::extract::RequestContext;
7
8/// Produces the key a request is rate-limited by (the "tracker").
9///
10/// The default is the client IP ([`ByIp`]). Implement this on a unit type to key
11/// by something else — a user id, an API key, a tenant — reusing dependency
12/// injection inside, since it has the full [`RequestContext`]:
13///
14/// ```ignore
15/// struct ByUser;
16/// impl ThrottleKey for ByUser {
17///     async fn throttle_key(ctx: &RequestContext) -> tork::Result<String> {
18///         Ok(CurrentUser::from_request(ctx).await?.id.to_string())
19///     }
20/// }
21/// ```
22pub trait ThrottleKey: Send + Sync + 'static {
23    /// Returns the key for `ctx`.
24    fn throttle_key(ctx: &RequestContext) -> impl Future<Output = Result<String>> + Send;
25}
26
27/// The default key strategy: the client IP address.
28///
29/// Falls back to `"unknown"` when no peer address is available (for example an
30/// in-process test request), so such requests share one counter rather than going
31/// unlimited.
32pub struct ByIp;
33
34impl ThrottleKey for ByIp {
35    fn throttle_key(ctx: &RequestContext) -> impl Future<Output = Result<String>> + Send {
36        let key = ctx
37            .peer_addr()
38            .map(|addr| addr.ip().to_string())
39            .unwrap_or_else(|| "unknown".to_string());
40        async move { Ok(key) }
41    }
42}