Added dns functionality

src/main.rs

@@ -7,6 +7,12 @@ use axum::{
 	routing::get,
 };
 use tera::Tera;
+use trust_dns_resolver::{
+	TokioAsyncResolver,
+	config::ResolverOpts,
+	config::ResolverConfig,
+	error::*,
+};
 
 use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
@@ -33,37 +39,75 @@ impl ToString for ResponseFormat {
 
 #[derive(serde::Deserialize, serde::Serialize)]
 struct IpQuery {
-	ip: Option<IpAddr>,
+	ip:     Option<IpAddr>,
 	format: Option<ResponseFormat>,
 }
 
+#[derive(serde::Deserialize, serde::Serialize)]
+struct DigQuery {
+	name:   String,
+	format: Option<ResponseFormat>
+}
+
+#[derive(serde::Deserialize, serde::Serialize)]
+struct MxRecord {
+	preference: u16,
+	exchange:   String,
+}
+
+#[derive(serde::Deserialize, serde::Serialize)]
+struct DigResult {
+	a:    Vec<IpAddr>,
+	aaaa: Vec<IpAddr>,
+	mx:   Vec<MxRecord>,
+}
+
 struct ServiceSharedState {
 	tera: Tera,
+	dns_resolver: TokioAsyncResolver,
 }
 
 #[tokio::main]
 async fn main() {
 	// Initalize Tera templates
 	// TODO: don't hardcode template directory
-	let tera = match Tera::new("templates/*.html") {
+	println!("Parsing Templates ...");
+	let res = Tera::new("templates/*.html");
+	let tera = match res {
 		Ok(t) => t,
 		Err(e) => {
 			println!("Template parsing error(s): {}", e);
 			::std::process::exit(1);
 		}
 	};
+	
+	// Initalize DNS resolver with os defaults
+	println!("Initalizing dns resolver ...");
+	let res = TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
+	let dns_resolver = match res {
+		Ok(resolver) => resolver,
+		Err(e) => {
+			println!("Error while setting up dns resolver: {e}");
+			::std::process::exit(1);
+		}
+	};
+	
 	// Initialize shared state
 	let shared_state = Arc::new(ServiceSharedState{
 		tera: tera,
+		dns_resolver: dns_resolver,
 	});
 	
 	// Initalize axum server
     let app = Router::new()
 	    .route("/", get(handle_default_route))
+	    .route("/dig", get(handle_dig_route))
     	.route("/hi", get(hello_world_handler))
     	.with_state(shared_state)
     	;
 
+	println!("Starting Server ...");
+
     axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
         .serve(app.into_make_service())
         .await
@@ -74,6 +118,33 @@ async fn hello_world_handler() -> &'static str {
 	"Hello, there, you, awesome creature!"
 }
 
+async fn simple_reverse_dns_lookup(
+		resolver: &TokioAsyncResolver,
+		address: &IpAddr,
+) -> Option<String> {
+	let revese_res = resolver.reverse_lookup(*address);
+	match revese_res.await {
+		Ok(lookup) => {
+			for name in lookup {
+				return Some(name.to_string())
+			}
+			None
+		}
+		Err(e) => {
+			let kind = e.kind();
+			match kind {
+				ResolveErrorKind::NoRecordsFound { .. } => {
+					//Ignore, that just happens …
+				}
+				_ => {
+					println!("Reverse lookup on {address} failed: {kind}");
+				}
+			}
+			None
+		}
+	}	
+}
+
 async fn handle_default_route(
 	Query(ip_query): Query<IpQuery>,
 	State(arc_state): State<Arc<ServiceSharedState>>,
@@ -84,9 +155,14 @@ async fn handle_default_route(
 	let format_string = format.to_string();
 
 	let state = Arc::clone(&arc_state);
+
+	// do reverse lookup
+	let hostname = simple_reverse_dns_lookup(&state.dns_resolver, &address);
+
 	let mut context = tera::Context::new();
 	context.insert("ip", &address);
 	context.insert("format", &format_string);
+	context.insert("hostname", &hostname.await);
 
 	match state.tera.render("index.html", &context) {
 		Ok(html) => Ok(Html(html)),
@@ -95,5 +171,73 @@ async fn handle_default_route(
 			Err(StatusCode::INTERNAL_SERVER_ERROR)
 		} 
 	}
+}
+
+async fn handle_dig_route(
+	Query(dig_query): Query<DigQuery>,
+	State(arc_state): State<Arc<ServiceSharedState>>,
+) -> Result<Html<String>,StatusCode> {
+	
+	let state = Arc::clone(&arc_state);
+	let name = dig_query.name;
+
+	let ipv4_lookup_res = state.dns_resolver.ipv4_lookup(&name);
+	let ipv6_lookup_res = state.dns_resolver.ipv6_lookup(&name);
+	let mx_lookup_res = state.dns_resolver.mx_lookup(&name);
+
+	let mut dig_result = DigResult{
+		a: Vec::new(),
+		aaaa: Vec::new(),
+		mx: Vec::new(),
+	};
+
+	match ipv4_lookup_res.await {
+		Ok(lookup) => {
+			for address in lookup {
+				dig_result.a.push(std::net::IpAddr::V4(address));
+			}
+		}
+		Err(e) => {
+			println!("There was an error while looking A up {name}: {e}");
+		}
+	}
+	
+	match ipv6_lookup_res.await {
+		Ok(lookup) => {
+			for address in lookup {
+				dig_result.aaaa.push(std::net::IpAddr::V6(address));
+			}
+		}
+		Err(e) => {
+			println!("There was an error while looking AAAA up {name}: {e}");
+		}
+	}
+
+	match mx_lookup_res.await {
+		Ok(lookup) => {
+			for mx in lookup {
+				dig_result.mx.push(MxRecord{
+					preference: mx.preference(),
+					exchange:   mx.exchange().to_string(),
+				});
+			}
+		}
+		Err(e) => {
+			println!("There was an error while looking MX up {name}: {e}");
+		}
+	}
+
+
+	let mut context = tera::Context::new();
+	context.insert("dig_query", &name);
+	context.insert("dig_result", &dig_result);
+
+	match state.tera.render("dig.html", &context) {
+		Ok(html) => Ok(Html(html)),
+		Err(e) => {
+			println!("There was an error while rendering index.html: {e}");
+			Err(StatusCode::INTERNAL_SERVER_ERROR)
+		} 
+	}
 	
 }