#[cfg(feature = "lazy")] #[macro_use] extern crate lazy_static; use std::cell::RefCell; use std::fs::File; use std::io::{self, Read, Seek, SeekFrom}; use std::net::IpAddr; use std::{fmt, str}; mod db; mod error; pub use error::{Error, Result}; mod owned; #[doc(hidden)] pub use db::DB_PATH; pub use owned::{OwnedIp2Region, OwnedIpInfo}; #[cfg(feature = "lazy")] use db::DB_BYTES; #[cfg(feature = "lazy")] pub use owned::{memory_search, memory_search_ip}; const INDEX_BLOCK_LENGTH: u32 = 12; const TOTAL_HEADER_LENGTH: usize = 8192; thread_local!(static BUF: RefCell<[u8;256]> = RefCell::new([0;256])); thread_local!(static BUF_BTREE: RefCell<[u8;TOTAL_HEADER_LENGTH]> = RefCell::new([0;TOTAL_HEADER_LENGTH])); #[allow(non_snake_case)] #[derive(Debug, Default, Clone, PartialEq)] pub struct IpInfo<'a> { pub city_id: u32, pub country: &'a str, pub region: &'a str, pub province: &'a str, pub city: &'a str, pub ISP: &'a str, } impl<'a> IpInfo<'a> { fn new(city_id: u32, fields: &[&'a str]) -> Self { if fields.len() < 5 { panic!(format!("invlid fields: {:?}", fields)); } IpInfo { country: fields[0], region: fields[1], province: fields[2], city: fields[3], ISP: fields[4], city_id, } } pub fn to_owned(&self) -> OwnedIpInfo { OwnedIpInfo { city_id: self.city_id, country: self.country.to_owned(), region: self.region.to_owned(), province: self.province.to_owned(), city: self.city.to_owned(), ISP: self.ISP.to_owned(), } } } impl<'a> fmt::Display for IpInfo<'a> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!( f, "{}|{}|{}|{}|{}|{}", self.city_id, self.country, self.region, self.province, self.city, self.ISP ) } } fn get_ip_info(city_id: u32, line: &[u8]) -> Result { let str = str::from_utf8(line)?; let fields = str.split('|').collect::>(); Ok(IpInfo::new(city_id, &fields[..])) } fn get_u32(bytes: &[u8], offset: u32) -> u32 { let offset = offset as usize; let tmp = (bytes[offset] as i64) & 0x000000FF | ((bytes[offset + 1] as i64) << 8) & 0x0000FF00 | ((bytes[offset + 2] as i64) << 16) & 0x00FF0000 | ((bytes[offset + 3] as i64) << 24) & 0xFF000000; tmp as u32 } fn ip2u32(ip: &IpAddr) -> Result { if ip.is_ipv6() { return Err(Error::UnsupportIpv6); } if ip.is_unspecified() { return Err(Error::IpIsUnspecified); } if ip.is_loopback() { return Err(Error::IpIsLoopback); } if ip.is_multicast() { return Err(Error::IpIsMulticast); } match ip { IpAddr::V4(v4) => { if v4.is_private() { return Err(Error::IpIsPrivate); } let mut sum: u32 = 0; for (i, n) in v4.octets().iter().enumerate() { sum += (*n as u32) << 24 - 8 * i; } return Ok(sum); } IpAddr::V6(_v6) => unreachable!(), } } pub struct Ip2Region { // db file handler db_file: File, //header block info header_sip: Vec, header_ptr: Vec, header_len: u32, // super block index info first_index_ptr: u32, last_index_ptr: u32, total_blocks: u32, } impl Ip2Region { pub fn new(path: &str) -> io::Result { let file = File::open(path)?; Ok(Ip2Region { db_file: file, header_sip: Vec::new(), header_ptr: Vec::new(), header_len: 0, first_index_ptr: 0, last_index_ptr: 0, total_blocks: 0, }) } pub fn to_owned(&mut self) -> Result { OwnedIp2Region::new2(&mut self.db_file).map_err(Error::Io) } pub fn binary_search>(&mut self, ip_str: S) -> Result { let ip = ip_str.as_ref().parse::()?; self.binary_search_ip(&ip) } pub fn binary_search_ip(&mut self, ip_addr: &IpAddr) -> Result { BUF.with(|buf| { let mut buf = buf.borrow_mut(); if self.total_blocks == 0 { self.db_file.seek(SeekFrom::Start(0))?; self.db_file.read_exact(&mut buf[..8])?; self.first_index_ptr = get_u32(&buf[..8], 0); self.last_index_ptr = get_u32(&buf[..8], 4); self.total_blocks = (self.last_index_ptr - self.first_index_ptr) / INDEX_BLOCK_LENGTH + 1; } let ip = ip2u32(ip_addr)?; let mut h = self.total_blocks; let (mut data_ptr, mut l) = (0u32, 0u32); while l <= h { let m = (l + h) >> 1; let p = self.first_index_ptr + m * INDEX_BLOCK_LENGTH; self.db_file.seek(SeekFrom::Start(p as u64))?; self.db_file .read_exact(&mut buf[0..INDEX_BLOCK_LENGTH as usize])?; let sip = get_u32(&buf[..INDEX_BLOCK_LENGTH as usize], 0); if ip < sip { h = m - 1; } else { let eip = get_u32(&buf[..INDEX_BLOCK_LENGTH as usize], 4); if ip > eip { l = m + 1; } else { data_ptr = get_u32(&buf[..INDEX_BLOCK_LENGTH as usize], 8); break; } } } if data_ptr == 0 { Err(Error::NotFound)?; } let data_len = (data_ptr >> 24) & 0xff; data_ptr = data_ptr & 0x00FFFFFF; self.db_file.seek(SeekFrom::Start(data_ptr as u64))?; self.db_file.read_exact(&mut buf[0..data_len as usize])?; get_ip_info( get_u32(&buf[..data_len as usize], 0), &buf[4..data_len as usize], ).map(|i| i.to_owned()) }) } pub fn btree_search>(&mut self, ip_str: S) -> Result { let ip = ip_str.as_ref().parse::()?; self.btree_search_ip(&ip) } pub fn btree_search_ip(&mut self, ip_addr: &IpAddr) -> Result { BUF_BTREE.with(|buf| { let mut buf = buf.borrow_mut(); if self.header_len == 0 { self.db_file.seek(SeekFrom::Start(8))?; self.db_file.read_exact(&mut buf[0..TOTAL_HEADER_LENGTH])?; let (mut i, mut idx) = (0, 0); while i < TOTAL_HEADER_LENGTH { let sip = get_u32(&buf[0..TOTAL_HEADER_LENGTH], i as u32); let idx_ptr = get_u32(&buf[0..TOTAL_HEADER_LENGTH], i as u32 + 4); if idx_ptr == 0 { break; } self.header_sip.push(sip); self.header_ptr.push(idx_ptr); i += 8; idx += 1; } self.header_len = idx } let ip = ip2u32(ip_addr)?; let mut h = self.header_len; let (mut sptr, mut eptr, mut l) = (0u32, 0u32, 0u32); while l <= h { let m = (l + h) >> 1; if m < self.header_len { if ip == self.header_sip[m as usize] { if m > 0 { sptr = self.header_ptr[m as usize - 1]; eptr = self.header_ptr[m as usize]; } else { sptr = self.header_ptr[m as usize]; eptr = self.header_ptr[m as usize + 1]; } break; } if ip < self.header_sip[m as usize] { if m == 0 { sptr = self.header_ptr[m as usize]; eptr = self.header_ptr[m as usize + 1]; break; } else if ip > self.header_sip[m as usize - 1] { sptr = self.header_ptr[m as usize - 1]; eptr = self.header_ptr[m as usize]; break; } h = m - 1 } else { if m == self.header_len - 1 { println!("m/hl: {}/{}", m, self.header_len); sptr = self.header_ptr[m as usize - 1]; eptr = self.header_ptr[m as usize]; break; } else if ip <= self.header_sip[m as usize + 1] { sptr = self.header_ptr[m as usize]; eptr = self.header_ptr[m as usize + 1]; break; } l = m + 1 } } } if sptr == 0 { Err(Error::NotFound)?; } let block_len = eptr - sptr; self.db_file.seek(SeekFrom::Start(sptr as u64))?; let buf_size = (block_len + INDEX_BLOCK_LENGTH) as usize; self.db_file.read_exact(&mut buf[..buf_size])?; let mut data_ptr = 0; h = block_len / INDEX_BLOCK_LENGTH; l = 0; while l <= h { let m = (l + h) >> 1; let p = m * INDEX_BLOCK_LENGTH; let sip = get_u32(&buf[..buf_size], p); if ip < sip { h = m - 1; } else { let eip = get_u32(&buf[..buf_size], p + 4); if ip > eip { l = m + 1; } else { data_ptr = get_u32(&buf[..buf_size], p + 8); break; } } } if data_ptr == 0 { Err(Error::NotFound)?; } let data_len = (data_ptr >> 24) & 0xff; data_ptr = data_ptr & 0x00FFFFFF; self.db_file.seek(SeekFrom::Start(data_ptr as u64))?; self.db_file.read_exact(&mut buf[0..data_len as usize])?; get_ip_info( get_u32(&buf[..data_len as usize], 0), &buf[4..data_len as usize], ).map(|i| i.to_owned()) }) } } // cargo test --features lazy -- --nocapture #[cfg(test)] mod tests { use super::*; #[test] fn it_works() { let mut ip2 = Ip2Region::new(DB_PATH).unwrap(); let ip2o = ip2.to_owned().unwrap(); for ip in vec![ "117.136.105.202", "47.95.47.253", "127.0.0.1", "10.0.0.1", "1.1.1.1", ] { #[cfg(feature = "lazy")] { println!("lzay__: {:?}", memory_search(ip)); if ip2o.memory_search(ip).is_ok() { assert_eq!(ip2o.memory_search(ip).unwrap(), memory_search(ip).unwrap()); } else { assert!(memory_search(ip).is_err()); } } println!("memory: {:?}", ip2o.memory_search(ip)); println!("binary: {:?}", ip2.binary_search(ip)); println!("btree : {:?}", ip2.btree_search(ip)); if ip2o.memory_search(ip).is_ok() { assert_eq!( ip2o.memory_search(ip).unwrap().to_owned(), ip2.binary_search(ip).unwrap() ); assert_eq!( ip2o.memory_search(ip).unwrap().to_owned(), ip2.btree_search(ip).unwrap() ); } else { assert!(ip2.binary_search(ip).is_err()); assert!(ip2.btree_search(ip).is_err()); } println!(); } } }