1use anyhow::Context;
2use itertools::Itertools;
3use maxminddb::Reader;
4use std::cell::RefCell;
5use std::collections::HashMap;
6use std::net::IpAddr;
7use std::path::Path;
8use std::rc::Rc;
9use std::str::FromStr;
10
11#[derive(Debug, Clone, Default)]
12pub struct GeoIpCity {
13 latitude: Option<f64>,
14 longitude: Option<f64>,
15 accuracy_radius: Option<u16>,
16 city: Option<String>,
17 subdivision: Option<String>,
18 subdivision_code: Option<String>,
19 country: Option<String>,
20 country_code: Option<String>,
21 continent: Option<String>,
22}
23
24impl GeoIpCity {
25 pub fn short_name(&self) -> String {
26 [
27 self.city.as_ref(),
28 self.subdivision_code.as_ref(),
29 self.country_code.as_ref(),
30 ]
31 .into_iter()
32 .flatten()
33 .join(", ")
34 }
35
36 pub fn long_name(&self) -> String {
37 [
38 self.city.as_ref(),
39 self.subdivision.as_ref(),
40 self.country.as_ref(),
41 self.continent.as_ref(),
42 ]
43 .into_iter()
44 .flatten()
45 .join(", ")
46 }
47
48 pub fn location(&self) -> String {
49 format!(
50 "{}, {} (~{}km)",
51 self.latitude.unwrap_or_default(),
52 self.longitude.unwrap_or_default(),
53 self.accuracy_radius.unwrap_or_default(),
54 )
55 }
56
57 pub const fn coordinates(&self) -> Option<(f64, f64, u16)> {
58 match (self.latitude, self.longitude, self.accuracy_radius) {
59 (Some(lat), Some(long), Some(radius)) => Some((lat, long, radius)),
60 _ => None,
61 }
62 }
63}
64
65mod ipinfo {
66 use serde::{Deserialize, Serialize};
67 use serde_with::serde_as;
68
69 #[serde_as]
79 #[derive(Debug, Serialize, Deserialize)]
80 pub struct IpInfoGeoIp {
81 #[serde(default)]
83 #[serde_as(as = "serde_with::NoneAsEmptyString")]
84 pub latitude: Option<String>,
85 #[serde(default)]
87 #[serde_as(as = "serde_with::NoneAsEmptyString")]
88 pub longitude: Option<String>,
89 #[serde(default)]
91 #[serde_as(as = "serde_with::NoneAsEmptyString")]
92 pub radius: Option<String>,
93 #[serde(default)]
95 #[serde_as(as = "serde_with::NoneAsEmptyString")]
96 pub city: Option<String>,
97 #[serde(default)]
99 #[serde_as(as = "serde_with::NoneAsEmptyString")]
100 pub region: Option<String>,
101 #[serde(default)]
103 #[serde_as(as = "serde_with::NoneAsEmptyString")]
104 pub postal_code: Option<String>,
105 #[serde(default)]
107 #[serde_as(as = "serde_with::NoneAsEmptyString")]
108 pub country: Option<String>,
109 #[serde(default)]
111 #[serde_as(as = "serde_with::NoneAsEmptyString")]
112 pub country_name: Option<String>,
113 #[serde(default)]
115 #[serde_as(as = "serde_with::NoneAsEmptyString")]
116 pub continent_name: Option<String>,
117 }
118
119 #[cfg(test)]
120 mod tests {
121 use super::*;
122
123 #[test]
124 fn test_empty() {
125 let json = "{}";
126 let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
127 assert_eq!(None, value.latitude);
128 assert_eq!(None, value.longitude);
129 assert_eq!(None, value.radius);
130 assert_eq!(None, value.city);
131 assert_eq!(None, value.region);
132 assert_eq!(None, value.postal_code);
133 assert_eq!(None, value.country.as_deref());
134 assert_eq!(None, value.country_name.as_deref());
135 assert_eq!(None, value.continent_name.as_deref());
136 }
137
138 #[test]
139 fn test_country_asn_db_format() {
140 let json = r#"
141 {
142 "start_ip": "40.96.54.192",
143 "end_ip": "40.96.54.255",
144 "country": "JP",
145 "country_name": "Japan",
146 "continent": "AS",
147 "continent_name": "Asia",
148 "asn": "AS8075",
149 "as_name": "Microsoft Corporation",
150 "as_domain": "microsoft.com"
151 }
152 "#;
153 let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
154 assert_eq!(None, value.latitude);
155 assert_eq!(None, value.longitude);
156 assert_eq!(None, value.radius);
157 assert_eq!(None, value.city);
158 assert_eq!(None, value.region);
159 assert_eq!(None, value.postal_code);
160 assert_eq!(Some("JP"), value.country.as_deref());
161 assert_eq!(Some("Japan"), value.country_name.as_deref());
162 assert_eq!(Some("Asia"), value.continent_name.as_deref());
163 }
164
165 #[test]
166 fn test_extended_db_format() {
167 let json = r#"
168 {
169 "start_ip": "60.127.10.249",
170 "end_ip": "60.127.10.249",
171 "join_key": "60.127.0.0",
172 "city": "Yokohama",
173 "region": "Kanagawa",
174 "country": "JP",
175 "latitude": "35.43333",
176 "longitude": "139.65",
177 "postal_code": "220-8588",
178 "timezone": "Asia/Tokyo",
179 "geoname_id": "1848354",
180 "radius": "500"
181 }
182 "#;
183 let value: IpInfoGeoIp = serde_json::from_str(json).unwrap();
184 assert_eq!(Some("35.43333"), value.latitude.as_deref());
185 assert_eq!(Some("139.65"), value.longitude.as_deref());
186 assert_eq!(Some("500"), value.radius.as_deref());
187 assert_eq!(Some("Yokohama"), value.city.as_deref());
188 assert_eq!(Some("Kanagawa"), value.region.as_deref());
189 assert_eq!(Some("220-8588"), value.postal_code.as_deref());
190 assert_eq!(Some("JP"), value.country.as_deref());
191 assert_eq!(None, value.country_name.as_deref());
192 assert_eq!(None, value.continent_name.as_deref());
193 }
194 }
195}
196
197impl From<ipinfo::IpInfoGeoIp> for GeoIpCity {
198 fn from(value: ipinfo::IpInfoGeoIp) -> Self {
199 Self {
200 latitude: value.latitude.and_then(|val| f64::from_str(&val).ok()),
201 longitude: value.longitude.and_then(|val| f64::from_str(&val).ok()),
202 accuracy_radius: value.radius.and_then(|val| u16::from_str(&val).ok()),
203 city: value.city,
204 subdivision: value.region,
205 subdivision_code: value.postal_code,
206 country: value.country_name,
207 country_code: value.country,
208 continent: value.continent_name,
209 }
210 }
211}
212
213impl From<(maxminddb::geoip2::City<'_>, &str)> for GeoIpCity {
214 fn from((value, locale): (maxminddb::geoip2::City<'_>, &str)) -> Self {
215 let city = value
216 .city
217 .as_ref()
218 .and_then(|city| city.names.as_ref())
219 .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
220 .map(ToString::to_string);
221 let subdivision = value
222 .subdivisions
223 .as_ref()
224 .and_then(|c| c.first())
225 .and_then(|c| c.names.as_ref())
226 .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
227 .map(ToString::to_string);
228 let subdivision_code = value
229 .subdivisions
230 .as_ref()
231 .and_then(|c| c.first())
232 .and_then(|c| c.iso_code.as_ref())
233 .map(ToString::to_string);
234 let country = value
235 .country
236 .as_ref()
237 .and_then(|country| country.names.as_ref())
238 .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
239 .map(ToString::to_string);
240 let country_code = value
241 .country
242 .as_ref()
243 .and_then(|country| country.iso_code.as_ref())
244 .map(ToString::to_string);
245 let continent = value
246 .continent
247 .as_ref()
248 .and_then(|continent| continent.names.as_ref())
249 .and_then(|names| names.get(locale).or_else(|| names.get(FALLBACK_LOCALE)))
250 .map(ToString::to_string);
251 let latitude = value
252 .location
253 .as_ref()
254 .and_then(|location| location.latitude);
255 let longitude = value
256 .location
257 .as_ref()
258 .and_then(|location| location.longitude);
259 let accuracy_radius = value
260 .location
261 .as_ref()
262 .and_then(|location| location.accuracy_radius);
263 Self {
264 latitude,
265 longitude,
266 accuracy_radius,
267 city,
268 subdivision,
269 subdivision_code,
270 country,
271 country_code,
272 continent,
273 }
274 }
275}
276
277const FALLBACK_LOCALE: &str = "en";
290
291type Cache = RefCell<HashMap<IpAddr, Rc<GeoIpCity>>>;
293
294#[derive(Debug)]
296pub struct GeoIpLookup {
297 reader: Option<Reader<Vec<u8>>>,
298 cache: Cache,
299 locale: String,
300}
301
302impl GeoIpLookup {
303 pub fn from_file<P: AsRef<Path>>(path: P, locale: String) -> anyhow::Result<Self> {
305 let reader = maxminddb::Reader::open_readfile(path.as_ref())
306 .context(format!("{}", path.as_ref().display()))?;
307 Ok(Self {
308 reader: Some(reader),
309 cache: RefCell::new(HashMap::new()),
310 locale,
311 })
312 }
313
314 pub fn empty() -> Self {
316 Self {
317 reader: None,
318 cache: RefCell::new(HashMap::new()),
319 locale: FALLBACK_LOCALE.to_string(),
320 }
321 }
322
323 pub fn lookup(&self, addr: IpAddr) -> anyhow::Result<Option<Rc<GeoIpCity>>> {
327 if let Some(reader) = &self.reader {
328 if let Some(geo) = self.cache.borrow().get(&addr) {
329 return Ok(Some(geo.clone()));
330 }
331 let city_data = if reader.metadata.database_type.starts_with("ipinfo") {
332 GeoIpCity::from(reader.lookup::<ipinfo::IpInfoGeoIp>(addr)?)
333 } else {
334 GeoIpCity::from((
335 reader.lookup::<maxminddb::geoip2::City<'_>>(addr)?,
336 self.locale.as_ref(),
337 ))
338 };
339 let geo = self.cache.borrow_mut().insert(addr, Rc::new(city_data));
340 Ok(geo)
341 } else {
342 Ok(None)
343 }
344 }
345}