From d202ebb14eff19054897ce8f8aef5a9dd588ef38 Mon Sep 17 00:00:00 2001
From: Slatian <baschdel@disroot.org>
Date: Thu, 23 Feb 2023 00:58:38 +0100
Subject: [PATCH] Added IDN support

---
 Cargo.lock               |  1 +
 Cargo.toml               |  1 +
 src/idna.rs              | 76 ++++++++++++++++++++++++++++++++++++++++
 src/main.rs              | 30 ++++++++++++----
 src/templating_engine.rs |  4 +--
 5 files changed, 104 insertions(+), 8 deletions(-)
 create mode 100644 src/idna.rs

diff --git a/Cargo.lock b/Cargo.lock
index 82a3e72..c64b475 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1190,6 +1190,7 @@ dependencies = [
  "axum",
  "axum-client-ip",
  "clap",
+ "idna 0.3.0",
  "lazy_static",
  "maxminddb",
  "regex",
diff --git a/Cargo.toml b/Cargo.toml
index 2aa8a08..ab2231b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,7 @@ authors = ["Slatian <baschdel@disroot.org>"]
 axum = { version = "0.6", features = ["macros", "headers"] }
 axum-client-ip = "0.4"
 clap = { version = "4", features = ["derive"] }
+idna = "0.3"
 lazy_static = "1.4.0"
 regex = "1.7"
 serde = { version = "1", features = ["derive"] }
diff --git a/src/idna.rs b/src/idna.rs
new file mode 100644
index 0000000..b3e40e6
--- /dev/null
+++ b/src/idna.rs
@@ -0,0 +1,76 @@
+/*
+ * The prupose of this module is to analyse a given string
+ * for being a IDNA domain name and decodes it,
+ * it is also able to return an encoded form if given unicode.
+ * The result is supposed to be used in a template
+ */
+
+use serde::{Deserialize, Serialize};
+use ::idna;
+
+#[derive(Deserialize, Serialize, Copy, Default, Clone, PartialEq)]
+#[serde(rename_all="lowercase")]
+pub enum NameType {
+	Ascii,
+	#[default]
+	Unicode,
+	IDNA,
+}
+
+// Note, that the
+#[derive(Deserialize, Serialize, Default, Clone)]
+pub struct IdnaName {
+	pub unicode: String,
+	// if null the unicode version only contains ascii range chars,
+	// not neccessary to encode
+	pub idna: Option<String>,
+	pub original_was: NameType,
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub decoder_error: Option<String>,
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub encoder_error: Option<String>,
+}
+
+impl IdnaName {
+	pub fn from_string(s: &String) -> Self {
+		let mut original_was = NameType::Unicode;
+		let unicode: String;
+		let decoder_error;
+		if s.starts_with("xn--") && s.is_ascii() {
+			original_was = NameType::IDNA;
+			let (uc, ures) = idna::domain_to_unicode(s);
+			unicode = uc;
+			decoder_error = ures.map_or_else(|e| Some(e.to_string()), |_| None);
+		} else {
+			unicode = s.clone();
+			decoder_error = None;
+		};
+		let (idna, encoder_error) = match idna::domain_to_ascii_strict(s) {
+			Ok(idna) => {
+				if &idna != s || original_was == NameType::IDNA{
+					(Some(idna), None)
+				} else {
+					original_was = NameType::Ascii;
+					(None, None)
+				}
+			},
+			Err(e) => {
+				(None, Some(e.to_string()))
+			}
+		};
+
+		IdnaName {
+			unicode: unicode,
+			idna: idna,
+			original_was: original_was,
+
+			decoder_error: decoder_error,
+			encoder_error: encoder_error,
+		}
+	}
+
+	pub fn was_ascii(&self) -> bool {
+		return self.original_was == NameType::Ascii;
+	}
+}
+
diff --git a/src/main.rs b/src/main.rs
index 1c0f0c9..f833bd6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -35,11 +35,15 @@ mod geoip;
 mod ipinfo;
 mod simple_dns;
 mod templating_engine;
+mod idna;
 
-use crate::geoip::QueryAsn;
-use crate::geoip::QueryLocation;
-use geoip::AsnResult;
-use geoip::LocationResult;
+use crate::geoip::{
+	QueryAsn,
+	QueryLocation,
+	AsnResult,
+	LocationResult,
+};
+use crate::idna::IdnaName;
 
 use crate::templating_engine::{
 	View,
@@ -69,6 +73,13 @@ pub struct IpResult {
 	ip_info:  AddressInfo,
 }
 
+#[derive(serde::Deserialize, serde::Serialize, Default, Clone)]
+pub struct DigResult {
+	records: simple_dns::DnsLookupResult,
+	#[serde(skip_serializing_if = "IdnaName::was_ascii")]
+	idna: IdnaName,
+}
+
 struct ServiceSharedState {
 	templating_engine: templating_engine::Engine,
 	dns_resolver:      TokioAsyncResolver,
@@ -467,11 +478,18 @@ async fn handle_dig_request(
 async fn get_dig_result(
 	dig_query: &String,
 	state:     &ServiceSharedState,
-) -> simple_dns::DnsLookupResult {
+) -> DigResult {
 	let name = &dig_query.trim().trim_end_matches(".").to_string();
 	if match_domain_hidden_list(&name, &state.config.dns.hidden_suffixes) {
 		Default::default()
 	} else {
-		simple_dns::lookup(&state.dns_resolver, name, true).await
+		let idna_name = IdnaName::from_string(&name);
+		DigResult {
+			records: simple_dns::lookup(
+				&state.dns_resolver,
+				&(idna_name.idna.clone().unwrap_or(name.to_owned())+"."),
+				true).await,
+			idna: idna_name,
+		}
 	}
 }
diff --git a/src/templating_engine.rs b/src/templating_engine.rs
index 54f3e11..624dca1 100644
--- a/src/templating_engine.rs
+++ b/src/templating_engine.rs
@@ -13,7 +13,7 @@ use axum::{
 use tera::Tera;
 use toml::Table;
 
-use crate::simple_dns;
+use crate::DigResult;
 use crate::IpResult;
 
 /* Response format */
@@ -52,7 +52,7 @@ pub struct TemplateSettings {
 #[serde(untagged)]
 pub enum View {
 	Asn { asn: u32 },
-	Dig { query: String, result: simple_dns::DnsLookupResult },
+	Dig { query: String, result: DigResult },
 	Index { result: IpResult, user_agent: Option<String> },
 	Ip { result: IpResult },
 	Message{ title: String, message: String },