1use std::cmp::Ordering;
4use std::net::{IpAddr, Ipv6Addr};
5
6use crate::decoder;
7use crate::error::MaxMindDbError;
8use crate::reader::Reader;
9use crate::result::{LookupResult, LookupSource, NetworkKind};
10
11#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
33pub struct WithinOptions {
34 include_aliased_networks: bool,
36 include_networks_without_data: bool,
38 skip_empty_values: bool,
40}
41
42impl WithinOptions {
43 #[must_use]
52 pub fn include_aliased_networks(mut self) -> Self {
53 self.include_aliased_networks = true;
54 self
55 }
56
57 #[must_use]
63 pub fn include_networks_without_data(mut self) -> Self {
64 self.include_networks_without_data = true;
65 self
66 }
67
68 #[must_use]
73 pub fn skip_empty_values(mut self) -> Self {
74 self.skip_empty_values = true;
75 self
76 }
77}
78
79#[derive(Debug)]
80pub(crate) struct WithinNode {
81 pub(crate) node: usize,
82 pub(crate) ip_int: IpInt,
83 pub(crate) prefix_len: usize,
84}
85
86#[derive(Debug)]
111pub struct Within<'de, S: AsRef<[u8]>> {
112 pub(crate) reader: &'de Reader<S>,
113 pub(crate) node_count: usize,
114 pub(crate) has_ipv4_subtree: bool,
115 pub(crate) stack: Vec<WithinNode>,
116 pub(crate) options: WithinOptions,
117}
118
119#[derive(Debug, Clone, Copy, PartialEq, Eq)]
120pub(crate) enum IpInt {
121 V4(u32),
122 V6(u128),
123}
124
125impl IpInt {
126 pub(crate) fn new(ip_addr: IpAddr) -> Self {
127 match ip_addr {
128 IpAddr::V4(v4) => IpInt::V4(v4.into()),
129 IpAddr::V6(v6) => IpInt::V6(v6.into()),
130 }
131 }
132
133 #[inline(always)]
134 pub(crate) fn get_bit(&self, index: usize) -> bool {
135 match self {
136 IpInt::V4(ip) => (ip >> (31 - index)) & 1 == 1,
137 IpInt::V6(ip) => (ip >> (127 - index)) & 1 == 1,
138 }
139 }
140
141 #[inline(always)]
142 pub(crate) fn set_bit(&mut self, index: usize) {
143 match self {
144 IpInt::V4(ip) => *ip |= 1 << (31 - index),
145 IpInt::V6(ip) => *ip |= 1 << (127 - index),
146 }
147 }
148
149 pub(crate) fn bit_count(&self) -> usize {
150 match self {
151 IpInt::V4(_) => 32,
152 IpInt::V6(_) => 128,
153 }
154 }
155
156 pub(crate) fn is_ipv4_in_ipv6(&self) -> bool {
157 match self {
158 IpInt::V4(_) => false,
159 IpInt::V6(ip) => *ip <= 0xFFFFFFFF,
160 }
161 }
162}
163
164impl<'de, S: AsRef<[u8]>> Iterator for Within<'de, S> {
165 type Item = Result<LookupResult<'de, S>, MaxMindDbError>;
166
167 fn next(&mut self) -> Option<Self::Item> {
168 while let Some(current) = self.stack.pop() {
169 let bit_count = current.ip_int.bit_count();
170
171 if !self.options.include_aliased_networks
173 && self.reader.ipv4_start != 0
174 && current.node == self.reader.ipv4_start
175 && bit_count == 128
176 && !current.ip_int.is_ipv4_in_ipv6()
177 {
178 continue;
179 }
180
181 match current.node.cmp(&self.node_count) {
182 Ordering::Greater => {
183 let data_offset = match self.reader.resolve_data_pointer(current.node) {
186 Ok(offset) => offset,
187 Err(e) => return Some(Err(e)),
188 };
189
190 if self.options.skip_empty_values {
192 match self.is_empty_value_at(data_offset) {
193 Ok(true) => continue, Ok(false) => {} Err(e) => return Some(Err(e)),
196 }
197 }
198
199 let network_kind = self.network_kind(¤t);
200 let ip_addr = network_ip_addr(network_kind, current.ip_int);
201
202 return Some(Ok(LookupResult::new_found(
203 self.reader,
204 data_offset,
205 current.prefix_len as u8,
206 ip_addr,
207 LookupSource::Iter,
208 network_kind,
209 )));
210 }
211 Ordering::Equal => {
212 if self.options.include_networks_without_data {
214 let network_kind = self.network_kind(¤t);
215 let ip_addr = network_ip_addr(network_kind, current.ip_int);
216 return Some(Ok(LookupResult::new_not_found(
217 self.reader,
218 current.prefix_len as u8,
219 ip_addr,
220 LookupSource::Iter,
221 network_kind,
222 )));
223 }
224 }
226 Ordering::Less => {
227 let mut right_ip_int = current.ip_int;
230
231 if current.prefix_len < bit_count {
232 right_ip_int.set_bit(current.prefix_len);
233 }
234
235 self.push_child(current.node, 1, right_ip_int, current.prefix_len + 1);
236 self.push_child(current.node, 0, current.ip_int, current.prefix_len + 1);
238 }
239 }
240 }
241 None
242 }
243}
244
245impl<'de, S: AsRef<[u8]>> Within<'de, S> {
246 #[inline(always)]
247 fn network_kind(&self, node: &WithinNode) -> NetworkKind {
248 match node.ip_int {
249 IpInt::V4(_) if self.reader.metadata().ip_version == 6 && !self.has_ipv4_subtree => {
250 NetworkKind::V6
251 }
252 IpInt::V4(_) => NetworkKind::V4,
253 IpInt::V6(_)
254 if node.ip_int.is_ipv4_in_ipv6()
255 && self.has_ipv4_subtree
256 && node.prefix_len >= self.reader.ipv4_start_bit_depth =>
257 {
258 NetworkKind::V4InV6Subtree
259 }
260 IpInt::V6(_) => NetworkKind::V6,
261 }
262 }
263
264 fn push_child(
265 &mut self,
266 parent_node: usize,
267 direction: usize,
268 ip_int: IpInt,
269 prefix_len: usize,
270 ) {
271 let node = self.reader.read_node(parent_node, direction);
272 self.stack.push(WithinNode {
273 node,
274 ip_int,
275 prefix_len,
276 });
277 }
278
279 fn is_empty_value_at(&self, data_offset: usize) -> Result<bool, MaxMindDbError> {
281 let buf = &self.reader.buf.as_ref()[self.reader.pointer_base..];
282 let mut dec =
283 decoder::Decoder::new_with_limit(buf, data_offset, self.reader.data_section_len);
284 let (size, type_num) = dec.peek_type()?;
285 match type_num {
286 decoder::TYPE_MAP | decoder::TYPE_ARRAY => Ok(size == 0),
287 _ => Ok(false), }
289 }
290}
291
292#[inline(always)]
293fn network_ip_addr(network_kind: NetworkKind, ip_int: IpInt) -> IpAddr {
294 let ip_addr = ip_int_to_addr(&ip_int);
295 match (network_kind, ip_int, ip_addr) {
296 (NetworkKind::V6, IpInt::V6(ip), IpAddr::V4(_)) => IpAddr::V6(ip.into()),
300 (NetworkKind::V6, IpInt::V4(_), IpAddr::V4(_)) => IpAddr::V6(Ipv6Addr::UNSPECIFIED),
303 (_, _, ip_addr) => ip_addr,
304 }
305}
306
307#[inline(always)]
309pub(crate) fn ip_int_to_addr(ip_int: &IpInt) -> IpAddr {
310 match ip_int {
311 IpInt::V4(ip) => IpAddr::V4((*ip).into()),
312 IpInt::V6(ip) => {
313 if *ip <= 0xFFFFFFFF {
315 IpAddr::V4((*ip as u32).into())
316 } else {
317 IpAddr::V6((*ip).into())
318 }
319 }
320 }
321}
322
323#[cfg(test)]
324mod tests {
325 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
326
327 use crate::result::NetworkKind;
328
329 use super::{network_ip_addr, IpInt};
330
331 #[test]
332 fn network_ip_addr_preserves_low_ipv6_bits() {
333 assert_eq!(
334 network_ip_addr(NetworkKind::V6, IpInt::V6(1)),
335 IpAddr::V6(Ipv6Addr::from(1_u128))
336 );
337 }
338
339 #[test]
340 fn network_ip_addr_formats_ipv4_subtree_records_as_ipv4() {
341 assert_eq!(
342 network_ip_addr(NetworkKind::V4InV6Subtree, IpInt::V6(1)),
343 IpAddr::V4(Ipv4Addr::from(1_u32))
344 );
345 }
346}