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}