1#![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
75pub 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 InvalidInputError(String),
99 ResourceLimit(String),
102}
103
104impl From<io::Error> for MaxMindDBError {
105 fn from(err: io::Error) -> MaxMindDBError {
106 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
127impl 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
163pub struct DecodeLimits {
164 pub max_depth: usize,
170 pub max_values: usize,
174 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
196pub struct ReaderOptions {
197 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 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 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 let ip_net =
243 match bytes_and_prefix_to_net(¤t.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 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 }
268 Ordering::Less => {
269 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 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 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#[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
317impl Reader<Vec<u8>> {
334 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 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 pub fn from_source(buf: S) -> Result<Reader<S>, MaxMindDBError> {
393 Reader::from_source_with_options(buf, ReaderOptions::default())
394 }
395
396 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 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 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 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 let mut i = 0_usize;
557 while i < prefix_len {
558 if node >= node_count {
559 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 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 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 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 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 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
739fn 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 if prefix >= 96 && bytes[..12].iter().all(|&b| b == 0) {
767 (
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 _ => return Err(MaxMindDBError::InvalidNetworkError("invalid address".to_owned())),
787 };
788 IpNetwork::new(ip, pre).map_err(|e| MaxMindDBError::InvalidNetworkError(e.to_string()))
789}
790
791fn 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}