session_guard/intel/
maxmind.rs

1//! [`IpIntel`] backed by `MaxMind` `GeoLite2` / `GeoIP2` databases.
2
3use std::fmt;
4use std::net::IpAddr;
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7
8use arc_swap::ArcSwap;
9use maxminddb::{PathElement, Reader};
10
11use super::{IpInfo, IpIntel};
12
13/// Which of the two databases an error is about.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15#[non_exhaustive]
16pub enum MaxMindDatabase {
17    /// The City or Country database, used for the ISO country code.
18    Country,
19    /// The ASN database, used for the autonomous system number.
20    Asn,
21}
22
23impl fmt::Display for MaxMindDatabase {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        f.write_str(match self {
26            Self::Country => "country",
27            Self::Asn => "ASN",
28        })
29    }
30}
31
32/// Errors loading `MaxMind` databases.
33#[derive(Debug, thiserror::Error)]
34#[non_exhaustive]
35pub enum MaxMindError {
36    /// The database file could not be read.
37    #[error("cannot read MaxMind {database} database {}: {source}", path.display())]
38    Io {
39        /// Which database.
40        database: MaxMindDatabase,
41        /// The path that failed.
42        path: PathBuf,
43        /// The underlying I/O error.
44        #[source]
45        source: std::io::Error,
46    },
47    /// The bytes are not a valid MMDB file.
48    #[error("invalid MaxMind {database} database: {source}")]
49    Invalid {
50        /// Which database.
51        database: MaxMindDatabase,
52        /// The parser's error.
53        #[source]
54        source: Box<dyn std::error::Error + Send + Sync + 'static>,
55    },
56    /// A valid database of the wrong kind, e.g. an ASN database passed as the country database.
57    #[error("expected a MaxMind {database} database, found {found:?}")]
58    WrongDatabase {
59        /// Which slot the database was passed for.
60        database: MaxMindDatabase,
61        /// The database type from the file's metadata.
62        found: String,
63    },
64}
65
66/// Country and ASN lookups from `MaxMind` databases (`GeoLite2-City` / `GeoLite2-Country` /
67/// `GeoIP2-*` for the country, `GeoLite2-ASN` for the network).
68///
69/// Databases are read into memory and can be swapped at runtime with [`reload`](Self::reload)
70/// (for example after a `geoipupdate` run); lookups on the request path are lock-free. Unknown
71/// addresses and lookup errors yield empty fields.
72///
73/// ```no_run
74/// use session_guard::{Guard, GuardConfig, MaxMindIntel, MemoryStore};
75///
76/// let intel = MaxMindIntel::open(
77///     "/var/lib/GeoIP/GeoLite2-Country.mmdb",
78///     Some("/var/lib/GeoIP/GeoLite2-ASN.mmdb"),
79/// )?;
80/// let guard = Guard::builder(MemoryStore::new(), GuardConfig::default()).ip_intel(intel).build();
81/// # let _ = guard;
82/// # Ok::<(), session_guard::MaxMindError>(())
83/// ```
84pub struct MaxMindIntel {
85    databases: ArcSwap<Databases>,
86}
87
88struct Databases {
89    country: Reader<Vec<u8>>,
90    asn: Option<Reader<Vec<u8>>>,
91}
92
93impl Databases {
94    fn read(country: &Path, asn: Option<impl AsRef<Path>>) -> Result<Self, MaxMindError> {
95        let read = |database, path: &Path| {
96            std::fs::read(path).map_err(|source| MaxMindError::Io {
97                database,
98                path: path.to_owned(),
99                source,
100            })
101        };
102        let country = read(MaxMindDatabase::Country, country)?;
103        let asn = asn.map(|path| read(MaxMindDatabase::Asn, path.as_ref())).transpose()?;
104        Self::parse(country, asn)
105    }
106
107    fn parse(country: Vec<u8>, asn: Option<Vec<u8>>) -> Result<Self, MaxMindError> {
108        Ok(Self {
109            country: reader(MaxMindDatabase::Country, country)?,
110            asn: asn.map(|bytes| reader(MaxMindDatabase::Asn, bytes)).transpose()?,
111        })
112    }
113}
114
115fn reader(database: MaxMindDatabase, bytes: Vec<u8>) -> Result<Reader<Vec<u8>>, MaxMindError> {
116    let reader = Reader::from_source(bytes).map_err(|source| MaxMindError::Invalid {
117        database,
118        source: Box::new(source),
119    })?;
120    let kind = reader.metadata().database_type.to_ascii_lowercase();
121    let has_country = kind.contains("country") || kind.contains("city");
122    let has_asn = kind.contains("asn");
123    let fits = match database {
124        MaxMindDatabase::Country => has_country || !has_asn,
125        MaxMindDatabase::Asn => has_asn || !has_country,
126    };
127    if !fits {
128        return Err(MaxMindError::WrongDatabase {
129            database,
130            found: reader.metadata().database_type.clone(),
131        });
132    }
133    Ok(reader)
134}
135
136impl MaxMindIntel {
137    /// Reads the country database (City or Country edition) and, optionally, the ASN database.
138    ///
139    /// # Errors
140    ///
141    /// Fails when a file cannot be read, is not a valid MMDB file, or is the wrong kind of
142    /// database.
143    pub fn open(country_db: impl AsRef<Path>, asn_db: Option<impl AsRef<Path>>) -> Result<Self, MaxMindError> {
144        let databases = Databases::read(country_db.as_ref(), asn_db)?;
145        Ok(Self {
146            databases: ArcSwap::from_pointee(databases),
147        })
148    }
149
150    /// Like [`open`](Self::open), from database contents already in memory.
151    ///
152    /// # Errors
153    ///
154    /// Fails when the bytes are not a valid MMDB file or the wrong kind of database.
155    pub fn from_bytes(country_db: Vec<u8>, asn_db: Option<Vec<u8>>) -> Result<Self, MaxMindError> {
156        Ok(Self {
157            databases: ArcSwap::from_pointee(Databases::parse(country_db, asn_db)?),
158        })
159    }
160
161    /// Replaces both databases with freshly read files. In-flight lookups finish on the old
162    /// databases. On error the current databases stay in place.
163    ///
164    /// To reload a source the guard uses, wrap it in an `Arc`, pass a clone to
165    /// [`GuardBuilder::ip_intel`](crate::GuardBuilder::ip_intel) and keep the other.
166    ///
167    /// # Errors
168    ///
169    /// As for [`open`](Self::open).
170    pub fn reload(&self, country_db: impl AsRef<Path>, asn_db: Option<impl AsRef<Path>>) -> Result<(), MaxMindError> {
171        let databases = Databases::read(country_db.as_ref(), asn_db)?;
172        self.databases.store(Arc::new(databases));
173        Ok(())
174    }
175
176    /// Like [`reload`](Self::reload), from database contents already in memory.
177    ///
178    /// # Errors
179    ///
180    /// As for [`from_bytes`](Self::from_bytes).
181    pub fn reload_from_bytes(&self, country_db: Vec<u8>, asn_db: Option<Vec<u8>>) -> Result<(), MaxMindError> {
182        self.databases.store(Arc::new(Databases::parse(country_db, asn_db)?));
183        Ok(())
184    }
185}
186
187const COUNTRY: [PathElement<'static>; 2] = [PathElement::Key("country"), PathElement::Key("iso_code")];
188const REGISTERED_COUNTRY: [PathElement<'static>; 2] =
189    [PathElement::Key("registered_country"), PathElement::Key("iso_code")];
190const ASN: [PathElement<'static>; 1] = [PathElement::Key("autonomous_system_number")];
191
192impl IpIntel for MaxMindIntel {
193    fn lookup(&self, ip: IpAddr) -> IpInfo {
194        let ip = crate::client::canonical_ip(ip);
195        let databases = self.databases.load();
196        let country = databases.country.lookup(ip).ok().and_then(|found| {
197            [&COUNTRY[..], &REGISTERED_COUNTRY[..]]
198                .into_iter()
199                .find_map(|path| found.decode_path::<&str>(path).ok().flatten())
200                .map(str::to_owned)
201        });
202        let asn = databases
203            .asn
204            .as_ref()
205            .and_then(|reader| reader.lookup(ip).ok())
206            .and_then(|found| found.decode_path::<u32>(&ASN).ok().flatten());
207        IpInfo::new(country, asn)
208    }
209}
210
211impl fmt::Debug for MaxMindIntel {
212    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213        let databases = self.databases.load();
214        f.debug_struct("MaxMindIntel")
215            .field("country", &databases.country.metadata().database_type)
216            .field("asn", &databases.asn.as_ref().map(|r| &r.metadata().database_type))
217            .finish()
218    }
219}
220
221#[cfg(test)]
222#[path = "maxmind_tests.rs"]
223mod tests;
session_guard/intel/maxmind.rs

session_guard/intel/
maxmind.rs
