Working geoip lookup

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(