Written geosite protobuf parser and tests 4 it

src/geoparsers/v2ray/parsing.rs

@@ -0,0 +1,79 @@
+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 })
+}