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

fix: fixed a bug in the binary tree traversal, refactored other parts of the code

Browse Source
This commit is contained in:
Ishan Jain 2025-01-22 08:07:12 +05:30
parent ea61402270
commit 5531c2eb93
4 changed files with 136 additions and 122 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
geofw-common/src/lib.rs
View File

@ -17,20 +17,20 @@ pub enum ProgramParameters {
pub const BLOCK_MARKER: u32 = 0x00ffffff; pub const BLOCK_MARKER: u32 = 0x00ffffff;
#[derive(Copy, Clone)] #[derive(Copy, Clone)]
pub enum MaxmindDb { pub enum MaxmindDbType {
Country, Country,
Asn, Asn,
} }
impl Display for MaxmindDb { impl Display for MaxmindDbType {
fn fmt(&self, f: &mut Formatter) -> FmtResult { fn fmt(&self, f: &mut Formatter) -> FmtResult {
let val = match self { let val = match self {
MaxmindDb::Country => "GeoLite2-Country", MaxmindDbType::Country => "GeoLite2-Country",
MaxmindDb::Asn => "GeoLite2-ASN", MaxmindDbType::Asn => "GeoLite2-ASN",
}; };
write!(f, "{val}") write!(f, "{val}")
} }
} }
// impl aya_log_ebpf::macro_support::Formatter for MaxmindDb {}

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

@ -9,7 +9,7 @@ use aya_ebpf::{
}; };
use aya_log_ebpf::{debug, warn}; use aya_log_ebpf::{debug, warn};
use core::{mem, net::IpAddr}; use core::{mem, net::IpAddr};
use geofw_common::{MaxmindDb, ProgramParameters, BLOCK_MARKER}; use geofw_common::{MaxmindDbType, ProgramParameters, BLOCK_MARKER};
use network_types::{ use network_types::{
eth::{EthHdr, EtherType}, eth::{EthHdr, EtherType},
ip::{Ipv4Hdr, Ipv6Hdr}, ip::{Ipv4Hdr, Ipv6Hdr},
@ -60,10 +60,10 @@ fn filter_ip_packet(ctx: XdpContext) -> Result<u32, u32> {
let ip: *const Ipv4Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?; let ip: *const Ipv4Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?;
let source = unsafe { (*ip).src_addr() }; let source = unsafe { (*ip).src_addr() };
let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V4(source)) let result = should_block(&ctx, MaxmindDbType::Asn, &BLOCKED_ASN, IpAddr::V4(source))
|| should_block( || should_block(
&ctx, &ctx,
MaxmindDb::Country, MaxmindDbType::Country,
&BLOCKED_COUNTRY, &BLOCKED_COUNTRY,
IpAddr::V4(source), IpAddr::V4(source),
); );
@ -82,10 +82,10 @@ fn filter_ipv6_packet(ctx: XdpContext) -> Result<u32, u32> {
let ip: *const Ipv6Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?; let ip: *const Ipv6Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?;
let source = unsafe { (*ip).src_addr() }; let source = unsafe { (*ip).src_addr() };
let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V6(source)) let result = should_block(&ctx, MaxmindDbType::Asn, &BLOCKED_ASN, IpAddr::V6(source))
|| should_block( || should_block(
&ctx, &ctx,
MaxmindDb::Country, MaxmindDbType::Country,
&BLOCKED_COUNTRY, &BLOCKED_COUNTRY,
IpAddr::V6(source), IpAddr::V6(source),
); );
@ -101,22 +101,27 @@ fn filter_ipv6_packet(ctx: XdpContext) -> Result<u32, u32> {
} }
} }
pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr: IpAddr) -> bool { pub fn should_block(
let record_size = match db_name { ctx: &XdpContext,
MaxmindDb::Country => unsafe { db_type: MaxmindDbType,
map: &Array<u8>,
addr: IpAddr,
) -> bool {
let record_size = match db_type {
MaxmindDbType::Country => unsafe {
PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8)) PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8))
}, },
MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) }, MaxmindDbType::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) },
}; };
let Some(&record_size) = record_size else { let Some(&record_size) = record_size else {
return false; return false;
}; };
let node_count = match db_name { let node_count = match db_type {
MaxmindDb::Country => unsafe { MaxmindDbType::Country => unsafe {
PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8)) PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8))
}, },
MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) }, MaxmindDbType::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) },
}; };
let Some(&node_count) = node_count else { let Some(&node_count) = node_count else {
return false; return false;
@ -135,7 +140,7 @@ pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr:
}; };
while i >= 0 && node < node_count { while i >= 0 && node < node_count {
let bit = (ip & (1 << 127)) == 0; let left = (ip & (1 << 127)) == 0;
ip <<= 1; ip <<= 1;
let mut slice = [0; 8]; let mut slice = [0; 8];
@ -152,31 +157,21 @@ pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr:
} }
} }
} }
node = node_from_bytes(slice, bit, record_size as u16); node = node_from_bytes(slice, left, record_size as u16);
i -= 1; i -= 1;
} }
node == BLOCK_MARKER node == BLOCK_MARKER
} }
fn node_from_bytes(n: [u8; 8], bit: bool, record_size: u16) -> u32 { fn node_from_bytes(n: [u8; 8], left: bool, record_size: u16) -> u32 {
match record_size { match record_size {
28 => { 28 if left => u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]]),
if bit { 28 => u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]]),
u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]]) 24 if left => u32::from_be_bytes([0, n[0], n[1], n[2]]),
} else { 24 => u32::from_be_bytes([0, n[3], n[4], n[5]]),
u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]])
}
}
24 => {
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]])
}
}
// this should never reach // This should never run unless we are using 32bit dbs
_ => 0, _ => 0,
} }
} }

54
geofw/src/main.rs
View File

@ -8,7 +8,7 @@ use aya::{
}; };
use flate2::bufread::GzDecoder; use flate2::bufread::GzDecoder;
use fxhash::FxHashSet; use fxhash::FxHashSet;
use geofw_common::{MaxmindDb, ProgramParameters}; use geofw_common::{MaxmindDbType, ProgramParameters};
use log::{debug, info, warn}; use log::{debug, info, warn};
use maxmind::{Data, ProcessedDb}; use maxmind::{Data, ProcessedDb};
use serde_derive::{Deserialize, Serialize}; use serde_derive::{Deserialize, Serialize};
@ -16,6 +16,7 @@ use std::{
fs::File, fs::File,
io::{BufReader, ErrorKind, Read, Write}, io::{BufReader, ErrorKind, Read, Write},
path::PathBuf, path::PathBuf,
time::Instant,
}; };
use tar::Archive; use tar::Archive;
use tokio::{signal, time}; use tokio::{signal, time};
@ -81,12 +82,12 @@ fn read_config(path: &str) -> Result<Config, String> {
} }
} }
fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result<ProcessedDb, String> { fn fetch_geoip_db(config: &Config, db_type: MaxmindDbType) -> Result<ProcessedDb, String> {
let mut unpack_path = PathBuf::new(); let mut unpack_path = PathBuf::new();
unpack_path.push(&config.db.path); unpack_path.push(&config.db.path);
unpack_path.push(format!("{}.mmdb", db_name)); unpack_path.push(format!("{}.mmdb", db_type));
let url = format!("https://download.maxmind.com/app/geoip_download?edition_id={}&license_key={}&suffix=tar.gz", db_name, config.db.maxmind_key); let url = format!("https://download.maxmind.com/app/geoip_download?edition_id={}&license_key={}&suffix=tar.gz", db_type, config.db.maxmind_key);
info!("path = {:?} fetching db from = {}", unpack_path, url); info!("path = {:?} fetching db from = {}", unpack_path, url);
@ -130,10 +131,10 @@ fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result<ProcessedDb, St
} }
}; };
let db = maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())?; let db = maxmind::MaxmindDb::from_file(&unpack_path.to_string_lossy())?;
match db_name { match db_type {
MaxmindDb::Country => Ok(db.consume(|data| -> bool { MaxmindDbType::Country => Ok(db.consume(|data| -> bool {
let Some(Data::Map(country)) = data.get("country".as_bytes()) else { let Some(Data::Map(country)) = data.get("country".as_bytes()) else {
return false; return false;
}; };
@ -143,7 +144,7 @@ fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result<ProcessedDb, St
config.source_countries.contains(&iso_code.to_string()) config.source_countries.contains(&iso_code.to_string())
})), })),
MaxmindDb::Asn => Ok(db.consume(|data| -> bool { MaxmindDbType::Asn => Ok(db.consume(|data| -> bool {
let Some(Data::U32(asn)) = data.get("autonomous_system_number".as_bytes()) else { let Some(Data::U32(asn)) = data.get("autonomous_system_number".as_bytes()) else {
return false; return false;
}; };
@ -194,17 +195,17 @@ async fn main() -> anyhow::Result<()> {
_ = interval.tick() => { _ = interval.tick() => {
info!("updating DB"); info!("updating DB");
match update_geoip_map(&config, &mut ebpf, MaxmindDb::Country, "BLOCKED_COUNTRY") { match update_geoip_map(&config, &mut ebpf, MaxmindDbType::Country, "BLOCKED_COUNTRY") {
Ok(_) => (), Ok(_) => (),
Err(e) => { Err(e) => {
warn!("error in updating map {} = {}", MaxmindDb::Country, e); warn!("error in updating map {} = {}", MaxmindDbType::Country, e);
} }
} }
match update_geoip_map(&config, &mut ebpf, MaxmindDb::Asn, "BLOCKED_ASN") { match update_geoip_map(&config, &mut ebpf, MaxmindDbType::Asn, "BLOCKED_ASN") {
Ok(_) => (), Ok(_) => (),
Err(e) => { Err(e) => {
warn!("error in updating map {} = {}", MaxmindDb::Asn, e); warn!("error in updating map {} = {}", MaxmindDbType::Asn, e);
} }
} }
} }
@ -217,35 +218,38 @@ async fn main() -> anyhow::Result<()> {
fn update_geoip_map( fn update_geoip_map(
config: &Config, config: &Config,
ebpf: &mut Ebpf, ebpf: &mut Ebpf,
db_name: MaxmindDb, db_type: MaxmindDbType,
map_name: &str, map_name: &str,
) -> Result<(), String> { ) -> Result<(), String> {
info!("updating maps db_name = {db_name} map_name = {map_name}"); info!("updating maps db_type = {db_type} map_name = {map_name}");
let mut map = Array::try_from(ebpf.map_mut(map_name).expect("error in getting map")) let mut map = Array::try_from(ebpf.map_mut(map_name).expect("error in getting map"))
.expect("error in processing map"); .expect("error in processing map");
let result = fetch_geoip_db(config, db_name)?; let result = fetch_geoip_db(config, db_type)?;
info!(
"set map = {map_name} up to the location = {} record_size = {} node_count = {}",
result.db.len(),
result.record_size,
result.node_count
);
let t = Instant::now();
for (i, v) in result.db.into_iter().enumerate() { for (i, v) in result.db.into_iter().enumerate() {
map.set(i as u32, v, 0).map_err(|e| e.to_string())?; map.set(i as u32, v, 0).map_err(|e| e.to_string())?;
} }
info!(
"updated map = {} record_size = {} node_count = {} est_size = {} time_taken = {:?}",
map_name,
result.record_size,
result.node_count,
result.record_size as u64 * result.node_count as u64,
t.elapsed()
);
let mut map: HashMap<&mut MapData, u8, u32> = HashMap::try_from( let mut map: HashMap<&mut MapData, u8, u32> = HashMap::try_from(
ebpf.map_mut("PARAMETERS") ebpf.map_mut("PARAMETERS")
.expect("error in getting parameter map"), .expect("error in getting parameter map"),
) )
.expect("error in processing parameter map"); .expect("error in processing parameter map");
match db_name { match db_type {
MaxmindDb::Country => { MaxmindDbType::Country => {
map.insert( map.insert(
ProgramParameters::CountryNodeCount as u8, ProgramParameters::CountryNodeCount as u8,
result.node_count, result.node_count,
@ -259,7 +263,7 @@ fn update_geoip_map(
) )
.expect("error in writing country record size to map"); .expect("error in writing country record size to map");
} }
MaxmindDb::Asn => { MaxmindDbType::Asn => {
map.insert(ProgramParameters::AsnNodeCount as u8, result.node_count, 0) map.insert(ProgramParameters::AsnNodeCount as u8, result.node_count, 0)
.expect("error in writing country node count to map"); .expect("error in writing country node count to map");
map.insert( map.insert(

139
geofw/src/maxmind.rs
View File

@ -2,6 +2,7 @@ use core::str;
use fxhash::FxHashMap; use fxhash::FxHashMap;
use geofw_common::BLOCK_MARKER; use geofw_common::BLOCK_MARKER;
use std::{ use std::{
collections::VecDeque,
fmt::{Debug, Display, Formatter, Result as FmtResult}, fmt::{Debug, Display, Formatter, Result as FmtResult},
fs::File, fs::File,
io::Read, io::Read,
@ -12,7 +13,7 @@ const METADATA_SECTION_START: &[u8] = &[
0xab, 0xcd, 0xef, 0x4d, 0x61, 0x78, 0x4d, 0x69, 0x6e, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0xab, 0xcd, 0xef, 0x4d, 0x61, 0x78, 0x4d, 0x69, 0x6e, 0x64, 0x2e, 0x63, 0x6f, 0x6d,
]; ];
pub struct MaxmindDB { pub struct MaxmindDb {
pub metadata: Metadata, pub metadata: Metadata,
pub data: Vec<u8>, pub data: Vec<u8>,
} }
@ -81,13 +82,13 @@ impl Display for Data<'_> {
} }
} }
impl Debug for MaxmindDB { impl Debug for MaxmindDb {
fn fmt(&self, f: &mut Formatter) -> FmtResult { fn fmt(&self, f: &mut Formatter) -> FmtResult {
f.write_fmt(format_args!("{:?}", self.metadata)) f.write_fmt(format_args!("{:?}", self.metadata))
} }
} }
impl MaxmindDB { impl MaxmindDb {
pub fn from_file(path: &str) -> Result<Self, String> { pub fn from_file(path: &str) -> Result<Self, String> {
let mut data = vec![]; let mut data = vec![];
let mut file = File::open(path).map_err(|e| format!("error in opening file: {}", e))?; let mut file = File::open(path).map_err(|e| format!("error in opening file: {}", e))?;
@ -131,40 +132,30 @@ impl MaxmindDB {
map map
} }
fn node_from_bytes(n: &[u8], bit: bool, record_size: u16) -> u32 { fn node_from_bytes(n: &[u8], left: bool, record_size: u16) -> u32 {
match record_size { match record_size {
28 => { 28 if left => u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]]),
if bit { 28 => u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]]),
u32::from_be_bytes([(n[3] & 0b1111_0000) >> 4, n[0], n[1], n[2]]) 24 if left => u32::from_be_bytes([0, n[0], n[1], n[2]]),
} else { 24 => u32::from_be_bytes([0, n[3], n[4], n[5]]),
u32::from_be_bytes([n[3] & 0b0000_1111, n[4], n[5], n[6]])
}
}
24 => {
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]])
}
}
_ => unreachable!(), _ => unreachable!(),
} }
} }
fn write_over_node_bytes(n: &mut [u8], bit: u128, record_size: u16, val: u32) { fn write_over_node_bytes(n: &mut [u8], left: bool, record_size: u16, val: u32) {
let val = val.to_be_bytes(); let val = val.to_be_bytes();
match record_size { match record_size {
28 if bit == 0 => { 28 if left => {
n[0..=2].copy_from_slice(&val[1..=3]); n[0..=2].copy_from_slice(&val[1..=3]);
n[3] = (n[3] & 0b0000_1111) | (val[0] << 4); n[3] = (n[3] & 0b0000_1111) | (val[0] << 4);
} }
28 if bit == 1 => { 28 => {
n[4..=6].copy_from_slice(&val[1..=3]); n[4..=6].copy_from_slice(&val[1..=3]);
n[3] = (n[3] & 0b1111_0000) | (val[0] & 0b0000_1111); n[3] = (n[3] & 0b1111_0000) | (val[0] & 0b0000_1111);
} }
24 if bit == 0 => n[0..=2].copy_from_slice(&val[1..=3]), 24 if left => n[0..=2].copy_from_slice(&val[1..=3]),
24 if bit == 1 => n[3..=5].copy_from_slice(&val[1..=3]), 24 => n[3..=5].copy_from_slice(&val[1..=3]),
_ => unreachable!(), _ => unreachable!(),
} }
} }
@ -189,10 +180,10 @@ impl MaxmindDB {
}; };
while i >= 0 && node < self.metadata.node_count { while i >= 0 && node < self.metadata.node_count {
let bit = (ip & (1 << i)) == 0; let left = (ip & (1 << i)) == 0;
let n = &self.data[node as usize * node_size..(node as usize * node_size) + node_size]; 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); node = Self::node_from_bytes(n, left, self.metadata.record_size);
i -= 1; i -= 1;
} }
@ -207,50 +198,50 @@ impl MaxmindDB {
} }
} }
pub fn consume(mut self, should_block: impl Fn(FxHashMap<&[u8], Data>) -> bool) -> ProcessedDb { pub fn consume(
let mut stack = vec![]; mut self,
should_block: impl Fn(&FxHashMap<&[u8], Data>) -> bool,
) -> ProcessedDb {
let mut stack = VecDeque::new();
let node_size = self.metadata.record_size as usize * 2 / 8; let node_size = self.metadata.record_size as usize * 2 / 8;
stack.push((0, 0)); stack.push_back((0, 0, false));
while let Some((node, parent, bit)) = stack.pop_front() {
if node == BLOCK_MARKER {
continue;
}
if node >= self.metadata.node_count {
let ds_offset = node - self.metadata.node_count;
let (data, _) =
self.read_data(self.metadata.data_section_start + ds_offset as usize - 16);
let Data::Map(data) = data else {
unreachable!()
};
if should_block(&data) {
// Mark the parent of this node as non existent
let node = parent;
Self::write_over_node_bytes(
&mut self.data
[node as usize * node_size..(node as usize * node_size) + node_size],
bit,
self.metadata.record_size,
BLOCK_MARKER,
);
}
continue;
}
while let Some((node, position)) = stack.pop() {
let n = let n =
&mut self.data[node as usize * node_size..(node as usize * node_size) + node_size]; &mut self.data[node as usize * node_size..(node as usize * node_size) + node_size];
let node_1 = Self::node_from_bytes(n, false, 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); let node_2 = Self::node_from_bytes(n, true, self.metadata.record_size);
if position < 128 && node_1 < self.metadata.node_count { stack.push_back((node_1, node, false));
stack.push((node_1, position + 1)); stack.push_back((node_2, node, true));
}
if position < 128 && node_2 < self.metadata.node_count {
stack.push((node_2, position + 1));
}
let data_section_offset = if node_1 != BLOCK_MARKER && node_1 > self.metadata.node_count
{
node_1 - self.metadata.node_count
} else if node_2 != BLOCK_MARKER && node_2 > self.metadata.node_count {
node_2 - self.metadata.node_count
} else {
continue;
};
let (data, _) = self
.read_data(self.metadata.data_section_start + data_section_offset as usize - 16);
let Data::Map(data) = data else {
unreachable!()
};
if should_block(data) {
// Mark this node as non existent
Self::write_over_node_bytes(
&mut self.data
[node as usize * node_size..(node as usize * node_size) + node_size],
0,
self.metadata.record_size,
BLOCK_MARKER,
);
}
} }
// Trim database to only contain the binary tree // Trim database to only contain the binary tree
@ -443,3 +434,27 @@ impl MaxmindDB {
(data_type, length, read + r) (data_type, length, read + r)
} }
} }
impl ProcessedDb {
pub fn lookup(&self, addr: IpAddr) -> bool {
let node_size = self.record_size as usize * 2 / 8;
let mut node = 0;
let mut ip = match addr {
IpAddr::V4(a) => a.to_bits() as u128,
IpAddr::V6(a) => a.to_bits(),
};
let mut i = 0;
while i < 128 && node < self.node_count {
let left = (ip & (1 << 127)) == 0;
ip <<= 1;
let n = &self.db[node as usize * node_size..(node as usize * node_size) + node_size];
node = MaxmindDb::node_from_bytes(n, left, self.record_size);
i += 1;
}
node == BLOCK_MARKER
}
}