1use super::IpAddressConfig;
2use std::{collections::HashMap, net::IpAddr};
3
4impl IpAddressConfig {
5 #[must_use]
11 pub fn resolve_ip(&self, headers: &HashMap<String, String>) -> Option<String> {
12 if self.disable_ip_tracking {
13 return None;
14 }
15 let proxies = self
16 .trusted_proxies
17 .iter()
18 .filter_map(|value| Network::parse(value))
19 .collect::<Vec<_>>();
20 self.headers
21 .iter()
22 .find_map(|name| {
23 headers
24 .iter()
25 .find(|(key, _)| key.eq_ignore_ascii_case(name))
26 .and_then(|(_, value)| self.resolve_header(value, &proxies))
27 })
28 .or_else(|| self.localhost_fallback.then(|| "127.0.0.1".to_owned()))
29 }
30
31 fn resolve_header(&self, value: &str, proxies: &[Network]) -> Option<String> {
32 let hops = value
33 .split(',')
34 .map(|hop| {
35 hop.trim_matches(|character: char| {
36 (character.is_whitespace() && character != '\u{85}') || character == '\u{feff}'
37 })
38 })
39 .filter(|hop| !hop.is_empty())
40 .collect::<Vec<_>>();
41 if proxies.is_empty() {
42 return if hops.len() == 1 {
43 Address::parse(hops.first()?).map(|ip| ip.normalize(self.ipv6_subnet))
44 } else {
45 None
46 };
47 }
48 for hop in hops.iter().rev() {
49 let ip = Address::parse(hop)?;
50 if !proxies.iter().any(|proxy| proxy.contains(&ip.bytes)) {
51 return Some(ip.normalize(self.ipv6_subnet));
52 }
53 }
54 None
55 }
56}
57
58struct Address {
59 bytes: Vec<u8>,
60 ipv6_groups: Option<Vec<String>>,
61}
62
63impl Address {
64 fn parse(value: &str) -> Option<Self> {
65 match value.parse::<IpAddr>().ok()? {
66 IpAddr::V4(ip) => Some(Self {
67 bytes: ip.octets().to_vec(),
68 ipv6_groups: None,
69 }),
70 IpAddr::V6(ip) => {
71 if let Some(mapped) = ip.to_ipv4_mapped().filter(|_| {
75 value.contains('.') || value.split(':').any(|group| group == "ffff")
76 }) {
77 return Some(Self {
78 bytes: mapped.octets().to_vec(),
79 ipv6_groups: None,
80 });
81 }
82 let groups = expand_ipv6(value);
83 let bytes = groups
84 .iter()
85 .flat_map(|group| segment(group).to_be_bytes())
86 .collect();
87 Some(Self {
88 bytes,
89 ipv6_groups: Some(groups),
90 })
91 }
92 }
93 }
94
95 fn normalize(self, prefix: f64) -> String {
96 let Some(groups) = self.ipv6_groups else {
97 return self
98 .bytes
99 .iter()
100 .map(u8::to_string)
101 .collect::<Vec<_>>()
102 .join(".");
103 };
104 if prefix.is_nan() || prefix >= 128.0 {
105 return groups.join(":").to_lowercase();
106 }
107 let mut bits = prefix.floor().max(0.0);
108 groups
109 .into_iter()
110 .map(|group| {
111 let retained = (0..16_u16)
112 .filter(|bit| f64::from(*bit) < bits)
113 .fold(0_u16, |mask, bit| mask | (1 << (15 - bit)));
114 bits -= 16.0;
115 if retained == u16::MAX {
116 group.to_lowercase()
117 } else {
118 format!("{:04x}", segment(&group) & retained)
119 }
120 })
121 .collect::<Vec<_>>()
122 .join(":")
123 }
124}
125
126fn expand_ipv6(value: &str) -> Vec<String> {
127 let raw = if let Some((left, right)) = value.split_once("::") {
128 let mut left = left
129 .split(':')
130 .filter(|group| !group.is_empty())
131 .collect::<Vec<_>>();
132 let right = right
133 .split(':')
134 .filter(|group| !group.is_empty())
135 .collect::<Vec<_>>();
136 left.extend(std::iter::repeat_n("0", 8 - left.len() - right.len()));
137 left.extend(right);
138 left
139 } else {
140 value.split(':').collect()
141 };
142 raw.into_iter()
143 .map(|group| format!("{group:0>4}"))
144 .collect()
145}
146
147fn segment(group: &str) -> u16 {
148 let prefix = group.split_once('.').map_or(group, |(first, _)| first);
149 u16::from_str_radix(prefix, 16).unwrap_or_default()
150}
151
152struct Network {
153 bytes: Vec<u8>,
154 prefix: u8,
155}
156
157impl Network {
158 fn parse(value: &str) -> Option<Self> {
159 let (ip, prefix) = value
160 .rsplit_once('/')
161 .map_or((value, None), |(ip, prefix)| (ip, Some(prefix)));
162 let address = Address::parse(ip)?;
163 let maximum = u8::try_from(address.bytes.len() * 8).ok()?;
164 let prefix = if let Some(prefix) = prefix {
165 if prefix.is_empty() || !prefix.bytes().all(|byte| byte.is_ascii_digit()) {
166 return None;
167 }
168 prefix
169 .parse::<u8>()
170 .ok()
171 .filter(|prefix| *prefix <= maximum)?
172 } else {
173 maximum
174 };
175 Some(Self {
176 bytes: address.bytes,
177 prefix,
178 })
179 }
180
181 fn contains(&self, bytes: &[u8]) -> bool {
182 if bytes.len() != self.bytes.len() {
183 return false;
184 }
185 let mut remaining = self.prefix;
186 for (ip, network) in bytes.iter().zip(&self.bytes) {
187 if remaining == 0 {
188 break;
189 }
190 let mask = u8::MAX << (8 - remaining.min(8));
191 if ip & mask != network & mask {
192 return false;
193 }
194 remaining = remaining.saturating_sub(8);
195 }
196 true
197 }
198}
199
200#[cfg(test)]
202mod tests {
203 use super::*;
204
205 #[test]
206 fn forwarded_chains_select_the_first_untrusted_hop_from_the_right() {
207 for (proxies, chain, expected) in [
208 (vec![], "198.51.100.7", Some("198.51.100.7")),
209 (vec![], "198.51.100.7, 10.0.0.1", None),
210 (vec!["invalid"], "198.51.100.7, 10.0.0.1", None),
211 (
212 vec!["10.0.0.0/8"],
213 "203.0.113.99, 198.51.100.7, 10.0.0.1",
214 Some("198.51.100.7"),
215 ),
216 (vec!["10.0.0.0/8"], "10.0.0.2, 10.0.0.1", None),
217 (
218 vec!["10.0.0.0/8"],
219 "198.51.100.7, malformed, 10.0.0.1",
220 None,
221 ),
222 ] {
223 let config = IpAddressConfig {
224 trusted_proxies: proxies.into_iter().map(str::to_owned).collect(),
225 localhost_fallback: false,
226 ..Default::default()
227 };
228 let headers = HashMap::from([("X-Forwarded-For".into(), chain.into())]);
229 assert_eq!(config.resolve_ip(&headers).as_deref(), expected, "{chain}");
230 }
231 }
232
233 #[test]
234 fn malformed_header_falls_back_but_tracking_opt_out_never_does() {
235 let mut config = IpAddressConfig {
236 headers: vec!["x-forwarded-for".into(), "x-real-ip".into()],
237 trusted_proxies: vec!["10.0.0.0/8".into()],
238 localhost_fallback: true,
239 ..Default::default()
240 };
241 let headers = HashMap::from([
242 (
243 "x-forwarded-for".into(),
244 "198.51.100.7, bad-hop, 10.0.0.1".into(),
245 ),
246 ("x-real-ip".into(), "203.0.113.8".into()),
247 ]);
248 assert_eq!(config.resolve_ip(&headers).as_deref(), Some("203.0.113.8"));
249 assert_eq!(
250 config.resolve_ip(&HashMap::new()).as_deref(),
251 Some("127.0.0.1")
252 );
253 config.disable_ip_tracking = true;
254 assert_eq!(config.resolve_ip(&headers), None);
255 assert_eq!(config.resolve_ip(&HashMap::new()), None);
256 }
257
258 #[test]
259 fn ipv6_proxy_trust_is_checked_before_grouping_the_client_address() {
260 let config = IpAddressConfig {
261 trusted_proxies: vec!["2001:db8:abcd:12::1/128".into()],
262 ipv6_subnet: 64.0,
263 localhost_fallback: false,
264 ..Default::default()
265 };
266 let headers = HashMap::from([(
267 "x-forwarded-for".into(),
268 "203.0.113.99, 2001:db8:abcd:12::2, 2001:db8:abcd:12::1".into(),
269 )]);
270 assert_eq!(
271 config.resolve_ip(&headers).as_deref(),
272 Some("2001:0db8:abcd:0012:0000:0000:0000:0000")
273 );
274 }
275}
276