From a5f0837bb53e16b3b8e529e9193d2fb8f8bb965b Mon Sep 17 00:00:00 2001 From: gongzhengyang Date: Tue, 20 Dec 2022 09:56:51 +0800 Subject: [PATCH] feat: add ToUIntIP for types ip value --- binding/rust/src/ip_value.rs | 61 ++++++++++++++++++++++++++ binding/rust/src/main.rs | 12 ++--- binding/rust/src/{lib.rs => search.rs} | 59 ++++++++++++++++++------- 3 files changed, 112 insertions(+), 20 deletions(-) create mode 100644 binding/rust/src/ip_value.rs rename binding/rust/src/{lib.rs => search.rs} (57%) diff --git a/binding/rust/src/ip_value.rs b/binding/rust/src/ip_value.rs new file mode 100644 index 0000000..0a11f92 --- /dev/null +++ b/binding/rust/src/ip_value.rs @@ -0,0 +1,61 @@ +use std::error::Error; +use std::net::Ipv4Addr; +use std::str::FromStr; + +pub trait ToUIntIP { + fn to_u32_ip(&self) -> Result>; +} + +impl ToUIntIP for u32 { + fn to_u32_ip(&self) -> Result> { + Ok(self.to_owned()) + } +} + +impl ToUIntIP for &str { + fn to_u32_ip(&self) -> Result> { + if let Ok(ip_addr) = Ipv4Addr::from_str(self) { + return Ok(u32::from(ip_addr)); + } + Ok(self.parse::()?) + } +} + +impl ToUIntIP for Ipv4Addr { + fn to_u32_ip(&self) -> Result> { + Ok(u32::from(*self)) + } +} + +#[cfg(test)] +mod test_ip { + use super::*; + + #[test] + fn test_ip_str_2_u32() { + let ip_str = "1.1.1.1"; + let result = ip_str.to_u32_ip().unwrap(); + assert_eq!(result, 1 << 24 | 1 << 16 | 1 << 8 | 1); + } + + #[test] + fn test_ip_u32_str() { + let ip = "12"; + let result = ip.to_u32_ip().unwrap(); + assert_eq!(result, 12); + } + + #[test] + fn test_ip_u32() { + let ip: u32 = 33; + let result = ip.to_u32_ip().unwrap(); + assert_eq!(result, 33); + } + + #[test] + fn test_ip_addr() { + let ip = Ipv4Addr::from_str("0.0.3.12").unwrap(); + let result = ip.to_u32_ip().unwrap(); + assert_eq!(result, 3 << 8 | 12) + } +} diff --git a/binding/rust/src/main.rs b/binding/rust/src/main.rs index 86682d1..023c20c 100644 --- a/binding/rust/src/main.rs +++ b/binding/rust/src/main.rs @@ -1,8 +1,10 @@ -mod lib; +mod ip_value; +mod search; fn main() { - let filepath = "../../data/ip2region.xdb"; - let searcher = lib::Searcher::new(filepath).expect("load file error"); - let result = searcher.search_by_ip("1.2.165.128"); - println!("{:?}", result); + println!(""); + // let filepath = "../../data/ip2region.xdb"; + // let searcher = lib::Searcher::new(filepath).expect("load file error"); + // let result = searcher.search_by_ip("1.2.165.128"); + // println!("{:?}", result); } diff --git a/binding/rust/src/lib.rs b/binding/rust/src/search.rs similarity index 57% rename from binding/rust/src/lib.rs rename to binding/rust/src/search.rs index 12762a0..e2f496b 100644 --- a/binding/rust/src/lib.rs +++ b/binding/rust/src/search.rs @@ -1,6 +1,9 @@ +use std::error::Error; use std::fs::File; use std::io::Read; use std::net::Ipv4Addr; +use std::str::FromStr; +use crate::ip_value::ToUIntIP; const HEADER_INFO_LENGTH: u32 = 256; // const VECTOR_INDEX_ROWS: u32 = 256; @@ -13,22 +16,22 @@ pub struct Searcher { } impl Searcher { - pub fn new(filepath: &'static str) -> Result> { + pub fn new(filepath: &'static str) -> Result> { let mut f = File::open(filepath)?; let mut buffer = Vec::new(); f.read_to_end(&mut buffer)?; Ok(Self { buffer }) } - pub fn search_by_ip(&self, ip: &'static str) -> Result> { - let ip = ip.parse::().unwrap_or_else(|_| { - let ip = ip - .parse::() - .expect("ip is not a valid ip or valid int"); - Ipv4Addr::from(ip) - }); + pub fn search_by_ip(&self, ip: T) -> Result> + where + T: ToUIntIP + { + let changed_value = ip.to_u32_ip()?; + self.search_by_ip_u32(changed_value) + } - let ip = u32::from(ip); + pub fn search_by_ip_u32(&self, ip: u32) -> Result> { let il0 = (ip >> 24) & 0xFF; let il1 = (ip >> 16) & 0xFF; let idx = VECTOR_INDEX_SIZE * (il0 * VECTOR_INDEX_COLS + il1); @@ -80,13 +83,39 @@ fn get_u32(bytes: &[u8], offset: usize) -> u32 { mod tests { use super::*; + ///test all types find correct #[test] fn test_search_by_ip() { - let filepath = "../../data/ip2region.xdb"; - let searcher = Searcher::new(filepath).expect("load file error"); - let result = searcher.search_by_ip("2.0.0.0"); - println!("{:?}", result); + let searcher = Searcher::new(get_xdb_filepath()).expect("load file error"); + searcher.search_by_ip("2.0.0.0").unwrap(); + searcher.search_by_ip("32").unwrap(); + searcher.search_by_ip(32).unwrap(); + searcher.search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap(); + } + + fn get_xdb_filepath() -> &'static str { + "../../data/ip2region.xdb" + } + + /// test find ip correct use the file ip.test.txt in ../../data + #[test] + fn test_random_choose_ip() { + let searcher = Searcher::new(get_xdb_filepath()).unwrap(); + 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 = searcher.search_by_ip(value).unwrap(); + assert_eq!(result.as_str(), ip_test_line[2]) + } + } + } } - -fn main() {}