ishan/geofw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: make optimizations for v4

Browse Source
This commit is contained in:
Ishan Jain 2025-01-22 04:15:50 +05:30
parent 8de845441c
commit ea61402270
3 changed files with 33 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
geofw-ebpf/src/main.rs
View File

@ -124,14 +124,18 @@ pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr:
let node_size = record_size as usize * 2 / 8;
let mut node = 0;
let mut i = 128;
let mut ip = match addr {
IpAddr::V4(a) => a.to_bits() as u128,
IpAddr::V4(a) => {
i = 32;
node = 96;
(a.to_bits() as u128) << 96
}
IpAddr::V6(a) => a.to_bits(),
};
let mut i = 0;
while i < 128 && node < node_count {
let bit = ip & (1 << 127);
while i >= 0 && node < node_count {
let bit = (ip & (1 << 127)) == 0;
ip <<= 1;
let mut slice = [0; 8];
@ -148,24 +152,24 @@ pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr:
}
}
}
node = node_from_bytes(slice, if bit > 0 { 1 } else { 0 }, record_size as u16);
i += 1;
node = node_from_bytes(slice, bit, record_size as u16);
i -= 1;
}
node == BLOCK_MARKER
}
fn node_from_bytes(n: [u8; 8], bit: u8, record_size: u16) -> u32 {
fn node_from_bytes(n: [u8; 8], bit: bool, record_size: u16) -> u32 {
match record_size {
28 => {
if bit == 0 {
if bit {
u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]])
} else {
u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]])
}
}
24 => {
if bit == 0 {
if bit {
u32::from_be_bytes([0, n[0], n[1], n[2]])
} else {
u32::from_be_bytes([0, n[3], n[4], n[5]])

2
geofw/src/main.rs
View File

@ -132,8 +132,6 @@ fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result<ProcessedDb, St
let db = maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())?;
info!("downloaded {}", db_name);
match db_name {
MaxmindDb::Country => Ok(db.consume(|data| -> bool {
let Some(Data::Map(country)) = data.get("country".as_bytes()) else {

32
geofw/src/maxmind.rs
View File

@ -131,17 +131,17 @@ impl MaxmindDB {
map
}
fn node_from_bytes(n: &[u8], bit: u128, record_size: u16) -> u32 {
fn node_from_bytes(n: &[u8], bit: bool, record_size: u16) -> u32 {
match record_size {
28 => {
if bit == 0 {
if bit {
u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]])
} else {
u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]])
}
}
24 => {
if bit == 0 {
if bit {
u32::from_be_bytes([0, n[0], n[1], n[2]])
} else {
u32::from_be_bytes([0, n[3], n[4], n[5]])
@ -173,19 +173,27 @@ impl MaxmindDB {
pub fn lookup(&self, addr: IpAddr) -> Option<Data> {
let node_size = self.metadata.record_size as usize * 2 / 8;
let mut node = 0;
let mut i = 0i8;
let mut ip = match addr {
IpAddr::V4(a) => a.to_bits() as u128,
IpAddr::V6(a) => a.to_bits(),
IpAddr::V4(a) => {
node = 96;
i = 31;
a.to_bits() as u128
}
IpAddr::V6(a) => {
node = 0;
i = 127;
a.to_bits()
}
};
let mut i = 0;
while i < 128 && node < self.metadata.node_count {
let bit = ip & (1 << 127);
ip <<= 1;
while i >= 0 && node < self.metadata.node_count {
let bit = (ip & (1 << i)) == 0;
let n = &self.data[node as usize * node_size..(node as usize * node_size) + node_size];
node = Self::node_from_bytes(n, bit, self.metadata.record_size);
i += 1;
i -= 1;
}
if node == self.metadata.node_count {
@ -207,8 +215,8 @@ impl MaxmindDB {
while let Some((node, position)) = stack.pop() {
let n =
&mut self.data[node as usize * node_size..(node as usize * node_size) + node_size];
let node_1 = Self::node_from_bytes(n, 0, self.metadata.record_size);
let node_2 = Self::node_from_bytes(n, 1, self.metadata.record_size);
let node_1 = Self::node_from_bytes(n, false, self.metadata.record_size);
let node_2 = Self::node_from_bytes(n, true, self.metadata.record_size);
if position < 128 && node_1 < self.metadata.node_count {
stack.push((node_1, position + 1));