spacegate_kernel/extension/
original_ip_addr.rs

1use std::{net::IpAddr, str::FromStr};
2
3use crate::{extractor::OptionalExtract, Extract, SgRequestExt};
4
5#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
6/// Extract original ip address from request
7///
8/// # Panics
9/// ⚠ **WARNING** ⚠
10///
11/// If peer addr is not settled, it will panic when there's no original ip information from headers.
12pub struct OriginalIpAddr(pub IpAddr);
13
14impl OriginalIpAddr {
15    pub fn into_inner(self) -> IpAddr {
16        self.0
17    }
18}
19
20impl std::ops::Deref for OriginalIpAddr {
21    type Target = IpAddr;
22
23    fn deref(&self) -> &Self::Target {
24        &self.0
25    }
26}
27impl OptionalExtract for OriginalIpAddr {
28    fn extract(req: &hyper::Request<crate::SgBody>) -> Option<Self> {
29        if let Some(ip) = req.extensions().get::<OriginalIpAddr>().cloned() {
30            return Some(ip);
31        }
32        const X_FORWARDED_FOR: &str = "x-forwarded-for";
33        const X_REAL_IP: &str = "x-real-ip";
34        fn header_to_ip(header: &hyper::header::HeaderValue) -> Option<IpAddr> {
35            let s = header.to_str().ok()?;
36            let ip = IpAddr::from_str(s).ok()?;
37            Some(ip)
38        }
39        let ip = req
40            .headers()
41            .get(X_REAL_IP)
42            .and_then(header_to_ip)
43            .or_else(|| req.headers().get_all(X_FORWARDED_FOR).iter().next().and_then(header_to_ip))
44            .or_else(|| req.extensions().get::<crate::extension::PeerAddr>().map(|peer| peer.0.ip()))?;
45        Some(Self(ip))
46    }
47}
48impl Extract for OriginalIpAddr {
49    /// # Panics
50    /// if peer addr is not settled, it will panic when there's no original ip information from headers.
51    fn extract(req: &hyper::Request<crate::SgBody>) -> Self {
52        req.extract::<Option<OriginalIpAddr>>().expect("peer addr is not settled")
53    }
54}