rama_net/rate/
key.rs

1use core::fmt::Debug;
2use core::hash::Hash;
3use std::net::IpAddr;
4
5use rama_core::error::BoxError;
6use rama_core::extensions::ExtensionsRef;
7
8use crate::address::ip::ipnet::IpNet;
9
10/// A key to rate limit on: the bucket-map key of a
11/// [`KeyedRatePolicy`](super::KeyedRatePolicy).
12pub trait RateKey: Hash + Eq + Clone + Debug + Send + Sync + 'static {
13    /// Telemetry attributes identifying this key.
14    #[cfg(feature = "opentelemetry")]
15    fn attributes(
16        &self,
17    ) -> impl Iterator<Item = rama_core::telemetry::opentelemetry::KeyValue> + '_ {
18        core::iter::empty()
19    }
20}
21
22impl RateKey for IpAddr {
23    #[cfg(feature = "opentelemetry")]
24    fn attributes(
25        &self,
26    ) -> impl Iterator<Item = rama_core::telemetry::opentelemetry::KeyValue> + '_ {
27        use rama_core::telemetry::opentelemetry::{KeyValue, semantic_conventions};
28        core::iter::once(KeyValue::new(
29            semantic_conventions::attribute::CLIENT_ADDRESS,
30            self.to_string(),
31        ))
32    }
33}
34
35impl RateKey for String {}
36impl RateKey for u64 {}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41    use rama_core::extensions::Extensions;
42
43    fn ip(s: &str) -> IpAddr {
44        s.parse().unwrap()
45    }
46
47    #[test]
48    fn ipv4_mapped_ipv6_collapses_to_ipv4() {
49        let key = ClientIpRateKey::new();
50        // the dual-stack peer form and the plain v4 form must key identically
51        assert_eq!(key.key_for(ip("::ffff:203.0.113.5")), ip("203.0.113.5"));
52        assert_eq!(
53            key.key_for(ip("::ffff:203.0.113.5")),
54            key.key_for(ip("203.0.113.5"))
55        );
56    }
57
58    #[test]
59    fn ipv6_aggregates_to_prefix() {
60        let key = ClientIpRateKey::new(); // default /64
61        // two addresses in the same /64 share one bucket ...
62        assert_eq!(
63            key.key_for(ip("2001:db8:1:2::1")),
64            key.key_for(ip("2001:db8:1:2:ffff:ffff:ffff:ffff"))
65        );
66        assert_eq!(key.key_for(ip("2001:db8:1:2::1")), ip("2001:db8:1:2::"));
67        // ... a different /64 does not
68        assert_ne!(
69            key.key_for(ip("2001:db8:1:2::1")),
70            key.key_for(ip("2001:db8:1:3::1"))
71        );
72    }
73
74    #[test]
75    fn ipv6_prefix_128_keys_exact_address() {
76        let key = ClientIpRateKey::new().with_ipv6_prefix(128);
77        assert_eq!(key.key_for(ip("2001:db8:1:2::1")), ip("2001:db8:1:2::1"));
78        assert_ne!(
79            key.key_for(ip("2001:db8:1:2::1")),
80            key.key_for(ip("2001:db8:1:2::2"))
81        );
82    }
83
84    #[test]
85    fn ipv4_is_never_aggregated() {
86        let key = ClientIpRateKey::new().with_ipv6_prefix(1);
87        assert_eq!(key.key_for(ip("203.0.113.5")), ip("203.0.113.5"));
88    }
89
90    #[test]
91    fn ipv6_prefix_is_clamped() {
92        assert_eq!(ClientIpRateKey::new().with_ipv6_prefix(0).ipv6_prefix, 1);
93        assert_eq!(
94            ClientIpRateKey::new().with_ipv6_prefix(200).ipv6_prefix,
95            128
96        );
97    }
98
99    #[test]
100    fn extractor_reads_and_canonicalises_client_ip() {
101        use crate::address::SocketAddress;
102        use crate::stream::SocketInfo;
103
104        let ext = Extensions::new();
105        ext.insert(SocketInfo::new(
106            None,
107            SocketAddress::new(ip("::ffff:203.0.113.5"), 0),
108        ));
109        let got = ClientIpRateKey::new().rate_key(&ext).unwrap();
110        assert_eq!(got, Some(ip("203.0.113.5")));
111    }
112}
113
114/// Derives the [`RateKey`] of an input for a
115/// [`KeyedRatePolicy`](super::KeyedRatePolicy).
116///
117/// `Ok(None)` means the key cannot be derived for this input (e.g. no
118/// client IP is known); how that is handled is up to the policy
119/// ([`KeyedRatePolicy::with_missing_key_allowed`](super::KeyedRatePolicy)).
120///
121/// Any `Fn(&Input) -> Result<Option<K>, BoxError>` is an extractor.
122pub trait InputToRateKey<Input>: Send + Sync + 'static {
123    /// The key type produced by this extractor.
124    type Key: RateKey;
125
126    /// Derive the rate key from the given input.
127    fn rate_key(&self, input: &Input) -> Result<Option<Self::Key>, BoxError>;
128}
129
130impl<Input, K, F> InputToRateKey<Input> for F
131where
132    F: Fn(&Input) -> Result<Option<K>, BoxError> + Send + Sync + 'static,
133    K: RateKey,
134{
135    type Key = K;
136
137    fn rate_key(&self, input: &Input) -> Result<Option<Self::Key>, BoxError> {
138        (self)(input)
139    }
140}
141
142/// The usual IPv6 end-site allocation, used as the default aggregation
143/// prefix so a client cannot dodge its bucket by rotating within its /64.
144const DEFAULT_IPV6_PREFIX: u8 = 64;
145
146/// An [`InputToRateKey`] extractor keying on the client IP address,
147/// resolved via [`client_ip`](crate::client_ip::client_ip):
148/// [`Forwarded`](crate::forwarded::Forwarded) information (populated by
149/// e.g. forwarded-header or PROXY-protocol layers) wins over the
150/// transport peer address ([`SocketInfo`](crate::stream::SocketInfo)).
151/// Only populate `Forwarded` from a trusted proxy boundary: accepting a
152/// client-supplied forwarding header lets that client choose and rotate its
153/// own rate key.
154///
155/// The resolved address is canonicalised before keying: IPv4-mapped IPv6
156/// peers (`::ffff:a.b.c.d`) collapse to their IPv4 form, and IPv6 clients
157/// are aggregated to [`with_ipv6_prefix`](Self::with_ipv6_prefix) (default
158/// `/64`). Without this a single client keys to `2^64` distinct buckets and
159/// per-client limiting is a no-op against exactly the clients most able to
160/// abuse it.
161///
162/// The aggregation prefix and the policy's `max_keys` must be sized together:
163/// one routed `/48` contains 65 536 `/64` keys, equal to the default
164/// [`KeyedRatePolicy`](super::KeyedRatePolicy) capacity. Aggregate more broadly
165/// when one client or tenant may legitimately control many `/64` networks.
166#[derive(Debug, Clone, Copy)]
167#[non_exhaustive]
168pub struct ClientIpRateKey {
169    ipv6_prefix: u8,
170}
171
172impl ClientIpRateKey {
173    /// Create a new [`ClientIpRateKey`], aggregating IPv6 clients to `/64`.
174    #[must_use]
175    pub const fn new() -> Self {
176        Self {
177            ipv6_prefix: DEFAULT_IPV6_PREFIX,
178        }
179    }
180
181    rama_utils::macros::generate_set_and_with! {
182        /// Aggregate IPv6 client addresses to this prefix length (clamped
183        /// to `1..=128`) before keying; `128` keys on the exact address.
184        /// IPv4 clients are always keyed on the exact address. Choose this
185        /// together with [`KeyedRatePolicy::set_max_keys`](super::KeyedRatePolicy::set_max_keys):
186        /// a broader prefix consumes fewer cache entries but groups more
187        /// clients into one budget.
188        pub fn ipv6_prefix(mut self, prefix: u8) -> Self {
189            self.ipv6_prefix = prefix.clamp(1, 128);
190            self
191        }
192    }
193
194    /// Canonicalise and aggregate a resolved client IP into its bucket key.
195    fn key_for(self, ip: IpAddr) -> IpAddr {
196        match ip.to_canonical() {
197            IpAddr::V6(v6) if self.ipv6_prefix < 128 => {
198                IpNet::new(IpAddr::V6(v6), self.ipv6_prefix)
199                    .map(|net| net.trunc().addr())
200                    .unwrap_or(IpAddr::V6(v6))
201            }
202            canonical => canonical,
203        }
204    }
205}
206
207impl Default for ClientIpRateKey {
208    fn default() -> Self {
209        Self::new()
210    }
211}
212
213impl<Input> InputToRateKey<Input> for ClientIpRateKey
214where
215    Input: ExtensionsRef + Send + Sync + 'static,
216{
217    type Key = IpAddr;
218
219    fn rate_key(&self, input: &Input) -> Result<Option<Self::Key>, BoxError> {
220        Ok(crate::client_ip::client_ip(input).map(|ip| self.key_for(ip)))
221    }
222}
rama_net/rate/key.rs

rama_net/rate/
key.rs
