1use std::net::IpAddr;
8
9use rainier_http::{ClientIp, Request, Response};
10
11use crate::pipeline::{Middleware, Next};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct Cidr {
16 address: IpAddr,
17 prefix: u8,
18}
19
20impl Cidr {
21 pub fn parse(value: &str) -> Option<Self> {
25 let value = value.trim();
26 let (address, prefix) = match value.split_once('/') {
27 Some((address, prefix)) => (address, Some(prefix)),
28 None => (value, None),
29 };
30
31 let address: IpAddr = address.parse().ok()?;
32 let full = if address.is_ipv4() { 32 } else { 128 };
33
34 let prefix = match prefix {
35 Some(prefix) => {
36 let prefix: u8 = prefix.parse().ok()?;
37 if prefix > full {
38 return None;
39 }
40 prefix
41 }
42 None => full,
43 };
44
45 Some(Self { address, prefix })
46 }
47
48 pub fn contains(&self, candidate: IpAddr) -> bool {
50 match (self.address, candidate) {
51 (IpAddr::V4(network), IpAddr::V4(candidate)) => {
52 matches(&network.octets(), &candidate.octets(), self.prefix)
53 }
54 (IpAddr::V6(network), IpAddr::V6(candidate)) => {
55 matches(&network.octets(), &candidate.octets(), self.prefix)
56 }
57 _ => false,
61 }
62 }
63}
64
65fn matches(network: &[u8], candidate: &[u8], prefix: u8) -> bool {
67 let whole = (prefix / 8) as usize;
68 if network[..whole] != candidate[..whole] {
69 return false;
70 }
71
72 let remaining = prefix % 8;
73 if remaining == 0 {
74 return true;
75 }
76
77 let mask = 0xffu8 << (8 - remaining);
78 network[whole] & mask == candidate[whole] & mask
79}
80
81#[derive(Debug, Clone, Default)]
83pub enum Trusted {
84 #[default]
86 None,
87 All,
94 These(Vec<Cidr>),
96}
97
98impl Trusted {
99 pub fn accepts(&self, peer: IpAddr) -> bool {
101 match self {
102 Trusted::None => false,
103 Trusted::All => true,
104 Trusted::These(networks) => networks.iter().any(|network| network.contains(peer)),
105 }
106 }
107
108 pub fn is_known_hop(&self, address: IpAddr) -> bool {
116 match self {
117 Trusted::None | Trusted::All => false,
118 Trusted::These(networks) => networks.iter().any(|network| network.contains(address)),
119 }
120 }
121}
122
123#[derive(Debug, Clone, Default)]
136pub struct TrustProxies {
137 trusted: Trusted,
138 header: String,
139}
140
141impl TrustProxies {
142 pub fn new() -> Self {
144 Self { trusted: Trusted::None, header: "x-forwarded-for".to_string() }
145 }
146
147 pub fn all() -> Self {
149 Self { trusted: Trusted::All, ..Self::new() }
150 }
151
152 pub fn these(networks: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
154 let parsed = networks
155 .into_iter()
156 .filter_map(|network| {
157 let network = network.as_ref();
158 let parsed = Cidr::parse(network);
159 if parsed.is_none() {
160 tracing::warn!(%network, "ignoring an unparseable trusted proxy");
161 }
162 parsed
163 })
164 .collect();
165
166 Self { trusted: Trusted::These(parsed), ..Self::new() }
167 }
168
169 pub fn header(mut self, header: impl Into<String>) -> Self {
171 self.header = header.into().to_ascii_lowercase();
172 self
173 }
174
175 pub fn trusts_anything(&self) -> bool {
177 !matches!(self.trusted, Trusted::None)
178 }
179
180 pub fn resolve(&self, request: &Request, peer: IpAddr) -> IpAddr {
191 if !self.trusted.accepts(peer) {
192 return peer;
193 }
194
195 let Some(header) = request.header(&self.header) else {
196 return peer;
197 };
198
199 header
200 .split(',')
201 .rev()
202 .filter_map(|entry| entry.trim().parse::<IpAddr>().ok())
203 .find(|address| !self.trusted.is_known_hop(*address))
204 .unwrap_or(peer)
205 }
206}
207
208#[async_trait::async_trait]
209impl Middleware for TrustProxies {
210 async fn handle(&self, mut request: Request, next: Next) -> Response {
211 if let Some(peer) = request.ip() {
212 let client = self.resolve(&request, peer);
213 if client != peer {
214 request.extensions_mut().insert(ClientIp(client));
215 }
216 }
217 next.run(request).await
218 }
219
220 fn name(&self) -> &'static str {
221 "TrustProxies"
222 }
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228
229 fn ip(value: &str) -> IpAddr {
230 value.parse().unwrap()
231 }
232
233 fn request(chain: &str) -> Request {
234 Request::builder().header("x-forwarded-for", chain).build()
235 }
236
237 #[test]
238 fn a_bare_address_is_a_single_host() {
239 let cidr = Cidr::parse("192.168.1.7").unwrap();
240
241 assert!(cidr.contains(ip("192.168.1.7")));
242 assert!(!cidr.contains(ip("192.168.1.8")));
243 }
244
245 #[test]
246 fn a_prefix_matches_its_network() {
247 let cidr = Cidr::parse("10.0.0.0/8").unwrap();
248
249 assert!(cidr.contains(ip("10.0.0.1")));
250 assert!(cidr.contains(ip("10.255.255.255")));
251 assert!(!cidr.contains(ip("11.0.0.1")));
252 }
253
254 #[test]
255 fn a_prefix_that_is_not_a_whole_byte_still_works() {
256 let cidr = Cidr::parse("172.16.0.0/12").unwrap();
257
258 assert!(cidr.contains(ip("172.16.0.1")));
259 assert!(cidr.contains(ip("172.31.255.255")));
260 assert!(!cidr.contains(ip("172.32.0.1")), "just outside the /12");
261 assert!(!cidr.contains(ip("172.15.255.255")));
262 }
263
264 #[test]
265 fn ipv6_works_too() {
266 assert!(Cidr::parse("::1").unwrap().contains(ip("::1")));
267 assert!(Cidr::parse("fd00::/8").unwrap().contains(ip("fd12:3456::1")));
268 assert!(!Cidr::parse("fd00::/8").unwrap().contains(ip("fe80::1")));
269 }
270
271 #[test]
272 fn families_do_not_cross() {
273 assert!(!Cidr::parse("10.0.0.0/8").unwrap().contains(ip("::ffff:10.0.0.1")));
274 assert!(!Cidr::parse("::/0").unwrap().contains(ip("10.0.0.1")));
275 }
276
277 #[test]
278 fn rubbish_does_not_parse() {
279 for bad in ["", "nonsense", "10.0.0.0/33", "10.0.0.0/-1", "::1/129", "10.0.0.0/x"] {
280 assert!(Cidr::parse(bad).is_none(), "{bad}");
281 }
282 }
283
284 #[test]
285 fn trusting_nobody_ignores_the_header() {
286 let middleware = TrustProxies::new();
287
288 assert!(!middleware.trusts_anything());
289 assert_eq!(middleware.resolve(&request("1.2.3.4"), ip("10.0.0.9")), ip("10.0.0.9"));
290 }
291
292 #[test]
293 fn an_untrusted_peer_is_believed_over_its_own_header() {
294 let middleware = TrustProxies::these(["10.0.0.0/8"]);
296
297 assert_eq!(middleware.resolve(&request("1.2.3.4"), ip("203.0.113.9")), ip("203.0.113.9"));
298 }
299
300 #[test]
301 fn a_trusted_peer_yields_the_client() {
302 let middleware = TrustProxies::these(["10.0.0.0/8"]);
303
304 assert_eq!(middleware.resolve(&request("203.0.113.9"), ip("10.0.0.1")), ip("203.0.113.9"));
305 }
306
307 #[test]
308 fn trusted_hops_are_walked_past_from_the_right() {
309 let middleware = TrustProxies::these(["10.0.0.0/8"]);
311 let chain = request("203.0.113.9, 10.0.0.5, 10.0.0.6");
312
313 assert_eq!(middleware.resolve(&chain, ip("10.0.0.1")), ip("203.0.113.9"));
314 }
315
316 #[test]
317 fn a_client_cannot_forge_an_address_by_prepending_one() {
318 let middleware = TrustProxies::these(["10.0.0.0/8"]);
322 let chain = request("1.2.3.4, 203.0.113.9, 10.0.0.5");
323
324 assert_eq!(
325 middleware.resolve(&chain, ip("10.0.0.1")),
326 ip("203.0.113.9"),
327 "the rightmost untrusted entry is the only one our own proxy vouched for"
328 );
329 }
330
331 #[test]
332 fn a_chain_of_only_trusted_hops_falls_back_to_the_peer() {
333 let middleware = TrustProxies::these(["10.0.0.0/8"]);
334 let chain = request("10.0.0.5, 10.0.0.6");
335
336 assert_eq!(middleware.resolve(&chain, ip("10.0.0.1")), ip("10.0.0.1"));
337 }
338
339 #[test]
340 fn rubbish_entries_are_skipped() {
341 let middleware = TrustProxies::these(["10.0.0.0/8"]);
342 let chain = request("203.0.113.9, unknown, 10.0.0.5");
343
344 assert_eq!(middleware.resolve(&chain, ip("10.0.0.1")), ip("203.0.113.9"));
345 }
346
347 #[test]
348 fn no_header_leaves_the_peer_alone() {
349 let middleware = TrustProxies::these(["10.0.0.0/8"]);
350
351 assert_eq!(middleware.resolve(&Request::builder().build(), ip("10.0.0.1")), ip("10.0.0.1"));
352 }
353
354 #[test]
355 fn trusting_everything_believes_the_chain() {
356 let middleware = TrustProxies::all();
357
358 assert!(middleware.trusts_anything());
359 assert_eq!(middleware.resolve(&request("203.0.113.9"), ip("10.0.0.1")), ip("203.0.113.9"));
360 }
361
362 #[test]
363 fn trusting_everything_takes_the_rightmost_entry() {
364 let middleware = TrustProxies::all();
368 let chain = request("1.2.3.4, 203.0.113.9");
369
370 assert_eq!(middleware.resolve(&chain, ip("10.0.0.1")), ip("203.0.113.9"));
371 }
372
373 #[test]
374 fn an_alternative_header_can_be_read() {
375 let middleware = TrustProxies::all().header("CF-Connecting-IP");
376 let request = Request::builder().header("cf-connecting-ip", "203.0.113.9").build();
377
378 assert_eq!(middleware.resolve(&request, ip("10.0.0.1")), ip("203.0.113.9"));
379 }
380
381 #[tokio::test]
382 async fn the_middleware_rewrites_the_recorded_address() {
383 use crate::pipeline::Pipeline;
384
385 let request = Request::builder()
386 .header("x-forwarded-for", "203.0.113.9, 10.0.0.5")
387 .build()
388 .with_extension(ClientIp(ip("10.0.0.1")));
389
390 let response = Pipeline::new()
391 .through(TrustProxies::these(["10.0.0.0/8"]))
392 .then(|request: Request| async move {
393 Response::text(request.ip().map(|ip| ip.to_string()).unwrap_or_default())
394 })
395 .run(request)
396 .await;
397
398 let body = response.into_http().into_body().collect().await.unwrap();
399 assert_eq!(body, "203.0.113.9");
400 }
401
402 #[tokio::test]
403 async fn a_request_with_no_recorded_address_passes_through() {
404 use crate::pipeline::Pipeline;
405
406 let response = Pipeline::new()
407 .through(TrustProxies::all())
408 .then(|request: Request| async move {
409 Response::text(format!("{}", request.ip().is_none()))
410 })
411 .run(request("203.0.113.9"))
412 .await;
413
414 let body = response.into_http().into_body().collect().await.unwrap();
415 assert_eq!(body, "true");
416 }
417}