ipnet_trie/lib.rs
1//! This crate provides storage and retrieval of IPv4 and IPv6 network prefixes.
2//!
3//! Currently, it uses `ipnet` crate, that provides IP network data structure and
4//! `prefix-trie` as backend, that provides fast lookup times, and a small memory footprint.
5//! Backend can be changed in future releases.
6//!
7//! ## Examples
8//!
9//! ```rust
10//! use std::net::{IpAddr, Ipv6Addr};
11//! use ipnet::{IpNet, Ipv6Net};
12//! use ipnet_trie::IpnetTrie;
13//!
14//! let mut table = IpnetTrie::new();
15//! let network = IpNet::from(Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap());
16//! let ip_address = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0x1);
17//!
18//! assert_eq!(table.insert(network, "foo"), None);
19//! // Get value for network from table
20//! assert_eq!(table.longest_match(&IpNet::from(ip_address.to_canonical())), Some((network, &"foo")));
21//! ```
22
23#![doc(
24 html_logo_url = "https://raw.githubusercontent.com/bgpkit/assets/main/logos/icon-transparent.png",
25 html_favicon_url = "https://raw.githubusercontent.com/bgpkit/assets/main/logos/favicon.ico"
26)]
27
28#[cfg(feature = "export")]
29mod export;
30
31use ipnet::{IpNet, Ipv4Net, Ipv6Net};
32use prefix_trie::PrefixMap;
33
34/// The number of unique IP addresses covered by prefixes in the trie.
35///
36/// Returned by [`IpnetTrie::ip_count`].
37///
38/// `ipv4` uses `u64` and can represent the entire IPv4 address space (2³²).
39/// `ipv6` is `Option<u128>`: `Some(count)` for values up to 2¹²⁸−1, and `None`
40/// when the entire IPv6 address space (2¹²⁸ addresses) is covered — a value too
41/// large to store in a `u128`.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub struct IpCount {
44 /// Number of unique IPv4 addresses covered (maximum 2³²).
45 pub ipv4: u64,
46 /// Number of unique IPv6 addresses covered, or `None` if the entire
47 /// IPv6 space (2¹²⁸) is covered.
48 pub ipv6: Option<u128>,
49}
50
51/// Table holding IPv4 and IPv6 network prefixes with value.
52#[derive(Default)]
53pub struct IpnetTrie<T> {
54 ipv4: PrefixMap<Ipv4Net, T>,
55 ipv6: PrefixMap<Ipv6Net, T>,
56}
57
58impl<T> Clone for IpnetTrie<T>
59where
60 T: Clone,
61{
62 fn clone(&self) -> Self {
63 Self {
64 ipv4: self.ipv4.clone(),
65 ipv6: self.ipv6.clone(),
66 }
67 }
68}
69
70/// Splits a source IP network into multiple IP networks based on a target IP network.
71///
72/// It makes sure the returning IP networks are non-overlapping and does not include the target prefix.
73///
74/// # Arguments
75///
76/// * `source` - The source IP network to split.
77/// * `target` - The target IP network used for splitting.
78///
79/// # Returns
80///
81/// A vector containing the split IP networks.
82/// ```
83/// use std::net::{IpAddr, Ipv4Addr};
84/// use ipnet::{IpNet, Ipv4Net};
85/// use ipnet_trie::exclude_prefix;
86///
87/// let source: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 22).unwrap());
88/// let target: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap());
89/// let split_networks = exclude_prefix(source, target);
90/// assert_eq!(split_networks.len(), 2);
91///
92/// let source: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap());
93/// let target: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap());
94/// let split_networks = exclude_prefix(source, target);
95/// assert_eq!(split_networks.len(), 0);
96///
97/// let source: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 23).unwrap());
98/// let target: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap());
99/// let split_networks = exclude_prefix(source, target);
100/// assert_eq!(split_networks.len(), 1);
101/// assert_ne!(split_networks[0], source);
102///
103/// let source: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap());
104/// let target: IpNet = IpNet::V4(Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 23).unwrap());
105/// let split_networks = exclude_prefix(source, target);
106/// assert_eq!(split_networks[0], source);
107/// assert_eq!(split_networks.len(), 1);
108/// ```
109pub fn exclude_prefix(source: IpNet, target: IpNet) -> Vec<IpNet> {
110 let new_prefixes = match source.contains(&target) {
111 true => {
112 // target_prefix is covered by sub_prefix, split it!
113 source
114 .subnets(target.prefix_len())
115 .unwrap()
116 .into_iter()
117 .filter(|p| *p != target)
118 .collect()
119 }
120 false => {
121 // target_prefix is not covered by sub_prefix, keep it as is
122 vec![source]
123 }
124 };
125
126 IpNet::aggregate(&new_prefixes)
127}
128
129impl<T> IpnetTrie<T> {
130 /// Constructs a new, empty `IpnetTrie<T>`.
131 pub fn new() -> Self {
132 Self {
133 ipv4: PrefixMap::new(),
134 ipv6: PrefixMap::new(),
135 }
136 }
137
138 /// Returns the number of elements in the table. First value is number of IPv4 networks and second is number of IPv6 networks.
139 pub fn len(&self) -> (usize, usize) {
140 (self.ipv4.iter().count(), self.ipv6.iter().count())
141 }
142
143 /// Returns `true` if table is empty.
144 pub fn is_empty(&self) -> bool {
145 self.ipv4.iter().next().is_none() && self.ipv6.iter().next().is_none()
146 }
147
148 /// Insert a value for the `IpNet`. If prefix existed previously, the old value is returned.
149 ///
150 /// # Examples
151 ///
152 /// ```
153 /// use ipnet_trie::IpnetTrie;
154 /// use ipnet::Ipv6Net;
155 /// use std::net::Ipv6Addr;
156 ///
157 /// let mut table = IpnetTrie::new();
158 /// let network = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
159 ///
160 /// assert_eq!(table.insert(network, "foo"), None);
161 /// // Insert duplicate
162 /// assert_eq!(table.insert(network, "bar"), Some("foo"));
163 /// // Value is replaced
164 /// assert_eq!(table.insert(network, "null"), Some("bar"));
165 /// ```
166 pub fn insert<N: Into<IpNet>>(&mut self, network: N, data: T) -> Option<T> {
167 match network.into() {
168 IpNet::V4(ipv4_network) => self.ipv4.insert(ipv4_network, data),
169 IpNet::V6(ipv6_network) => self.ipv6.insert(ipv6_network, data),
170 }
171 }
172
173 /// Remove a `IpNet` from table. If prefix existed, the value is returned.
174 ///
175 /// # Examples
176 ///
177 /// ```
178 /// use ipnet_trie::IpnetTrie;
179 /// use std::net::Ipv6Addr;
180 /// use ipnet::Ipv6Net;
181 ///
182 /// let mut table = IpnetTrie::new();
183 /// let network = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
184 ///
185 /// assert_eq!(table.insert(network, "foo"), None);
186 /// // Remove network from table
187 /// assert_eq!(table.remove(network), Some("foo"));
188 /// // Network is removed
189 /// assert_eq!(table.exact_match(network), None);
190 /// ```
191 pub fn remove<N: Into<IpNet>>(&mut self, network: N) -> Option<T> {
192 match network.into() {
193 IpNet::V4(ipv4_network) => self.ipv4.remove(&ipv4_network),
194 IpNet::V6(ipv6_network) => self.ipv6.remove(&ipv6_network),
195 }
196 }
197
198 /// Get pointer to value from table based on exact network match.
199 /// If network is not in table, `None` is returned.
200 ///
201 /// # Examples
202 ///
203 /// ```
204 /// use ipnet_trie::IpnetTrie;
205 /// use std::net::Ipv6Addr;
206 /// use ipnet::Ipv6Net;
207 ///
208 /// let mut table = IpnetTrie::new();
209 /// let network_a = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
210 /// let network_b = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 128).unwrap();
211 ///
212 /// assert_eq!(table.insert(network_a, "foo"), None);
213 /// // Get value for network from trie
214 /// assert_eq!(table.exact_match(network_a), Some(&"foo"));
215 /// // Network B doesn not exist in trie
216 /// assert_eq!(table.exact_match(network_b), None);
217 /// ```
218 pub fn exact_match<N: Into<IpNet>>(&self, network: N) -> Option<&T> {
219 match network.into() {
220 IpNet::V4(ipv4_network) => self.ipv4.get(&ipv4_network),
221 IpNet::V6(ipv6_network) => self.ipv6.get(&ipv6_network),
222 }
223 }
224
225 /// Get mutable pointer to value from table based on exact network match.
226 /// If network is not in table, `None` is returned.
227 ///
228 /// # Examples
229 ///
230 /// ```
231 /// use ipnet_trie::IpnetTrie;
232 /// use std::net::Ipv6Addr;
233 /// use ipnet::Ipv6Net;
234 ///
235 /// let mut table = IpnetTrie::new();
236 /// let network_a = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
237 /// let network_b = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 128).unwrap();
238 ///
239 /// assert_eq!(table.insert(network_a, "foo"), None);
240 /// // Get value for network from trie
241 /// assert_eq!(table.exact_match_mut(network_a), Some(&mut "foo"));
242 /// // Network B does not exist in trie
243 /// assert_eq!(table.exact_match(network_b), None);
244 /// ```
245 pub fn exact_match_mut<N: Into<IpNet>>(&mut self, network: N) -> Option<&mut T> {
246 match network.into() {
247 IpNet::V4(ipv4_network) => self.ipv4.get_mut(&ipv4_network),
248 IpNet::V6(ipv6_network) => self.ipv6.get_mut(&ipv6_network),
249 }
250 }
251
252 /// Find most specific IP network in table that contains given IP address. If no network in table contains
253 /// given IP address, `None` is returned.
254 ///
255 /// # Examples
256 ///
257 /// ```
258 /// use ipnet_trie::IpnetTrie;
259 /// use ipnet::{IpNet, Ipv6Net};
260 /// use std::net::{IpAddr, Ipv6Addr};
261 ///
262 /// let mut table = IpnetTrie::new();
263 /// let network = IpNet::new(IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0)), 64).unwrap();
264 /// let ip_address = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0x1);
265 ///
266 /// assert_eq!(table.insert(network, "foo"), None);
267 /// // Get value for network from table
268 /// assert_eq!(table.longest_match(&IpNet::from(ip_address.to_canonical())), Some((network, &"foo")));
269 /// ```
270 pub fn longest_match(&self, ipnet: &IpNet) -> Option<(IpNet, &T)> {
271 match ipnet {
272 IpNet::V4(net) => self
273 .longest_match_ipv4(net)
274 .map(|(net, data)| (IpNet::V4(*net), data)),
275 IpNet::V6(net) => self
276 .longest_match_ipv6(net)
277 .map(|(net, data)| (IpNet::V6(*net), data)),
278 }
279 }
280
281 /// Find most specific IP network in table that contains given IP address. If no network in table contains
282 /// given IP address, `None` is returned.
283 ///
284 /// # Examples
285 ///
286 /// ```
287 /// use ipnet_trie::IpnetTrie;
288 /// use ipnet::{IpNet, Ipv6Net};
289 /// use std::net::{IpAddr, Ipv6Addr};
290 ///
291 /// let mut table = IpnetTrie::new();
292 /// let network = IpNet::new(IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0)), 64).unwrap();
293 /// let ip_address = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0x1);
294 ///
295 /// assert_eq!(table.insert(network, "foo"), None);
296 /// // Get value for network from table
297 /// assert_eq!(table.longest_match_mut(&IpNet::from(ip_address.to_canonical())), Some((network, &mut "foo")));
298 /// ```
299 pub fn longest_match_mut(&mut self, ipnet: &IpNet) -> Option<(IpNet, &mut T)> {
300 match ipnet {
301 IpNet::V4(net) => self
302 .longest_match_ipv4_mut(net)
303 .map(|(net, data)| (IpNet::V4(*net), data)),
304 IpNet::V6(net) => self
305 .longest_match_ipv6_mut(net)
306 .map(|(net, data)| (IpNet::V6(*net), data)),
307 }
308 }
309
310 /// Specific version of `longest_match` for IPv4 address.
311 #[inline]
312 pub fn longest_match_ipv4(&self, net: &Ipv4Net) -> Option<(&Ipv4Net, &T)> {
313 self.ipv4.get_lpm(net)
314 }
315
316 /// Specific version of `longest_match` for IPv6 address.
317 #[inline]
318 pub fn longest_match_ipv6(&self, net: &Ipv6Net) -> Option<(&Ipv6Net, &T)> {
319 self.ipv6.get_lpm(net)
320 }
321
322 /// Specific version of `longest_match` for IPv4 address.
323 #[inline]
324 pub fn longest_match_ipv4_mut(&mut self, net: &Ipv4Net) -> Option<(&Ipv4Net, &mut T)> {
325 self.ipv4.get_lpm_mut(net)
326 }
327
328 /// Specific version of `longest_match` for IPv6 address.
329 #[inline]
330 pub fn longest_match_ipv6_mut(&mut self, net: &Ipv6Net) -> Option<(&Ipv6Net, &mut T)> {
331 self.ipv6.get_lpm_mut(net)
332 }
333
334 /// Find all IP networks in table that contains given IP address.
335 /// Returns iterator of `IpNet` and reference to value.
336 ///
337 /// # Examples
338 ///
339 /// ```
340 /// use ipnet_trie::IpnetTrie;
341 /// use ipnet::{IpNet, Ipv6Net};
342 /// use std::net::{IpAddr, Ipv6Addr};
343 ///
344 /// let mut table = IpnetTrie::new();
345 /// let network = IpNet::new(IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0)), 64).unwrap();
346 /// let ip_address = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0x1);
347 ///
348 /// assert_eq!(table.insert(network, "foo"), None);
349 /// // Get value for network from table
350 /// assert_eq!(table.matches(&IpNet::from(ip_address.to_canonical())).len(), 1);
351 /// ```
352 pub fn matches(&self, ipnet: &IpNet) -> Vec<(IpNet, &T)> {
353 match ipnet {
354 IpNet::V4(net) => self
355 .matches_ipv4(net)
356 .into_iter()
357 .map(|(net, data)| (IpNet::V4(*net), data))
358 .collect(),
359 IpNet::V6(net) => self
360 .matches_ipv6(net)
361 .into_iter()
362 .map(|(net, data)| (IpNet::V6(*net), data))
363 .collect(),
364 }
365 }
366
367 /// Specific version of `matches` for IPv4 address.
368 pub fn matches_ipv4(&self, net: &Ipv4Net) -> Vec<(&Ipv4Net, &T)> {
369 match self.ipv4.get_spm(net) {
370 None => vec![],
371 Some((shortest, _)) => self.ipv4.children(*shortest).collect(),
372 }
373 }
374
375 /// Specific version of `matches` for IPv6 address.
376 pub fn matches_ipv6(&self, net: &Ipv6Net) -> Vec<(&Ipv6Net, &T)> {
377 match self.ipv6.get_spm(net) {
378 None => vec![],
379 Some((shortest, _)) => self.ipv6.children(*shortest).collect(),
380 }
381 }
382
383 /// Iterator for all networks in table, first are iterated IPv4 and then IPv6 networks. Order is not guaranteed.
384 ///
385 /// # Examples
386 ///
387 /// ```
388 /// use ipnet_trie::IpnetTrie;
389 /// use ipnet::{IpNet, Ipv4Net, Ipv6Net};
390 /// use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
391 ///
392 /// let mut table: IpnetTrie<&str> = IpnetTrie::new();
393 /// let network_a = Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap();
394 /// assert_eq!(table.insert(network_a, "foo"), None);
395 /// let network_b = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
396 /// assert_eq!(table.insert(network_b, "foo"), None);
397 ///
398 /// let mut iterator = table.iter();
399 /// assert_eq!(iterator.next(), Some((IpNet::V4(network_a), &"foo")));
400 /// assert_eq!(iterator.next(), Some((IpNet::V6(network_b), &"foo")));
401 /// assert_eq!(iterator.next(), None);
402 /// ```
403 pub fn iter(&self) -> impl Iterator<Item = (IpNet, &T)> {
404 self.iter_ipv4()
405 .map(|(network, data)| (IpNet::V4(*network), data))
406 .chain(
407 self.iter_ipv6()
408 .map(|(network, data)| (IpNet::V6(*network), data)),
409 )
410 }
411
412 /// Mutable iterator for all networks in table, first are iterated IPv4 and then IPv6 networks. Order is not guaranteed.
413 ///
414 /// # Examples
415 ///
416 /// ```
417 /// use ipnet_trie::IpnetTrie;
418 /// use ipnet::{IpNet, Ipv4Net, Ipv6Net};
419 /// use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
420 ///
421 /// let mut table: IpnetTrie<&str> = IpnetTrie::new();
422 /// let network_a = Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap();
423 /// assert_eq!(table.insert(network_a, "foo"), None);
424 /// let network_b = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
425 /// assert_eq!(table.insert(network_b, "foo"), None);
426 ///
427 /// let mut iterator = table.iter_mut();
428 /// for (network, value) in iterator {
429 /// *value = "bar";
430 /// }
431 ///
432 /// assert_eq!(table.exact_match(network_a), Some(&"bar"));
433 /// assert_eq!(table.exact_match(network_b), Some(&"bar"));
434 /// ```
435 pub fn iter_mut(&mut self) -> impl Iterator<Item = (IpNet, &mut T)> {
436 self.ipv4
437 .iter_mut()
438 .map(|(net, data)| (IpNet::from(*net), data))
439 .chain(
440 self.ipv6
441 .iter_mut()
442 .map(|(net, data)| (IpNet::from(*net), data)),
443 )
444 }
445
446 /// Iterator for all IPv4 networks in table. Order is not guaranteed.
447 pub fn iter_ipv4(&self) -> impl Iterator<Item = (&Ipv4Net, &T)> {
448 self.ipv4.iter()
449 }
450
451 /// Iterator for all IPv6 networks in table. Order is not guaranteed.
452 pub fn iter_ipv6(&self) -> impl Iterator<Item = (&Ipv6Net, &T)> {
453 self.ipv6.iter()
454 }
455
456 /// Retains only the elements specified by the predicate.
457 ///
458 /// In other words, remove all pairs `(k, v)` such that `f(ip_network, &mut v)` returns `false`.
459 ///
460 /// # Examples
461 ///
462 /// ```
463 /// use ipnet_trie::IpnetTrie;
464 /// use ipnet::{IpNet, Ipv4Net, Ipv6Net};
465 /// use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
466 ///
467 /// let mut table: IpnetTrie<&str> = IpnetTrie::new();
468 /// let network_a = Ipv4Net::new(Ipv4Addr::new(192, 168, 0, 0), 24).unwrap();
469 /// assert_eq!(table.insert(network_a, "foo"), None);
470 /// let network_b = Ipv6Net::new(Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0), 64).unwrap();
471 /// assert_eq!(table.insert(network_b, "foo"), None);
472 ///
473 /// // Keep just IPv4 networks
474 /// table.retain(|network, _| network.network().is_ipv4());
475 ///
476 /// assert_eq!(table.exact_match(network_a), Some(&"foo"));
477 /// assert_eq!(table.exact_match(network_b), None);
478 /// ```
479 pub fn retain<F>(&mut self, mut f: F)
480 where
481 F: FnMut(IpNet, &mut T) -> bool,
482 {
483 let mut to_delete = vec![];
484 for (network, data) in self.iter_mut() {
485 if !f(network, data) {
486 to_delete.push(network);
487 }
488 }
489 for network in to_delete {
490 self.remove(network);
491 }
492 }
493
494 /// Count the number of unique IPv4 and IPv6 addresses in the trie.
495 ///
496 /// Returns an [`IpCount`] struct. The `ipv4` field uses `u64` and can represent
497 /// the entire IPv4 address space (2³²). The `ipv6` field is `Option<u128>`:
498 /// `Some(count)` for any value up to 2¹²⁸−1, and `None` when the entire IPv6
499 /// address space (2¹²⁸ addresses) is covered, since that value cannot be
500 /// represented in a `u128`.
501 ///
502 /// ```rust
503 /// use std::str::FromStr;
504 /// use ipnet::{Ipv4Net, Ipv6Net};
505 /// use ipnet_trie::{IpnetTrie, IpCount};
506 ///
507 /// let mut table = IpnetTrie::new();
508 /// table.insert(Ipv4Net::from_str("192.0.2.129/25").unwrap(), 1);
509 /// table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
510 /// table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
511 /// table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
512 /// assert_eq!(table.ip_count(), IpCount { ipv4: 256, ipv6: Some(0) });
513 ///
514 /// table.insert(Ipv4Net::from_str("198.51.100.0/25").unwrap(), 1);
515 /// table.insert(Ipv4Net::from_str("198.51.100.64/26").unwrap(), 1);
516 /// assert_eq!(table.ip_count(), IpCount { ipv4: 384, ipv6: Some(0) });
517 ///
518 /// table.insert(Ipv4Net::from_str("198.51.100.65/26").unwrap(), 1);
519 /// assert_eq!(table.ip_count(), IpCount { ipv4: 384, ipv6: Some(0) });
520 ///
521 /// table.insert(Ipv6Net::from_str("2001:DB80::/48").unwrap(), 1);
522 /// assert_eq!(table.ip_count(), IpCount { ipv4: 384, ipv6: Some(2_u128.pow(80)) });
523 /// table.insert(Ipv6Net::from_str("2001:DB80::/49").unwrap(), 1);
524 /// assert_eq!(table.ip_count(), IpCount { ipv4: 384, ipv6: Some(2_u128.pow(80)) });
525 /// table.insert(Ipv6Net::from_str("2001:DB81::/48").unwrap(), 1);
526 /// assert_eq!(table.ip_count(), IpCount { ipv4: 384, ipv6: Some(2_u128.pow(81)) });
527 ///
528 /// // Full IPv4 space (0.0.0.0/0) — 2^32 addresses, fits in u64
529 /// let mut v4_full = IpnetTrie::new();
530 /// v4_full.insert(Ipv4Net::from_str("0.0.0.0/0").unwrap(), 1);
531 /// assert_eq!(v4_full.ip_count(), IpCount { ipv4: 1u64 << 32, ipv6: Some(0) });
532 ///
533 /// // Full IPv6 space (::/0) — 2^128 addresses, cannot fit in u128
534 /// let mut v6_full = IpnetTrie::new();
535 /// v6_full.insert(Ipv6Net::from_str("::/0").unwrap(), 1);
536 /// assert_eq!(v6_full.ip_count(), IpCount { ipv4: 0, ipv6: None });
537 /// ```
538 pub fn ip_count(&self) -> IpCount {
539 let (root_ipv4_prefixes, root_ipv6_prefixes) = self.get_aggregated_prefixes();
540
541 let mut ipv4_space: u64 = 0;
542 for prefix in root_ipv4_prefixes {
543 ipv4_space += 1u64 << (32 - prefix.prefix_len() as u32);
544 }
545
546 let mut ipv6_space: u128 = 0;
547 let mut ipv6_full = false;
548 for prefix in root_ipv6_prefixes {
549 let host_bits = 128 - prefix.prefix_len() as u32;
550 // 2^128 cannot be represented in u128; checked_shl returns None when
551 // host_bits >= 128, and checked_add catches any intermediate overflow.
552 match 1u128
553 .checked_shl(host_bits)
554 .and_then(|c| ipv6_space.checked_add(c))
555 {
556 Some(v) => ipv6_space = v,
557 None => {
558 ipv6_full = true;
559 break;
560 }
561 }
562 }
563
564 IpCount {
565 ipv4: ipv4_space,
566 ipv6: if ipv6_full { None } else { Some(ipv6_space) },
567 }
568 }
569
570 /// Retrieves the aggregated prefixes for both IPv4 and IPv6 from the given data.
571 ///
572 /// # Returns
573 ///
574 /// A tuple containing two vectors. The first vector contains the aggregated IPv4 prefixes,
575 /// and the second vector contains the aggregated IPv6 prefixes.
576 pub fn get_aggregated_prefixes(&self) -> (Vec<Ipv4Net>, Vec<Ipv6Net>) {
577 // get a Vector of all prefixes for IPv4 and IPv6
578 let mut all_prefixes = self
579 .ipv4
580 .iter()
581 .map(|(net, _data)| IpNet::from(*net))
582 .collect::<Vec<IpNet>>();
583 all_prefixes.extend(self.ipv6.iter().map(|(net, _data)| IpNet::from(*net)));
584
585 // get the aggregated prefixes
586 let aggregated_prefixes = IpNet::aggregate(&all_prefixes);
587
588 // split aggregated_prefixes into IPv4 and IPv6 prefixes
589 let mut ipv4_prefixes = Vec::new();
590 let mut ipv6_prefixes = Vec::new();
591 for prefix in aggregated_prefixes {
592 match prefix {
593 IpNet::V4(net) => ipv4_prefixes.push(net),
594 IpNet::V6(net) => ipv6_prefixes.push(net),
595 }
596 }
597 (ipv4_prefixes, ipv6_prefixes)
598 }
599
600 /// Find the difference between two prefix tries, returning two vectors of IpNets, one for
601 /// added prefixes, and one for removed prefixes.
602 ///
603 /// - added prefixes: all prefixes in other that are not in self
604 /// - removed prefixes: all prefixes in self that are not in other
605 pub fn diff(&self, other: &Self) -> (Vec<IpNet>, Vec<IpNet>) {
606 let mut added = IpnetTrie::<bool>::new();
607 let mut removed = IpnetTrie::<bool>::new();
608
609 // Find added prefixes: all prefixes in self that are not in other
610 // Method: build a trie using all prefixes in self, then remove all prefixes in other on the trie.
611 // The remaining prefixes are the added prefixes.
612 let (self_ipv4_prefixes, self_ipv6_prefixes) = self.get_aggregated_prefixes();
613 let (other_ipv4_prefixes, other_ipv6_prefixes) = other.get_aggregated_prefixes();
614
615 let mut self_ipv4_map: PrefixMap<Ipv4Net, bool> = PrefixMap::new();
616 for prefix in &self_ipv4_prefixes {
617 self_ipv4_map.insert(*prefix, true);
618 }
619 let mut other_ipv4_map: PrefixMap<Ipv4Net, bool> = PrefixMap::new();
620 for prefix in &other_ipv4_prefixes {
621 other_ipv4_map.insert(*prefix, true);
622 }
623 // check added prefixes in other
624 for v4_prefix in &other_ipv4_prefixes {
625 if self_ipv4_map.get_lpm(v4_prefix).is_some() {
626 // Prefix is covered by some super-prefix in self, nothing added
627 continue;
628 }
629
630 // Prefix is not covered by some super-prefix in self, there might be some overlapping sub-prefixes.
631 // get non-overlapping sub-prefixes
632 let sub_prefixes = IpNet::aggregate(
633 &self_ipv4_map
634 .children(*v4_prefix)
635 .map(|(p, _)| IpNet::from(*p))
636 .collect::<Vec<IpNet>>(),
637 );
638
639 if sub_prefixes.is_empty() {
640 // Self-trie does not have any sub-prefixes of the given trie
641 added.insert(*v4_prefix, true);
642 } else {
643 // Self-trie has sub-prefixes of the given trie, in other words, the other-trie
644 // added a new covering super-prefix comparing to the self-trie.
645
646 let mut target_prefixes: Vec<IpNet> = vec![(*v4_prefix).into()];
647 for sub_prefix in sub_prefixes {
648 let mut new_prefixes = vec![];
649 for target_prefix in target_prefixes {
650 // make sure none of the new prefixes overlap with the sub-prefix prefix
651 new_prefixes.extend(exclude_prefix(target_prefix, sub_prefix));
652 }
653
654 target_prefixes = IpNet::aggregate(&new_prefixes);
655 }
656 for target_prefix in target_prefixes {
657 added.insert(target_prefix, true);
658 }
659 }
660 }
661 // check deleted prefixes in other
662 for v4_prefix in &self_ipv4_prefixes {
663 if other_ipv4_map.get_lpm(v4_prefix).is_some() {
664 continue;
665 }
666 // get non-overlapping sub-prefixes
667 let sub_prefixes = IpNet::aggregate(
668 &other_ipv4_map
669 .children(*v4_prefix)
670 .map(|(p, _)| IpNet::from(*p))
671 .collect::<Vec<IpNet>>(),
672 );
673
674 if sub_prefixes.is_empty() {
675 // Self-trie does not have any sub-prefixes of the given trie
676 removed.insert(*v4_prefix, true);
677 } else {
678 // Self-trie has sub-prefixes of the given trie, in other words, the other-trie
679 // added a new covering super-prefix comparing to the self-trie.
680
681 let mut target_prefixes: Vec<IpNet> = vec![(*v4_prefix).into()];
682 for sub_prefix in sub_prefixes {
683 let mut new_prefixes = vec![];
684 for target_prefix in target_prefixes {
685 // make sure none of the new prefixes overlap with the sub-prefix prefix
686 new_prefixes.extend(exclude_prefix(target_prefix, sub_prefix));
687 }
688
689 target_prefixes = IpNet::aggregate(&new_prefixes);
690 }
691 for target_prefix in target_prefixes {
692 removed.insert(target_prefix, true);
693 }
694 }
695 }
696
697 let mut self_ipv6_map: PrefixMap<Ipv6Net, bool> = PrefixMap::new();
698 for prefix in &self_ipv6_prefixes {
699 self_ipv6_map.insert(*prefix, true);
700 }
701 let mut other_ipv6_map: PrefixMap<Ipv6Net, bool> = PrefixMap::new();
702 for prefix in &other_ipv6_prefixes {
703 other_ipv6_map.insert(*prefix, true);
704 }
705 // check added prefixes in other
706 for v6_prefix in &other_ipv6_prefixes {
707 if self_ipv6_map.get_lpm(v6_prefix).is_some() {
708 // Prefix is covered by some super-prefix in self, nothing added
709 continue;
710 }
711
712 // Prefix is not covered by some super-prefix in self, there might be some overlapping sub-prefixes.
713 // get non-overlapping sub-prefixes
714 let sub_prefixes = IpNet::aggregate(
715 &self_ipv6_map
716 .children(*v6_prefix)
717 .map(|(p, _)| IpNet::from(*p))
718 .collect::<Vec<IpNet>>(),
719 );
720
721 if sub_prefixes.is_empty() {
722 // Self-trie does not have any sub-prefixes of the given trie
723 added.insert(*v6_prefix, true);
724 } else {
725 // Self-trie has sub-prefixes of the given trie, in other words, the other-trie
726 // added a new covering super-prefix comparing to the self-trie.
727
728 let mut target_prefixes: Vec<IpNet> = vec![(*v6_prefix).into()];
729 for sub_prefix in sub_prefixes {
730 let mut new_prefixes = vec![];
731 for target_prefix in target_prefixes {
732 // make sure none of the new prefixes overlap with the sub-prefix prefix
733 new_prefixes.extend(exclude_prefix(target_prefix, sub_prefix));
734 }
735
736 target_prefixes = IpNet::aggregate(&new_prefixes);
737 }
738 for target_prefix in target_prefixes {
739 added.insert(target_prefix, true);
740 }
741 }
742 }
743 // check deleted prefixes in other
744 for v6_prefix in &self_ipv6_prefixes {
745 if other_ipv6_map.get_lpm(v6_prefix).is_some() {
746 continue;
747 }
748 // get non-overlapping sub-prefixes
749 let sub_prefixes = IpNet::aggregate(
750 &other_ipv6_map
751 .children(*v6_prefix)
752 .map(|(p, _)| IpNet::from(*p))
753 .collect::<Vec<IpNet>>(),
754 );
755
756 if sub_prefixes.is_empty() {
757 // Self-trie does not have any sub-prefixes of the given trie
758 removed.insert(*v6_prefix, true);
759 } else {
760 // Self-trie has sub-prefixes of the given trie, in other words, the other-trie
761 // added a new covering super-prefix comparing to the self-trie.
762
763 let mut target_prefixes: Vec<IpNet> = vec![(*v6_prefix).into()];
764 for sub_prefix in sub_prefixes {
765 let mut new_prefixes = vec![];
766 for target_prefix in target_prefixes {
767 // make sure none of the new prefixes overlap with the sub-prefix prefix
768 new_prefixes.extend(exclude_prefix(target_prefix, sub_prefix));
769 }
770
771 target_prefixes = IpNet::aggregate(&new_prefixes);
772 }
773 for target_prefix in target_prefixes {
774 removed.insert(target_prefix, true);
775 }
776 }
777 }
778
779 (
780 added.iter().map(|(p, _)| p).collect(),
781 removed.iter().map(|(p, _)| p).collect(),
782 )
783 }
784}
785
786#[cfg(test)]
787mod tests {
788 use crate::IpnetTrie;
789 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
790 use std::net::{Ipv4Addr, Ipv6Addr};
791 use std::str::FromStr;
792
793 #[test]
794 fn insert_ipv4_ipv6() {
795 let mut table = IpnetTrie::new();
796 table.insert(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 0), 16).unwrap(), 1);
797 table.insert(
798 Ipv6Net::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 128).unwrap(),
799 1,
800 );
801 }
802
803 #[test]
804 fn exact_match_ipv4() {
805 let mut table = IpnetTrie::new();
806 table.insert(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 0), 16).unwrap(), 1);
807 let m = table.exact_match(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 0), 16).unwrap());
808 assert_eq!(m, Some(&1));
809 let m = table.exact_match(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 0), 17).unwrap());
810 assert_eq!(m, None);
811 let m = table.exact_match(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 0), 15).unwrap());
812 assert_eq!(m, None);
813 }
814
815 #[test]
816 fn exact_match_ipv6() {
817 let mut table = IpnetTrie::new();
818 table.insert(
819 Ipv6Net::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 127).unwrap(),
820 1,
821 );
822 let m =
823 table.exact_match(Ipv6Net::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 127).unwrap());
824 assert_eq!(m, Some(&1));
825 let m =
826 table.exact_match(Ipv6Net::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 128).unwrap());
827 assert_eq!(m, None);
828 let m =
829 table.exact_match(Ipv6Net::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 126).unwrap());
830 assert_eq!(m, None);
831 }
832
833 #[test]
834 fn test_ip_count() {
835 use crate::IpCount;
836
837 let mut table = IpnetTrie::new();
838 table.insert(Ipv4Net::from_str("192.0.2.129/25").unwrap(), 1);
839 table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
840 table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
841 table.insert(Ipv4Net::from_str("192.0.2.0/24").unwrap(), 1);
842 assert_eq!(
843 table.ip_count(),
844 IpCount {
845 ipv4: 256,
846 ipv6: Some(0)
847 }
848 );
849
850 table.insert(Ipv4Net::from_str("198.51.100.0/25").unwrap(), 1);
851 table.insert(Ipv4Net::from_str("198.51.100.64/26").unwrap(), 1);
852 assert_eq!(
853 table.ip_count(),
854 IpCount {
855 ipv4: 384,
856 ipv6: Some(0)
857 }
858 );
859
860 table.insert(Ipv4Net::from_str("198.51.100.65/26").unwrap(), 1);
861 assert_eq!(
862 table.ip_count(),
863 IpCount {
864 ipv4: 384,
865 ipv6: Some(0)
866 }
867 );
868
869 table.insert(Ipv6Net::from_str("2001:DB80::/48").unwrap(), 1);
870 assert_eq!(
871 table.ip_count(),
872 IpCount {
873 ipv4: 384,
874 ipv6: Some(2_u128.pow(80))
875 }
876 );
877 table.insert(Ipv6Net::from_str("2001:DB80::/49").unwrap(), 1);
878 assert_eq!(
879 table.ip_count(),
880 IpCount {
881 ipv4: 384,
882 ipv6: Some(2_u128.pow(80))
883 }
884 );
885 }
886
887 #[test]
888 fn test_ip_count_full_space() {
889 use crate::IpCount;
890
891 // Full IPv4 space: 0.0.0.0/0 — 2^32 addresses, fits in u64
892 let mut v4_full = IpnetTrie::new();
893 v4_full.insert(Ipv4Net::from_str("0.0.0.0/0").unwrap(), 1);
894 assert_eq!(
895 v4_full.ip_count(),
896 IpCount {
897 ipv4: 1u64 << 32,
898 ipv6: Some(0)
899 }
900 );
901
902 // Full IPv6 space: ::/0 — 2^128 addresses, cannot fit in u128
903 let mut v6_full = IpnetTrie::new();
904 v6_full.insert(Ipv6Net::from_str("::/0").unwrap(), 1);
905 assert_eq!(
906 v6_full.ip_count(),
907 IpCount {
908 ipv4: 0,
909 ipv6: None
910 }
911 );
912
913 // Two /1 prefixes also cover all of IPv6 (2^127 + 2^127 = 2^128)
914 let mut v6_two_halves = IpnetTrie::new();
915 v6_two_halves.insert(Ipv6Net::from_str("::/1").unwrap(), 1);
916 v6_two_halves.insert(Ipv6Net::from_str("8000::/1").unwrap(), 1);
917 assert_eq!(
918 v6_two_halves.ip_count(),
919 IpCount {
920 ipv4: 0,
921 ipv6: None
922 }
923 );
924
925 // Both full spaces at once
926 let mut both = IpnetTrie::new();
927 both.insert(Ipv4Net::from_str("0.0.0.0/0").unwrap(), 1);
928 both.insert(Ipv6Net::from_str("::/0").unwrap(), 1);
929 assert_eq!(
930 both.ip_count(),
931 IpCount {
932 ipv4: 1u64 << 32,
933 ipv6: None
934 }
935 );
936 }
937
938 #[test]
939 fn test_comparison() {
940 let mut trie_1 = IpnetTrie::new();
941 trie_1.insert(Ipv4Net::from_str("192.168.0.0/23").unwrap(), 1);
942 trie_1.insert(Ipv4Net::from_str("192.168.2.0/24").unwrap(), 1);
943
944 let mut trie_2 = IpnetTrie::new();
945 trie_2.insert(Ipv4Net::from_str("192.168.2.0/24").unwrap(), 1);
946
947 let (_added, removed) = trie_1.diff(&trie_2);
948 assert_eq!(removed.len(), 1);
949 assert_eq!(
950 removed[0],
951 IpNet::V4(Ipv4Net::from_str("192.168.0.0/23").unwrap())
952 );
953
954 trie_2.insert(Ipv4Net::from_str("192.168.0.0/24").unwrap(), 1);
955 let (_added, removed) = trie_1.diff(&trie_2);
956 assert_eq!(removed.len(), 1);
957 assert_eq!(
958 removed[0],
959 IpNet::from(Ipv4Net::from_str("192.168.1.0/24").unwrap())
960 );
961
962 trie_2.insert(Ipv4Net::from_str("192.168.3.0/24").unwrap(), 1);
963 let (added, removed) = trie_1.diff(&trie_2);
964 assert_eq!(removed.len(), 1);
965 assert_eq!(
966 removed[0],
967 IpNet::from(Ipv4Net::from_str("192.168.1.0/24").unwrap())
968 );
969 assert_eq!(added.len(), 1);
970 assert_eq!(
971 added[0],
972 IpNet::from(Ipv4Net::from_str("192.168.3.0/24").unwrap())
973 );
974
975 trie_2.insert(Ipv6Net::from_str("2001:DB80::/48").unwrap(), 1);
976 let (added, removed) = trie_1.diff(&trie_2);
977 assert_eq!(removed.len(), 1);
978 assert_eq!(
979 removed[0],
980 IpNet::from(Ipv4Net::from_str("192.168.1.0/24").unwrap())
981 );
982 assert_eq!(added.len(), 2);
983 assert_eq!(
984 added[0],
985 IpNet::from(Ipv4Net::from_str("192.168.3.0/24").unwrap())
986 );
987 assert_eq!(
988 added[1],
989 IpNet::from(Ipv6Net::from_str("2001:DB80::/48").unwrap())
990 );
991 }
992}