use std::cmp::Ordering;
use std::net::IpAddr;
use crate::decoder;
use crate::error::MaxMindDbError;
use crate::reader::Reader;
use crate::result::{LookupResult, LookupSource, NetworkKind};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct WithinOptions {
include_aliased_networks: bool,
include_networks_without_data: bool,
skip_empty_values: bool,
}
impl WithinOptions {
#[must_use]
pub fn include_aliased_networks(mut self) -> Self {
self.include_aliased_networks = true;
self
}
#[must_use]
pub fn include_networks_without_data(mut self) -> Self {
self.include_networks_without_data = true;
self
}
#[must_use]
pub fn skip_empty_values(mut self) -> Self {
self.skip_empty_values = true;
self
}
}
#[derive(Debug)]
pub(crate) struct WithinNode {
pub(crate) node: usize,
pub(crate) ip_int: IpInt,
pub(crate) prefix_len: usize,
}
#[derive(Debug)]
pub struct Within<'de, S: AsRef<[u8]>> {
pub(crate) reader: &'de Reader<S>,
pub(crate) node_count: usize,
pub(crate) stack: Vec<WithinNode>,
pub(crate) options: WithinOptions,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum IpInt {
V4(u32),
V6(u128),
}
impl IpInt {
pub(crate) fn new(ip_addr: IpAddr) -> Self {
match ip_addr {
IpAddr::V4(v4) => IpInt::V4(v4.into()),
IpAddr::V6(v6) => IpInt::V6(v6.into()),
}
}
#[inline(always)]
pub(crate) fn get_bit(&self, index: usize) -> bool {
match self {
IpInt::V4(ip) => (ip >> (31 - index)) & 1 == 1,
IpInt::V6(ip) => (ip >> (127 - index)) & 1 == 1,
}
}
#[inline(always)]
pub(crate) fn set_bit(&mut self, index: usize) {
match self {
IpInt::V4(ip) => *ip |= 1 << (31 - index),
IpInt::V6(ip) => *ip |= 1 << (127 - index),
}
}
pub(crate) fn bit_count(&self) -> usize {
match self {
IpInt::V4(_) => 32,
IpInt::V6(_) => 128,
}
}
pub(crate) fn is_ipv4_in_ipv6(&self) -> bool {
match self {
IpInt::V4(_) => false,
IpInt::V6(ip) => *ip <= 0xFFFFFFFF,
}
}
}
impl<'de, S: AsRef<[u8]>> Iterator for Within<'de, S> {
type Item = Result<LookupResult<'de, S>, MaxMindDbError>;
fn next(&mut self) -> Option<Self::Item> {
while let Some(current) = self.stack.pop() {
let bit_count = current.ip_int.bit_count();
if !self.options.include_aliased_networks
&& self.reader.ipv4_start != 0
&& current.node == self.reader.ipv4_start
&& bit_count == 128
&& !current.ip_int.is_ipv4_in_ipv6()
{
continue;
}
match current.node.cmp(&self.node_count) {
Ordering::Greater => {
let ip_addr = ip_int_to_addr(¤t.ip_int);
let data_offset = match self.reader.resolve_data_pointer(current.node) {
Ok(offset) => offset,
Err(e) => return Some(Err(e)),
};
if self.options.skip_empty_values {
match self.is_empty_value_at(data_offset) {
Ok(true) => continue, Ok(false) => {} Err(e) => return Some(Err(e)),
}
}
let network_kind = match current.ip_int {
IpInt::V4(_) => NetworkKind::V4,
IpInt::V6(_)
if current.ip_int.is_ipv4_in_ipv6()
&& self.reader.has_ipv4_subtree()
&& current.prefix_len >= self.reader.ipv4_start_bit_depth =>
{
NetworkKind::V4InV6Subtree
}
IpInt::V6(_) => NetworkKind::V6,
};
return Some(Ok(LookupResult::new_found(
self.reader,
data_offset,
current.prefix_len as u8,
ip_addr,
LookupSource::Iter,
network_kind,
)));
}
Ordering::Equal => {
if self.options.include_networks_without_data {
let ip_addr = ip_int_to_addr(¤t.ip_int);
let network_kind = match current.ip_int {
IpInt::V4(_) => NetworkKind::V4,
IpInt::V6(_)
if current.ip_int.is_ipv4_in_ipv6()
&& self.reader.has_ipv4_subtree()
&& current.prefix_len >= self.reader.ipv4_start_bit_depth =>
{
NetworkKind::V4InV6Subtree
}
IpInt::V6(_) => NetworkKind::V6,
};
return Some(Ok(LookupResult::new_not_found(
self.reader,
current.prefix_len as u8,
ip_addr,
LookupSource::Iter,
network_kind,
)));
}
}
Ordering::Less => {
let mut right_ip_int = current.ip_int;
if current.prefix_len < bit_count {
right_ip_int.set_bit(current.prefix_len);
}
self.push_child(current.node, 1, right_ip_int, current.prefix_len + 1);
self.push_child(current.node, 0, current.ip_int, current.prefix_len + 1);
}
}
}
None
}
}
impl<'de, S: AsRef<[u8]>> Within<'de, S> {
fn push_child(
&mut self,
parent_node: usize,
direction: usize,
ip_int: IpInt,
prefix_len: usize,
) {
let node = self.reader.read_node(parent_node, direction);
self.stack.push(WithinNode {
node,
ip_int,
prefix_len,
});
}
fn is_empty_value_at(&self, data_offset: usize) -> Result<bool, MaxMindDbError> {
let buf = &self.reader.buf.as_ref()[self.reader.pointer_base..];
let mut dec = decoder::Decoder::new(buf, data_offset);
let (size, type_num) = dec.peek_type()?;
match type_num {
decoder::TYPE_MAP | decoder::TYPE_ARRAY => Ok(size == 0),
_ => Ok(false), }
}
}
pub(crate) fn ip_int_to_addr(ip_int: &IpInt) -> IpAddr {
match ip_int {
IpInt::V4(ip) => IpAddr::V4((*ip).into()),
IpInt::V6(ip) => {
if *ip <= 0xFFFFFFFF {
IpAddr::V4((*ip as u32).into())
} else {
IpAddr::V6((*ip).into())
}
}
}
}