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 = 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 city: maxminddb::geoip2::City = match reader.lookup(ip_address) {
49            Ok(Some(x)) => x,
50            Ok(None) => return Ok(None),
51            Err(e) => return Err(e),
52        };
53
54        Ok(Some(Geo {
55            country_code: Annotated::from(
56                city.country
57                    .as_ref()
58                    .and_then(|country| Some(country.iso_code.as_ref()?.to_string())),
59            ),
60            city: Annotated::from(
61                city.city
62                    .as_ref()
63                    .and_then(|city| Some(city.names.as_ref()?.get("en")?.to_string())),
64            ),
65            subdivision: Annotated::from(city.subdivisions.as_ref().and_then(|subdivisions| {
66                subdivisions.first().and_then(|subdivision| {
67                    subdivision.names.as_ref().and_then(|subdivision_names| {
68                        subdivision_names
69                            .get("en")
70                            .map(|subdivision_name| subdivision_name.to_string())
71                    })
72                })
73            })),
74            region: Annotated::from(
75                city.country
76                    .as_ref()
77                    .and_then(|country| Some(country.names.as_ref()?.get("en")?.to_string())),
78            ),
79            ..Default::default()
80        }))
81    }
82
83    /// Like [`Self::try_lookup`], but swallows errors.
84    pub fn lookup(&self, ip_address: IpAddr) -> Option<Geo> {
85        self.try_lookup(ip_address).ok().flatten()
86    }
87}
88
89impl fmt::Debug for GeoIpLookup {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        f.debug_struct("GeoIpLookup").finish()
92    }
93}