radixtarget_rust/
utils.rs

1// utils.rs: Utility functions for radix trees
2use ipnet::{IpNet, Ipv4Net, Ipv6Net};
3use std::collections::hash_map::DefaultHasher;
4use std::hash::{Hash, Hasher};
5use std::net::IpAddr;
6
7/// Hash a value to u64 using the default hasher
8pub fn hash_u64<T: Hash + ?Sized>(value: &T) -> u64 {
9    let mut hasher = DefaultHasher::new();
10    value.hash(&mut hasher);
11    hasher.finish()
12}
13
14/// Convert an IP network to a vector of bits for radix tree traversal
15pub fn ipnet_to_bits(net: &IpNet) -> Vec<u8> {
16    let (addr, prefix) = match net {
17        IpNet::V4(n) => (n.network().octets().to_vec(), net.prefix_len()),
18        IpNet::V6(n) => (n.network().octets().to_vec(), net.prefix_len()),
19    };
20    let mut bits = Vec::with_capacity(prefix as usize);
21    for byte in addr {
22        for i in (0..8).rev() {
23            if bits.len() == prefix as usize {
24                return bits;
25            }
26            bits.push((byte >> i) & 1);
27        }
28    }
29    bits
30}
31
32/// Canonicalize an IP network by ensuring it uses the network address
33pub fn canonicalize_ipnet(network: &IpNet) -> IpNet {
34    match network {
35        IpNet::V4(n) => IpNet::V4(Ipv4Net::new(n.network(), n.prefix_len()).unwrap()),
36        IpNet::V6(n) => IpNet::V6(Ipv6Net::new(n.network(), n.prefix_len()).unwrap()),
37    }
38}
39
40/// Normalize a DNS hostname to canonical form (IDNA + lowercase)
41pub fn normalize_dns(hostname: &str) -> String {
42    idna::domain_to_ascii(hostname)
43        .unwrap_or_else(|_| hostname.to_string())
44        .to_lowercase()
45}
46
47/// Host size key function for sorting - equivalent to Python's host_size_key
48/// Returns (priority, string_repr) where priority is:
49/// - For IP networks: negative number of addresses (so bigger networks come first)
50/// - For DNS names: positive length (so shorter domains come first)
51pub fn host_size_key(host: &str) -> (i64, String) {
52    // Try to parse as IP network first
53    if let Ok(ipnet) = host.parse::<IpNet>() {
54        let num_addresses = match ipnet {
55            IpNet::V4(net) => (32 - net.prefix_len()) as u64,
56            IpNet::V6(net) => (128 - net.prefix_len()) as u64,
57        };
58        // Format as network address with prefix length (like Python's ipaddress module)
59        (
60            -(num_addresses as i64),
61            format!("{}/{}", ipnet.network(), ipnet.prefix_len()),
62        )
63    } else if let Ok(ipaddr) = host.parse::<IpAddr>() {
64        // Single IP address - use original string to avoid re-formatting
65        if ipaddr.is_ipv4() {
66            (-1, format!("{}/32", host))
67        } else {
68            (-1, format!("{}/128", host))
69        }
70    } else {
71        // DNS name - normalize and return length
72        let canonical = normalize_dns(host);
73        (canonical.len() as i64, canonical)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn test_host_size_key_single_ipv4() {
83        assert_eq!(host_size_key("1.2.3.4"), (-1, "1.2.3.4/32".to_string()));
84    }
85
86    #[test]
87    fn test_host_size_key_single_ipv6() {
88        assert_eq!(host_size_key("::1"), (-1, "::1/128".to_string()));
89    }
90
91    #[test]
92    fn test_host_size_key_ipv4_networks() {
93        // /24 network = 32-24 = 8 host bits
94        assert_eq!(host_size_key("1.2.3.0/24"), (-8, "1.2.3.0/24".to_string()));
95
96        // /28 network = 32-28 = 4 host bits
97        assert_eq!(host_size_key("1.2.3.0/28"), (-4, "1.2.3.0/28".to_string()));
98
99        // /30 network = 32-30 = 2 host bits
100        assert_eq!(host_size_key("1.2.3.4/30"), (-2, "1.2.3.4/30".to_string()));
101    }
102
103    #[test]
104    fn test_host_size_key_ipv6_networks() {
105        // /64 network = 128-64 = 64 host bits (network address is normalized)
106        assert_eq!(host_size_key("::1/64"), (-64, "::/64".to_string()));
107
108        // /120 network = 128-120 = 8 host bits (network address is normalized)
109        assert_eq!(host_size_key("::1/120"), (-8, "::/120".to_string()));
110    }
111
112    #[test]
113    fn test_host_size_key_dns_names() {
114        assert_eq!(
115            host_size_key("evilcorp.com"),
116            (12, "evilcorp.com".to_string())
117        );
118        assert_eq!(
119            host_size_key("www.evilcorp.com"),
120            (16, "www.evilcorp.com".to_string())
121        );
122        assert_eq!(
123            host_size_key("api.www.evilcorp.com"),
124            (20, "api.www.evilcorp.com".to_string())
125        );
126    }
127
128    #[test]
129    fn test_host_size_key_dns_normalization() {
130        // Test IDNA normalization and lowercasing
131        assert_eq!(
132            host_size_key("EXAMPLE.COM"),
133            (11, "example.com".to_string())
134        );
135        assert_eq!(
136            host_size_key("Example.Com"),
137            (11, "example.com".to_string())
138        );
139    }
140
141    #[test]
142    fn test_normalize_dns_case_insensitive() {
143        assert_eq!(normalize_dns("EXAMPLE.COM"), "example.com");
144        assert_eq!(
145            normalize_dns("MiXeD.CaSe.DoMaIn.CoM"),
146            "mixed.case.domain.com"
147        );
148    }
149
150    #[test]
151    fn test_normalize_dns_unicode_punycode() {
152        assert_eq!(normalize_dns("café.com"), "xn--caf-dma.com");
153        assert_eq!(normalize_dns("日本.jp"), "xn--wgv71a.jp");
154    }
155}
radixtarget_rust/utils.rs

radixtarget_rust/
utils.rs
