diff --git a/geofw-common/src/lib.rs b/geofw-common/src/lib.rs index 25fdf5d..7850d10 100644 --- a/geofw-common/src/lib.rs +++ b/geofw-common/src/lib.rs @@ -1,4 +1,7 @@ #![no_std] + +use core::fmt::{Display, Formatter, Result as FmtResult}; + pub enum ProgramParameters { CountryNodeCount = 1, CountryRecordSize = 2, @@ -12,3 +15,22 @@ pub enum ProgramParameters { // or since for this projet, I am only working with 24 bit dbs // the value is set to 0x00ffffff pub const BLOCK_MARKER: u32 = 0x00ffffff; + +#[derive(Copy, Clone)] +pub enum MaxmindDb { + Country, + Asn, +} + +impl Display for MaxmindDb { + fn fmt(&self, f: &mut Formatter) -> FmtResult { + let val = match self { + MaxmindDb::Country => "GeoLite2-Country", + MaxmindDb::Asn => "GeoLite2-ASN", + }; + + write!(f, "{val}") + } +} + +// impl aya_log_ebpf::macro_support::Formatter for MaxmindDb {} diff --git a/geofw-ebpf/src/main.rs b/geofw-ebpf/src/main.rs index 53b342a..b827df5 100644 --- a/geofw-ebpf/src/main.rs +++ b/geofw-ebpf/src/main.rs @@ -7,9 +7,9 @@ use aya_ebpf::{ maps::{Array, HashMap}, programs::XdpContext, }; -use aya_log_ebpf::{info, warn}; -use core::{mem, net::IpAddr, usize}; -use geofw_common::{ProgramParameters, BLOCK_MARKER}; +use aya_log_ebpf::{debug, warn}; +use core::{mem, net::IpAddr}; +use geofw_common::{MaxmindDb, ProgramParameters, BLOCK_MARKER}; use network_types::{ eth::{EthHdr, EtherType}, ip::{Ipv4Hdr, Ipv6Hdr}, @@ -60,11 +60,16 @@ fn filter_ip_packet(ctx: XdpContext) -> Result { let ip: *const Ipv4Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?; let source = unsafe { (*ip).src_addr() }; - let result = should_block(&ctx, "asn", &BLOCKED_ASN, IpAddr::V4(source)) - || should_block(&ctx, "country", &BLOCKED_COUNTRY, IpAddr::V4(source)); + let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V4(source)) + || should_block( + &ctx, + MaxmindDb::Country, + &BLOCKED_COUNTRY, + IpAddr::V4(source), + ); if result { - info!(&ctx, "ipv4 source = {} result = {}", source, result as u8); + debug!(&ctx, "ipv4 source = {} blocked = {}", source, result as u8); Ok(xdp_action::XDP_DROP) } else { @@ -77,11 +82,16 @@ fn filter_ipv6_packet(ctx: XdpContext) -> Result { let ip: *const Ipv6Hdr = ptr_at(&ctx, EthHdr::LEN).ok_or(xdp_action::XDP_PASS)?; let source = unsafe { (*ip).src_addr() }; - let result = should_block(&ctx, "asn", &BLOCKED_ASN, IpAddr::V6(source)) - || should_block(&ctx, "country", &BLOCKED_COUNTRY, IpAddr::V6(source)); + let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V6(source)) + || should_block( + &ctx, + MaxmindDb::Country, + &BLOCKED_COUNTRY, + IpAddr::V6(source), + ); if result { - info!(&ctx, "ipv6 source = {} result = {}", source, result as u8); + debug!(&ctx, "ipv6 source = {} blocked = {}", source, result as u8); Ok(xdp_action::XDP_DROP) } else { @@ -91,35 +101,24 @@ fn filter_ipv6_packet(ctx: XdpContext) -> Result { } } -pub fn should_block(ctx: &XdpContext, map_name: &str, map: &Array, addr: IpAddr) -> bool { - let record_size = match map_name { - "country" => unsafe { PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8)) }, - "asn" => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) }, - - _ => { - warn!(ctx, "unknown map: {}", map_name); - return false; - } +pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array, addr: IpAddr) -> bool { + let record_size = match db_name { + MaxmindDb::Country => unsafe { + PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8)) + }, + MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) }, }; let Some(&record_size) = record_size else { - // warn!(ctx, "error in geting record size: {}", map_name); - return false; }; - let node_count = match map_name { - "country" => unsafe { PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8)) }, - "asn" => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) }, - - _ => { - warn!(ctx, "unknown map: {}", map_name); - - return false; - } + let node_count = match db_name { + MaxmindDb::Country => unsafe { + PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8)) + }, + MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) }, }; let Some(&node_count) = node_count else { - // warn!(ctx, "error in getting node count: {}", map_name); - return false; }; @@ -142,9 +141,8 @@ pub fn should_block(ctx: &XdpContext, map_name: &str, map: &Array, addr: IpA None => { warn!( ctx, - "error in getting item at position node * node + i as i32 = {} map = {}", - node * node + i as u32, - map_name + "error in reading position = {}", + node * node_size as u32 + i as u32, ); return false; } diff --git a/geofw/src/main.rs b/geofw/src/main.rs index d63d533..47928f3 100644 --- a/geofw/src/main.rs +++ b/geofw/src/main.rs @@ -8,8 +8,8 @@ use aya::{ }; use flate2::bufread::GzDecoder; use fxhash::FxHashSet; -use geofw_common::ProgramParameters; -use log::{debug, error, info, warn}; +use geofw_common::{MaxmindDb, ProgramParameters}; +use log::{debug, info, warn}; use maxmind::{Data, ProcessedDb}; use serde_derive::{Deserialize, Serialize}; use std::{ @@ -56,9 +56,6 @@ impl Default for Db { } } -const COUNTRY_DB: &str = "GeoLite2-Country"; -const ASN_DB: &str = "GeoLite2-ASN"; - fn read_config(path: &str) -> Result { match File::open(path) { Ok(mut f) => { @@ -84,7 +81,7 @@ fn read_config(path: &str) -> Result { } } -fn fetch_geoip_db(config: &Config, db_name: &str) -> Result { +fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> Result { let mut unpack_path = PathBuf::new(); unpack_path.push(&config.db.path); unpack_path.push(format!("{}.mmdb", db_name)); @@ -138,7 +135,7 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result info!("downloaded {}", db_name); match db_name { - COUNTRY_DB => Ok(db.consume(|data| -> bool { + MaxmindDb::Country => Ok(db.consume(|data| -> bool { let Some(Data::Map(country)) = data.get("country".as_bytes()) else { return false; }; @@ -148,15 +145,13 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result config.source_countries.contains(&iso_code.to_string()) })), - ASN_DB => Ok(db.consume(|data| -> bool { + MaxmindDb::Asn => Ok(db.consume(|data| -> bool { let Some(Data::U32(asn)) = data.get("autonomous_system_number".as_bytes()) else { return false; }; config.source_asn.contains(asn) })), - - _ => Err("unknown db".to_string()), } } @@ -201,17 +196,17 @@ async fn main() -> anyhow::Result<()> { _ = interval.tick() => { info!("updating DB"); - match update_geoip_map(&config, &mut ebpf, COUNTRY_DB, "BLOCKED_COUNTRY") { + match update_geoip_map(&config, &mut ebpf, MaxmindDb::Country, "BLOCKED_COUNTRY") { Ok(_) => (), Err(e) => { - warn!("error in updating map {} = {}", COUNTRY_DB, e); + warn!("error in updating map {} = {}", MaxmindDb::Country, e); } } - match update_geoip_map(&config, &mut ebpf, ASN_DB, "BLOCKED_ASN") { + match update_geoip_map(&config, &mut ebpf, MaxmindDb::Asn, "BLOCKED_ASN") { Ok(_) => (), Err(e) => { - warn!("error in updating map {} = {}", ASN_DB, e); + warn!("error in updating map {} = {}", MaxmindDb::Asn, e); } } } @@ -224,7 +219,7 @@ async fn main() -> anyhow::Result<()> { fn update_geoip_map( config: &Config, ebpf: &mut Ebpf, - db_name: &str, + db_name: MaxmindDb, map_name: &str, ) -> Result<(), String> { info!("updating maps db_name = {db_name} map_name = {map_name}"); @@ -252,7 +247,7 @@ fn update_geoip_map( .expect("error in processing parameter map"); match db_name { - COUNTRY_DB => { + MaxmindDb::Country => { map.insert( ProgramParameters::CountryNodeCount as u8, result.node_count, @@ -266,7 +261,7 @@ fn update_geoip_map( ) .expect("error in writing country record size to map"); } - ASN_DB => { + MaxmindDb::Asn => { map.insert(ProgramParameters::AsnNodeCount as u8, result.node_count, 0) .expect("error in writing country node count to map"); map.insert( @@ -276,8 +271,6 @@ fn update_geoip_map( ) .expect("error in writing country record size to map"); } - - _ => unreachable!(), } Ok(()) diff --git a/geofw/src/maxmind.rs b/geofw/src/maxmind.rs index ac0666e..f79d305 100644 --- a/geofw/src/maxmind.rs +++ b/geofw/src/maxmind.rs @@ -1,4 +1,5 @@ -use fxhash::{FxHashMap, FxHashSet}; +use core::str; +use fxhash::FxHashMap; use geofw_common::BLOCK_MARKER; use std::{ fmt::{Debug, Display, Formatter, Result as FmtResult}, @@ -50,7 +51,7 @@ pub struct ProcessedDb { impl Display for Data<'_> { fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { match self { - Data::String(s) => f.write_str(&String::from_utf8_lossy(s)), + Data::String(s) => write!(f, "{}", unsafe { str::from_utf8_unchecked(s) }), Data::Double(s) => write!(f, "{s}"), Data::Bytes(s) => write!(f, "{s:?}"), Data::U16(s) => write!(f, "{s}"),