From 0d81bae32a1e2d8986ec67f90c02e2a83d1da79f Mon Sep 17 00:00:00 2001 From: Namilsk Date: Tue, 28 Apr 2026 17:29:31 +0300 Subject: [PATCH] Refactor: Fully migrated parser to full `prost` crate, auto-fmt; Fix: Remove no more needed tests --- src/geoparsers/geoip2.rs | 4 +- src/geoparsers/geosite.rs | 40 +++++++++++++++ src/geoparsers/mod.rs | 1 + src/geoparsers/v2ray/mod.rs | 2 +- src/geoparsers/v2ray/parsing.rs | 88 +++++---------------------------- src/geoparsers/v2ray/types.rs | 12 +---- src/sniffing/headers.rs | 6 +-- tests/v2ray_geosite.rs | 18 +++++-- 8 files changed, 75 insertions(+), 96 deletions(-) create mode 100644 src/geoparsers/geosite.rs diff --git a/src/geoparsers/geoip2.rs b/src/geoparsers/geoip2.rs index 1fd6ab1..75ab15a 100644 --- a/src/geoparsers/geoip2.rs +++ b/src/geoparsers/geoip2.rs @@ -8,9 +8,10 @@ use std::net::IpAddr; // TODO: V2Ray protobuf parsing && Test 4 ts /// Interface enum for `dst_addr` info +#[allow(dead_code)] #[derive(Debug, Deserialize)] pub enum RouteType { - /// GeoSite MMDB type, like `category-ads-all` + /// GeoSite type, like `category-ads-all` GeoSite(String), /// Result with GeoCode like "RU" GeoIp(String), @@ -18,6 +19,7 @@ pub enum RouteType { } /// Routing actions +#[allow(dead_code)] #[derive(Debug, Deserialize)] pub enum RouteAction { #[serde(alias = "block")] diff --git a/src/geoparsers/geosite.rs b/src/geoparsers/geosite.rs new file mode 100644 index 0000000..4357bdc --- /dev/null +++ b/src/geoparsers/geosite.rs @@ -0,0 +1,40 @@ +use crate::geoparsers::v2ray::parsing::decode_geosite_stream; +use crate::geoparsers::v2ray::types::{GeoSite, GeoSiteList}; + +use std::fs; +use std::io::{BufReader, Read}; + +pub struct GeoSiteService { + index: GeoSiteList, +} + +impl GeoSiteService { + // TODO: Make more smart memory mapping; geosite files can be > 70MB + pub fn new(path: &str) -> Result> { + let file = fs::File::open(path)?; + let reader = BufReader::new(file); + let geosite_list = decode_geosite_stream(reader)?; + + Ok(Self { + index: geosite_list, + }) + } + + // Idk but i think it can work + pub fn lookup(&self, value: &str) -> Option<&GeoSite> { + self.index + .entry + .iter() + .find(|site| site.domain.iter().any(|d| d.value == value)) + } + + /// Returns the number of GeoSite entries in the list + pub fn len(&self) -> usize { + self.index.entry.len() + } + + /// Returns true if the GeoSite list is empty + pub fn is_empty(&self) -> bool { + self.index.entry.is_empty() + } +} diff --git a/src/geoparsers/mod.rs b/src/geoparsers/mod.rs index c3cfece..b157aa5 100644 --- a/src/geoparsers/mod.rs +++ b/src/geoparsers/mod.rs @@ -1,3 +1,4 @@ pub mod geoip2; +pub mod geosite; pub mod toml; pub mod v2ray; diff --git a/src/geoparsers/v2ray/mod.rs b/src/geoparsers/v2ray/mod.rs index 971be55..14f07ff 100644 --- a/src/geoparsers/v2ray/mod.rs +++ b/src/geoparsers/v2ray/mod.rs @@ -1,2 +1,2 @@ -pub mod parsing; +pub(crate) mod parsing; pub mod types; diff --git a/src/geoparsers/v2ray/parsing.rs b/src/geoparsers/v2ray/parsing.rs index 4f0bbba..11042c0 100644 --- a/src/geoparsers/v2ray/parsing.rs +++ b/src/geoparsers/v2ray/parsing.rs @@ -1,79 +1,15 @@ -use crate::geoparsers::v2ray::types::{Domain, GeoSite, GeoSiteList}; +use crate::geoparsers::v2ray::types::GeoSiteList; use prost::Message; -use prost::bytes::Buf; -use std::fs; +use std::io::Read; -pub struct GeoSiteService { - index: GeoSiteList, -} - -impl GeoSiteService { - // TODO: Make more smart memory mapping; geosite files can be > 70MB - pub fn new(path: &str) -> Result> { - let bytes = fs::read(path)?; - let geosite_list = decode_geosite_stream(&bytes)?; - - Ok(Self { - index: geosite_list, - }) - } - - // Idk but i think it can work - pub fn lookup(&self, value: &str) -> Option<&GeoSite> { - self.index - .entry - .iter() - .find(|site| site.domain.iter().any(|d| d.value == value)) - } - - /// Returns the number of GeoSite entries in the list - pub fn len(&self) -> usize { - self.index.entry.len() - } - - /// Returns true if the GeoSite list is empty - pub fn is_empty(&self) -> bool { - self.index.entry.is_empty() - } -} - -/// Decode a stream of length-delimited GeoSite messages -/// `geosite.dat` ts is not one protobuf-message, stream of length-delimited messages -/// so we need ts helper -fn decode_geosite_stream(bytes: &[u8]) -> Result> { - let mut buf = bytes; - let mut entries = Vec::new(); - - while buf.has_remaining() { - // Read tag (0x0a field 1, wire type 2) - let tag = buf.get_u8(); - if tag != 0x0a { - return Err(format!("Unexpected tag: {:#04x}", tag).into()); - } - // varint - let mut len = 0usize; - let mut shift = 0; - loop { - if !buf.has_remaining() { - return Err("Unexpected end of buffer while reading varint".into()); - } - let b = buf.get_u8(); - len |= ((b & 0x7f) as usize) << shift; - if b & 0x80 == 0 { - break; - } - shift += 7; - if shift >= 70 { - return Err("Varint too long".into()); - } - } - - let entry_bytes = &buf[..len]; - let site = GeoSite::decode(entry_bytes)?; - entries.push(site); - - buf.advance(len); - } - - Ok(GeoSiteList { entry: entries }) +/// Decodes full geosite protobuf; TODO: MMAP +pub fn decode_geosite_stream( + mut reader: R, +) -> Result> { + let mut buf = Vec::new(); + reader.read_to_end(&mut buf)?; + + let list = GeoSiteList::decode(&buf[..])?; + + Ok(list) } diff --git a/src/geoparsers/v2ray/types.rs b/src/geoparsers/v2ray/types.rs index d7c0436..b82fbbf 100644 --- a/src/geoparsers/v2ray/types.rs +++ b/src/geoparsers/v2ray/types.rs @@ -33,17 +33,7 @@ pub mod domain { } } /// Type of domain value. - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] + #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum Type { /// The value is used as is. diff --git a/src/sniffing/headers.rs b/src/sniffing/headers.rs index bbe58de..c68eb50 100644 --- a/src/sniffing/headers.rs +++ b/src/sniffing/headers.rs @@ -69,9 +69,9 @@ impl fmt::Display for PacketInfo { // y:y:y:y:y:y:y:y = 8 hexademical; y = segment, pair of 2 u8 big endian write!( f, - // FIXME: fe80:0:0:0:93a:245e:daac:7a75 -> ff12:0:0:0:0:0:0:8384 - // we should drop zeroes in display: fe80::93a:245e:daac:7a75 -> ff12::8384 - // NOTE: fe80:0:0:93a:0:245e and fe80:0:93a:0:0:245e can NOT both be compressed to fe80::93a::245e by obvious reasons + // FIXME: fe80:0:0:0:93a:245e:daac:7a75 -> ff12:0:0:0:0:0:0:8384 + // we should drop zeroes in display: fe80::93a:245e:daac:7a75 -> ff12::8384 + // NOTE: fe80:0:0:93a:0:245e and fe80:0:93a:0:0:245e can NOT both be compressed to fe80::93a::245e by obvious reasons "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x} port:{} -> {:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x} port:{} {:?} is dns? {:?}", src_ip[0], src_ip[1], diff --git a/tests/v2ray_geosite.rs b/tests/v2ray_geosite.rs index d400839..261309a 100644 --- a/tests/v2ray_geosite.rs +++ b/tests/v2ray_geosite.rs @@ -1,5 +1,5 @@ -use nsc::geoparsers::v2ray::parsing::GeoSiteService; -use nsc::geoparsers::v2ray::types::Domain; +use nsc::geoparsers::geosite::GeoSiteService; +use nsc::geoparsers::v2ray::types::{Domain, domain}; use std::fs; use std::path::PathBuf; @@ -7,11 +7,11 @@ fn download_geosite() -> Result> { let tmp_dir = std::env::temp_dir().join("seccontrol_test"); fs::create_dir_all(&tmp_dir)?; - let geosite_path = tmp_dir.join("geosite.dat"); + let geosite_path = tmp_dir.join("dlc.dat"); if !geosite_path.exists() { // Use v2fly domain-list-community which has standard protobuf format - let url = "https://github.com/v2fly/domain-list-community/releases/latest/download/dlc.dat"; + let url = "https://github.com/v2fly/domain-list-community/releases/download/20260330021310/dlc.dat"; let response = ureq::get(url).call()?; let mut file = fs::File::create(&geosite_path)?; let mut reader = response.into_reader(); @@ -41,6 +41,16 @@ fn geosite_service_creation() { fn lookup_existing_domain() { let service = get_geosite_service().expect("Failed to create service"); + let domain = Domain { + r#type: nsc::geoparsers::v2ray::types::domain::Type::Full as i32, + value: "google.com".to_string(), + attribute: vec![], + }; + + let result = service.lookup(domain.value.as_str()); + assert!(result.is_some(), "Should return GeoSite categs"); + println!("{:?}", result); + assert!(!service.is_empty(), "Service should have entries"); println!("Loaded {} GeoSite entries", service.len()); }