use std::error::Error; use std::fmt::Display; use std::fs::File; use std::io::Read; use std::path::Path; use once_cell::sync::OnceCell; use crate::ToUIntIP; const HEADER_INFO_LENGTH: usize = 256; const VECTOR_INDEX_COLS: usize = 256; const VECTOR_INDEX_SIZE: usize = 8; const SEGMENT_INDEX_SIZE: usize = 14; const VECTOR_INDEX_LENGTH: usize = 512 * 1024; const XDB_FILEPATH_ENV: &str = "XDB_FILEPATH"; static CACHE: OnceCell> = OnceCell::new(); /// check https://mp.weixin.qq.com/s/ndjzu0BgaeBmDOCw5aqHUg for details pub fn search_by_ip(ip: T) -> Result> where T: ToUIntIP + Display, { let ip = ip.to_u32_ip()?; let il0 = ((ip >> 24) & 0xFF) as usize; let il1 = ((ip >> 16) & 0xFF) as usize; let idx = VECTOR_INDEX_SIZE * (il0 * VECTOR_INDEX_COLS + il1); let start_point = idx; let vector_cache = get_vector_index_cache(); let start_ptr = get_block_by_size(vector_cache, start_point, 4); let end_ptr = get_block_by_size(vector_cache, start_point + 4, 4); let mut left: usize = 0; let mut right: usize = (end_ptr - start_ptr) / SEGMENT_INDEX_SIZE; while left <= right { let mid = (left + right) >> 1; let offset = start_ptr + mid * SEGMENT_INDEX_SIZE; let buffer_ip_value = &get_full_cache()[offset..offset + SEGMENT_INDEX_SIZE]; let start_ip = get_block_by_size(buffer_ip_value, 0, 4); if ip < (start_ip as u32) { right = mid - 1; } else if ip > (get_block_by_size(buffer_ip_value, 4, 4) as u32) { left = mid + 1; } else { let data_length = get_block_by_size(buffer_ip_value, 8, 2); let data_offset = get_block_by_size(buffer_ip_value, 10, 4); let result = String::from_utf8( get_full_cache()[data_offset..(data_offset + data_length)].to_vec(), ); return Ok(result?); } } Err("not matched".into()) } /// it will check ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb fn default_detect_xdb_file() -> Result> { let prefix = "../".to_owned(); for recurse in 1..4 { let filepath = prefix.repeat(recurse) + "data/ip2region.xdb"; if Path::new(filepath.as_str()).exists() { return Ok(filepath); } } Err("default filepath not find the xdb file, so you must set xdb_filepath".into()) } #[inline] pub fn get_block_by_size(bytes: &[u8], offset: usize, length: usize) -> usize { let mut result: usize = 0; for (index, value) in bytes[offset..offset + length].iter().enumerate() { result += usize::from(*value) << (index << 3); } result } pub fn searcher_init(xdb_filepath: Option) { let xdb_filepath = xdb_filepath.unwrap_or_else(|| default_detect_xdb_file().unwrap()); std::env::set_var(XDB_FILEPATH_ENV, xdb_filepath); CACHE.get_or_init(load_file); } pub fn get_vector_index_cache() -> &'static [u8] { let full_cache: &'static Vec = get_full_cache(); &full_cache[HEADER_INFO_LENGTH..(HEADER_INFO_LENGTH + VECTOR_INDEX_LENGTH)] } fn load_file() -> Vec { let xdb_filepath = std::env::var("XDB_FILEPATH").unwrap_or_else(|_| default_detect_xdb_file().unwrap()); tracing::debug!("load xdb searcher file at {} ", xdb_filepath); let mut f = File::open(xdb_filepath).expect("file open error"); let mut buffer = Vec::new(); f.read_to_end(&mut buffer).expect("load file error"); buffer } pub fn get_full_cache() -> &'static Vec { CACHE.get_or_init(load_file) } #[cfg(test)] mod tests { use std::fs::File; use std::io::Read; use std::net::Ipv4Addr; use std::str::FromStr; use std::thread; use super::*; ///test all types find correct #[test] fn test_multi_type_ip() { searcher_init(None); search_by_ip("2.0.0.0").unwrap(); search_by_ip("32").unwrap(); search_by_ip(4294408949).unwrap(); search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap(); } #[test] fn test_match_all_ip_correct() { searcher_init(None); let mut file = File::open("../../../data/ip.test.txt").unwrap(); let mut contents = String::new(); file.read_to_string(&mut contents).unwrap(); for line in contents.split("\n") { if !line.contains("|") { continue; } let ip_test_line = line.splitn(3, "|").collect::>(); let start_ip = Ipv4Addr::from_str(ip_test_line[0]).unwrap(); let end_ip = Ipv4Addr::from_str(ip_test_line[1]).unwrap(); for value in u32::from(start_ip)..u32::from(end_ip) + 1 { let result = search_by_ip(value).unwrap(); assert_eq!(result.as_str(), ip_test_line[2]) } } } #[test] fn test_multi_thread_only_load_xdb_once() { searcher_init(None); let handle = thread::spawn(|| { let result = search_by_ip("2.2.2.2").unwrap(); println!("ip search in spawn: {result}"); }); let r = search_by_ip("1.1.1.1").unwrap(); println!("ip search in main thread: {r}"); handle.join().unwrap(); } #[test] fn test_multi_searcher_init() { for _ in 0..5 { thread::spawn(|| { searcher_init(None); }); } searcher_init(None); searcher_init(Some(String::from("test"))); search_by_ip(123).unwrap(); } }