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