radixtarget_rust/
ip.rs

1use crate::node::{BaseNode, IPNode};
2use crate::utils::{canonicalize_ipnet, ipnet_to_bits};
3use ipnet::IpNet;
4use std::collections::HashSet;
5
6#[derive(Debug, Clone)]
7pub struct IpRadixTree {
8    pub root: IPNode,
9    pub acl_mode: bool,
10}
11
12impl IpRadixTree {
13    pub fn new(acl_mode: bool) -> Self {
14        IpRadixTree {
15            root: IPNode::new(),
16            acl_mode,
17        }
18    }
19    pub fn insert(&mut self, network: IpNet) -> Option<String> {
20        let canonical = canonicalize_ipnet(&network);
21
22        // If ACL mode is enabled, check if the network is already covered by the tree
23        if self.acl_mode && self.get(&canonical).is_some() {
24            return None; // Skip insertion if already covered
25        }
26
27        let mut node = &mut self.root;
28        let bits = ipnet_to_bits(&canonical);
29        for &bit in &bits {
30            node = node
31                .children
32                .entry(u64::from(bit))
33                .or_insert_with(|| Box::new(IPNode::new()));
34        }
35        node.network = Some(canonical);
36
37        // If ACL mode is enabled, clear children of the inserted node
38        if self.acl_mode {
39            node.clear();
40        }
41
42        Some(canonical.to_string())
43    }
44    pub fn get(&self, net: &IpNet) -> Option<String> {
45        let canonical = canonicalize_ipnet(net);
46        let mut node = &self.root;
47        let bits = ipnet_to_bits(&canonical);
48        let mut best: Option<&IpNet> = None;
49        if let Some(n) = &node.network
50            && n.contains(&canonical.network())
51            && n.prefix_len() <= canonical.prefix_len()
52        {
53            best = Some(n);
54        }
55        for &bit in &bits {
56            if let Some(child) = node.children.get(&u64::from(bit)) {
57                node = child;
58                if let Some(n) = &node.network
59                    && n.contains(&canonical.network())
60                    && n.prefix_len() <= canonical.prefix_len()
61                {
62                    best = Some(n);
63                }
64            } else {
65                break;
66            }
67        }
68        best.map(|n| n.to_string())
69    }
70    pub fn delete(&mut self, network: IpNet) -> bool {
71        let canonical = canonicalize_ipnet(&network);
72        Self::delete_rec(&mut self.root, &ipnet_to_bits(&canonical), 0, &canonical)
73    }
74    fn delete_rec(node: &mut IPNode, bits: &[u8], depth: usize, network: &IpNet) -> bool {
75        if depth == bits.len() {
76            if node.network.as_ref() == Some(network) {
77                node.network = None;
78                return true;
79            }
80            return false;
81        }
82        let bit = bits[depth];
83        if let Some(child) = node.children.get_mut(&u64::from(bit)) {
84            let deleted = Self::delete_rec(child, bits, depth + 1, network);
85            if child.children.is_empty() && child.network.is_none() {
86                node.children.remove(&u64::from(bit));
87            }
88            return deleted;
89        }
90        false
91    }
92    pub fn prune(&mut self) -> usize {
93        self.root.prune()
94    }
95
96    /// Get all networks stored in the tree
97    pub fn hosts(&self) -> HashSet<String> {
98        self.root.all_hosts()
99    }
100    /// Defrag the entire tree, merging mergeable networks. Returns (cleaned, new) as sets of strings, with overlap removed.
101    pub fn defrag(&mut self) -> (HashSet<String>, HashSet<String>) {
102        let mut cleaned = HashSet::new();
103        let mut new = HashSet::new();
104        loop {
105            let pairs = self.root.mergeable_pairs();
106            if pairs.is_empty() {
107                break;
108            }
109            for (left_net, right_net, supernet) in pairs {
110                cleaned.insert(left_net.to_string());
111                cleaned.insert(right_net.to_string());
112                new.insert(supernet.to_string());
113                self.delete(left_net);
114                self.delete(right_net);
115                self.insert(supernet);
116            }
117        }
118        // Remove any overlap between cleaned and new
119        let overlap: HashSet<_> = cleaned.intersection(&new).cloned().collect();
120        for k in &overlap {
121            cleaned.remove(k);
122            new.remove(k);
123        }
124        (cleaned, new)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use crate::utils::{canonicalize_ipnet, ipnet_to_bits};
132    use ipnet::IpNet;
133    use std::str::FromStr;
134
135    /// Test basic insertion and lookup for both IPv4 and IPv6 networks.
136    /// Ensures that inserted networks are found, and unrelated addresses are not.
137    #[test]
138    fn test_insert_and_get_ipv4_and_ipv6() {
139        let mut tree = IpRadixTree::new(false);
140        let net_v4 = IpNet::from_str("8.8.8.0/24").unwrap();
141        let net_v6 = IpNet::from_str("dead::/64").unwrap();
142        let hash_v4 = tree.insert(net_v4);
143        assert_eq!(hash_v4, Some(net_v4.to_string()), "insert(8.8.8.0/24) hash");
144        let hash_v6 = tree.insert(net_v6);
145        assert_eq!(hash_v6, Some(net_v6.to_string()), "insert(dead::/64) hash");
146        assert_eq!(
147            tree.get(&IpNet::from_str("8.8.8.8/32").unwrap()),
148            Some(net_v4.to_string())
149        );
150        assert_eq!(
151            tree.get(&IpNet::from_str("dead::beef/128").unwrap()),
152            Some(net_v6.to_string())
153        );
154        assert_eq!(tree.get(&IpNet::from_str("1.1.1.1/32").unwrap()), None);
155        assert_eq!(tree.get(&IpNet::from_str("cafe::beef/128").unwrap()), None);
156    }
157
158    /// Test insertion of a network with host bits set (e.g., 192.168.1.42/24).
159    /// The tree should sanitize the network to the correct base address and match all addresses in the subnet.
160    #[test]
161    fn test_insert_network_with_host_bits() {
162        let mut tree = IpRadixTree::new(false);
163        // 192.168.1.42/24 should be sanitized to 192.168.1.0/24
164        let net = IpNet::from_str("192.168.1.0/24").unwrap();
165        let hash = tree.insert(net);
166        assert_eq!(hash, Some(net.to_string()), "insert(192.168.1.0/24) hash");
167        assert_eq!(
168            tree.get(&IpNet::from_str("192.168.1.1/32").unwrap()),
169            Some(net.to_string())
170        );
171        assert_eq!(tree.get(&IpNet::from_str("192.168.2.1/32").unwrap()), None);
172    }
173
174    /// Test insertion and lookup of a single IP address as a /32 network.
175    /// Ensures only the exact address matches.
176    #[test]
177    fn test_insert_single_ip() {
178        let mut tree = IpRadixTree::new(false);
179        let net = IpNet::from_str("10.0.0.1/32").unwrap();
180        let hash = tree.insert(net);
181        assert_eq!(hash, Some(net.to_string()), "insert(10.0.0.1/32) hash");
182        assert_eq!(
183            tree.get(&IpNet::from_str("10.0.0.1/32").unwrap()),
184            Some(net.to_string())
185        );
186        assert_eq!(tree.get(&IpNet::from_str("10.0.0.2/32").unwrap()), None);
187    }
188
189    /// Test overlapping networks and longest-prefix match logic.
190    /// Ensures the most specific (longest prefix) network is returned for a given address.
191    #[test]
192    fn test_overlapping_and_longest_prefix() {
193        let mut tree = IpRadixTree::new(false);
194        let net1 = IpNet::from_str("10.0.0.0/8").unwrap();
195        let net2 = IpNet::from_str("10.1.0.0/16").unwrap();
196        let net3 = IpNet::from_str("10.1.2.0/24").unwrap();
197        let hash3 = tree.insert(net3);
198        assert_eq!(hash3, Some(net3.to_string()), "insert(10.1.2.0/24) hash");
199        let hash2 = tree.insert(net2);
200        assert_eq!(hash2, Some(net2.to_string()), "insert(10.1.0.0/16) hash");
201        let hash1 = tree.insert(net1);
202        assert_eq!(hash1, Some(net1.to_string()), "insert(10.0.0.0/8) hash");
203        assert_eq!(
204            tree.get(&IpNet::from_str("10.1.2.3/32").unwrap()),
205            Some(net3.to_string())
206        );
207        assert_eq!(
208            tree.get(&IpNet::from_str("10.1.2.3/30").unwrap()),
209            Some(net3.to_string())
210        );
211        assert_eq!(
212            tree.get(&IpNet::from_str("10.1.3.3/32").unwrap()),
213            Some(net2.to_string())
214        );
215        assert_eq!(
216            tree.get(&IpNet::from_str("10.1.3.3/30").unwrap()),
217            Some(net2.to_string())
218        );
219        assert_eq!(
220            tree.get(&IpNet::from_str("10.2.2.2/32").unwrap()),
221            Some(net1.to_string())
222        );
223        assert_eq!(
224            tree.get(&IpNet::from_str("10.2.2.2/30").unwrap()),
225            Some(net1.to_string())
226        );
227    }
228
229    /// Test IPv6 longest-prefix match logic.
230    /// Ensures the most specific IPv6 network is returned for a given address.
231    #[test]
232    fn test_ipv6_longest_prefix() {
233        let mut tree = IpRadixTree::new(false);
234        // Insert overlapping networks in various orders, with and without host bits
235        let net1 = IpNet::from_str("2001:db8::/32").unwrap();
236        let net2 = IpNet::from_str("2001:db8:abcd::/48").unwrap();
237        let net3 = IpNet::from_str("2001:db8:abcd:1234::/64").unwrap();
238        let net4 = IpNet::from_str("2001:db8:abcd:1234:5678::/80").unwrap();
239        // Insert out of order and with host bits
240        let hash4 = tree.insert(IpNet::from_str("2001:db8:abcd:1234:5678:9abc::1/80").unwrap()); // net4
241        assert_eq!(
242            hash4,
243            Some(net4.to_string()),
244            "insert(2001:db8:abcd:1234:5678::/80) hash"
245        );
246        let hash3 = tree.insert(IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()); // net3
247        assert_eq!(
248            hash3,
249            Some(net3.to_string()),
250            "insert(2001:db8:abcd:1234::/64) hash"
251        );
252        let hash2 = tree.insert(IpNet::from_str("2001:db8:abcd::/48").unwrap()); // net2
253        assert_eq!(
254            hash2,
255            Some(net2.to_string()),
256            "insert(2001:db8:abcd::/48) hash"
257        );
258        let hash1 = tree.insert(IpNet::from_str("2001:db8::1/32").unwrap()); // net1
259        assert_eq!(hash1, Some(net1.to_string()), "insert(2001:db8::/32) hash");
260
261        // Queries for most specific match
262        // Should match net4
263        assert_eq!(
264            tree.get(&IpNet::from_str("2001:db8:abcd:1234:5678:9abc::dead/128").unwrap()),
265            Some(net4.to_string())
266        );
267        // Should match net3
268        assert_eq!(
269            tree.get(&IpNet::from_str("2001:db8:abcd:1234::beef/128").unwrap()),
270            Some(net3.to_string())
271        );
272        // Should match net2
273        assert_eq!(
274            tree.get(&IpNet::from_str("2001:db8:abcd::cafe/128").unwrap()),
275            Some(net2.to_string())
276        );
277        // Should match net1
278        assert_eq!(
279            tree.get(&IpNet::from_str("2001:db8::1234/128").unwrap()),
280            Some(net1.to_string())
281        );
282
283        // Queries with host bits and different prefix lengths
284        assert_eq!(
285            tree.get(&IpNet::from_str("2001:db8:abcd:1234:5678:9abc::dead/81").unwrap()),
286            Some(net4.to_string())
287        );
288        assert_eq!(
289            tree.get(&IpNet::from_str("2001:db8:abcd:1234::beef/65").unwrap()),
290            Some(net3.to_string())
291        );
292        assert_eq!(
293            tree.get(&IpNet::from_str("2001:db8:abcd::cafe/49").unwrap()),
294            Some(net2.to_string())
295        );
296        assert_eq!(
297            tree.get(&IpNet::from_str("2001:db8::1234/33").unwrap()),
298            Some(net1.to_string())
299        );
300
301        // Query outside all networks
302        assert_eq!(
303            tree.get(&IpNet::from_str("2001:dead::1/128").unwrap()),
304            None
305        );
306        assert_eq!(tree.get(&IpNet::from_str("3000::/32").unwrap()), None);
307    }
308
309    /// Test deletion combined with querying by network and host.
310    /// Ensures that after deleting a network, the correct fallback (less specific) parent is returned, or None if no match remains.
311    #[test]
312    fn test_deletion_and_querying() {
313        let mut tree = IpRadixTree::new(false);
314        // Insert overlapping networks
315        let net1 = IpNet::from_str("192.168.0.0/16").unwrap();
316        let net2 = IpNet::from_str("192.168.1.0/24").unwrap();
317        let net3 = IpNet::from_str("192.168.1.128/25").unwrap();
318        let hash1 = tree.insert(net1);
319        assert_eq!(hash1, Some(net1.to_string()), "insert(192.168.0.0/16) hash");
320        let hash2 = tree.insert(net2);
321        assert_eq!(hash2, Some(net2.to_string()), "insert(192.168.1.0/24) hash");
322        let hash3 = tree.insert(net3);
323        assert_eq!(
324            hash3,
325            Some(net3.to_string()),
326            "insert(192.168.1.128/25) hash"
327        );
328        // Query before deletion
329        assert_eq!(
330            tree.get(&IpNet::from_str("192.168.1.128/25").unwrap()),
331            Some(net3.to_string())
332        ); // Most specific: /25
333        assert_eq!(
334            tree.get(&IpNet::from_str("192.168.1.129/32").unwrap()),
335            Some(net3.to_string())
336        );
337        assert_eq!(
338            tree.get(&IpNet::from_str("192.168.1.0/24").unwrap()),
339            Some(net2.to_string())
340        );
341        // Delete the most specific network
342        assert!(tree.delete(IpNet::from_str("192.168.1.128/25").unwrap()));
343        // Now queries in 192.168.1.128/25 should fall back to /24
344        assert_eq!(
345            tree.get(&IpNet::from_str("192.168.1.128/25").unwrap()),
346            Some(net2.to_string())
347        );
348        assert_eq!(
349            tree.get(&IpNet::from_str("192.168.1.129/32").unwrap()),
350            Some(net2.to_string())
351        );
352        // Delete the /24
353        assert!(tree.delete(IpNet::from_str("192.168.1.0/24").unwrap()));
354        // Now queries in 192.168.1.0/24 should fall back to /16
355        assert_eq!(
356            tree.get(&IpNet::from_str("192.168.1.129/32").unwrap()),
357            Some(net1.to_string())
358        );
359        // Delete the /16
360        assert!(tree.delete(IpNet::from_str("192.168.0.0/16").unwrap()));
361        // Now nothing should match
362        assert_eq!(
363            tree.get(&IpNet::from_str("192.168.1.129/32").unwrap()),
364            None
365        );
366
367        // IPv6 case
368        let netv6_1 = IpNet::from_str("2001:db8::/32").unwrap();
369        let netv6_2 = IpNet::from_str("2001:db8:abcd::/48").unwrap();
370        let netv6_3 = IpNet::from_str("2001:db8:abcd:1234::/64").unwrap();
371        let hashv6_1 = tree.insert(netv6_1);
372        assert_eq!(
373            hashv6_1,
374            Some(netv6_1.to_string()),
375            "insert(2001:db8::/32) hash"
376        );
377        let hashv6_2 = tree.insert(netv6_2);
378        assert_eq!(
379            hashv6_2,
380            Some(netv6_2.to_string()),
381            "insert(2001:db8:abcd::/48) hash"
382        );
383        let hashv6_3 = tree.insert(netv6_3);
384        assert_eq!(
385            hashv6_3,
386            Some(netv6_3.to_string()),
387            "insert(2001:db8:abcd:1234::/64) hash"
388        );
389        // Query before deletion
390        assert_eq!(
391            tree.get(&IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()),
392            Some(netv6_3.to_string())
393        ); // Most specific: /64
394        assert_eq!(
395            tree.get(&IpNet::from_str("2001:db8:abcd::/48").unwrap()),
396            Some(netv6_2.to_string())
397        ); // Most specific: /48
398        assert_eq!(
399            tree.get(&IpNet::from_str("2001:db8::/32").unwrap()),
400            Some(netv6_1.to_string())
401        ); // Most specific: /32
402        // Delete the most specific
403        assert!(tree.delete(IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()));
404        assert_eq!(
405            tree.get(&IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()),
406            Some(netv6_2.to_string())
407        );
408        // Delete the /48
409        assert!(tree.delete(IpNet::from_str("2001:db8:abcd::/48").unwrap()));
410        assert_eq!(
411            tree.get(&IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()),
412            Some(netv6_1.to_string())
413        );
414        // Delete the /32
415        assert!(tree.delete(IpNet::from_str("2001:db8::/32").unwrap()));
416        assert_eq!(
417            tree.get(&IpNet::from_str("2001:db8:abcd:1234::/64").unwrap()),
418            None
419        );
420    }
421
422    /// Test insertion and lookup with custom data types as values.
423    /// Ensures the tree can store and retrieve arbitrary data.
424    #[test]
425    fn test_insert_with_custom_data_types() {
426        let mut tree = IpRadixTree::new(false);
427        let net = IpNet::from_str("8.8.8.0/24").unwrap();
428        let hash = tree.insert(net);
429        assert_eq!(hash, Some(net.to_string()), "insert(8.8.8.0/24) hash");
430        assert_eq!(
431            tree.get(&IpNet::from_str("8.8.8.8/32").unwrap()),
432            Some(net.to_string())
433        );
434    }
435
436    /// Test insertion and lookup of a zero-prefix (default) IPv4 network.
437    /// Ensures all IPv4 addresses match the default route.
438    #[test]
439    fn test_insert_and_get_zero_prefix() {
440        let mut tree = IpRadixTree::new(false);
441        let net = IpNet::from_str("0.0.0.0/0").unwrap();
442        let hash = tree.insert(net);
443        assert_eq!(hash, Some(net.to_string()), "insert(0.0.0.0/0) hash");
444        assert_eq!(
445            tree.get(&IpNet::from_str("1.2.3.4/32").unwrap()),
446            Some(net.to_string())
447        );
448        assert_eq!(
449            tree.get(&IpNet::from_str("255.255.255.255/32").unwrap()),
450            Some(net.to_string())
451        );
452    }
453
454    /// Test insertion and lookup of a zero-prefix (default) IPv6 network.
455    /// Ensures all IPv6 addresses match the default route.
456    #[test]
457    fn test_insert_and_get_ipv6_zero_prefix() {
458        let mut tree = IpRadixTree::new(false);
459        let net = IpNet::from_str("::/0").unwrap();
460        let hash = tree.insert(net);
461        assert_eq!(hash, Some(net.to_string()), "insert(::/0) hash");
462        assert_eq!(
463            tree.get(&IpNet::from_str("::1/128").unwrap()),
464            Some(net.to_string())
465        );
466        assert_eq!(
467            tree.get(&IpNet::from_str("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/128").unwrap()),
468            Some(net.to_string())
469        );
470    }
471
472    /// Test insertion and lookup of both IPv4 and IPv6 single-address networks in the same tree.
473    /// Ensures both types can coexist and be queried independently.
474    #[test]
475    fn test_insert_and_get_multiple_types() {
476        let mut tree = IpRadixTree::new(false);
477        let net4 = IpNet::from_str("1.2.3.4/32").unwrap();
478        let net6 = IpNet::from_str("dead::beef/128").unwrap();
479        let hash4 = tree.insert(net4);
480        assert_eq!(hash4, Some(net4.to_string()), "insert(1.2.3.4/32) hash");
481        let hash6 = tree.insert(net6);
482        assert_eq!(hash6, Some(net6.to_string()), "insert(dead::beef/128) hash");
483        assert_eq!(
484            tree.get(&IpNet::from_str("1.2.3.4/32").unwrap()),
485            Some(net4.to_string())
486        );
487        assert_eq!(
488            tree.get(&IpNet::from_str("dead::beef/128").unwrap()),
489            Some(net6.to_string())
490        );
491    }
492
493    /// Test insertion of an IPv6 network with host bits set.
494    /// Ensures the tree sanitizes the network and matches all addresses in the subnet, but not outside.
495    #[test]
496    fn test_insert_and_get_with_host_bits_ipv6() {
497        let mut tree = IpRadixTree::new(false);
498        // Insert with host bits set; should normalize to dead:beef::0/120
499        let inserted = IpNet::from_str("dead:beef::1/120").unwrap();
500        let normalized = IpNet::from_str("dead:beef::0/120").unwrap();
501        let hash = tree.insert(inserted);
502        assert_eq!(
503            hash,
504            Some(normalized.to_string()),
505            "insert(dead:beef::0/120) hash"
506        );
507        // Query with addresses within the /120
508        assert_eq!(
509            tree.get(&IpNet::from_str("dead:beef::2/128").unwrap()),
510            Some(normalized.to_string())
511        );
512        assert_eq!(
513            tree.get(&IpNet::from_str("dead:beef::ff/128").unwrap()),
514            Some(normalized.to_string())
515        );
516        // Query with a network within the /120 (should normalize query)
517        assert_eq!(
518            tree.get(&IpNet::from_str("dead:beef::2/121").unwrap()),
519            Some(normalized.to_string())
520        );
521        // Query with an address outside the /120
522        assert_eq!(
523            tree.get(&IpNet::from_str("dead:beef::100/128").unwrap()),
524            None
525        );
526        // Delete with host bits set (should normalize and delete the /120)
527        assert!(tree.delete(IpNet::from_str("dead:beef::ff/120").unwrap()));
528        // After deletion, nothing in the /120 should match
529        assert_eq!(
530            tree.get(&IpNet::from_str("dead:beef::2/128").unwrap()),
531            None
532        );
533        assert_eq!(
534            tree.get(&IpNet::from_str("dead:beef::ff/128").unwrap()),
535            None
536        );
537        assert_eq!(
538            tree.get(&IpNet::from_str("dead:beef::2/121").unwrap()),
539            None
540        );
541    }
542
543    /// Test querying by network string for both IPv4 and IPv6.
544    /// Ensures that querying by network returns the correct data, matching the Python API.
545    #[test]
546    fn test_query_by_network() {
547        let mut tree = IpRadixTree::new(false);
548        let net_v4 = IpNet::from_str("8.8.8.0/24").unwrap();
549        let net_v6 = IpNet::from_str("dead::/64").unwrap();
550        let hash_v4 = tree.insert(net_v4);
551        assert_eq!(hash_v4, Some(net_v4.to_string()), "insert(8.8.8.0/24) hash");
552        let hash_v6 = tree.insert(net_v6);
553        assert_eq!(hash_v6, Some(net_v6.to_string()), "insert(dead::/64) hash");
554        // Query by network string
555        assert_eq!(
556            tree.get(&IpNet::from_str("8.8.8.0/24").unwrap()),
557            Some(net_v4.to_string())
558        );
559        assert_eq!(
560            tree.get(&IpNet::from_str("dead::/64").unwrap()),
561            Some(net_v6.to_string())
562        );
563        // Query by non-existent network
564        assert_eq!(tree.get(&IpNet::from_str("8.8.9.0/24").unwrap()), None);
565        assert_eq!(tree.get(&IpNet::from_str("cafe::/64").unwrap()), None);
566    }
567
568    /// Test querying by overlapping networks, ensuring longest-prefix match for network queries.
569    #[test]
570    fn test_query_by_overlapping_networks() {
571        let mut tree = IpRadixTree::new(false);
572        let net1 = IpNet::from_str("10.0.0.0/8").unwrap();
573        let net2 = IpNet::from_str("10.1.0.0/16").unwrap();
574        let net3 = IpNet::from_str("10.1.2.0/24").unwrap();
575        let hash1 = tree.insert(net1);
576        assert_eq!(hash1, Some(net1.to_string()), "insert(10.0.0.0/8) hash");
577        let hash2 = tree.insert(net2);
578        assert_eq!(hash2, Some(net2.to_string()), "insert(10.1.0.0/16) hash");
579        let hash3 = tree.insert(net3);
580        assert_eq!(hash3, Some(net3.to_string()), "insert(10.1.2.0/24) hash");
581        // Query by network string
582        assert_eq!(
583            tree.get(&IpNet::from_str("10.1.2.0/24").unwrap()),
584            Some(net3.to_string())
585        );
586        assert_eq!(
587            tree.get(&IpNet::from_str("10.1.0.0/16").unwrap()),
588            Some(net2.to_string())
589        );
590        assert_eq!(
591            tree.get(&IpNet::from_str("10.0.0.0/8").unwrap()),
592            Some(net1.to_string())
593        );
594        // Query by less specific network (should match the most specific containing network)
595        assert_eq!(
596            tree.get(&IpNet::from_str("10.1.2.0/26").unwrap()),
597            Some(net3.to_string())
598        );
599        assert_eq!(
600            tree.get(&IpNet::from_str("10.1.0.0/22").unwrap()),
601            Some(net2.to_string())
602        );
603        assert_eq!(
604            tree.get(&IpNet::from_str("10.1.0.0/12").unwrap()),
605            Some(net1.to_string())
606        );
607        // Query by a network not contained by any inserted network
608        assert_eq!(tree.get(&IpNet::from_str("10.0.0.0/7").unwrap()), None);
609        assert_eq!(tree.get(&IpNet::from_str("11.0.0.0/8").unwrap()), None);
610    }
611
612    /// Test querying by child networks (subnets) that aren't an exact match.
613    /// Ensures the most specific (longest prefix) parent is returned for overlapping networks.
614    #[test]
615    fn test_query_by_child_networks_longest_prefix() {
616        let mut tree = IpRadixTree::new(false);
617        // Insert overlapping networks
618        let net1 = IpNet::from_str("192.168.0.0/16").unwrap();
619        let net2 = IpNet::from_str("192.168.1.0/24").unwrap();
620        let net3 = IpNet::from_str("192.168.1.128/25").unwrap();
621        let hash1 = tree.insert(net1);
622        assert_eq!(hash1, Some(net1.to_string()), "insert(192.168.0.0/16) hash");
623        let hash2 = tree.insert(net2);
624        assert_eq!(hash2, Some(net2.to_string()), "insert(192.168.1.0/24) hash");
625        let hash3 = tree.insert(net3);
626        assert_eq!(
627            hash3,
628            Some(net3.to_string()),
629            "insert(192.168.1.128/25) hash"
630        );
631        // Query by a child network that is a subnet of both /16 and /24, but not an exact match
632        assert_eq!(
633            tree.get(&IpNet::from_str("192.168.1.128/26").unwrap()),
634            Some(net3.to_string())
635        ); // Most specific: /25
636        assert_eq!(
637            tree.get(&IpNet::from_str("192.168.1.0/25").unwrap()),
638            Some(net2.to_string())
639        ); // Most specific: /24
640        assert_eq!(
641            tree.get(&IpNet::from_str("192.168.2.0/24").unwrap()),
642            Some(net1.to_string())
643        ); // Only /16 matches
644        // Query by a network that doesn't match any
645        assert_eq!(tree.get(&IpNet::from_str("10.0.0.0/8").unwrap()), None);
646
647        // IPv6 case
648        let netv6_1 = IpNet::from_str("2001:db8::/32").unwrap();
649        let netv6_2 = IpNet::from_str("2001:db8:abcd::/48").unwrap();
650        let netv6_3 = IpNet::from_str("2001:db8:abcd:1234::/64").unwrap();
651        let hashv6_1 = tree.insert(netv6_1);
652        assert_eq!(
653            hashv6_1,
654            Some(netv6_1.to_string()),
655            "insert(2001:db8::/32) hash"
656        );
657        let hashv6_2 = tree.insert(netv6_2);
658        assert_eq!(
659            hashv6_2,
660            Some(netv6_2.to_string()),
661            "insert(2001:db8:abcd::/48) hash"
662        );
663        let hashv6_3 = tree.insert(netv6_3);
664        assert_eq!(
665            hashv6_3,
666            Some(netv6_3.to_string()),
667            "insert(2001:db8:abcd:1234::/64) hash"
668        );
669        assert_eq!(
670            tree.get(&IpNet::from_str("2001:db8:abcd:1234::/80").unwrap()),
671            Some(netv6_3.to_string())
672        ); // Most specific: /64
673        assert_eq!(
674            tree.get(&IpNet::from_str("2001:db8:abcd::/56").unwrap()),
675            Some(netv6_2.to_string())
676        ); // Most specific: /48
677        assert_eq!(
678            tree.get(&IpNet::from_str("2001:db8::/40").unwrap()),
679            Some(netv6_1.to_string())
680        ); // Most specific: /32
681        assert_eq!(tree.get(&IpNet::from_str("2001:dead::/32").unwrap()), None); // No match
682    }
683
684    /// Test that IPv4 and IPv6 addresses with the same bits are treated as overlapping in a single IpNet-based tree.
685    /// This is expected, as the tree does not distinguish between address families.
686    #[test]
687    fn test_ipv4_ipv6_same_bits_overlap() {
688        let mut tree = IpRadixTree::new(false);
689        // 1.0.0.0/30 (IPv4) and 100::/30 (IPv6) have the same first 30 bits
690        let net4 = IpNet::from_str("1.0.0.0/30").unwrap();
691        let net6 = IpNet::from_str("100::/30").unwrap();
692        // Assert that the bit vectors for both networks are identical for the first 30 bits
693        let expected_bits = vec![
694            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
695            0,
696        ];
697        assert_eq!(
698            ipnet_to_bits(&net4),
699            expected_bits,
700            "net4 bits did not match expected"
701        );
702        assert_eq!(
703            ipnet_to_bits(&net6),
704            expected_bits,
705            "net6 bits did not match expected"
706        );
707
708        let hash4 = tree.insert(net4);
709        assert_eq!(hash4, Some(net4.to_string()), "insert(1.0.0.0/30) hash");
710        assert_eq!(
711            tree.get(&IpNet::from_str("1.0.0.1/30").unwrap()),
712            Some(net4.to_string())
713        );
714
715        let hash6 = tree.insert(net6);
716        assert_eq!(hash6, Some(net6.to_string()), "insert(100::/30) hash");
717        assert_eq!(
718            tree.get(&IpNet::from_str("100::1/30").unwrap()),
719            Some(net6.to_string())
720        );
721
722        // Oops, we clobbered our IPv4 network. This is expected, and the reason why we maintain 2 separate trees.
723        assert_eq!(tree.get(&IpNet::from_str("1.0.0.1/30").unwrap()), None);
724        assert_eq!(
725            tree.get(&IpNet::from_str("100::1/30").unwrap()),
726            Some(net6.to_string())
727        );
728    }
729
730    #[test]
731    fn test_mergeable_pairs_two_children() {
732        let net1 = IpNet::from_str("192.168.0.0/24").unwrap();
733        let net2 = IpNet::from_str("192.168.1.0/24").unwrap();
734        let supernet = IpNet::from_str("192.168.0.0/23").unwrap();
735        let mut parent = IPNode::new();
736        let mut child1 = Box::new(IPNode::new());
737        child1.network = Some(net1);
738        let mut child2 = Box::new(IPNode::new());
739        child2.network = Some(net2);
740        parent.children.insert(0, child1);
741        parent.children.insert(1, child2);
742
743        let pairs = parent.mergeable_pairs();
744        assert_eq!(pairs.len(), 1);
745        let (left, right, merged) = &pairs[0];
746        // The pair should be the two /24s, and the merged should be the /23
747        assert!(
748            (left == &net1 && right == &net2) || (left == &net2 && right == &net1),
749            "Pair should be the two /24s"
750        );
751        assert_eq!(merged, &supernet, "Merged should be the /23");
752    }
753
754    #[test]
755    fn test_ipradix_defrag_merge() {
756        let mut tree = IpRadixTree::new(false);
757        let net1 = IpNet::from_str("192.168.0.0/24").unwrap();
758        let net2 = IpNet::from_str("192.168.1.0/24").unwrap();
759        let supernet = IpNet::from_str("192.168.0.0/23").unwrap();
760        tree.insert(net1);
761        tree.insert(net2);
762        // Before defrag, lookups should return the /24s
763        let ip1 = IpNet::from_str("192.168.0.1/32").unwrap();
764        let ip2 = IpNet::from_str("192.168.1.1/32").unwrap();
765        assert_eq!(tree.get(&ip1), Some(net1.to_string()));
766        assert_eq!(tree.get(&ip2), Some(net2.to_string()));
767        // Defrag
768        let (cleaned, new) = tree.defrag();
769        // The cleaned set should contain the two /24s, the new set should contain the /23
770        assert!(cleaned.contains(&net1.to_string()));
771        assert!(cleaned.contains(&net2.to_string()));
772        assert!(new.contains(&supernet.to_string()));
773        // After defrag, lookups should return the /23
774        assert_eq!(tree.get(&ip1), Some(supernet.to_string()));
775        assert_eq!(tree.get(&ip2), Some(supernet.to_string()));
776    }
777
778    #[test]
779    fn test_clear_method() {
780        use crate::node::BaseNode;
781
782        let mut tree = IpRadixTree::new(false);
783
784        // Insert networks in random order
785        let mut networks = vec![
786            "10.0.0.0/8",
787            "10.1.0.0/16",
788            "10.1.1.0/24",
789            "10.1.2.0/24",
790            "10.1.1.128/25",
791            "10.1.1.192/26",
792            "10.1.1.224/27",
793            "10.1.1.240/28",
794            "192.168.0.0/16",
795            "192.168.1.0/24",
796        ];
797
798        // Shuffle randomly
799        use rand::seq::SliceRandom;
800        use rand::thread_rng;
801        networks.shuffle(&mut thread_rng());
802
803        for net_str in &networks {
804            let net = IpNet::from_str(net_str).unwrap();
805            tree.insert(net);
806        }
807
808        // Verify all networks are present
809        for net_str in &networks {
810            let net = IpNet::from_str(net_str).unwrap();
811            assert!(
812                tree.get(&net).is_some(),
813                "Network {} should be present",
814                net_str
815            );
816        }
817
818        // Find the node for "10.1.1.0/24" and clear it
819        let target_net = IpNet::from_str("10.1.1.0/24").unwrap();
820        let canonical = canonicalize_ipnet(&target_net);
821        let mut node = &mut tree.root;
822        let bits = ipnet_to_bits(&canonical);
823        for &bit in &bits {
824            node = node
825                .children
826                .get_mut(&u64::from(bit))
827                .expect("Node should exist");
828        }
829
830        // Clear the 10.1.1.0/24 node (should clear its children)
831        let cleared_hosts = node.clear();
832
833        // Should have cleared: 10.1.1.128/25, 10.1.1.192/26, 10.1.1.224/27, 10.1.1.240/28
834        let expected_cleared = vec![
835            "10.1.1.128/25",
836            "10.1.1.192/26",
837            "10.1.1.224/27",
838            "10.1.1.240/28",
839        ];
840
841        assert_eq!(
842            cleared_hosts.len(),
843            expected_cleared.len(),
844            "Should have cleared {} networks, got {}: {:?}",
845            expected_cleared.len(),
846            cleared_hosts.len(),
847            cleared_hosts
848        );
849
850        // Check that all expected networks were cleared
851        for expected in &expected_cleared {
852            assert!(
853                cleared_hosts.contains(&expected.to_string()),
854                "Should have cleared {}",
855                expected
856            );
857        }
858
859        // Verify the cleared networks are no longer accessible or fall back to parent
860        for cleared in &expected_cleared {
861            let net = IpNet::from_str(cleared).unwrap();
862            let result = tree.get(&net);
863            assert!(
864                result.is_none() || result == Some("10.1.1.0/24".to_string()),
865                "Cleared network {} should not be accessible or should fall back to parent",
866                cleared
867            );
868        }
869
870        // Verify that 10.1.1.0/24 itself is still accessible
871        assert_eq!(tree.get(&target_net), Some("10.1.1.0/24".to_string()));
872
873        // Verify that unrelated networks are still accessible
874        let unrelated = vec![
875            "10.0.0.0/8",
876            "10.1.0.0/16",
877            "10.1.2.0/24",
878            "192.168.0.0/16",
879            "192.168.1.0/24",
880        ];
881        for net_str in &unrelated {
882            let net = IpNet::from_str(net_str).unwrap();
883            assert_eq!(
884                tree.get(&net),
885                Some(net_str.to_string()),
886                "Unrelated network {} should still be accessible",
887                net_str
888            );
889        }
890    }
891
892    #[test]
893    fn test_acl_mode_skip_existing() {
894        let mut tree = IpRadixTree::new(true);
895
896        // First insertion should succeed
897        let net = IpNet::from_str("192.168.1.0/24").unwrap();
898        let result1 = tree.insert(net);
899        assert_eq!(result1, Some("192.168.1.0/24".to_string()));
900
901        // Second insertion of same network should return None
902        let result2 = tree.insert(net);
903        assert_eq!(result2, None);
904
905        // Different network should still work
906        let other_net = IpNet::from_str("10.0.0.0/8").unwrap();
907        let result3 = tree.insert(other_net);
908        assert_eq!(result3, Some("10.0.0.0/8".to_string()));
909
910        // Verify both networks are accessible
911        assert_eq!(tree.get(&net), Some("192.168.1.0/24".to_string()));
912        assert_eq!(tree.get(&other_net), Some("10.0.0.0/8".to_string()));
913    }
914
915    #[test]
916    fn test_acl_mode_skip_children() {
917        let mut tree = IpRadixTree::new(true);
918
919        // Insert parent network first
920        let parent = IpNet::from_str("192.168.0.0/16").unwrap();
921        assert_eq!(tree.insert(parent), Some("192.168.0.0/16".to_string()));
922        assert_eq!(tree.get(&parent), Some("192.168.0.0/16".to_string()));
923
924        // Insert child network should return None (already covered by parent)
925        let child = IpNet::from_str("192.168.1.0/24").unwrap();
926        assert_eq!(tree.insert(child), None);
927
928        // Get child network should return parent
929        assert_eq!(tree.get(&child), Some("192.168.0.0/16".to_string()));
930
931        // Test with IPv6 as well
932        let parent_v6 = IpNet::from_str("2001:db8::/32").unwrap();
933        assert_eq!(tree.insert(parent_v6), Some("2001:db8::/32".to_string()));
934
935        let child_v6 = IpNet::from_str("2001:db8:abcd::/48").unwrap();
936        assert_eq!(tree.insert(child_v6), None);
937        assert_eq!(tree.get(&child_v6), Some("2001:db8::/32".to_string()));
938    }
939
940    #[test]
941    fn test_acl_mode_clear_children() {
942        let mut tree = IpRadixTree::new(true);
943
944        // Insert child networks first
945        let child1 = IpNet::from_str("192.168.1.0/24").unwrap();
946        let child2 = IpNet::from_str("192.168.2.0/24").unwrap();
947        let child3 = IpNet::from_str("192.168.3.0/24").unwrap();
948
949        tree.insert(child1);
950        tree.insert(child2);
951        tree.insert(child3);
952
953        // Verify children are accessible
954        assert_eq!(tree.get(&child1), Some("192.168.1.0/24".to_string()));
955        assert_eq!(tree.get(&child2), Some("192.168.2.0/24".to_string()));
956        assert_eq!(tree.get(&child3), Some("192.168.3.0/24".to_string()));
957
958        // Insert parent network - should clear children
959        let parent = IpNet::from_str("192.168.0.0/16").unwrap();
960        let result = tree.insert(parent);
961        assert_eq!(result, Some("192.168.0.0/16".to_string()));
962
963        // Parent should be accessible
964        assert_eq!(tree.get(&parent), Some("192.168.0.0/16".to_string()));
965
966        // Children should now fall back to parent
967        assert_eq!(tree.get(&child1), Some("192.168.0.0/16".to_string()));
968        assert_eq!(tree.get(&child2), Some("192.168.0.0/16".to_string()));
969        assert_eq!(tree.get(&child3), Some("192.168.0.0/16".to_string()));
970
971        // Test with IPv6 as well
972        let child_v6_1 = IpNet::from_str("2001:db8:abcd::/48").unwrap();
973        let child_v6_2 = IpNet::from_str("2001:db8:beef::/48").unwrap();
974
975        tree.insert(child_v6_1);
976        tree.insert(child_v6_2);
977
978        assert_eq!(
979            tree.get(&child_v6_1),
980            Some("2001:db8:abcd::/48".to_string())
981        );
982        assert_eq!(
983            tree.get(&child_v6_2),
984            Some("2001:db8:beef::/48".to_string())
985        );
986
987        // Insert parent IPv6 network
988        let parent_v6 = IpNet::from_str("2001:db8::/32").unwrap();
989        let result_v6 = tree.insert(parent_v6);
990        assert_eq!(result_v6, Some("2001:db8::/32".to_string()));
991
992        // Children should fall back to parent
993        assert_eq!(tree.get(&child_v6_1), Some("2001:db8::/32".to_string()));
994        assert_eq!(tree.get(&child_v6_2), Some("2001:db8::/32".to_string()));
995    }
996}
radixtarget_rust/ip.rs

radixtarget_rust/
ip.rs
