469 lines
15 KiB
Rust
469 lines
15 KiB
Rust
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(bytes: &[u8], offset: usize, base: usize) -> Result<(&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");
|