79 lines
2.2 KiB
Rust
79 lines
2.2 KiB
Rust
use crate::geoparsers::v2ray::types::{Domain, GeoSite, GeoSiteList};
|
|
use prost::bytes::Buf;
|
|
use prost::Message;
|
|
use std::fs;
|
|
|
|
pub struct GeoSiteService {
|
|
index: GeoSiteList,
|
|
}
|
|
|
|
impl GeoSiteService {
|
|
// TODO: Make more smart memory mapping; geosite files can be > 70MB
|
|
pub fn new(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
|
|
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<GeoSiteList, Box<dyn std::error::Error>> {
|
|
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 })
|
|
}
|