xdb_parse/
utils.rs

1use std::{
2    fs::File,
3    io::Read,
4    net::{IpAddr, Ipv4Addr, Ipv6Addr},
5    path::PathBuf,
6    str::FromStr,
7};
8
9use crate::{
10    TOTAL_HEADER_SIZE, IPV4_SEGMENT_INDEX_BLOCK_SIZE, IPV6_SEGMENT_INDEX_BLOCK_SIZE,
11    TOTAL_VECTOR_INDEX_SIZE, VECTOR_COL_SIZE, VECTOR_INDEX_BLOCK_SIZE,
12    error::XdbError,
13    xdb::{Header, SegmentIndex},
14};
15///ip database file load to memery
16pub fn load_file(path: PathBuf) -> Result<Vec<u8>, XdbError> {
17    let mut data = File::open(path)?;
18    let mut buf = Vec::new();
19    data.read_to_end(&mut buf)?;
20    Ok(buf)
21}
22///A trivial function
23pub fn print_ip(header: Header, data: &[u8], ip_version: &str) -> Result<(), XdbError> {
24
25    if ip_version == "ipv4" {
26        let count = (header.index_end_address - header.index_base_address + 1)
27            / IPV4_SEGMENT_INDEX_BLOCK_SIZE;
28
29        for i in 0..count {
30            let index_start_addr =
31                header.index_base_address + i as u32 * IPV4_SEGMENT_INDEX_BLOCK_SIZE;
32            let index_end_addr = index_start_addr + IPV4_SEGMENT_INDEX_BLOCK_SIZE;
33            let segment_block = SegmentIndex::<u32>::try_parse(
34                &data[(index_start_addr as usize)..(index_end_addr as usize)],
35            )?;
36            let ret = String::from_utf8(
37                data[segment_block.data_ptr as usize
38                    ..(segment_block.data_ptr + segment_block.data_len as u32) as usize]
39                    .to_vec(),
40            )?;
41            let mut res = String::new();
42            res.push_str(Ipv4Addr::from(segment_block.ip_start).to_string().as_str());
43            res.push_str("|");
44            res.push_str(Ipv4Addr::from(segment_block.ip_end).to_string().as_str());
45            res.push_str("|");
46            res.push_str(ret.as_str());
47
48            println!("{}", res);
49        }
50    } else if ip_version == "ipv6" {
51        let count = (header.index_end_address - header.index_base_address + 1)
52            / IPV6_SEGMENT_INDEX_BLOCK_SIZE;
53        for i in 0..count {
54            let index_start_addr =
55                header.index_base_address + i as u32 * IPV6_SEGMENT_INDEX_BLOCK_SIZE;
56            let index_end_addr = index_start_addr + IPV6_SEGMENT_INDEX_BLOCK_SIZE;
57            let segment_block = SegmentIndex::<u128>::try_parse(
58                &data[(index_start_addr as usize)..(index_end_addr as usize)],
59            )?;
60            let ret = String::from_utf8(
61                data[segment_block.data_ptr as usize
62                    ..(segment_block.data_ptr + segment_block.data_len as u32) as usize]
63                    .to_vec(),
64            )?;
65            let mut res = String::new();
66            res.push_str(Ipv6Addr::from(segment_block.ip_start).to_string().as_str());
67            res.push_str("|");
68            res.push_str(Ipv6Addr::from(segment_block.ip_end).to_string().as_str());
69            res.push_str("|");
70            res.push_str(ret.as_str());
71        }
72    } else {
73        return Err(XdbError::InvalidIPVersion("Invalid IP version".into()));
74    }
75    Ok(())
76}
77
78///search ipv4 or ipv6
79pub fn search_ip(ip: &str, data: &[u8]) -> Result<String, XdbError> {
80    let ip = parse_ip_or_number(ip)?;
81    match ip {
82        IpAddr::V4(ipv4) => {
83            let ip = ipv4.to_bits();
84            let il0 = ((ip >> 24) & 0xFF) as usize;
85            let il1 = ((ip >> 16) & 0xFF) as usize;
86
87            let idx = il0 as u32 * VECTOR_COL_SIZE as u32 * VECTOR_INDEX_BLOCK_SIZE as u32
88                + il1 as u32 * VECTOR_INDEX_BLOCK_SIZE as u32;
89            let vec_segment: &[u8] =
90                &data[TOTAL_HEADER_SIZE..(TOTAL_HEADER_SIZE + TOTAL_VECTOR_INDEX_SIZE)];
91            let slice = &vec_segment[(idx as usize)..(idx as usize + VECTOR_INDEX_BLOCK_SIZE)];
92            let start_ptr = u32::from_le_bytes(slice[0..4].try_into()?);
93            let end_ptr = u32::from_le_bytes(slice[4..8].try_into()?);
94
95            let mut left: usize = 0;
96            let mut right: usize =
97                ((end_ptr - start_ptr) / IPV4_SEGMENT_INDEX_BLOCK_SIZE as u32) as usize;
98
99            while left <= right {
100                let mid = (left + right) / 2;
101                let offset = start_ptr as usize + mid * IPV4_SEGMENT_INDEX_BLOCK_SIZE as usize;
102                let ip_value = &data[offset..offset + IPV4_SEGMENT_INDEX_BLOCK_SIZE as usize];
103                let start_ip = u32::from_le_bytes(ip_value[0..4].try_into()?);
104                if ip < (start_ip as u32) {
105                    right = mid - 1;
106                } else if ip > (u32::from_le_bytes(ip_value[4..8].try_into()?)) {
107                    left = mid + 1;
108                } else {
109                    let data_length = u16::from_le_bytes(ip_value[8..10].try_into()?);
110                    let data_offset = u32::from_le_bytes(ip_value[10..14].try_into()?);
111                    let result = String::from_utf8(
112                        data[data_offset as usize..(data_offset as usize + data_length as usize)]
113                            .to_vec(),
114                    )?;
115                    return Ok(result);
116                }
117            }
118            //now get vector segment data
119        }
120        IpAddr::V6(ipv6) => {
121            let ip = ipv6.to_bits();
122            // For IPv6, use the first two bytes (16 bits) similar to IPv4
123            let il0 = ((ip >> 120) & 0xFF) as usize; // First byte (bits 120-127)
124            let il1 = ((ip >> 112) & 0xFF) as usize; // Second byte (bits 112-119)
125            let vec_segment: &[u8] =
126                &data[TOTAL_HEADER_SIZE..(TOTAL_HEADER_SIZE + TOTAL_VECTOR_INDEX_SIZE)];
127            // Calculate index (same formula as IPv4)
128            let idx = il0 as u32 * VECTOR_COL_SIZE as u32 * VECTOR_INDEX_BLOCK_SIZE as u32
129                + il1 as u32 * VECTOR_INDEX_BLOCK_SIZE as u32;
130            let slice = &vec_segment[(idx as usize)..(idx as usize + VECTOR_INDEX_BLOCK_SIZE)];
131            let start_ptr = u32::from_le_bytes(slice[0..4].try_into()?);
132            let end_ptr = u32::from_le_bytes(slice[4..8].try_into()?);
133
134            let mut left: usize = 0;
135            let mut right: usize =
136                ((end_ptr - start_ptr) / IPV6_SEGMENT_INDEX_BLOCK_SIZE as u32) as usize;
137
138            while left <= right {
139                let mid = (left + right) / 2;
140                let offset = start_ptr as usize + mid * IPV6_SEGMENT_INDEX_BLOCK_SIZE as usize;
141                let ip_value = &data[offset..offset + IPV6_SEGMENT_INDEX_BLOCK_SIZE as usize];
142
143                // Parse u128 start and end IP addresses
144                let start_ip = u128::from_le_bytes(ip_value[0..16].try_into()?);
145                let end_ip = u128::from_le_bytes(ip_value[16..32].try_into()?);
146
147                if ip < start_ip {
148                    right = mid - 1;
149                } else if ip > end_ip {
150                    left = mid + 1;
151                } else {
152                    let data_length = u16::from_le_bytes(ip_value[32..34].try_into()?);
153                    let data_offset = u32::from_le_bytes(ip_value[34..38].try_into()?);
154                    let result = String::from_utf8(
155                        data[data_offset as usize..(data_offset as usize + data_length as usize)]
156                            .to_vec(),
157                    )?;
158                    return Ok(result);
159                }
160            }
161        }
162    };
163    Ok("".into())
164}
165
166///Numeric string ip parse to standard ip format
167fn parse_ip_or_number(ip: &str) -> Result<IpAddr, XdbError> {
168    // 1. normal ip byparse
169    if let Ok(ip_addr) = IpAddr::from_str(ip) {
170        return Ok(ip_addr);
171    }
172
173    // 2. parse to IPv4
174    if let Ok(num) = ip.parse::<u32>() {
175        return Ok(IpAddr::V4(Ipv4Addr::from(num)));
176    }
177
178    // 3. parse to IPv6
179    if let Ok(num) = ip.parse::<u128>() {
180        return Ok(IpAddr::V6(Ipv6Addr::from(num)));
181    }
182    Err(XdbError::InvalidIP("InvalidIP or IP error".into()))
183}
184
185#[cfg(test)]
186mod tests {
187    use std::{sync::Arc, thread, time::Instant};
188
189    use anyhow::Result;
190
191    use crate::utils::{load_file, parse_ip_or_number, search_ip};
192
193    #[test]
194    fn test_multi_thread_only_load_xdb_once() -> Result<()> {
195        /*
196        ────────────
197            Nextest run ID a4d3234a-55bf-4ce9-a5aa-d6244d36ade8 with nextest profile: default
198                Starting 1 test across 2 binaries (1 test skipped)
199                Running [ 00:00:00] 0/1: 0 running, 0 passed, 0 skipped
200                PASS [   0.276s] xdb-parse utils::tests::test_multi_thread_only_load_xdb_once
201            Nextest run ID c1650311-801c-4628-83c5-fe1a46469326 with nextest profile: default
202                Starting 1 test across 2 binaries (1 test skipped)
203                Running [ 00:00:00] 0/1: 0 running, 0 passed, 0 skipped
204                PASS [   0.021s] xdb-parse utils::tests::test_multi_thread_only_load_xdb_once
205         */
206        let start = Instant::now();
207        let path = "./assets/ip2region_v6.xdb";
208        let data = load_file(path.into())?;
209        let data = Arc::new(data);
210        // let path = "./assets/ip2region_v6.xdb";
211
212        let data_clone = Arc::clone(&data);
213        let handle = thread::spawn(move || {
214            // let result = search_ip("2.2.2.2",&get_cache()).unwrap();
215            // let result =search_ip("2001:0db8:85a3:0000:0000:8a2e:0370:7334", &get_cache()).unwrap();
216           search_ip("2408:8352:da10:1ad:c283:c9ff:fec6:4046", &data_clone).unwrap();
217            // println!("ip search in spawn: {result}");
218        });
219        // let r = search_ip("2001:0db8:85a3:0000:0000:8a2e:0370:7334", &get_cache()).unwrap();
220        search_ip("2408:8352:da10:1ad:c283:c9ff:fec6:4046", &data).unwrap();
221        // println!("ip search in main thread: {r}");
222        let time = start.elapsed();
223        println!("use time:{:?}", time);
224        handle.join().unwrap();
225        Ok(())
226    }
227
228    #[test]
229    fn test_use_once_cell_search_ipv4() -> Result<()> {
230        /*
231        ---- utils::tests::test_use_once_cell_search_ipv4 stdout ----
232            use time:5.3644ms-ret:美国|佛罗里达|0|康卡斯特
233         */
234        let start = Instant::now();
235        let path = "./assets/ip2region_v4.xdb";
236        let data = load_file(path.into())?;
237
238        let ret = search_ip("73.24.63.66", &data)?;
239        let time = start.elapsed();
240        println!("use time:{:?}-ret:{}", time, ret);
241        Ok(())
242    }
243    #[test]
244    fn test_use_once_cell_search_ipv6() -> Result<()> {
245        /*
246        ---- utils::tests::test_use_once_cell_search_ipv6 stdout ----
247                use time:251.6989ms-ret:美国|亚拉巴马州|杰斐逊|专线用户
248         */
249        let start = Instant::now();
250        let path = "./assets/ip2region_v6.xdb";
251
252        // xdb_init(Some(path.into())).unwrap();
253        let data = load_file(path.into())?;
254        let ret = search_ip("2001:0db8:85a3:0000:0000:8a2e:0370:7334", &data)?;
255        let time = start.elapsed();
256        println!("use time:{:?}-ret:{}", time, ret);
257        Ok(())
258    }
259
260    #[test]
261    fn test_parse_ip() -> Result<()> {
262        println!("{:?}", parse_ip_or_number("192.168.1.1")?); // V4
263        println!("{:?}", parse_ip_or_number("2400:3200::1")?); // V6
264        println!("{:?}", parse_ip_or_number("3232235776")?); // V4
265        Ok(())
266    }
267}
xdb_parse/utils.rs

xdb_parse/
utils.rs
