tower_rate_limiter/limiter/key_extractor.rs
1//! Request-facing client-key extraction interfaces.
2
3use std::{
4 fmt::{self, Display},
5 net::IpAddr,
6 sync::Arc,
7};
8
9use http::Request;
10
11use super::error::RateLimitError;
12
13/// Extract a client key synchronously from a request.
14///
15/// A Rate Limiter owns one extractor. Normalize client identity once at this request boundary and
16/// return the resulting displayable key; do not chain multiple extractors inside the middleware.
17pub trait KeyExtractor: Clone {
18 /// The type of the key.
19 type Key: Display;
20
21 /// Extract the key from the request.
22 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, RateLimitError>;
23}
24
25/// Extract an IP client key from only the socket peer.
26///
27/// Use this extractor for direct connections or whenever the transport peer itself should identify
28/// the caller. It never reads forwarding Headers. Choose one built-in IP extractor for the
29/// application's network topology; do not layer this extractor with another one.
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct IpKeyExtractor;
32
33impl IpKeyExtractor {
34 /// Construct the default extractor.
35 pub const fn new() -> Self {
36 Self
37 }
38}
39
40impl KeyExtractor for IpKeyExtractor {
41 type Key = IpAddr;
42
43 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, RateLimitError> {
44 http_extract::extract_socket_ip(request).ok_or_else(|| {
45 RateLimitError::Key(
46 String::from("socket_ip_unavailable"),
47 String::from("request extensions do not contain a socket ip address"),
48 )
49 })
50 }
51}
52
53/// Extract a client IP key from supported client-IP headers, falling back to the socket IP.
54///
55/// Use this extractor only when a platform or deployment boundary already guarantees that every
56/// accepted client-IP Header is trustworthy, such as an application reachable only through a proxy
57/// that removes or overwrites those Headers. This extractor does not validate the socket peer. If no
58/// client-IP Header is present, it falls back to an Axum `ConnectInfo<SocketAddr>` or a generic
59/// `SocketAddr` request extension. Choose one built-in IP extractor for the application's network
60/// topology; do not layer this extractor with another one.
61#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
62pub struct ClientIpKeyExtractor;
63
64impl ClientIpKeyExtractor {
65 /// Construct the header-aware client IP extractor.
66 pub const fn new() -> Self {
67 Self
68 }
69}
70
71impl KeyExtractor for ClientIpKeyExtractor {
72 type Key = IpAddr;
73
74 /// Extract the client IP from the request, falling back to the socket IP.
75 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, RateLimitError> {
76 http_extract::extract_proxy_client_ip(request)
77 .map_err(|error| RateLimitError::Key(String::from("invalid_client_ip"), error.to_string()))?
78 .ok_or_else(|| {
79 RateLimitError::Key(
80 String::from("client_ip_unavailable"),
81 String::from("request does not contain a client or socket IP address"),
82 )
83 })
84 }
85}
86
87/// Extract a client IP key only when the socket peer satisfies an application trust policy.
88///
89/// Use this extractor when trusted proxies and direct or untrusted peers may reach the same
90/// application. It requires the socket peer and validates that peer with the synchronous policy
91/// supplied to [`Self::new`] before reading any forwarding Header. Choose one built-in IP extractor
92/// for the application's network topology; do not layer this extractor with another one.
93///
94/// An untrusted peer always uses its socket IP and all forwarding Headers are ignored. A trusted
95/// peer uses the same Header order and strict parsing as [`ClientIpKeyExtractor`], falling back to
96/// the peer when no supported Header is present.
97///
98/// The policy establishes which transport peers may assert a client address; Header parsing does
99/// not authenticate the value. Applications must still ensure every trusted proxy removes or
100/// overwrites each supported client-IP Header.
101#[derive(Clone)]
102pub struct TrustedProxyClientIpKeyExtractor {
103 is_trusted_proxy: Arc<dyn Fn(IpAddr) -> bool + Send + Sync>,
104}
105
106impl fmt::Debug for TrustedProxyClientIpKeyExtractor {
107 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
108 formatter
109 .debug_struct("TrustedProxyClientIpKeyExtractor")
110 .finish_non_exhaustive()
111 }
112}
113
114impl TrustedProxyClientIpKeyExtractor {
115 /// Construct an extractor with an application-defined trusted-peer policy.
116 pub fn new<F>(is_trusted_proxy: F) -> Self
117 where
118 F: Fn(IpAddr) -> bool + Send + Sync + 'static,
119 {
120 Self {
121 is_trusted_proxy: Arc::new(is_trusted_proxy),
122 }
123 }
124}
125
126impl KeyExtractor for TrustedProxyClientIpKeyExtractor {
127 type Key = IpAddr;
128
129 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, RateLimitError> {
130 let peer = http_extract::extract_socket_ip(request).ok_or_else(|| {
131 RateLimitError::Key(
132 String::from("socket_ip_unavailable"),
133 String::from("request extensions do not contain a socket ip address"),
134 )
135 })?;
136
137 if !(self.is_trusted_proxy)(peer) {
138 return Ok(peer);
139 }
140
141 http_extract::extract_client_ip(request.headers())
142 .map(|client_ip| client_ip.unwrap_or(peer))
143 .map_err(|error| RateLimitError::Key(String::from("invalid_client_ip"), error.to_string()))
144 }
145}