relay_event_normalization/
geo.rs

1use std::fmt;
2use std::net::IpAddr;
3use std::path::Path;
4use std::sync::Arc;
5
6use relay_event_schema::protocol::Geo;
7use relay_protocol::Annotated;
8
9#[cfg(feature = "mmap")]
10type ReaderType = maxminddb::Mmap;
11
12#[cfg(not(feature = "mmap"))]
13type ReaderType = Vec<u8>;
14
15/// An error in the `GeoIpLookup`.
16pub type GeoIpError = maxminddb::MaxMindDbError;
17
18/// A geo ip lookup helper based on maxmind db files.
19///
20/// The helper is internally reference counted and can be cloned cheaply.
21#[derive(Clone, Default)]
22pub struct GeoIpLookup(Option<Arc<maxminddb::Reader<ReaderType>>>);
23
24impl GeoIpLookup {
25    /// Opens a maxminddb file by path.
26    pub fn open<P>(path: P) -> Result<Self, GeoIpError>
27    where
28        P: AsRef<Path>,
29    {
30        #[cfg(feature = "mmap")]
31        let reader = unsafe { maxminddb::Reader::open_mmap(path)? };
32        #[cfg(not(feature = "mmap"))]
33        let reader = maxminddb::Reader::open_readfile(path)?;
34        Ok(GeoIpLookup(Some(Arc::new(reader))))
35    }
36
37    /// Creates a new [`GeoIpLookup`] instance without any data loaded.
38    pub fn empty() -> Self {
39        Self(None)
40    }
41
42    /// Looks up an IP address.
43    pub fn try_lookup(&self, ip_address: IpAddr) -> Result<Option<Geo>, GeoIpError> {
44        let Some(reader) = self.0.as_ref() else {
45            return Ok(None);
46        };
47
48        let Some(city) = reader
49            .lookup(ip_address)?
50            .decode::<maxminddb::geoip2::City>()?
51        else {
52            return Ok(None);
53        };
54
55        Ok(Some(Geo {
56            country_code: Annotated::from(city.country.iso_code.map(String::from)),
57            city: Annotated::from(city.city.names.english.map(String::from)),
58            subdivision: Annotated::from(
59                city.subdivisions
60                    .first()
61                    .and_then(|subdivision| subdivision.names.english.map(String::from)),
62            ),
63            region: Annotated::from(city.country.names.english.map(String::from)),
64            ..Default::default()
65        }))
66    }
67
68    /// Like [`Self::try_lookup`], but swallows errors.
69    pub fn lookup(&self, ip_address: IpAddr) -> Option<Geo> {
70        self.try_lookup(ip_address).ok().flatten()
71    }
72}
73
74impl fmt::Debug for GeoIpLookup {
75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
76        f.debug_struct("GeoIpLookup").finish()
77    }
78}