fix: disable some logs, swapped strings in matches for enums

This commit is contained in:
Ishan Jain 2025-01-21 16:58:10 +05:30
parent 6d0e300e44
commit 8de845441c
4 changed files with 69 additions and 55 deletions

View File

@ -1,4 +1,7 @@
#![no_std] #![no_std]
use core::fmt::{Display, Formatter, Result as FmtResult};
pub enum ProgramParameters { pub enum ProgramParameters {
CountryNodeCount = 1, CountryNodeCount = 1,
CountryRecordSize = 2, CountryRecordSize = 2,
@ -12,3 +15,22 @@ pub enum ProgramParameters {
// or since for this projet, I am only working with 24 bit dbs // or since for this projet, I am only working with 24 bit dbs
// the value is set to 0x00ffffff // the value is set to 0x00ffffff
pub const BLOCK_MARKER: u32 = 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 {}

View File

@ -7,9 +7,9 @@ use aya_ebpf::{
maps::{Array, HashMap}, maps::{Array, HashMap},
programs::XdpContext, programs::XdpContext,
}; };
use aya_log_ebpf::{info, warn}; use aya_log_ebpf::{debug, warn};
use core::{mem, net::IpAddr, usize}; use core::{mem, net::IpAddr};
use geofw_common::{ProgramParameters, BLOCK_MARKER}; use geofw_common::{MaxmindDb, ProgramParameters, BLOCK_MARKER};
use network_types::{ use network_types::{
eth::{EthHdr, EtherType}, eth::{EthHdr, EtherType},
ip::{Ipv4Hdr, Ipv6Hdr}, ip::{Ipv4Hdr, Ipv6Hdr},
@ -60,11 +60,16 @@ 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, "asn", &BLOCKED_ASN, IpAddr::V4(source)) let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V4(source))
|| should_block(&ctx, "country", &BLOCKED_COUNTRY, IpAddr::V4(source)); || should_block(
&ctx,
MaxmindDb::Country,
&BLOCKED_COUNTRY,
IpAddr::V4(source),
);
if result { 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) Ok(xdp_action::XDP_DROP)
} else { } else {
@ -77,11 +82,16 @@ 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, "asn", &BLOCKED_ASN, IpAddr::V6(source)) let result = should_block(&ctx, MaxmindDb::Asn, &BLOCKED_ASN, IpAddr::V6(source))
|| should_block(&ctx, "country", &BLOCKED_COUNTRY, IpAddr::V6(source)); || should_block(
&ctx,
MaxmindDb::Country,
&BLOCKED_COUNTRY,
IpAddr::V6(source),
);
if result { 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) Ok(xdp_action::XDP_DROP)
} else { } else {
@ -91,35 +101,24 @@ fn filter_ipv6_packet(ctx: XdpContext) -> Result<u32, u32> {
} }
} }
pub fn should_block(ctx: &XdpContext, map_name: &str, map: &Array<u8>, addr: IpAddr) -> bool { pub fn should_block(ctx: &XdpContext, db_name: MaxmindDb, map: &Array<u8>, addr: IpAddr) -> bool {
let record_size = match map_name { let record_size = match db_name {
"country" => unsafe { PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8)) }, MaxmindDb::Country => unsafe {
"asn" => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) }, PARAMETERS.get(&(ProgramParameters::CountryRecordSize as u8))
},
_ => { MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnRecordSize as u8)) },
warn!(ctx, "unknown map: {}", map_name);
return false;
}
}; };
let Some(&record_size) = record_size else { let Some(&record_size) = record_size else {
// warn!(ctx, "error in geting record size: {}", map_name);
return false; return false;
}; };
let node_count = match map_name { let node_count = match db_name {
"country" => unsafe { PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8)) }, MaxmindDb::Country => unsafe {
"asn" => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) }, PARAMETERS.get(&(ProgramParameters::CountryNodeCount as u8))
},
_ => { MaxmindDb::Asn => unsafe { PARAMETERS.get(&(ProgramParameters::AsnNodeCount as u8)) },
warn!(ctx, "unknown map: {}", map_name);
return false;
}
}; };
let Some(&node_count) = node_count else { let Some(&node_count) = node_count else {
// warn!(ctx, "error in getting node count: {}", map_name);
return false; return false;
}; };
@ -142,9 +141,8 @@ pub fn should_block(ctx: &XdpContext, map_name: &str, map: &Array<u8>, addr: IpA
None => { None => {
warn!( warn!(
ctx, ctx,
"error in getting item at position node * node + i as i32 = {} map = {}", "error in reading position = {}",
node * node + i as u32, node * node_size as u32 + i as u32,
map_name
); );
return false; return false;
} }

View File

@ -8,8 +8,8 @@ use aya::{
}; };
use flate2::bufread::GzDecoder; use flate2::bufread::GzDecoder;
use fxhash::FxHashSet; use fxhash::FxHashSet;
use geofw_common::ProgramParameters; use geofw_common::{MaxmindDb, ProgramParameters};
use log::{debug, error, 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};
use std::{ 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<Config, String> { fn read_config(path: &str) -> Result<Config, String> {
match File::open(path) { match File::open(path) {
Ok(mut f) => { Ok(mut f) => {
@ -84,7 +81,7 @@ fn read_config(path: &str) -> Result<Config, String> {
} }
} }
fn fetch_geoip_db(config: &Config, db_name: &str) -> Result<ProcessedDb, String> { fn fetch_geoip_db(config: &Config, db_name: MaxmindDb) -> 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_name));
@ -138,7 +135,7 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result<ProcessedDb, String>
info!("downloaded {}", db_name); info!("downloaded {}", db_name);
match 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 { let Some(Data::Map(country)) = data.get("country".as_bytes()) else {
return false; return false;
}; };
@ -148,15 +145,13 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result<ProcessedDb, String>
config.source_countries.contains(&iso_code.to_string()) 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 { let Some(Data::U32(asn)) = data.get("autonomous_system_number".as_bytes()) else {
return false; return false;
}; };
config.source_asn.contains(asn) config.source_asn.contains(asn)
})), })),
_ => Err("unknown db".to_string()),
} }
} }
@ -201,17 +196,17 @@ async fn main() -> anyhow::Result<()> {
_ = interval.tick() => { _ = interval.tick() => {
info!("updating DB"); 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(_) => (), Ok(_) => (),
Err(e) => { 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(_) => (), Ok(_) => (),
Err(e) => { 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( fn update_geoip_map(
config: &Config, config: &Config,
ebpf: &mut Ebpf, ebpf: &mut Ebpf,
db_name: &str, db_name: MaxmindDb,
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_name = {db_name} map_name = {map_name}");
@ -252,7 +247,7 @@ fn update_geoip_map(
.expect("error in processing parameter map"); .expect("error in processing parameter map");
match db_name { match db_name {
COUNTRY_DB => { MaxmindDb::Country => {
map.insert( map.insert(
ProgramParameters::CountryNodeCount as u8, ProgramParameters::CountryNodeCount as u8,
result.node_count, result.node_count,
@ -266,7 +261,7 @@ fn update_geoip_map(
) )
.expect("error in writing country record size to map"); .expect("error in writing country record size to map");
} }
ASN_DB => { MaxmindDb::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(
@ -276,8 +271,6 @@ fn update_geoip_map(
) )
.expect("error in writing country record size to map"); .expect("error in writing country record size to map");
} }
_ => unreachable!(),
} }
Ok(()) Ok(())

View File

@ -1,4 +1,5 @@
use fxhash::{FxHashMap, FxHashSet}; use core::str;
use fxhash::FxHashMap;
use geofw_common::BLOCK_MARKER; use geofw_common::BLOCK_MARKER;
use std::{ use std::{
fmt::{Debug, Display, Formatter, Result as FmtResult}, fmt::{Debug, Display, Formatter, Result as FmtResult},
@ -50,7 +51,7 @@ pub struct ProcessedDb {
impl Display for Data<'_> { impl Display for Data<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
match self { 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::Double(s) => write!(f, "{s}"),
Data::Bytes(s) => write!(f, "{s:?}"), Data::Bytes(s) => write!(f, "{s:?}"),
Data::U16(s) => write!(f, "{s}"), Data::U16(s) => write!(f, "{s}"),