Skip to main content

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}