v0.28.0

src/geoip/mmdb.rs

@@ -0,0 +1,334 @@
+use anyhow::{bail, ensure, Context, Result};
+use std::{fs, net::IpAddr, path::Path};
+
+const METADATA_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
+const MAX_METADATA_BYTES: usize = 128 * 1024;
+const MAX_VALUES: usize = 65_536;
+const MAX_DEPTH: usize = 512;
+const MAX_STRING_BYTES: usize = 4096;
+
+#[derive(Debug)]
+pub(super) struct Database {
+    bytes: Vec<u8>,
+    node_count: u32,
+    record_size: u16,
+    ip_version: u16,
+    search_tree_size: usize,
+    data_section_start: usize,
+    metadata_offset: usize,
+    database_type: String,
+}
+
+#[derive(Clone, Copy)]
+struct Header {
+    kind: u8,
+    size: usize,
+    payload: usize,
+    next: usize,
+    pointer: Option<usize>,
+}
+
+#[derive(Default)]
+struct Budget {
+    values: usize,
+}
+
+impl Budget {
+    fn charge(&mut self) -> Result<()> {
+        self.values = self.values.saturating_add(1);
+        ensure!(self.values <= MAX_VALUES, "MMDB value traversal limit exceeded");
+        Ok(())
+    }
+}
+
+impl Database {
+    pub(super) fn load(path: &Path, max_bytes: u64) -> Result<Self> {
+        let metadata = fs::symlink_metadata(path)
+            .with_context(|| format!("open GeoIP MMDB {}", path.display()))?;
+        ensure!(metadata.is_file() && !metadata.file_type().is_symlink(), "GeoIP MMDB must be a regular file");
+        ensure!(metadata.len() <= max_bytes, "GeoIP MMDB exceeds configured maximum size");
+        let bytes = fs::read(path).with_context(|| format!("read GeoIP MMDB {}", path.display()))?;
+        let marker = metadata_marker(&bytes).context("MMDB metadata marker not found")?;
+        let metadata_offset = marker + METADATA_MARKER.len();
+        let base = metadata_offset;
+        let node_count = u32::try_from(read_uint_path(&bytes, metadata_offset, base, &["node_count"])?).context("MMDB node_count is too large")?;
+        let record_size = u16::try_from(read_uint_path(&bytes, metadata_offset, base, &["record_size"])?).context("MMDB record_size is too large")?;
+        let ip_version = u16::try_from(read_uint_path(&bytes, metadata_offset, base, &["ip_version"])?).context("MMDB ip_version is too large")?;
+        let major = read_uint_path(&bytes, metadata_offset, base, &["binary_format_major_version"])?;
+        let database_type = read_string_path(&bytes, metadata_offset, base, &["database_type"])?;
+        ensure!(node_count > 0, "MMDB node_count must be positive");
+        ensure!(matches!(record_size, 24 | 28 | 32), "unsupported MMDB record_size {record_size}");
+        ensure!(matches!(ip_version, 4 | 6), "invalid MMDB ip_version {ip_version}");
+        ensure!(major == 2, "unsupported MMDB major format version {major}");
+        let node_bytes = usize::from(record_size)
+            .checked_mul(2)
+            .and_then(|bits| bits.checked_div(8))
+            .context("MMDB node size overflow")?;
+        let search_tree_size = (node_count as usize)
+            .checked_mul(node_bytes)
+            .context("MMDB search tree size overflow")?;
+        let data_section_start = search_tree_size.checked_add(16).context("MMDB data section overflow")?;
+        ensure!(data_section_start <= marker, "MMDB search tree overlaps metadata");
+        ensure!(bytes.get(search_tree_size..data_section_start) == Some(&[0u8; 16][..]), "MMDB data-section separator is invalid");
+        Ok(Self {
+            bytes,
+            node_count,
+            record_size,
+            ip_version,
+            search_tree_size,
+            data_section_start,
+            metadata_offset,
+            database_type,
+        })
+    }
+
+    pub(super) fn database_type(&self) -> &str {
+        &self.database_type
+    }
+
+    pub(super) fn country(&self, ip: IpAddr) -> Option<String> {
+        let offset = self.lookup_offset(ip).ok().flatten()?;
+        for path in [
+            &["country", "iso_code"][..],
+            &["registered_country", "iso_code"][..],
+            &["country_code"][..],
+        ] {
+            if let Ok(code) = read_string_path(&self.bytes, offset, self.data_section_start, path) {
+                if let Some(code) = super::normalize_country_code(&code) {
+                    return Some(String::from_utf8_lossy(&code).into_owned());
+                }
+            }
+        }
+        let mut budget = Budget::default();
+        if let Ok(code) = read_string(&self.bytes, offset, self.data_section_start, 0, &mut budget) {
+            if let Some(code) = super::normalize_country_code(&code) {
+                return Some(String::from_utf8_lossy(&code).into_owned());
+            }
+        }
+        None
+    }
+
+    fn lookup_offset(&self, ip: IpAddr) -> Result<Option<usize>> {
+        let (bytes, bits): ([u8; 16], usize) = match (self.ip_version, ip) {
+            (4, IpAddr::V4(value)) => {
+                let mut out = [0u8; 16];
+                out[..4].copy_from_slice(&value.octets());
+                (out, 32)
+            }
+            (4, IpAddr::V6(_)) => return Ok(None),
+            (6, IpAddr::V4(value)) => {
+                let mut out = [0u8; 16];
+                out[12..].copy_from_slice(&value.octets());
+                (out, 128)
+            }
+            (6, IpAddr::V6(value)) => (value.octets(), 128),
+            _ => bail!("invalid MMDB IP version"),
+        };
+        let mut node = 0u32;
+        for bit_index in 0..bits {
+            if node >= self.node_count {
+                break;
+            }
+            let bit = (bytes[bit_index / 8] >> (7 - (bit_index % 8))) & 1;
+            node = self.read_node(node, bit as usize)?;
+        }
+        if node == self.node_count {
+            return Ok(None);
+        }
+        ensure!(node >= self.node_count.saturating_add(16), "invalid MMDB data pointer");
+        let offset = (node as usize)
+            .checked_sub(self.node_count as usize)
+            .and_then(|value| value.checked_add(self.search_tree_size))
+            .context("MMDB data pointer overflow")?;
+        ensure!(offset >= self.data_section_start && offset < self.metadata_offset - METADATA_MARKER.len(), "MMDB data pointer is out of bounds");
+        Ok(Some(offset))
+    }
+
+    fn read_node(&self, node: u32, side: usize) -> Result<u32> {
+        ensure!((node as usize) < self.node_count as usize, "MMDB node is out of range");
+        let base = match self.record_size {
+            24 => node as usize * 6,
+            28 => node as usize * 7,
+            32 => node as usize * 8,
+            _ => unreachable!(),
+        };
+        let value = match self.record_size {
+            24 => read_be(&self.bytes, base + side * 3, 3)? as u32,
+            28 => {
+                let middle = *self.bytes.get(base + 3).context("MMDB node is truncated")?;
+                if side == 0 {
+                    (((middle & 0xf0) as u32) << 20) | read_be(&self.bytes, base, 3)? as u32
+                } else {
+                    (((middle & 0x0f) as u32) << 24) | read_be(&self.bytes, base + 4, 3)? as u32
+                }
+            }
+            32 => read_be(&self.bytes, base + side * 4, 4)? as u32,
+            _ => unreachable!(),
+        };
+        Ok(value)
+    }
+}
+
+fn metadata_marker(bytes: &[u8]) -> Option<usize> {
+    let start = bytes.len().saturating_sub(MAX_METADATA_BYTES + METADATA_MARKER.len());
+    bytes[start..]
+        .windows(METADATA_MARKER.len())
+        .rposition(|window| window == METADATA_MARKER)
+        .map(|offset| start + offset)
+}
+
+fn read_uint_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<u64> {
+    let mut budget = Budget::default();
+    let value = find_path(bytes, offset, base, path, 0, &mut budget)?.context("MMDB metadata field missing")?;
+    read_uint(bytes, value, base, 0, &mut budget)
+}
+
+fn read_string_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<String> {
+    let mut budget = Budget::default();
+    let value = find_path(bytes, offset, base, path, 0, &mut budget)?.context("MMDB string field missing")?;
+    read_string(bytes, value, base, 0, &mut budget)
+}
+
+fn find_path(bytes: &[u8], offset: usize, base: usize, path: &[&str], depth: usize, budget: &mut Budget) -> Result<Option<usize>> {
+    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
+    budget.charge()?;
+    let header = header(bytes, offset, base)?;
+    if let Some(pointer) = header.pointer {
+        return find_path(bytes, pointer, base, path, depth + 1, budget);
+    }
+    if path.is_empty() {
+        return Ok(Some(offset));
+    }
+    ensure!(header.kind == 7, "MMDB selected value is not a map");
+    let mut cursor = header.payload;
+    for _ in 0..header.size {
+        let (key, next) = inline_string(bytes, cursor, base)?;
+        cursor = next;
+        let value_offset = cursor;
+        if key == path[0] {
+            if path.len() == 1 {
+                return Ok(Some(value_offset));
+            }
+            if let Some(found) = find_path(bytes, value_offset, base, &path[1..], depth + 1, budget)? {
+                return Ok(Some(found));
+            }
+        }
+        cursor = skip(bytes, value_offset, base, depth + 1, budget)?;
+    }
+    Ok(None)
+}
+
+fn read_string(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<String> {
+    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
+    budget.charge()?;
+    let h = header(bytes, offset, base)?;
+    if let Some(pointer) = h.pointer {
+        return read_string(bytes, pointer, base, depth + 1, budget);
+    }
+    ensure!(h.kind == 2, "MMDB value is not a string");
+    ensure!(h.size <= MAX_STRING_BYTES, "MMDB string exceeds the decode limit");
+    let data = bytes.get(h.payload..h.next).context("MMDB string is truncated")?;
+    Ok(std::str::from_utf8(data).context("MMDB string is not UTF-8")?.to_owned())
+}
+
+fn read_uint(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<u64> {
+    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
+    budget.charge()?;
+    let h = header(bytes, offset, base)?;
+    if let Some(pointer) = h.pointer {
+        return read_uint(bytes, pointer, base, depth + 1, budget);
+    }
+    ensure!(matches!(h.kind, 5 | 6 | 9), "MMDB value is not an unsigned integer");
+    ensure!(h.size <= 8, "MMDB integer is too wide");
+    read_be(bytes, h.payload, h.size)
+}
+
+fn inline_string<'a>(bytes: &'a [u8], offset: usize, base: usize) -> Result<(&'a str, usize)> {
+    let h = header(bytes, offset, base)?;
+    ensure!(h.pointer.is_none() && h.kind == 2, "MMDB map key is not an inline string");
+    let value = bytes.get(h.payload..h.next).context("MMDB map key is truncated")?;
+    Ok((std::str::from_utf8(value).context("MMDB map key is not UTF-8")?, h.next))
+}
+
+fn skip(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<usize> {
+    ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
+    budget.charge()?;
+    let h = header(bytes, offset, base)?;
+    if h.pointer.is_some() {
+        return Ok(h.next);
+    }
+    let mut cursor = h.payload;
+    match h.kind {
+        7 => {
+            for _ in 0..h.size {
+                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
+                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
+            }
+            Ok(cursor)
+        }
+        11 => {
+            for _ in 0..h.size {
+                cursor = skip(bytes, cursor, base, depth + 1, budget)?;
+            }
+            Ok(cursor)
+        }
+        14 => Ok(h.payload),
+        12 | 13 => bail!("unsupported deprecated MMDB data type {}", h.kind),
+        _ => Ok(h.next),
+    }
+}
+
+fn header(bytes: &[u8], offset: usize, base: usize) -> Result<Header> {
+    let control = *bytes.get(offset).context("MMDB value is truncated")?;
+    let primary = control >> 5;
+    let size_bits = (control & 0x1f) as usize;
+    if primary == 1 {
+        let class = (size_bits >> 3) & 0x03;
+        let low = size_bits & 0x07;
+        let count = class + 1;
+        let payload = offset + 1;
+        let raw = read_be(bytes, payload, count)? as usize;
+        let pointer = match class {
+            0 => (low << 8) | raw,
+            1 => ((low << 16) | raw).checked_add(2048).context("MMDB pointer overflow")?,
+            2 => ((low << 24) | raw).checked_add(526_336).context("MMDB pointer overflow")?,
+            3 => raw,
+            _ => unreachable!(),
+        };
+        let target = base.checked_add(pointer).context("MMDB pointer overflow")?;
+        ensure!(target < bytes.len(), "MMDB pointer is out of bounds");
+        return Ok(Header { kind: 1, size: 0, payload, next: payload + count, pointer: Some(target) });
+    }
+    let mut cursor = offset + 1;
+    let kind = if primary == 0 {
+        let extended = *bytes.get(cursor).context("MMDB extended type is truncated")?;
+        cursor += 1;
+        extended.checked_add(7).context("MMDB type overflow")?
+    } else {
+        primary
+    };
+    let size = match size_bits {
+        0..=28 => size_bits,
+        29 => 29 + *bytes.get(cursor).context("MMDB size is truncated")? as usize,
+        30 => 285 + read_be(bytes, cursor, 2)? as usize,
+        31 => 65_821 + read_be(bytes, cursor, 3)? as usize,
+        _ => unreachable!(),
+    };
+    cursor += match size_bits { 29 => 1, 30 => 2, 31 => 3, _ => 0 };
+    let next = match kind {
+        7 | 11 | 14 => cursor,
+        _ => cursor.checked_add(size).context("MMDB field size overflow")?,
+    };
+    ensure!(next <= bytes.len(), "MMDB field is out of bounds");
+    Ok(Header { kind, size, payload: cursor, next, pointer: None })
+}
+
+fn read_be(bytes: &[u8], offset: usize, count: usize) -> Result<u64> {
+    ensure!(count <= 8, "MMDB integer is too wide");
+    let end = offset.checked_add(count).context("MMDB read overflow")?;
+    let data = bytes.get(offset..end).context("MMDB value is truncated")?;
+    Ok(data.iter().fold(0u64, |value, byte| (value << 8) | u64::from(*byte)))
+}
+
+#[cfg(test)]
+include!("mmdb_tests.rs");