From 909966e4750eb9ce1b8f0304793f0608e2070587 Mon Sep 17 00:00:00 2001 From: Ishan Jain Date: Tue, 21 Jan 2025 16:09:41 +0530 Subject: [PATCH] fix: use closures in consume --- config.json | 2 +- geofw/src/main.rs | 99 ++++++++++++++++++++++++++++++-------------- geofw/src/maxmind.rs | 70 ++----------------------------- 3 files changed, 74 insertions(+), 97 deletions(-) diff --git a/config.json b/config.json index b5216df..9e47460 100644 --- a/config.json +++ b/config.json @@ -3,7 +3,7 @@ "maxmind_key": "", "//": "refresh every 24h", "refresh_interval": 86400, - "path": "/tmp/geofw" + "path": "/home/ishan/geofw/geofw" }, "interface": "enp6s18", "source_countries": [ diff --git a/geofw/src/main.rs b/geofw/src/main.rs index 931dfe9..cfe9ea7 100644 --- a/geofw/src/main.rs +++ b/geofw/src/main.rs @@ -9,18 +9,18 @@ use aya::{ use flate2::bufread::GzDecoder; use fxhash::FxHashSet; use geofw_common::ProgramParameters; -use log::{debug, info, warn}; -use maxmind::ProcessedDb; +use log::{debug, error, info, warn}; +use maxmind::{Data, ProcessedDb}; use serde_derive::{Deserialize, Serialize}; use std::{ fs::File, - io::{BufReader, Read}, + io::{BufReader, ErrorKind, Read}, path::PathBuf, }; use tar::Archive; use tokio::{signal, time}; -#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct Config { pub db: Db, pub interface: String, @@ -28,22 +28,52 @@ pub struct Config { pub source_asn: FxHashSet, } -#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] +impl Default for Config { + fn default() -> Self { + Self { + db: Default::default(), + interface: "enp1s0".to_string(), + source_countries: Default::default(), + source_asn: Default::default(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct Db { pub maxmind_key: String, pub refresh_interval: i64, pub path: String, } +impl Default for Db { + fn default() -> Self { + Self { + maxmind_key: "".to_string(), + refresh_interval: 86400, + path: "/tmp/geofw".to_string(), + } + } +} + const COUNTRY_DB: &str = "GeoLite2-Country"; const ASN_DB: &str = "GeoLite2-ASN"; fn read_config(path: &str) -> Result { - let mut f = File::open(path).map_err(|e| e.to_string())?; - let mut contents = vec![]; - f.read_to_end(&mut contents).map_err(|e| e.to_string())?; - - serde_json::from_slice(&contents).map_err(|e| e.to_string()) + match File::open(path) { + Ok(mut f) => { + let mut contents = vec![]; + f.read_to_end(&mut contents).map_err(|e| e.to_string())?; + serde_json::from_slice(&contents).map_err(|e| e.to_string()) + } + Err(e) if e.kind() == ErrorKind::NotFound => { + if let Err(e) = File::create(path) { + warn!("error in writing config to {}: {}", path, e); + } + Ok(Default::default()) + } + Err(e) => Err(format!("permission denied reading {}: {}", path, e)), + } } fn fetch_geoip_db(config: &Config, db_name: &str) -> Result { @@ -51,19 +81,15 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result unpack_path.push(&config.db.path); unpack_path.push(format!("{}.mmdb", db_name)); - info!("unpack path = {:?}", unpack_path); - let url = format!("https://download.maxmind.com/app/geoip_download?edition_id={}&license_key={}&suffix=tar.gz", db_name, config.db.maxmind_key); - info!("fetching db from = {}", url); + info!("path = {:?} fetching db from = {}", unpack_path, url); let response = ureq::get(&url).call(); - let db = match response { + match response { Ok(v) if v.status() != 200 => { warn!("response from maxmind is not 200 = {}", v.status()); - - maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())? } Ok(resp) => { let reader = resp.into_reader(); @@ -74,40 +100,53 @@ fn fetch_geoip_db(config: &Config, db_name: &str) -> Result .entries() .map_err(|e| format!("error in listing files in the archive: {}", e))?; - let mut db_entry = entries + let db_entry = entries .into_iter() .filter_map(|e| e.ok()) .filter_map(|entry| { - let path = match entry.path() { - Ok(v) => v, - Err(_) => return None, + let Ok(path) = entry.path() else { + return None; }; - if path.extension().is_none_or(|x| x != "mmdb") { return None; } Some(entry) }) - .next() - .unwrap(); + .next(); + + let Some(mut db_entry) = db_entry else { + return Err("error in finding mmdb file in the tarball".to_string()); + }; db_entry.unpack(&unpack_path).map_err(|e| e.to_string())?; - - maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())? } - Err(e) => { warn!("error in fetching db from maxmind: {}", e); - - maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())? } }; + let db = maxmind::MaxmindDB::from_file(&unpack_path.to_string_lossy())?; + info!("downloaded {}", db_name); match db_name { - COUNTRY_DB => Ok(db.consume_country_database(&config.source_countries)), - ASN_DB => Ok(db.consume_asn_database(&config.source_asn)), + COUNTRY_DB => Ok(db.consume(|data| -> bool { + let Some(Data::Map(country)) = data.get("country".as_bytes()) else { + return false; + }; + let Some(iso_code) = country.get("iso_code".as_bytes()) else { + return false; + }; + + config.source_countries.contains(&iso_code.to_string()) + })), + ASN_DB => 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()), } diff --git a/geofw/src/maxmind.rs b/geofw/src/maxmind.rs index 824191f..bb4e119 100644 --- a/geofw/src/maxmind.rs +++ b/geofw/src/maxmind.rs @@ -51,7 +51,7 @@ 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::Double(s) => f.write_str(&s.to_string()), + Data::Double(s) => write!(f, "{s}"), Data::Bytes(s) => f.write_fmt(format_args!("{:?}", s)), Data::U16(s) => f.write_str(&s.to_string()), Data::U32(s) => f.write_str(&s.to_string()), @@ -168,6 +168,7 @@ impl MaxmindDB { } } + #[allow(unused)] pub fn lookup(&self, addr: IpAddr) -> Option { let node_size = self.metadata.record_size as usize * 2 / 8; let mut node = 0; @@ -197,7 +198,7 @@ impl MaxmindDB { } } - pub fn consume_country_database(mut self, countries: &FxHashSet) -> ProcessedDb { + pub fn consume(mut self, should_block: impl Fn(FxHashMap<&[u8], Data>) -> bool) -> ProcessedDb { let mut stack = vec![]; let node_size = self.metadata.record_size as usize * 2 / 8; stack.push((0, 0)); @@ -230,71 +231,8 @@ impl MaxmindDB { let Data::Map(data) = data else { unreachable!() }; - let Some(Data::Map(country)) = data.get("country".as_bytes()) else { - continue; - }; - let Some(iso_code) = country.get("iso_code".as_bytes()) else { - unreachable!() - }; - if countries.contains(&iso_code.to_string()) { - // 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 - ProcessedDb { - node_count: self.metadata.node_count, - record_size: self.metadata.record_size, - db: self.data[..self.metadata.data_section_start].to_vec(), - } - } - - pub fn consume_asn_database(mut self, asns: &FxHashSet) -> ProcessedDb { - let mut stack = vec![]; - let node_size = self.metadata.record_size as usize * 2 / 8; - stack.push((0, 0)); - - 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); - - if position < 128 && node_1 < self.metadata.node_count { - stack.push((node_1, position + 1)); - } - 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!() - }; - let Some(Data::U32(asn)) = data.get("autonomous_system_number".as_bytes()) else { - continue; - }; - - if asns.contains(asn) { + if should_block(data) { // Mark this node as non existent Self::write_over_node_bytes( &mut self.data