Working geoip lookup

src/geoip.rs

@@ -7,20 +7,11 @@ use maxminddb;
 use maxminddb::geoip2;
 
 use std::collections::BTreeMap;
-
 use std::net::IpAddr;
-
+use std::path::Path;
 
 /* Datatypes */
 
-// TODO
-pub enum IsTheNameOf {
-	Continent,
-	Country,
-	Subdivision,
-	City,
-}
-
 #[derive(serde::Deserialize, serde::Serialize, Default)]
 pub struct NamedLocation {
 	iso_code: Option<String>,
@@ -53,15 +44,16 @@ pub struct AsnResult {
 	name: Option<String>,
 }
 
-struct MMDBCarrier {
-	mmdb: Option<maxminddb::Reader<Vec<u8>>>,
+pub struct MMDBCarrier {
+	pub mmdb: Option<maxminddb::Reader<Vec<u8>>>,
+	pub name: String,
 }
 
-trait QueryLocation {
+pub trait QueryLocation {
 	fn query_location_for_ip(&self, address: IpAddr, laguages: &Vec<String>) -> Option<LocationResult>;
 }
 
-trait QueryAsn {
+pub trait QueryAsn {
 	fn query_asn_for_ip(&self, address: IpAddr) -> Option<AsnResult>;
 }
 
@@ -128,7 +120,7 @@ pub fn geoip2_subdivision_to_named_location(item: geoip2::model::Subdivision, la
 
 impl QueryAsn for MMDBCarrier {
 	fn query_asn_for_ip(&self, address: IpAddr) -> Option<AsnResult> {
-		match self.mmdb {
+		match &self.mmdb {
 			Some(mmdb) => {
 				match  mmdb.lookup::<geoip2::Asn>(address) {
 					Ok(res) => {
@@ -150,7 +142,7 @@ impl QueryAsn for MMDBCarrier {
 
 impl QueryLocation for MMDBCarrier {
 	fn query_location_for_ip(&self, address: IpAddr, languages: &Vec<String>) -> Option<LocationResult> {
-		match self.mmdb {
+		match &self.mmdb {
 			Some(mmdb) => {
 				match  mmdb.lookup::<geoip2::City>(address) {
 					Ok(res) => {
@@ -168,7 +160,7 @@ impl QueryLocation for MMDBCarrier {
 								
 							subdivisions: match res.subdivisions {
 								Some(sds) => {
-									let subdivisions = Vec::new();
+									let mut subdivisions = Vec::new();
 									subdivisions.reserve_exact(sds.len());
 									for sd in sds {
 										subdivisions.push(geoip2_subdivision_to_named_location(sd, languages));
@@ -198,3 +190,27 @@ impl QueryLocation for MMDBCarrier {
 	}
 }
 
+impl MMDBCarrier {
+	pub fn load_database_from_path(&mut self, path: &Path) -> Result<(),maxminddb::MaxMindDBError> {
+		println!("Loading {} from '{}' ...", &self.name, path.display());
+		match maxminddb::Reader::open_readfile(path) {
+			Ok(reader) => {
+				let wording = if self.mmdb.is_some() {
+					"Replaced old"
+				} else {
+					"Loaded new"
+				};
+				self.mmdb = Some(reader);
+				println!("{} {} with new one.", wording, &self.name);
+				Ok(())
+			},
+			Err(e) => {
+				println!("Error while reading {}: {}", &self.name, &e);
+				if self.mmdb.is_some() {
+					println!("Not replacing old database.");
+				}
+				Err(e)
+			},
+		}
+	}
+}

src/main.rs

@@ -5,8 +5,7 @@ use axum::{
 	Router,
 	routing::get,
 };
-use maxminddb;
-use maxminddb::geoip2;
+
 use tera::Tera;
 use trust_dns_resolver::{
 	TokioAsyncResolver,
@@ -16,12 +15,14 @@ use trust_dns_resolver::{
 
 use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
-
+use std::path::Path;
 
 mod simple_dns;
 mod templating_engine;
 mod geoip;
 
+use crate::geoip::QueryAsn;
+use crate::geoip::QueryLocation;
 use geoip::AsnResult;
 use geoip::LocationResult;
 
@@ -44,15 +45,17 @@ pub struct DigQuery {
 pub struct IpResult {
 	hostname: Option<String>,
 	asn:      Option<AsnResult>,
+	location: Option<LocationResult>,
 }
 
 struct ServiceSharedState {
 	templating_engine: templating_engine::Engine,
 	dns_resolver:      TokioAsyncResolver,
-	mmdb_asn:          Option<maxminddb::Reader<Vec<u8>>>,
-	mmdb_location:     Option<maxminddb::Reader<Vec<u8>>>,
+	asn_db:            geoip::MMDBCarrier,
+	location_db:       geoip::MMDBCarrier,
 }
 
+
 #[tokio::main]
 async fn main() {
 	// Initalize Tera templates
@@ -69,20 +72,17 @@ async fn main() {
 
 	// Initalize GeoIP Database
 
-	println!("Opening GeoIP ASN Databse ...");
-	let mmdb_asn = maxminddb::Reader::open_readfile("mmdb/GeoLite2-ASN.mmdb");
-	match mmdb_asn {
-		Ok(_) => { /* NOP */ },
-		Err(ref e) => println!("Error while opening GeoIP ASN Databse: {e}"),
-	} 
-
-	println!("Opening GeoIP Location Databse ...");
-	let mmdb_location = maxminddb::Reader::open_readfile("mmdb/GeoLite2-City.mmdb");
-	match mmdb_location {
-		Ok(_) => { /* NOP */ },
-		Err(ref e) => println!("Error while opening GeoIP Location Databse: {e}"),
-	} 
+	let mut asn_db = geoip::MMDBCarrier {
+		mmdb: None,
+		name: "GeoIP ASN Database".to_string(),
+	};
+	asn_db.load_database_from_path(Path::new("mmdb/GeoLite2-ASN.mmdb")).ok();
 
+	let mut location_db = geoip::MMDBCarrier {
+		mmdb: None,
+		name: "GeoIP Location Database".to_string(),
+	};
+	location_db.load_database_from_path(Path::new("mmdb/GeoLite2-City.mmdb")).ok();
 	
 	// Initalize DNS resolver with os defaults
 	println!("Initalizing dns resolver ...");
@@ -96,13 +96,14 @@ async fn main() {
 	};
 	
 	// Initialize shared state
-	let shared_state = Arc::new(ServiceSharedState{
+	let shared_state = Arc::new(
+		ServiceSharedState{
 		templating_engine: templating_engine::Engine{
 			tera: tera,
 		},
 		dns_resolver: dns_resolver,
-		mmdb_asn: mmdb_asn.ok(),
-		mmdb_location: mmdb_location.ok(),
+		asn_db: asn_db,
+		location_db: location_db,
 	});
 	
 	// Initalize axum server
@@ -144,47 +145,17 @@ async fn handle_default_route(
 
 	// do reverse lookup
 	let hostname = simple_dns::reverse_lookup(&state.dns_resolver, &address);
-	
-	let asn_result = match &state.mmdb_asn {
-		Some(mmdb_asn) => {
-			match mmdb_asn.lookup::<geoip2::Asn>(address) {
-				Ok(res) => {
-					Some(AsnResult {
-						asn: res.autonomous_system_number,
-						name: res.autonomous_system_organization.map(ToString::to_string),
-					})
-				}
-				Err(e) => {
-					println!("Error while looking up ASN for {address}: {e}");
-					Default::default()
-				}
-			}
-		},
-		None => None,
-	};
 
-	let location_result = match &state.mmdb_location {
-		Some(mmdb_location) => {
-			match mmdb_location.lookup::<geoip2::City>(address) {
-				Ok(res) => {
-					Some(AsnResult {
-						asn: res.autonomous_system_number,
-						name: res.autonomous_system_organization.map(ToString::to_string),
-					})
-				}
-				Err(e) => {
-					println!("Error while looking up ASN for {address}: {e}");
-					Default::default()
-				}
-			}
-		},
-		None => None,
-	};
-	
+	// asn lookup
+	let asn_result = state.asn_db.query_asn_for_ip(address);
+
+	// location lookup
+	let location_result = state.location_db.query_location_for_ip(address, &vec!["en".to_string()]);
 
 	let result = IpResult{
 		hostname: hostname.await,
 		asn: asn_result,
+		location: location_result,
 	};
 
 	state.templating_engine.render_view(