blob: 0f897bf44b2945cc41c1841a8d9d62ab2add63b3 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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 })
}
|
