From bc7bf2ffa823953c253a5c303048a034e62de902 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Tue, 15 Oct 2024 15:14:47 +0200
Subject: [PATCH] dns-resolver: create HickoryResolver wrapper

---
 crates/dns-resolver/src/lib.rs      |  7 ++-----
 crates/dns-resolver/src/resolver.rs | 24 +++++++++++++++++++++---
 crates/mod-dns-resolver/src/lib.rs  |  2 +-
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/crates/dns-resolver/src/lib.rs b/crates/dns-resolver/src/lib.rs
index b2d0e61e..d5c3db36 100644
--- a/crates/dns-resolver/src/lib.rs
+++ b/crates/dns-resolver/src/lib.rs
@@ -13,9 +13,9 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex};
 use std::time::Instant;
 
 mod resolver;
-pub use resolver::Resolver;
 #[cfg(feature = "default-unbound")]
 pub use resolver::UnboundResolver;
+pub use resolver::{HickoryResolver, Resolver};
 
 static RESOLVER: LazyLock<ArcSwap<Resolver>> =
     LazyLock::new(|| ArcSwap::from_pointee(default_resolver()));
@@ -35,10 +35,7 @@ fn default_resolver() -> Resolver {
 
 #[cfg(not(feature = "default-unbound"))]
 fn default_resolver() -> Resolver {
-    Resolver::Tokio(
-        hickory_resolver::TokioAsyncResolver::tokio_from_system_conf()
-            .expect("Parsing /etc/resolv.conf failed"),
-    )
+    Resolver::Tokio(HickoryResolver::new().expect("Parsing /etc/resolv.conf failed"))
 }
 
 fn mx_cache_get(name: &Name) -> Option<Arc<MailExchanger>> {
diff --git a/crates/dns-resolver/src/resolver.rs b/crates/dns-resolver/src/resolver.rs
index 41725992..20dd277d 100644
--- a/crates/dns-resolver/src/resolver.rs
+++ b/crates/dns-resolver/src/resolver.rs
@@ -5,7 +5,7 @@ use hickory_resolver::proto::rr::DNSClass;
 use hickory_resolver::proto::rr::{RData, RecordType};
 use hickory_resolver::{IntoName, TokioAsyncResolver, TryParseIp};
 #[cfg(feature = "unbound")]
-use libunbound::{Context, AsyncContext};
+use libunbound::{AsyncContext, Context};
 use std::net::IpAddr;
 use std::time::{Duration, Instant};
 
@@ -72,8 +72,26 @@ impl From<AsyncContext> for UnboundResolver {
     }
 }
 
+pub struct HickoryResolver {
+    inner: TokioAsyncResolver,
+}
+
+impl HickoryResolver {
+    pub fn new() -> Result<Self, hickory_resolver::error::ResolveError> {
+        Ok(Self {
+            inner: TokioAsyncResolver::tokio_from_system_conf()?,
+        })
+    }
+}
+
+impl From<TokioAsyncResolver> for HickoryResolver {
+    fn from(inner: TokioAsyncResolver) -> Self {
+        Self { inner }
+    }
+}
+
 pub enum Resolver {
-    Tokio(TokioAsyncResolver),
+    Tokio(HickoryResolver),
     #[cfg(feature = "unbound")]
     Unbound(UnboundResolver),
 }
@@ -89,7 +107,7 @@ impl Resolver {
         rrtype: RecordType,
     ) -> anyhow::Result<Answer> {
         match self {
-            Self::Tokio(t) => match t.lookup(name, rrtype).await {
+            Self::Tokio(t) => match t.inner.lookup(name, rrtype).await {
                 Ok(result) => {
                     let expires = result.valid_until();
                     let records = result.iter().cloned().collect();
diff --git a/crates/mod-dns-resolver/src/lib.rs b/crates/mod-dns-resolver/src/lib.rs
index 860e8816..cc4d86ab 100644
--- a/crates/mod-dns-resolver/src/lib.rs
+++ b/crates/mod-dns-resolver/src/lib.rs
@@ -128,7 +128,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
 
             let resolver = TokioAsyncResolver::tokio(r_config, config.options);
 
-            dns_resolver::reconfigure_resolver(Resolver::Tokio(resolver));
+            dns_resolver::reconfigure_resolver(Resolver::Tokio(resolver.into()));
 
             Ok(())
         })?,