diff --git a/geofw-ebpf/src/main.rs b/geofw-ebpf/src/main.rs index b827df5..fef98d0 100644 --- a/geofw-ebpf/src/main.rs +++ b/geofw-ebpf/src/main.rs @@ -124,14 +124,18 @@ pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array, 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, 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]]) diff --git a/geofw/src/main.rs b/geofw/src/main.rs index 47928f3..702a566 100644 --- a/geofw/src/main.rs +++ b/geofw/src/main.rs @@ -132,8 +132,6 @@ fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result Ok(db.consume(|data| -> bool { let Some(Data::Map(country)) = data.get("country".as_bytes()) else { diff --git a/geofw/src/maxmind.rs b/geofw/src/maxmind.rs index f79d305..caa5271 100644 --- a/geofw/src/maxmind.rs +++ b/geofw/src/maxmind.rs @@ -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 { 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));