Skip to main content

maxminddb/
maxminddb.rs

1//! A reader for the MaxMind DB file format, including the GeoIP2 and GeoLite2
2//! databases published by MaxMind.
3//!
4//! # Usage
5//!
6//! Open a database and look up an address:
7//!
8//! ```no_run
9//! use std::net::IpAddr;
10//!
11//! use maxminddb::geoip2;
12//!
13//! # fn main() -> Result<(), maxminddb::MaxMindDBError> {
14//! let reader = maxminddb::Reader::open_readfile("GeoIP2-City.mmdb")?;
15//! let ip: IpAddr = "89.160.20.128".parse().unwrap();
16//! let city: geoip2::City<'_> = reader.lookup(ip)?;
17//! println!("{city:?}");
18//! # Ok(())
19//! # }
20//! ```
21//!
22//! [`Reader::from_source`] accepts any buffer that implements `AsRef<[u8]>`,
23//! which is useful when the database is already in memory or memory-mapped by
24//! the caller.
25//!
26//! To iterate over every network in the database, use [`Reader::within`]. The
27//! `examples/` directory contains `lookup`, `within`, and `write_database`
28//! examples.
29//!
30//! # Resource limits
31//!
32//! A MaxMind DB encodes values with pointers, and many pointers may share one
33//! target, so a small record can describe a very large or deeply nested
34//! structure when expanded. To keep hostile or corrupt databases from consuming
35//! unbounded time or memory, decoding is bounded by [`DecodeLimits`]:
36//!
37//! * a maximum nesting depth of `512` (maps, arrays, and followed pointers),
38//! * a maximum of `65_536` decoded values per operation, and
39//! * a maximum of `2 MiB` of string and byte payload per operation.
40//!
41//! These are the values recommended by the MaxMind DB specification. When a
42//! limit is exceeded, decoding fails with [`MaxMindDBError::ResourceLimit`]
43//! instead of panicking or exhausting memory. The limits apply to metadata
44//! decoding at open as well as to record lookups.
45//!
46//! Use [`Reader::from_source_with_options`] or
47//! [`Reader::open_readfile_with_options`] with a [`ReaderOptions`] value to
48//! configure the limits.
49//!
50//! A database that is a valid IPv4-only file but is looked up with an IPv6
51//! address or CIDR is rejected with [`MaxMindDBError::InvalidInputError`]; the
52//! address family must match the database.
53//!
54//! # Supported databases
55//!
56//! Both IPv4 and IPv6 databases are supported, as are the 24-, 28-, and 32-bit
57//! search tree record sizes. The [`geoip2`] module provides typed records for
58//! the common GeoIP2 and GeoLite2 database shapes. For other databases,
59//! deserialize into your own type or into [`Value`].
60#![deny(trivial_casts, trivial_numeric_casts, unused_import_braces)]
61
62use std::{
63    cmp::Ordering,
64    collections::BTreeMap,
65    fmt::{self, Display, Formatter},
66    io,
67    marker::PhantomData,
68    net::{IpAddr, Ipv4Addr, Ipv6Addr},
69    path::Path,
70};
71
72use ipnetwork::IpNetwork;
73use serde::{Deserialize, Serialize, de};
74
75// #[cfg(feature = "mmap")]
76// pub use memmap2::Mmap;
77// #[cfg(feature = "mmap")]
78// use memmap2::MmapOptions;
79// #[cfg(feature = "mmap")]
80// use std::fs::File;
81
82pub mod decoder;
83pub mod encoder;
84pub mod geoip2;
85pub mod writer;
86pub use encoder::Value;
87
88#[derive(Debug, PartialEq, Eq)]
89pub enum MaxMindDBError {
90    AddressNotFoundError(IpAddr),
91    InvalidDatabaseError(String),
92    IoError(String),
93    MapError(String),
94    DecodingError(String),
95    InvalidNetworkError(String),
96    /// The caller supplied input that is incompatible with the database, such as
97    /// an IPv6 address or CIDR for an IPv4-only database.
98    InvalidInputError(String),
99    /// Decoding the database record would have exceeded one of the resource
100    /// limits configured via [`ReaderOptions`], so it was stopped.
101    ResourceLimit(String),
102}
103
104impl From<io::Error> for MaxMindDBError {
105    fn from(err: io::Error) -> MaxMindDBError {
106        // clean up and clean up MaxMindDBError generally
107        MaxMindDBError::IoError(err.to_string())
108    }
109}
110
111impl Display for MaxMindDBError {
112    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), fmt::Error> {
113        match self {
114            MaxMindDBError::AddressNotFoundError(msg) => write!(fmt, "AddressNotFoundError: {msg}")?,
115            MaxMindDBError::InvalidDatabaseError(msg) => write!(fmt, "InvalidDatabaseError: {msg}")?,
116            MaxMindDBError::IoError(msg) => write!(fmt, "IoError: {msg}")?,
117            MaxMindDBError::MapError(msg) => write!(fmt, "MapError: {msg}")?,
118            MaxMindDBError::DecodingError(msg) => write!(fmt, "DecodingError: {msg}")?,
119            MaxMindDBError::InvalidNetworkError(msg) => write!(fmt, "InvalidNetworkError: {msg}")?,
120            MaxMindDBError::InvalidInputError(msg) => write!(fmt, "InvalidInputError: {msg}")?,
121            MaxMindDBError::ResourceLimit(msg) => write!(fmt, "ResourceLimit: {msg}")?,
122        }
123        Ok(())
124    }
125}
126
127// Use default implementation for `core::error::Error`
128impl core::error::Error for MaxMindDBError {}
129
130impl de::Error for MaxMindDBError {
131    fn custom<T: Display>(msg: T) -> Self {
132        MaxMindDBError::DecodingError(format!("{msg}"))
133    }
134}
135
136#[derive(Deserialize, Serialize, Clone, Debug)]
137pub struct Metadata {
138    pub binary_format_major_version: u16,
139    pub binary_format_minor_version: u16,
140    pub build_epoch: u64,
141    pub database_type: String,
142    pub description: BTreeMap<String, String>,
143    pub ip_version: u16,
144    pub languages: Vec<String>,
145    pub node_count: u32,
146    pub record_size: u16,
147}
148
149/// Resource limits applied while decoding a single data entry or the metadata
150/// section.
151///
152/// A MaxMind DB can encode a value as a pointer, and many pointers can share
153/// one target, so a small record can describe a structure that is huge, deeply
154/// nested, or effectively infinite when expanded. These limits bound the time
155/// and memory one decode operation may consume. When a limit is exceeded, the
156/// decode fails with [`MaxMindDBError::ResourceLimit`].
157///
158/// The defaults follow the recommendations in the MaxMind DB specification: a
159/// maximum nesting depth of `512`, a maximum of `65_536` decoded values, and a
160/// maximum of `2 MiB` of string and byte payload materialized per operation.
161/// The largest records MaxMind produces stay far below these values.
162#[derive(Debug, Clone, Copy, PartialEq, Eq)]
163pub struct DecodeLimits {
164    /// Maximum nesting depth. Each nested map, array, or followed pointer adds
165    /// one level. This also bounds pointer cycles and unbounded pointer chains.
166    ///
167    /// The decoder is recursive, so raising this value increases stack usage
168    /// proportionally. The default of `512` is safe on typical thread stacks.
169    pub max_depth: usize,
170    /// Maximum number of logical values decoded per operation. Every value,
171    /// map key, and map value counts once; following a pointer does not add to
172    /// the count.
173    pub max_values: usize,
174    /// Maximum number of string and byte payload bytes materialized per
175    /// operation. Each occurrence of a string or byte value is charged.
176    pub max_payload_bytes: usize,
177}
178
179impl Default for DecodeLimits {
180    fn default() -> Self {
181        DecodeLimits {
182            max_depth: 512,
183            max_values: 65_536,
184            max_payload_bytes: 2 * 1024 * 1024,
185        }
186    }
187}
188
189/// Options controlling how a [`Reader`] opens and decodes a database.
190///
191/// Use [`Reader::from_source_with_options`] or
192/// [`Reader::open_readfile_with_options`] to supply custom options. The
193/// convenience constructors use `ReaderOptions::default()`, which applies the
194/// recommended resource limits described on [`DecodeLimits`].
195#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
196pub struct ReaderOptions {
197    /// Resource limits applied to metadata and record decoding.
198    pub limits: DecodeLimits,
199}
200
201#[derive(Debug, Clone, Copy)]
202struct WithinNode {
203    node: usize,
204    ip_bytes: [u8; 16],
205    prefix_len: usize,
206}
207
208#[derive(Debug)]
209pub struct Within<'de, T: Deserialize<'de>, S: AsRef<[u8]>> {
210    reader: &'de Reader<S>,
211    node_count: usize,
212    /// Length in bytes of the addresses being iterated (`4` or `16`).
213    byte_len: usize,
214    stack: Vec<WithinNode>,
215    phantom: PhantomData<&'de T>,
216}
217
218#[derive(Debug)]
219pub struct WithinItem<T> {
220    pub ip_net: IpNetwork,
221    pub info: T,
222}
223
224impl<'de, T: Deserialize<'de>, S: AsRef<[u8]>> Iterator for Within<'de, T, S> {
225    type Item = Result<WithinItem<T>, MaxMindDBError>;
226
227    fn next(&mut self) -> Option<Self::Item> {
228        let bit_count = self.byte_len * 8;
229        while let Some(current) = self.stack.pop() {
230            // Skip networks that are aliases for the IPv4 network
231            if self.reader.ipv4_start != 0
232                && current.node == self.reader.ipv4_start
233                && bit_count == 128
234                && current.ip_bytes[..12].iter().any(|&b| b != 0)
235            {
236                continue;
237            }
238
239            match current.node.cmp(&self.node_count) {
240                Ordering::Greater => {
241                    // This is a data node, emit it and we're done (until the following next call)
242                    let ip_net =
243                        match bytes_and_prefix_to_net(&current.ip_bytes[..self.byte_len], current.prefix_len as u8) {
244                            Ok(ip_net) => ip_net,
245                            Err(e) => return Some(Err(e)),
246                        };
247                    // TODO: should this block become a helper method on reader?
248                    let rec = match self.reader.resolve_data_pointer(current.node) {
249                        Ok(rec) => rec,
250                        Err(e) => return Some(Err(e)),
251                    };
252                    let mut decoder = decoder::Decoder::with_limits(
253                        &self.reader.buf.as_ref()[self.reader.pointer_base..],
254                        rec,
255                        self.reader.limits,
256                    );
257                    return match T::deserialize(&mut decoder) {
258                        Ok(info) => Some(Ok(WithinItem {
259                            ip_net,
260                            info,
261                        })),
262                        Err(e) => Some(Err(e)),
263                    };
264                }
265                Ordering::Equal => {
266                    // Dead end, nothing to do
267                }
268                Ordering::Less => {
269                    // A well-formed search tree cannot be deeper than the
270                    // address width. A corrupt or cyclic tree would otherwise
271                    // descend forever and grow the stack without bound.
272                    if current.prefix_len >= bit_count {
273                        return Some(Err(MaxMindDBError::InvalidDatabaseError(
274                            "search tree is deeper than the address size".to_owned(),
275                        )));
276                    }
277                    // In order traversal of our children
278                    // right/1-bit
279                    let mut right_ip_bytes = current.ip_bytes;
280                    right_ip_bytes[current.prefix_len >> 3] |= 1 << ((bit_count - current.prefix_len - 1) % 8);
281                    let node = match self.reader.read_node(current.node, 1) {
282                        Ok(node) => node,
283                        Err(e) => return Some(Err(e)),
284                    };
285                    self.stack.push(WithinNode {
286                        node,
287                        ip_bytes: right_ip_bytes,
288                        prefix_len: current.prefix_len + 1,
289                    });
290                    // left/0-bit
291                    let node = match self.reader.read_node(current.node, 0) {
292                        Ok(node) => node,
293                        Err(e) => return Some(Err(e)),
294                    };
295                    self.stack.push(WithinNode {
296                        node,
297                        ip_bytes: current.ip_bytes,
298                        prefix_len: current.prefix_len + 1,
299                    });
300                }
301            }
302        }
303        None
304    }
305}
306
307/// A reader for the MaxMind DB format. The lifetime of the decoded values is tied to the lifetime of the underlying buffer holding the contents of the database file.
308#[derive(Debug)]
309pub struct Reader<S: AsRef<[u8]>> {
310    buf: S,
311    pub metadata: Metadata,
312    ipv4_start: usize,
313    pointer_base: usize,
314    limits: DecodeLimits,
315}
316
317// #[cfg(feature = "mmap")]
318// impl<'de> Reader<Mmap> {
319//     /// Open a MaxMind DB database file by memory mapping it.
320//     ///
321//     /// # Example
322//     ///
323//     /// ```
324//     /// let reader = maxminddb::Reader::open_mmap("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
325//     /// ```
326//     pub fn open_mmap<P: AsRef<Path>>(database: P) -> Result<Reader<Mmap>, MaxMindDBError> {
327//         let file_read = File::open(database)?;
328//         let mmap = unsafe { MmapOptions::new().map(&file_read) }?;
329//         Reader::from_source(mmap)
330//     }
331// }
332
333impl Reader<Vec<u8>> {
334    /// Open a MaxMind DB database file by loading it into memory.
335    ///
336    /// Uses [`ReaderOptions::default`] for resource limits. Use
337    /// [`Reader::open_readfile_with_options`] to configure them.
338    ///
339    /// # Example
340    ///
341    /// ```
342    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
343    /// ```
344    ///
345    /// # Errors
346    ///
347    /// Returns [`MaxMindDBError::IoError`] if the file cannot be read, or an
348    /// error from [`Reader::from_source`] if its contents are not a valid
349    /// database.
350    pub fn open_readfile<P: AsRef<Path>>(database: P) -> Result<Reader<Vec<u8>>, MaxMindDBError> {
351        Reader::open_readfile_with_options(database, ReaderOptions::default())
352    }
353
354    /// Open a MaxMind DB database file by loading it into memory, using
355    /// `options` for resource limits.
356    ///
357    /// # Errors
358    ///
359    /// Returns [`MaxMindDBError::IoError`] if the file cannot be read, or an
360    /// error from [`Reader::from_source_with_options`] if its contents are not
361    /// a valid database or exceed the configured limits.
362    pub fn open_readfile_with_options<P: AsRef<Path>>(
363        database: P,
364        options: ReaderOptions,
365    ) -> Result<Reader<Vec<u8>>, MaxMindDBError> {
366        use std::fs;
367
368        let buf: Vec<u8> = fs::read(&database)?;
369        Reader::from_source_with_options(buf, options)
370    }
371}
372
373impl<'de, S: AsRef<[u8]>> Reader<S> {
374    /// Open a MaxMind DB database from anything that implements AsRef<[u8]>
375    ///
376    /// Uses [`ReaderOptions::default`] for resource limits. Use
377    /// [`Reader::from_source_with_options`] to configure them.
378    ///
379    /// # Example
380    ///
381    /// ```
382    /// use std::fs;
383    /// let buf = fs::read("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
384    /// let reader = maxminddb::Reader::from_source(buf).unwrap();
385    /// ```
386    ///
387    /// # Errors
388    ///
389    /// Returns [`MaxMindDBError::InvalidDatabaseError`] if the buffer is not a
390    /// valid MaxMind DB file, and [`MaxMindDBError::ResourceLimit`] if decoding
391    /// the metadata exceeds the configured limits.
392    pub fn from_source(buf: S) -> Result<Reader<S>, MaxMindDBError> {
393        Reader::from_source_with_options(buf, ReaderOptions::default())
394    }
395
396    /// Open a MaxMind DB database from anything that implements `AsRef<[u8]>`,
397    /// using `options` for resource limits.
398    ///
399    /// # Errors
400    ///
401    /// Returns [`MaxMindDBError::InvalidDatabaseError`] if the buffer is not a
402    /// valid MaxMind DB file (including unsupported format versions, invalid
403    /// record sizes or IP versions, and a search tree that does not fit in the
404    /// buffer), and [`MaxMindDBError::ResourceLimit`] if decoding the metadata
405    /// exceeds the configured limits.
406    pub fn from_source_with_options(buf: S, options: ReaderOptions) -> Result<Reader<S>, MaxMindDBError> {
407        let data_section_separator_size = 16;
408
409        let metadata_start = find_metadata_start(buf.as_ref())?;
410        let mut type_decoder = decoder::Decoder::with_limits(&buf.as_ref()[metadata_start..], 0, options.limits);
411        let metadata = Metadata::deserialize(&mut type_decoder)?;
412        validate_metadata(&metadata)?;
413
414        let search_tree_size = (metadata.node_count as usize)
415            .checked_mul(metadata.record_size as usize)
416            .map(|v| v / 4)
417            .ok_or_else(|| MaxMindDBError::InvalidDatabaseError("search tree size overflow".to_owned()))?;
418        if search_tree_size > metadata_start.saturating_sub(data_section_separator_size) {
419            return Err(MaxMindDBError::InvalidDatabaseError(
420                "the MaxMind DB file's search tree is larger than the file".to_owned(),
421            ));
422        }
423
424        let mut reader = Reader {
425            buf,
426            pointer_base: search_tree_size + data_section_separator_size,
427            metadata,
428            ipv4_start: 0,
429            limits: options.limits,
430        };
431        reader.ipv4_start = reader.find_ipv4_start()?;
432
433        Ok(reader)
434    }
435
436    /// Lookup the socket address in the opened MaxMind DB
437    ///
438    /// Example:
439    ///
440    /// ```
441    /// use maxminddb::geoip2;
442    /// use std::net::IpAddr;
443    /// use std::str::FromStr;
444    ///
445    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
446    ///
447    /// let ip: IpAddr = FromStr::from_str("89.160.20.128").unwrap();
448    /// let city: geoip2::City = reader.lookup(ip).unwrap();
449    /// print!("{:?}", city);
450    /// ```
451    ///
452    /// # Errors
453    ///
454    /// Returns [`MaxMindDBError::AddressNotFoundError`] if the address has no
455    /// record, [`MaxMindDBError::InvalidInputError`] for an IPv6 address in an
456    /// IPv4-only database, and a decoding error (including
457    /// [`MaxMindDBError::ResourceLimit`]) if the record cannot be decoded.
458    pub fn lookup<T>(&'de self, address: IpAddr) -> Result<T, MaxMindDBError>
459    where
460        T: Deserialize<'de>,
461    {
462        self.lookup_prefix(address).map(|(v, _)| v)
463    }
464
465    /// Lookup the socket address in the opened MaxMind DB
466    ///
467    /// Example:
468    ///
469    /// ```
470    /// use maxminddb::geoip2;
471    /// use std::net::IpAddr;
472    /// use std::str::FromStr;
473    ///
474    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
475    ///
476    /// let ip: IpAddr = "89.160.20.128".parse().unwrap();
477    /// let (city, prefix_len) = reader.lookup_prefix::<geoip2::City>(ip).unwrap();
478    /// print!("{:?}, prefix length: {}", city, prefix_len);
479    /// ```
480    ///
481    /// # Errors
482    ///
483    /// Returns the same errors as [`Reader::lookup`]. The prefix length of the
484    /// matched network is returned alongside the value.
485    pub fn lookup_prefix<T>(&'de self, address: IpAddr) -> Result<(T, usize), MaxMindDBError>
486    where
487        T: Deserialize<'de>,
488    {
489        if self.metadata.ip_version == 4 && address.is_ipv6() {
490            return Err(MaxMindDBError::InvalidInputError(
491                "IPv6 address cannot be looked up in an IPv4-only database".to_owned(),
492            ));
493        }
494        let mut ip_bytes_buf = [0_u8; 16];
495        let ip_bytes = ip_to_bytes(address, &mut ip_bytes_buf);
496        let (pointer, prefix_len) = self.find_address_in_tree(ip_bytes)?;
497        if pointer == 0 {
498            return Err(MaxMindDBError::AddressNotFoundError(address));
499        }
500
501        let rec = self.resolve_data_pointer(pointer)?;
502        let mut decoder = decoder::Decoder::with_limits(&self.buf.as_ref()[self.pointer_base..], rec, self.limits);
503
504        T::deserialize(&mut decoder).map(|v| (v, prefix_len))
505    }
506
507    /// Iterate over blocks of IP networks in the opened MaxMind DB
508    ///
509    /// Example:
510    ///
511    /// ```
512    /// use ipnetwork::IpNetwork;
513    /// use maxminddb::{geoip2, Within};
514    ///
515    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
516    ///
517    /// let ip_net = IpNetwork::V6("::/0".parse().unwrap());
518    /// let mut iter: Within<geoip2::City, _> = reader.within(ip_net).unwrap();
519    /// while let Some(next) = iter.next() {
520    ///     let item = next.unwrap();
521    ///     println!("ip_net={}, city={:?}", item.ip_net, item.info);
522    /// }
523    /// ```
524    ///
525    /// # Errors
526    ///
527    /// Returns [`MaxMindDBError::InvalidInputError`] if `cidr` is an IPv6 CIDR
528    /// and the database is IPv4-only. Problems found while walking the search
529    /// tree are yielded as `Err` items from the returned iterator rather than
530    /// panicking.
531    pub fn within<T>(&'de self, cidr: IpNetwork) -> Result<Within<'de, T, S>, MaxMindDBError>
532    where
533        T: Deserialize<'de>,
534    {
535        let ip_address = cidr.network();
536        if self.metadata.ip_version == 4 && ip_address.is_ipv6() {
537            return Err(MaxMindDBError::InvalidInputError(
538                "IPv6 CIDR cannot be iterated in an IPv4-only database".to_owned(),
539            ));
540        }
541        let prefix_len = cidr.prefix() as usize;
542        let byte_len = match ip_address {
543            IpAddr::V4(_) => 4,
544            IpAddr::V6(_) => 16,
545        };
546        let mut ip_bytes_buf = [0_u8; 16];
547        let ip_bytes = ip_to_bytes(ip_address, &mut ip_bytes_buf);
548        let bit_count = byte_len * 8;
549
550        let mut node = self.start_node(bit_count);
551        let node_count = self.metadata.node_count as usize;
552
553        let mut stack: Vec<WithinNode> = Vec::with_capacity(bit_count - prefix_len);
554
555        // Traverse down the tree to the level that matches the cidr mark
556        let mut i = 0_usize;
557        while i < prefix_len {
558            if node >= node_count {
559                // We've hit a dead end before we exhausted our prefix
560                break;
561            }
562            let bit = 1 & (ip_bytes[i >> 3] >> (7 - (i % 8))) as usize;
563            node = self.read_node(node, bit)?;
564            i += 1;
565        }
566
567        if node < node_count {
568            // Ok, now anything that's below node in the tree is "within", start with the node we
569            // traversed to as our to be processed stack.
570            let mut node_ip = [0_u8; 16];
571            node_ip[..byte_len].copy_from_slice(ip_bytes);
572            stack.push(WithinNode {
573                node,
574                ip_bytes: node_ip,
575                prefix_len,
576            });
577        }
578        // else the stack will be empty and we'll be returning an iterator that visits nothing,
579        // which makes sense.
580
581        let within: Within<T, S> = Within {
582            reader: self,
583            node_count,
584            byte_len,
585            stack,
586            phantom: PhantomData,
587        };
588
589        Ok(within)
590    }
591
592    fn find_address_in_tree(&self, ip_address: &[u8]) -> Result<(usize, usize), MaxMindDBError> {
593        // Dispatch on the record size once so the per-bit node reads below do
594        // not repeat the branch on the hot lookup path.
595        match self.metadata.record_size {
596            24 => self.find_address_in_tree_sized::<24>(ip_address),
597            28 => self.find_address_in_tree_sized::<28>(ip_address),
598            32 => self.find_address_in_tree_sized::<32>(ip_address),
599            s => Err(MaxMindDBError::InvalidDatabaseError(format!("unknown record size: {s:?}"))),
600        }
601    }
602
603    #[inline]
604    fn find_address_in_tree_sized<const RECORD_SIZE: usize>(
605        &self,
606        ip_address: &[u8],
607    ) -> Result<(usize, usize), MaxMindDBError> {
608        let bit_count = ip_address.len() * 8;
609        let mut node = self.start_node(bit_count);
610
611        let node_count = self.metadata.node_count as usize;
612        let mut prefix_len = bit_count;
613
614        for i in 0..bit_count {
615            if node >= node_count {
616                prefix_len = i;
617                break;
618            }
619            let bit = 1 & (ip_address[i >> 3] >> (7 - (i % 8)));
620
621            node = self.read_node_sized::<RECORD_SIZE>(node, bit as usize)?;
622        }
623        match node_count {
624            n if n == node => Ok((0, prefix_len)),
625            n if node > n => Ok((node, prefix_len)),
626            _ => {
627                // All bits consumed but we are at a node (exact prefix match).
628                // Read the left child to get the data pointer. For an exact prefix
629                // match the writer sets both children to the same data pointer.
630                let data_ptr = self.read_node_sized::<RECORD_SIZE>(node, 0)?;
631                if data_ptr > node_count {
632                    Ok((data_ptr, prefix_len))
633                } else {
634                    Err(MaxMindDBError::InvalidDatabaseError("invalid node in search tree".to_owned()))
635                }
636            }
637        }
638    }
639
640    fn start_node(&self, length: usize) -> usize {
641        if length == 128 { 0 } else { self.ipv4_start }
642    }
643
644    fn find_ipv4_start(&self) -> Result<usize, MaxMindDBError> {
645        if self.metadata.ip_version != 6 {
646            return Ok(0);
647        }
648
649        // We are looking up an IPv4 address in an IPv6 tree. Skip over the
650        // first 96 nodes.
651        let mut node: usize = 0_usize;
652        for _ in 0_u8..96 {
653            if node >= self.metadata.node_count as usize {
654                break;
655            }
656            node = self.read_node(node, 0)?;
657        }
658        Ok(node)
659    }
660
661    fn read_node(&self, node_number: usize, index: usize) -> Result<usize, MaxMindDBError> {
662        match self.metadata.record_size {
663            24 => self.read_node_sized::<24>(node_number, index),
664            28 => self.read_node_sized::<28>(node_number, index),
665            32 => self.read_node_sized::<32>(node_number, index),
666            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
667                "unknown record size: \
668                 {s:?}"
669            ))),
670        }
671    }
672
673    #[inline(always)]
674    fn read_node_sized<const RECORD_SIZE: usize>(
675        &self,
676        node_number: usize,
677        index: usize,
678    ) -> Result<usize, MaxMindDBError> {
679        fn oob() -> MaxMindDBError {
680            MaxMindDBError::InvalidDatabaseError("search tree node offset beyond end of file".to_owned())
681        }
682        let buf = self.buf.as_ref();
683        let base_offset = node_number
684            .checked_mul(RECORD_SIZE)
685            .map(|v| v / 4)
686            .ok_or_else(|| MaxMindDBError::InvalidDatabaseError("search tree node offset overflow".to_owned()))?;
687        let end_of = |start: usize, len: usize| start.checked_add(len).filter(|&end| end <= buf.len());
688
689        let val = if RECORD_SIZE == 24 {
690            let offset = base_offset.checked_add(index * 3).ok_or_else(oob)?;
691            let end = end_of(offset, 3).ok_or_else(oob)?;
692            to_usize(0, &buf[offset..end])
693        } else if RECORD_SIZE == 28 {
694            let mid = base_offset.checked_add(3).ok_or_else(oob)?;
695            if mid >= buf.len() {
696                return Err(oob());
697            }
698            let mut middle = buf[mid];
699            if index != 0 {
700                middle &= 0x0F
701            } else {
702                middle = (0xF0 & middle) >> 4
703            }
704            let offset = base_offset.checked_add(index * 4).ok_or_else(oob)?;
705            let end = end_of(offset, 3).ok_or_else(oob)?;
706            to_usize(middle, &buf[offset..end])
707        } else {
708            let offset = base_offset.checked_add(index * 4).ok_or_else(oob)?;
709            let end = end_of(offset, 4).ok_or_else(oob)?;
710            to_usize(0, &buf[offset..end])
711        };
712        Ok(val)
713    }
714
715    fn resolve_data_pointer(&self, pointer: usize) -> Result<usize, MaxMindDBError> {
716        let node_count = self.metadata.node_count as usize;
717        if pointer <= node_count {
718            return Err(MaxMindDBError::AddressNotFoundError(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))));
719        }
720        let separator_size = 16;
721        if pointer < node_count + separator_size {
722            return Err(MaxMindDBError::InvalidDatabaseError(format!(
723                "data pointer value {pointer} falls in the invalid range (node_count + 1 to node_count + 15)"
724            )));
725        }
726        let resolved = pointer - node_count - separator_size;
727
728        match self.pointer_base.checked_add(resolved) {
729            Some(end) if end < self.buf.as_ref().len() => Ok(resolved),
730            _ => Err(MaxMindDBError::InvalidDatabaseError(
731                "the MaxMind DB file's search tree \
732                 is corrupt"
733                    .to_owned(),
734            )),
735        }
736    }
737}
738
739// I haven't moved all patterns of this form to a generic function as
740// the FromPrimitive trait is unstable
741fn to_usize(base: u8, bytes: &[u8]) -> usize {
742    bytes.iter().fold(base as usize, |acc, &b| (acc << 8) | b as usize)
743}
744
745fn ip_to_bytes(address: IpAddr, buf: &mut [u8; 16]) -> &[u8] {
746    match address {
747        IpAddr::V4(a) => {
748            buf[..4].copy_from_slice(&a.octets());
749            &buf[..4]
750        }
751        IpAddr::V6(a) => {
752            buf.copy_from_slice(&a.octets());
753            &buf[..16]
754        }
755    }
756}
757
758#[allow(clippy::many_single_char_names)]
759fn bytes_and_prefix_to_net(bytes: &[u8], prefix: u8) -> Result<IpNetwork, MaxMindDBError> {
760    let (ip, pre) = match bytes.len() {
761        4 => (IpAddr::V4(Ipv4Addr::new(bytes[0], bytes[1], bytes[2], bytes[3])), prefix),
762        16 => {
763            // Only a v4-mapped prefix of at least 96 bits is treated as IPv4.
764            // Lower IPv6 networks (for example `::/64`) are kept as IPv6 rather
765            // than underflowing `prefix - 96`.
766            if prefix >= 96 && bytes[..12].iter().all(|&b| b == 0) {
767                // It's actually v4, but in v6 form, convert would be nice if ipnetwork had this
768                // logic.
769                (
770                    IpAddr::V4(Ipv4Addr::new(bytes[12], bytes[13], bytes[14], bytes[15])),
771                    prefix - 96,
772                )
773            } else {
774                let a = u16::from(bytes[0]) << 8 | u16::from(bytes[1]);
775                let b = u16::from(bytes[2]) << 8 | u16::from(bytes[3]);
776                let c = u16::from(bytes[4]) << 8 | u16::from(bytes[5]);
777                let d = u16::from(bytes[6]) << 8 | u16::from(bytes[7]);
778                let e = u16::from(bytes[8]) << 8 | u16::from(bytes[9]);
779                let f = u16::from(bytes[10]) << 8 | u16::from(bytes[11]);
780                let g = u16::from(bytes[12]) << 8 | u16::from(bytes[13]);
781                let h = u16::from(bytes[14]) << 8 | u16::from(bytes[15]);
782                (IpAddr::V6(Ipv6Addr::new(a, b, c, d, e, f, g, h)), prefix)
783            }
784        }
785        // This should never happen
786        _ => return Err(MaxMindDBError::InvalidNetworkError("invalid address".to_owned())),
787    };
788    IpNetwork::new(ip, pre).map_err(|e| MaxMindDBError::InvalidNetworkError(e.to_string()))
789}
790
791/// Rejects metadata that does not describe a readable database.
792fn validate_metadata(metadata: &Metadata) -> Result<(), MaxMindDBError> {
793    if metadata.binary_format_major_version != 2 {
794        return Err(MaxMindDBError::InvalidDatabaseError(format!(
795            "unsupported MaxMind DB binary format major version: {}",
796            metadata.binary_format_major_version
797        )));
798    }
799    if metadata.ip_version != 4 && metadata.ip_version != 6 {
800        return Err(MaxMindDBError::InvalidDatabaseError(format!(
801            "unsupported IP version in MaxMind DB metadata: {}",
802            metadata.ip_version
803        )));
804    }
805    if !matches!(metadata.record_size, 24 | 28 | 32) {
806        return Err(MaxMindDBError::InvalidDatabaseError(format!(
807            "unsupported record size in MaxMind DB metadata: {}",
808            metadata.record_size
809        )));
810    }
811    if metadata.node_count == 0 {
812        return Err(MaxMindDBError::InvalidDatabaseError(
813            "the MaxMind DB file has no search tree nodes".to_owned(),
814        ));
815    }
816    Ok(())
817}
818
819fn find_metadata_start(buf: &[u8]) -> Result<usize, MaxMindDBError> {
820    const METADATA_START_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
821
822    memchr::memmem::rfind(buf, METADATA_START_MARKER)
823        .map(|x| x + METADATA_START_MARKER.len())
824        .ok_or_else(|| MaxMindDBError::InvalidDatabaseError("Could not find MaxMind DB metadata in file.".to_owned()))
825}
826
827#[cfg(test)]
828mod reader_test;
829
830#[cfg(test)]
831mod robustness_test;
832
833#[cfg(test)]
834mod test_support;
835
836#[cfg(test)]
837mod alloc_test;
838
839#[cfg(test)]
840mod tests {
841    use std::net::{IpAddr, Ipv4Addr};
842
843    use super::MaxMindDBError;
844
845    #[test]
846    fn test_error_display() {
847        assert_eq!(
848            format!(
849                "{}",
850                MaxMindDBError::AddressNotFoundError(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
851            ),
852            "AddressNotFoundError: 127.0.0.1".to_owned(),
853        );
854        assert_eq!(
855            format!("{}", MaxMindDBError::InvalidDatabaseError("something went wrong".to_owned())),
856            "InvalidDatabaseError: something went wrong".to_owned(),
857        );
858        assert_eq!(
859            format!("{}", MaxMindDBError::IoError("something went wrong".to_owned())),
860            "IoError: something went wrong".to_owned(),
861        );
862        assert_eq!(
863            format!("{}", MaxMindDBError::MapError("something went wrong".to_owned())),
864            "MapError: something went wrong".to_owned(),
865        );
866        assert_eq!(
867            format!("{}", MaxMindDBError::DecodingError("something went wrong".to_owned())),
868            "DecodingError: something went wrong".to_owned(),
869        );
870    }
871}