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

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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 {}

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

@ -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<u32, u32> {
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<u32, u32> {
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<u32, u32> {
}
}
pub fn should_block(ctx: &XdpContext, map_name: &str, map: &Array<u8>, 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<u8>, 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<u8>, 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;
}

31
geofw/src/main.rs
View File

@ -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<Config, String> {
match File::open(path) {
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();
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<ProcessedDb, String>
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<ProcessedDb, 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 {
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(())

5
geofw/src/maxmind.rs
View File

@ -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}"),