Made the geoip database reloadable

This commit is contained in:
Slatian 2023-03-26 14:27:16 +02:00
parent 51f27be997
commit 6a57780490
4 changed files with 42 additions and 24 deletions

1
Cargo.lock generated
View File

@ -372,6 +372,7 @@ dependencies = [
"idna 0.3.0", "idna 0.3.0",
"lazy_static", "lazy_static",
"maxminddb", "maxminddb",
"parking_lot",
"regex", "regex",
"serde", "serde",
"tera", "tera",

View File

@ -13,6 +13,7 @@ clap = { version = "4", features = ["derive"] }
governor = "0.5" governor = "0.5"
idna = "0.3" idna = "0.3"
lazy_static = "1.4.0" lazy_static = "1.4.0"
parking_lot = "0.12"
regex = "1.7" regex = "1.7"
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }

View File

@ -6,6 +6,8 @@
use maxminddb; use maxminddb;
use maxminddb::geoip2; use maxminddb::geoip2;
use parking_lot::RwLock;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::net::IpAddr; use std::net::IpAddr;
use std::path::Path; use std::path::Path;
@ -47,8 +49,9 @@ pub struct AsnResult {
} }
pub struct MMDBCarrier { pub struct MMDBCarrier {
pub mmdb: Option<maxminddb::Reader<Vec<u8>>>, pub mmdb: RwLock<Option<maxminddb::Reader<Vec<u8>>>>,
pub name: String, pub name: String,
pub path: Option<String>,
} }
pub trait QueryLocation { pub trait QueryLocation {
@ -122,7 +125,8 @@ pub fn geoip2_subdivision_to_named_location(item: geoip2::city::Subdivision, lan
impl QueryAsn for MMDBCarrier { impl QueryAsn for MMDBCarrier {
fn query_asn_for_ip(&self, address: &IpAddr) -> Option<AsnResult> { fn query_asn_for_ip(&self, address: &IpAddr) -> Option<AsnResult> {
match &self.mmdb { let mmdb = self.mmdb.read();
match &*mmdb {
Some(mmdb) => { Some(mmdb) => {
match mmdb.lookup::<geoip2::Asn>(*address) { match mmdb.lookup::<geoip2::Asn>(*address) {
Ok(res) => { Ok(res) => {
@ -144,7 +148,8 @@ impl QueryAsn for MMDBCarrier {
impl QueryLocation for MMDBCarrier { impl QueryLocation for MMDBCarrier {
fn query_location_for_ip(&self, address: &IpAddr, languages: &Vec<&String>) -> Option<LocationResult> { fn query_location_for_ip(&self, address: &IpAddr, languages: &Vec<&String>) -> Option<LocationResult> {
match &self.mmdb { let mmdb = self.mmdb.read();
match &*mmdb {
Some(mmdb) => { Some(mmdb) => {
match mmdb.lookup::<geoip2::City>(*address) { match mmdb.lookup::<geoip2::City>(*address) {
Ok(res) => { Ok(res) => {
@ -210,22 +215,38 @@ impl QueryLocation for MMDBCarrier {
} }
impl MMDBCarrier { impl MMDBCarrier {
pub fn load_database_from_path(&mut self, path: &Path) -> Result<(),maxminddb::MaxMindDBError> { pub fn new(name: String, path: Option<String>) -> MMDBCarrier {
MMDBCarrier {
mmdb: RwLock::new(None),
name: name,
path: path,
}
}
pub fn reload_database(&self) -> Result<(),maxminddb::MaxMindDBError> {
match &self.path {
Some(path) => self.load_database_from_path(Path::new(&path)),
None => Ok(()),
}
}
pub fn load_database_from_path(&self, path: &Path) -> Result<(),maxminddb::MaxMindDBError> {
let mut mmdb = self.mmdb.write();
println!("Loading {} from '{}' ...", &self.name, path.display()); println!("Loading {} from '{}' ...", &self.name, path.display());
match maxminddb::Reader::open_readfile(path) { match maxminddb::Reader::open_readfile(path) {
Ok(reader) => { Ok(reader) => {
let wording = if self.mmdb.is_some() { let wording = if mmdb.is_some() {
"Replaced old" "Replaced old"
} else { } else {
"Loaded new" "Loaded new"
}; };
self.mmdb = Some(reader); *mmdb = Some(reader);
println!("{} {} with new one.", wording, &self.name); println!("{} {} with new one.", wording, &self.name);
Ok(()) Ok(())
}, },
Err(e) => { Err(e) => {
println!("Error while reading {}: {}", &self.name, &e); println!("Error while reading {}: {}", &self.name, &e);
if self.mmdb.is_some() { if mmdb.is_some() {
println!("Not replacing old database."); println!("Not replacing old database.");
} }
Err(e) Err(e)

View File

@ -30,7 +30,6 @@ use trust_dns_resolver::{
use std::fs; use std::fs;
use std::net::IpAddr; use std::net::IpAddr;
use std::sync::Arc; use std::sync::Arc;
use std::path::Path;
mod config; mod config;
mod geoip; mod geoip;
@ -205,23 +204,19 @@ async fn main() {
// Initalize GeoIP Database // Initalize GeoIP Database
let mut asn_db = geoip::MMDBCarrier { let asn_db = geoip::MMDBCarrier::new(
mmdb: None, "GeoIP ASN Database".to_string(),
name: "GeoIP ASN Database".to_string(), config.geoip.asn_database.clone()
}; );
match &config.geoip.asn_database {
Some(path) => { asn_db.load_database_from_path(Path::new(&path)).ok(); },
None => {},
}
let mut location_db = geoip::MMDBCarrier { asn_db.reload_database().ok();
mmdb: None,
name: "GeoIP Location Database".to_string(), let location_db = geoip::MMDBCarrier::new(
}; "GeoIP Location Database".to_string(),
match &config.geoip.location_database { config.geoip.location_database.clone()
Some(path) => { location_db.load_database_from_path(Path::new(&path)).ok(); }, );
None => {},
} location_db.reload_database().ok();
// Initalize DNS resolver with os defaults // Initalize DNS resolver with os defaults
println!("Initalizing dns resolver ..."); println!("Initalizing dns resolver ...");