From a9132e201d2b06b4adae0dddf3c3cff74eefda1c Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Tue, 15 Oct 2024 15:09:51 +0200
Subject: [PATCH] dns-resolver: create new UnboundResolver type

---
 crates/dns-resolver/src/lib.rs      |  8 +++----
 crates/dns-resolver/src/resolver.rs | 33 +++++++++++++++++++++++++----
 crates/mod-dns-resolver/src/lib.rs  |  4 ++--
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/crates/dns-resolver/src/lib.rs b/crates/dns-resolver/src/lib.rs
index 3b7dbde9..60c8c3b0 100644
--- a/crates/dns-resolver/src/lib.rs
+++ b/crates/dns-resolver/src/lib.rs
@@ -14,6 +14,8 @@ use std::time::Instant;
 
 mod resolver;
 pub use resolver::Resolver;
+#[cfg(feature = "default-unbound")]
+pub use resolver::UnboundResolver;
 
 static RESOLVER: LazyLock<ArcSwap<Resolver>> =
     LazyLock::new(|| ArcSwap::from_pointee(default_resolver()));
@@ -28,11 +30,7 @@ static IP_CACHE: LazyLock<StdMutex<LruCacheWithTtl<Name, Arc<Vec<IpAddr>>>>> =
 
 #[cfg(feature = "default-unbound")]
 fn default_resolver() -> Resolver {
-    // This resolves directly against the root
-    let context = libunbound::Context::new().unwrap();
-    // and enables DNSSEC
-    context.add_builtin_trust_anchors().unwrap();
-    Resolver::Unbound(context.into_async().unwrap())
+    Resolver::Unbound(UnboundResolver::new())
 }
 
 #[cfg(not(feature = "default-unbound"))]
diff --git a/crates/dns-resolver/src/resolver.rs b/crates/dns-resolver/src/resolver.rs
index d59425b6..059702ae 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::AsyncContext;
+use libunbound::{Context, AsyncContext};
 use std::net::IpAddr;
 use std::time::{Duration, Instant};
 
@@ -47,10 +47,35 @@ impl Answer {
     }
 }
 
+#[cfg(feature = "unbound")]
+pub struct UnboundResolver {
+    cx: AsyncContext,
+}
+
+#[cfg(feature = "unbound")]
+impl UnboundResolver {
+    pub fn new() -> Self {
+        // This resolves directly against the root
+        let context = Context::new().unwrap();
+        // and enables DNSSEC
+        context.add_builtin_trust_anchors().unwrap();
+        Self {
+            cx: context.into_async().unwrap(),
+        }
+    }
+}
+
+#[cfg(feature = "unbound")]
+impl From<AsyncContext> for UnboundResolver {
+    fn from(cx: AsyncContext) -> Self {
+        Self { cx }
+    }
+}
+
 pub enum Resolver {
     Tokio(TokioAsyncResolver),
     #[cfg(feature = "unbound")]
-    Unbound(AsyncContext),
+    Unbound(UnboundResolver),
 }
 
 impl Resolver {
@@ -99,10 +124,10 @@ impl Resolver {
                 },
             },
             #[cfg(feature = "unbound")]
-            Self::Unbound(ctx) => {
+            Self::Unbound(resolver) => {
                 let name = name.into_name()?;
                 let name = name.to_ascii();
-                let answer = ctx.resolve(&name, rrtype, DNSClass::IN).await?;
+                let answer = resolver.cx.resolve(&name, rrtype, DNSClass::IN).await?;
                 let mut records = vec![];
                 for r in answer.rdata() {
                     if let Ok(r) = r {
diff --git a/crates/mod-dns-resolver/src/lib.rs b/crates/mod-dns-resolver/src/lib.rs
index dfadb7a0..860e8816 100644
--- a/crates/mod-dns-resolver/src/lib.rs
+++ b/crates/mod-dns-resolver/src/lib.rs
@@ -1,6 +1,6 @@
 use anyhow::Context;
 use config::{any_err, get_or_create_sub_module, serialize_options};
-use dns_resolver::Resolver;
+use dns_resolver::{Resolver, UnboundResolver};
 use dns_resolver::{get_resolver, resolve_a_or_aaaa, MailExchanger};
 use hickory_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts};
 use hickory_resolver::{Name, TokioAsyncResolver};
@@ -185,7 +185,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
                 .context("make async resolver context")
                 .map_err(any_err)?;
 
-            dns_resolver::reconfigure_resolver(Resolver::Unbound(context));
+            dns_resolver::reconfigure_resolver(Resolver::Unbound(UnboundResolver::from(context)));
 
             Ok(())
         })?,