Skip to main content

mmdb_writer/
load.rs

1//! Loading an existing `.mmdb` file into a [`Writer`] so it can be extended and rewritten.
2//! Gated on the `load` feature (which pulls in the `maxminddb` reader).
3
4use std::path::Path;
5use std::time::{Duration, UNIX_EPOCH};
6
7use ipnet::IpNet;
8use maxminddb::{Reader, WithinOptions};
9
10use crate::error::Error;
11use crate::net::IpVersion;
12use crate::record_size::RecordSize;
13use crate::value::Value;
14use crate::writer::Writer;
15
16impl Writer {
17    /// Load an existing database from a file, ready to be extended.
18    ///
19    /// The returned [`Writer`] is seeded with every network in the file and the file's
20    /// metadata (database type, languages, descriptions, IP version, record size, and build
21    /// epoch). Reserved-network and aliasing options are not stored in the format, so they
22    /// take their defaults.
23    ///
24    /// # Errors
25    ///
26    /// Returns [`Error::Io`] if the file cannot be read, or [`Error::Load`] if it is not a
27    /// valid database.
28    pub fn load_from_path(path: impl AsRef<Path>) -> Result<Self, Error> {
29        let data = std::fs::read(path)?;
30        Self::load(&data)
31    }
32
33    /// Load an existing database from a byte buffer, ready to be extended.
34    ///
35    /// See [`load_from_path`](Self::load_from_path) for details.
36    ///
37    /// # Errors
38    ///
39    /// Returns [`Error::Load`] if `bytes` is not a valid database.
40    pub fn load(bytes: impl AsRef<[u8]>) -> Result<Self, Error> {
41        let reader = Reader::from_source(bytes.as_ref()).map_err(load_err)?;
42        let meta = reader.metadata();
43
44        let ip_version = if meta.ip_version == 4 {
45            IpVersion::V4
46        } else {
47            IpVersion::V6
48        };
49        let epoch = UNIX_EPOCH + Duration::from_secs(meta.build_epoch);
50
51        let description: Vec<(&str, &str)> = meta
52            .description
53            .iter()
54            .map(|(lang, text)| (lang.as_str(), text.as_str()))
55            .collect();
56        let mut writer = Writer::builder(meta.database_type.clone())
57            .ip_version(ip_version)
58            .languages(meta.languages.clone())
59            .description(&description)
60            .record_size(record_size_from(meta.record_size)?)
61            .build_epoch(epoch)
62            .build();
63
64        // Default options skip aliased networks (so IPv4 data is not yielded several times)
65        // and networks without data.
66        let items = reader
67            .networks(WithinOptions::default())
68            .map_err(load_err)?;
69        for item in items {
70            let item = item.map_err(load_err)?;
71            let Some(value) = item.decode::<Value>().map_err(load_err)? else {
72                continue;
73            };
74            let network = item.network().map_err(load_err)?;
75            let net = IpNet::new(network.ip(), network.prefix())
76                .map_err(|e| Error::Load(format!("reader produced an invalid network: {e}")))?;
77            writer.insert_value(net, value)?;
78        }
79
80        Ok(writer)
81    }
82}
83
84// Taken by value so it can be used directly as a `map_err` argument.
85#[allow(clippy::needless_pass_by_value)]
86fn load_err(e: maxminddb::MaxMindDbError) -> Error {
87    Error::Load(e.to_string())
88}
89
90fn record_size_from(bits: u16) -> Result<RecordSize, Error> {
91    match bits {
92        24 => Ok(RecordSize::Bits24),
93        28 => Ok(RecordSize::Bits28),
94        32 => Ok(RecordSize::Bits32),
95        other => Err(Error::Load(format!("unsupported record size {other}"))),
96    }
97}