xdb_parse/
search.rs

1//! IP geolocation lookup functions.
2//!
3//! Provides four search entry points with different trade-offs:
4//!
5//! | Function | Input | When to use |
6//! |----------|-------|-------------|
7//! | [`search_by_uint`] | `u32` | You have a raw IPv4 integer |
8//! | [`search_by_u128`] | `u128` | You have a raw IPv6 integer |
9//! | [`search_by_ipaddr`] | `IpAddr` | You already have a parsed `IpAddr` |
10//! | [`search_ip`] | `&str` | You have an IP string (dotted-decimal, colon-hex, or numeric) |
11//!
12//! All functions return `&str` borrowing from the loaded xdb data — zero heap allocation per query.
13
14use std::net::IpAddr;
15
16use crate::error::XdbError;
17use crate::ip::IntoIpAddr;
18use crate::{
19    IPV4_SEGMENT_INDEX_BLOCK_SIZE, IPV6_SEGMENT_INDEX_BLOCK_SIZE, TOTAL_HEADER_SIZE,
20    TOTAL_VECTOR_INDEX_SIZE, VECTOR_COL_SIZE, VECTOR_INDEX_BLOCK_SIZE,
21};
22
23/// Search by raw `u32` IPv4 address.
24///
25/// This is the fastest path — no string parsing, no `IpAddr` conversion.
26/// Returns a `&str` borrowing from the loaded xdb buffer.
27///
28/// # Example
29///
30/// ```no_run
31/// let data = xdb_parse::load_file("./assets/ip2region_v4.xdb".into())?;
32/// let ip: u32 = 0x4a7d_2b63; // 74.125.43.99
33/// let location = xdb_parse::search_by_uint(ip, &data)?;
34/// println!("{}", location);
35/// # Ok::<(), xdb_parse::error::XdbError>(())
36/// ```
37pub fn search_by_uint(ip: u32, data: &[u8]) -> Result<&str, XdbError> {
38    let il0 = ((ip >> 24) & 0xFF) as usize;
39    let il1 = ((ip >> 16) & 0xFF) as usize;
40
41    let idx = il0 as u32 * VECTOR_COL_SIZE as u32 * VECTOR_INDEX_BLOCK_SIZE as u32
42        + il1 as u32 * VECTOR_INDEX_BLOCK_SIZE as u32;
43
44    let vec_segment = &data[TOTAL_HEADER_SIZE..TOTAL_HEADER_SIZE + TOTAL_VECTOR_INDEX_SIZE];
45    let slice = &vec_segment[idx as usize..idx as usize + VECTOR_INDEX_BLOCK_SIZE];
46    let start_ptr = u32::from_le_bytes(slice[0..4].try_into()?);
47    let end_ptr = u32::from_le_bytes(slice[4..8].try_into()?);
48
49    let mut left: usize = 0;
50    let mut right: usize = ((end_ptr - start_ptr) / IPV4_SEGMENT_INDEX_BLOCK_SIZE) as usize;
51
52    while left < right {
53        let mid = (left + right) / 2;
54        let offset = start_ptr as usize + mid * IPV4_SEGMENT_INDEX_BLOCK_SIZE as usize;
55        let block: [u8; IPV4_SEGMENT_INDEX_BLOCK_SIZE as usize] =
56            data[offset..offset + IPV4_SEGMENT_INDEX_BLOCK_SIZE as usize].try_into()?;
57        let start_ip = read_u32(&block, 0);
58        if ip < start_ip {
59            right = mid;
60            continue;
61        }
62        let end_ip = read_u32(&block, 4);
63        if ip > end_ip {
64            left = mid + 1;
65        } else {
66            let data_len = read_u16(&block, 8) as usize;
67            let data_ptr = read_u32(&block, 10) as usize;
68            let bytes = &data[data_ptr..data_ptr + data_len];
69            return std::str::from_utf8(bytes).map_err(|_| XdbError::InvalidUtf8);
70        }
71    }
72    Ok("")
73}
74
75/// Search by raw `u128` IPv6 address.
76///
77/// This is the fastest path for IPv6 — no string parsing, no `IpAddr` conversion.
78/// Returns a `&str` borrowing from the loaded xdb buffer.
79///
80/// # Example
81///
82/// ```no_run
83/// let data = xdb_parse::load_file("./assets/ip2region_v6.xdb".into())?;
84/// let ip: u128 = 0x2001_0db8_85a3_0000_0000_8a2e_0370_7334;
85/// let location = xdb_parse::search_by_u128(ip, &data)?;
86/// println!("{}", location);
87/// # Ok::<(), xdb_parse::error::XdbError>(())
88/// ```
89pub fn search_by_u128(ip: u128, data: &[u8]) -> Result<&str, XdbError> {
90    let il0 = ((ip >> 120) & 0xFF) as usize;
91    let il1 = ((ip >> 112) & 0xFF) as usize;
92
93    let idx = il0 as u32 * VECTOR_COL_SIZE as u32 * VECTOR_INDEX_BLOCK_SIZE as u32
94        + il1 as u32 * VECTOR_INDEX_BLOCK_SIZE as u32;
95
96    let vec_segment = &data[TOTAL_HEADER_SIZE..TOTAL_HEADER_SIZE + TOTAL_VECTOR_INDEX_SIZE];
97    let slice = &vec_segment[idx as usize..idx as usize + VECTOR_INDEX_BLOCK_SIZE];
98    let start_ptr = u32::from_le_bytes(slice[0..4].try_into()?);
99    let end_ptr = u32::from_le_bytes(slice[4..8].try_into()?);
100
101    let mut left: usize = 0;
102    let mut right: usize = ((end_ptr - start_ptr) / IPV6_SEGMENT_INDEX_BLOCK_SIZE) as usize;
103
104    while left < right {
105        let mid = (left + right) / 2;
106        let offset = start_ptr as usize + mid * IPV6_SEGMENT_INDEX_BLOCK_SIZE as usize;
107        let block: [u8; IPV6_SEGMENT_INDEX_BLOCK_SIZE as usize] =
108            data[offset..offset + IPV6_SEGMENT_INDEX_BLOCK_SIZE as usize].try_into()?;
109        let start_ip = read_u128(&block, 0);
110        if ip < start_ip {
111            right = mid;
112            continue;
113        }
114        let end_ip = read_u128(&block, 16);
115        if ip > end_ip {
116            left = mid + 1;
117        } else {
118            let data_len = read_u16(&block, 32) as usize;
119            let data_ptr = read_u32(&block, 34) as usize;
120            let bytes = &data[data_ptr..data_ptr + data_len];
121            return std::str::from_utf8(bytes).map_err(|_| XdbError::InvalidUtf8);
122        }
123    }
124    Ok("")
125}
126
127/// Search by any IP-compatible input.
128///
129/// Accepts all types that implement [`IntoIpAddr`]:
130///
131/// | Type | Example |
132/// |------|---------|
133/// | `&str` / `String` | `"192.168.1.1"`, `"::1"`, `"3232235777"` |
134/// | `u32` | `0xC0A80101`, `3232235777` |
135/// | `u128` | `0x2001_0db8_...` |
136/// | `IpAddr` / `Ipv4Addr` / `Ipv6Addr` | from socket addresses |
137///
138/// # Example
139///
140/// ```no_run
141/// let data = xdb_parse::load_file("./assets/ip2region_v4.xdb".into())?;
142///
143/// // string
144/// let loc = xdb_parse::search_ip("73.24.63.66", &data)?;
145///
146/// // raw u32 (fastest)
147/// let loc = xdb_parse::search_ip(0x4918_3F42u32, &data)?;
148///
149/// // IpAddr from socket
150/// let addr = "73.24.63.66".parse::<std::net::IpAddr>()?;
151/// let loc = xdb_parse::search_ip(addr, &data)?;
152/// # Ok::<(), xdb_parse::error::XdbError>(())
153/// ```
154pub fn search_ip(ip: impl IntoIpAddr, data: &[u8]) -> Result<&str, XdbError> {
155    search_by_ipaddr(ip.into_ip_addr()?, data)
156}
157
158/// Search by [`std::net::IpAddr`].
159///
160/// Dispatches to [`search_by_uint`] for IPv4 or [`search_by_u128`] for IPv6.
161/// Use this when you already have a parsed `IpAddr` from another source
162/// (e.g. socket address).
163///
164/// # Example
165///
166/// ```no_run
167/// use std::net::IpAddr;
168///
169/// let data = xdb_parse::load_file("./assets/ip2region_v4.xdb".into())?;
170/// let ip: IpAddr = "73.24.63.66".parse()?;
171/// let location = xdb_parse::search_by_ipaddr(ip, &data)?;
172/// # Ok::<(), xdb_parse::error::XdbError>(())
173/// ```
174pub fn search_by_ipaddr(ip: IpAddr, data: &[u8]) -> Result<&str, XdbError> {
175    match ip {
176        IpAddr::V4(v4) => search_by_uint(v4.to_bits(), data),
177        IpAddr::V6(v6) => search_by_u128(v6.to_bits(), data),
178    }
179}
180
181#[inline(always)]
182fn read_u32<const N: usize>(block: &[u8; N], offset: usize) -> u32 {
183    u32::from_le_bytes(unsafe { *(block.as_ptr().add(offset) as *const [u8; 4]) })
184}
185
186#[inline(always)]
187fn read_u16<const N: usize>(block: &[u8; N], offset: usize) -> u16 {
188    u16::from_le_bytes(unsafe { *(block.as_ptr().add(offset) as *const [u8; 2]) })
189}
190
191#[inline(always)]
192fn read_u128<const N: usize>(block: &[u8; N], offset: usize) -> u128 {
193    u128::from_le_bytes(unsafe { *(block.as_ptr().add(offset) as *const [u8; 16]) })
194}
195
196#[cfg(test)]
197mod tests {
198    use std::sync::Arc;
199    use std::thread;
200    use std::time::Instant;
201
202    use anyhow::Result;
203
204    use super::*;
205    use crate::ip::parse_ip;
206    use crate::load_file;
207
208    #[test]
209    fn test_search_ipv4() -> Result<()> {
210        let path = "./assets/ip2region_v4.xdb";
211        let data = load_file(path.into())?;
212        let ret = search_ip("73.24.63.66", &data)?;
213        println!("result: {ret}");
214        assert!(!ret.is_empty());
215        Ok(())
216    }
217
218    #[test]
219    fn test_search_ipv6() -> Result<()> {
220        let path = "./assets/ip2region_v6.xdb";
221        let data = load_file(path.into())?;
222        let ret = search_ip("2001:0db8:85a3:0000:0000:8a2e:0370:7334", &data)?;
223        println!("result: {ret}");
224        assert!(!ret.is_empty());
225        Ok(())
226    }
227
228    #[test]
229    fn test_search_by_uint() -> Result<()> {
230        let path = "./assets/ip2region_v4.xdb";
231        let data = load_file(path.into())?;
232        let ret = search_by_uint(0x4a7d_2b63, &data)?; // 74.125.43.99
233        println!("result: {ret}");
234        assert!(!ret.is_empty());
235        Ok(())
236    }
237
238    #[test]
239    fn test_search_by_u128() -> Result<()> {
240        let path = "./assets/ip2region_v6.xdb";
241        let data = load_file(path.into())?;
242        let ip: u128 = 0x2001_0db8_85a3_0000_0000_8a2e_0370_7334;
243        let ret = search_by_u128(ip, &data)?;
244        println!("result: {ret}");
245        assert!(!ret.is_empty());
246        Ok(())
247    }
248
249    #[test]
250    fn test_multi_thread() -> Result<()> {
251        let start = Instant::now();
252        let path = "./assets/ip2region_v6.xdb";
253        let data = load_file(path.into())?;
254        let data = Arc::new(data);
255
256        let data_clone = Arc::clone(&data);
257        let handle = thread::spawn(move || {
258            search_ip("2408:8352:da10:1ad:c283:c9ff:fec6:4046", &data_clone).unwrap();
259        });
260        search_ip("2408:8352:da10:1ad:c283:c9ff:fec6:4046", &data).unwrap();
261        println!("use time:{:?}", start.elapsed());
262        handle.join().unwrap();
263        Ok(())
264    }
265
266    #[test]
267    fn test_parse_ip() -> Result<()> {
268        println!("{:?}", parse_ip("192.168.1.1")?);
269        println!("{:?}", parse_ip("2400:3200::1")?);
270        println!("{:?}", parse_ip("3232235776")?);
271        Ok(())
272    }
273}
xdb_parse/search.rs

xdb_parse/
search.rs
