#![deny(trivial_casts, trivial_numeric_casts, unused_import_braces)]
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::fmt::{self, Display, Formatter};
use std::io;
use std::marker::PhantomData;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::path::Path;
use ipnetwork::IpNetwork;
use serde::{de, Deserialize};
#[cfg(feature = "mmap")]
pub use memmap2::Mmap;
#[cfg(feature = "mmap")]
use memmap2::MmapOptions;
#[cfg(feature = "mmap")]
use std::fs::File;
#[derive(Debug, PartialEq, Eq)]
pub enum MaxMindDBError {
AddressNotFoundError(String),
InvalidDatabaseError(String),
IoError(String),
MapError(String),
DecodingError(String),
InvalidNetworkError(String),
}
impl From<io::Error> for MaxMindDBError {
fn from(err: io::Error) -> MaxMindDBError {
MaxMindDBError::IoError(err.to_string())
}
}
impl Display for MaxMindDBError {
fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), fmt::Error> {
match self {
MaxMindDBError::AddressNotFoundError(msg) => {
write!(fmt, "AddressNotFoundError: {msg}")?
}
MaxMindDBError::InvalidDatabaseError(msg) => {
write!(fmt, "InvalidDatabaseError: {msg}")?
}
MaxMindDBError::IoError(msg) => write!(fmt, "IoError: {msg}")?,
MaxMindDBError::MapError(msg) => write!(fmt, "MapError: {msg}")?,
MaxMindDBError::DecodingError(msg) => write!(fmt, "DecodingError: {msg}")?,
MaxMindDBError::InvalidNetworkError(msg) => write!(fmt, "InvalidNetworkError: {msg}")?,
}
Ok(())
}
}
impl std::error::Error for MaxMindDBError {}
impl de::Error for MaxMindDBError {
fn custom<T: Display>(msg: T) -> Self {
MaxMindDBError::DecodingError(format!("{msg}"))
}
}
#[derive(Deserialize, Debug)]
pub struct Metadata {
pub binary_format_major_version: u16,
pub binary_format_minor_version: u16,
pub build_epoch: u64,
pub database_type: String,
pub description: BTreeMap<String, String>,
pub ip_version: u16,
pub languages: Vec<String>,
pub node_count: u32,
pub record_size: u16,
}
#[derive(Debug)]
struct WithinNode {
node: usize,
ip_bytes: Vec<u8>,
prefix_len: usize,
}
#[derive(Debug)]
pub struct Within<'de, T: Deserialize<'de>, S: AsRef<[u8]>> {
reader: &'de Reader<S>,
node_count: usize,
stack: Vec<WithinNode>,
phantom: PhantomData<&'de T>,
}
#[derive(Debug)]
pub struct WithinItem<T> {
pub ip_net: IpNetwork,
pub info: T,
}
impl<'de, T: Deserialize<'de>, S: AsRef<[u8]>> Iterator for Within<'de, T, S> {
type Item = Result<WithinItem<T>, MaxMindDBError>;
fn next(&mut self) -> Option<Self::Item> {
while let Some(current) = self.stack.pop() {
let bit_count = current.ip_bytes.len() * 8;
if self.reader.ipv4_start != 0
&& current.node == self.reader.ipv4_start
&& bit_count == 128
&& current.ip_bytes[..12].iter().any(|&b| b != 0)
{
continue;
}
match current.node.cmp(&self.node_count) {
Ordering::Greater => {
let ip_net = match bytes_and_prefix_to_net(
¤t.ip_bytes,
current.prefix_len as u8,
) {
Ok(ip_net) => ip_net,
Err(e) => return Some(Err(e)),
};
let rec = match self.reader.resolve_data_pointer(current.node) {
Ok(rec) => rec,
Err(e) => return Some(Err(e)),
};
let mut decoder = decoder::Decoder::new(
&self.reader.buf.as_ref()[self.reader.pointer_base..],
rec,
);
return match T::deserialize(&mut decoder) {
Ok(info) => Some(Ok(WithinItem { ip_net, info })),
Err(e) => Some(Err(e)),
};
}
Ordering::Equal => {
}
Ordering::Less => {
let mut right_ip_bytes = current.ip_bytes.clone();
right_ip_bytes[current.prefix_len >> 3] |=
1 << ((bit_count - current.prefix_len - 1) % 8);
let node = match self.reader.read_node(current.node, 1) {
Ok(node) => node,
Err(e) => return Some(Err(e)),
};
self.stack.push(WithinNode {
node,
ip_bytes: right_ip_bytes,
prefix_len: current.prefix_len + 1,
});
let node = match self.reader.read_node(current.node, 0) {
Ok(node) => node,
Err(e) => return Some(Err(e)),
};
self.stack.push(WithinNode {
node,
ip_bytes: current.ip_bytes.clone(),
prefix_len: current.prefix_len + 1,
});
}
}
}
None
}
}
#[derive(Debug)]
pub struct Reader<S: AsRef<[u8]>> {
buf: S,
pub metadata: Metadata,
ipv4_start: usize,
pointer_base: usize,
}
#[cfg(feature = "mmap")]
impl<'de> Reader<Mmap> {
pub fn open_mmap<P: AsRef<Path>>(database: P) -> Result<Reader<Mmap>, MaxMindDBError> {
let file_read = File::open(database)?;
let mmap = unsafe { MmapOptions::new().map(&file_read) }?;
Reader::from_source(mmap)
}
}
impl Reader<Vec<u8>> {
pub fn open_readfile<P: AsRef<Path>>(database: P) -> Result<Reader<Vec<u8>>, MaxMindDBError> {
use std::fs;
let buf: Vec<u8> = fs::read(&database)?;
Reader::from_source(buf)
}
}
impl<'de, S: AsRef<[u8]>> Reader<S> {
pub fn from_source(buf: S) -> Result<Reader<S>, MaxMindDBError> {
let data_section_separator_size = 16;
let metadata_start = find_metadata_start(buf.as_ref())?;
let mut type_decoder = decoder::Decoder::new(&buf.as_ref()[metadata_start..], 0);
let metadata = Metadata::deserialize(&mut type_decoder)?;
let search_tree_size = (metadata.node_count as usize) * (metadata.record_size as usize) / 4;
let mut reader = Reader {
buf,
pointer_base: search_tree_size + data_section_separator_size,
metadata,
ipv4_start: 0,
};
reader.ipv4_start = reader.find_ipv4_start()?;
Ok(reader)
}
pub fn lookup<T>(&'de self, address: IpAddr) -> Result<T, MaxMindDBError>
where
T: Deserialize<'de>,
{
self.lookup_prefix(address).map(|(v, _)| v)
}
pub fn lookup_prefix<T>(&'de self, address: IpAddr) -> Result<(T, usize), MaxMindDBError>
where
T: Deserialize<'de>,
{
let ip_bytes = ip_to_bytes(address);
let (pointer, prefix_len) = self.find_address_in_tree(&ip_bytes)?;
if pointer == 0 {
return Err(MaxMindDBError::AddressNotFoundError(
"Address not found in database".to_owned(),
));
}
let rec = self.resolve_data_pointer(pointer)?;
let mut decoder = decoder::Decoder::new(&self.buf.as_ref()[self.pointer_base..], rec);
T::deserialize(&mut decoder).map(|v| (v, prefix_len))
}
pub fn within<T>(&'de self, cidr: IpNetwork) -> Result<Within<T, S>, MaxMindDBError>
where
T: Deserialize<'de>,
{
let ip_address = cidr.network();
let prefix_len = cidr.prefix() as usize;
let ip_bytes = ip_to_bytes(ip_address);
let bit_count = ip_bytes.len() * 8;
let mut node = self.start_node(bit_count);
let node_count = self.metadata.node_count as usize;
let mut stack: Vec<WithinNode> = Vec::with_capacity(bit_count - prefix_len);
let mut i = 0_usize;
while i < prefix_len {
let bit = 1 & (ip_bytes[i >> 3] >> (7 - (i % 8))) as usize;
node = self.read_node(node, bit)?;
if node >= node_count {
break;
}
i += 1;
}
if node < node_count {
stack.push(WithinNode {
node,
ip_bytes,
prefix_len,
});
}
let within: Within<T, S> = Within {
reader: self,
node_count,
stack,
phantom: PhantomData,
};
Ok(within)
}
fn find_address_in_tree(&self, ip_address: &[u8]) -> Result<(usize, usize), MaxMindDBError> {
let bit_count = ip_address.len() * 8;
let mut node = self.start_node(bit_count);
let node_count = self.metadata.node_count as usize;
let mut prefix_len = bit_count;
for i in 0..bit_count {
if node >= node_count {
prefix_len = i;
break;
}
let bit = 1 & (ip_address[i >> 3] >> (7 - (i % 8)));
node = self.read_node(node, bit as usize)?;
}
match node_count {
n if n == node => Ok((0, prefix_len)),
n if node > n => Ok((node, prefix_len)),
_ => Err(MaxMindDBError::InvalidDatabaseError(
"invalid node in search tree".to_owned(),
)),
}
}
fn start_node(&self, length: usize) -> usize {
if length == 128 {
0
} else {
self.ipv4_start
}
}
fn find_ipv4_start(&self) -> Result<usize, MaxMindDBError> {
if self.metadata.ip_version != 6 {
return Ok(0);
}
let mut node: usize = 0_usize;
for _ in 0_u8..96 {
if node >= self.metadata.node_count as usize {
break;
}
node = self.read_node(node, 0)?;
}
Ok(node)
}
fn read_node(&self, node_number: usize, index: usize) -> Result<usize, MaxMindDBError> {
let buf = self.buf.as_ref();
let base_offset = node_number * (self.metadata.record_size as usize) / 4;
let val = match self.metadata.record_size {
24 => {
let offset = base_offset + index * 3;
to_usize(0, &buf[offset..offset + 3])
}
28 => {
let mut middle = buf[base_offset + 3];
if index != 0 {
middle &= 0x0F
} else {
middle = (0xF0 & middle) >> 4
}
let offset = base_offset + index * 4;
to_usize(middle, &buf[offset..offset + 3])
}
32 => {
let offset = base_offset + index * 4;
to_usize(0, &buf[offset..offset + 4])
}
s => {
return Err(MaxMindDBError::InvalidDatabaseError(format!(
"unknown record size: \
{s:?}"
)))
}
};
Ok(val)
}
fn resolve_data_pointer(&self, pointer: usize) -> Result<usize, MaxMindDBError> {
let resolved = pointer - (self.metadata.node_count as usize) - 16;
if resolved > self.buf.as_ref().len() {
return Err(MaxMindDBError::InvalidDatabaseError(
"the MaxMind DB file's search tree \
is corrupt"
.to_owned(),
));
}
Ok(resolved)
}
}
fn to_usize(base: u8, bytes: &[u8]) -> usize {
bytes
.iter()
.fold(base as usize, |acc, &b| (acc << 8) | b as usize)
}
fn ip_to_bytes(address: IpAddr) -> Vec<u8> {
match address {
IpAddr::V4(a) => a.octets().to_vec(),
IpAddr::V6(a) => a.octets().to_vec(),
}
}
#[allow(clippy::many_single_char_names)]
fn bytes_and_prefix_to_net(bytes: &[u8], prefix: u8) -> Result<IpNetwork, MaxMindDBError> {
let (ip, pre) = match bytes.len() {
4 => (
IpAddr::V4(Ipv4Addr::new(bytes[0], bytes[1], bytes[2], bytes[3])),
prefix,
),
16 => {
if bytes[0] == 0
&& bytes[1] == 0
&& bytes[2] == 0
&& bytes[3] == 0
&& bytes[4] == 0
&& bytes[5] == 0
&& bytes[6] == 0
&& bytes[7] == 0
&& bytes[8] == 0
&& bytes[9] == 0
&& bytes[10] == 0
&& bytes[11] == 0
{
(
IpAddr::V4(Ipv4Addr::new(bytes[12], bytes[13], bytes[14], bytes[15])),
prefix - 96,
)
} else {
let a = u16::from(bytes[0]) << 8 | u16::from(bytes[1]);
let b = u16::from(bytes[2]) << 8 | u16::from(bytes[3]);
let c = u16::from(bytes[4]) << 8 | u16::from(bytes[5]);
let d = u16::from(bytes[6]) << 8 | u16::from(bytes[7]);
let e = u16::from(bytes[8]) << 8 | u16::from(bytes[9]);
let f = u16::from(bytes[10]) << 8 | u16::from(bytes[11]);
let g = u16::from(bytes[12]) << 8 | u16::from(bytes[13]);
let h = u16::from(bytes[14]) << 8 | u16::from(bytes[15]);
(IpAddr::V6(Ipv6Addr::new(a, b, c, d, e, f, g, h)), prefix)
}
}
_ => {
return Err(MaxMindDBError::InvalidNetworkError(
"invalid address".to_owned(),
))
}
};
IpNetwork::new(ip, pre).map_err(|e| MaxMindDBError::InvalidNetworkError(e.to_string()))
}
fn find_metadata_start(buf: &[u8]) -> Result<usize, MaxMindDBError> {
const METADATA_START_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
memchr::memmem::rfind(buf, METADATA_START_MARKER)
.map(|x| x + METADATA_START_MARKER.len())
.ok_or_else(|| {
MaxMindDBError::InvalidDatabaseError(
"Could not find MaxMind DB metadata in file.".to_owned(),
)
})
}
mod decoder;
pub mod geoip2;
#[cfg(test)]
mod reader_test;
#[cfg(test)]
mod tests {
use super::MaxMindDBError;
#[test]
fn test_error_display() {
assert_eq!(
format!(
"{}",
MaxMindDBError::AddressNotFoundError("something went wrong".to_owned())
),
"AddressNotFoundError: something went wrong".to_owned(),
);
assert_eq!(
format!(
"{}",
MaxMindDBError::InvalidDatabaseError("something went wrong".to_owned())
),
"InvalidDatabaseError: something went wrong".to_owned(),
);
assert_eq!(
format!(
"{}",
MaxMindDBError::IoError("something went wrong".to_owned())
),
"IoError: something went wrong".to_owned(),
);
assert_eq!(
format!(
"{}",
MaxMindDBError::MapError("something went wrong".to_owned())
),
"MapError: something went wrong".to_owned(),
);
assert_eq!(
format!(
"{}",
MaxMindDBError::DecodingError("something went wrong".to_owned())
),
"DecodingError: something went wrong".to_owned(),
);
}
}