diff --git a/crates/didbot-serve/src/rate_limit.rs b/crates/didbot-serve/src/rate_limit.rs index 9146e9d1..e270a49d 100644 --- a/crates/didbot-serve/src/rate_limit.rs +++ b/crates/didbot-serve/src/rate_limit.rs @@ -50,7 +50,7 @@ use std::collections::HashMap; use std::convert::Infallible; -use std::net::{IpAddr, SocketAddr}; +use std::net::{IpAddr, Ipv6Addr, SocketAddr}; use std::sync::Mutex; use std::time::{Duration, Instant}; @@ -233,6 +233,10 @@ impl Default for RateLimiter { /// `trust_forwarded_headers` being `true` with neither header present both /// answer the same `"unknown"`, the shared bucket every caller falls back to /// when this cannot tell them apart. +/// +/// An IPv6 address is keyed on its /64, because one host routinely holds a +/// whole /64. An IPv4-mapped IPv6 address is keyed as its IPv4 address. A +/// forwarded value that is not an IP address is keyed as it is. pub fn caller_key( headers: &HeaderMap, peer: Option, @@ -240,13 +244,27 @@ pub fn caller_key( ) -> String { if trust_forwarded_headers { if let Some(key) = forwarded_caller(headers) { - return key; + return match key.parse::() { + Ok(ip) => address_key(ip), + Err(_) => key, + }; } } - peer.map(|ip| ip.to_string()) + peer.map(address_key) .unwrap_or_else(|| "unknown".to_owned()) } +/// The budget key for one address: see [`caller_key`]. +fn address_key(ip: IpAddr) -> String { + match ip.to_canonical() { + IpAddr::V4(v4) => v4.to_string(), + IpAddr::V6(v6) => { + let prefix = Ipv6Addr::from_bits(v6.to_bits() & !(u128::MAX >> 64)); + format!("{prefix}/64") + } + } +} + /// The caller a reverse proxy's `X-Forwarded-For` (or `X-Real-IP`) header /// names, for a deployment that has opted into trusting it. `None` when /// neither header is present or readable. @@ -368,4 +386,24 @@ mod tests { let headers = HeaderMap::new(); assert_eq!(caller_key(&headers, None, true), "unknown"); } + + /// One host can pick any address in its /64, so a per-address key would + /// hand it an unlimited number of budgets. + #[test] + fn caller_key_shares_one_budget_across_an_ipv6_64() { + let headers = HeaderMap::new(); + let key = |ip: &str| caller_key(&headers, Some(ip.parse().unwrap()), false); + assert_eq!( + key("2001:db8:1:2::1"), + key("2001:db8:1:2:ffff:ffff:ffff:ffff") + ); + assert_ne!(key("2001:db8:1:2::1"), key("2001:db8:1:3::1")); + assert_eq!(key("::ffff:203.0.113.9"), "203.0.113.9"); + + let mut forwarded = HeaderMap::new(); + forwarded.insert("x-forwarded-for", "2001:db8:1:2::9".parse().unwrap()); + assert_eq!(caller_key(&forwarded, None, true), key("2001:db8:1:2::1")); + forwarded.insert("x-forwarded-for", "not-an-address".parse().unwrap()); + assert_eq!(caller_key(&forwarded, None, true), "not-an-address"); + } }