From fdcdad47d35d1022e5a967b4fbb241ba2d7950c7 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Mon, 21 Oct 2024 15:57:33 +0200
Subject: [PATCH] mod-dns-resolver: wire up test resolver setup

---
 crates/mod-dns-resolver/src/lib.rs | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/crates/mod-dns-resolver/src/lib.rs b/crates/mod-dns-resolver/src/lib.rs
index ea82beb4..823b6f30 100644
--- a/crates/mod-dns-resolver/src/lib.rs
+++ b/crates/mod-dns-resolver/src/lib.rs
@@ -1,7 +1,7 @@
 use anyhow::Context;
 use config::{any_err, get_or_create_sub_module, serialize_options};
 use dns_resolver::{
-    get_resolver, resolve_a_or_aaaa, HickoryResolver, MailExchanger, UnboundResolver,
+    get_resolver, resolve_a_or_aaaa, HickoryResolver, MailExchanger, TestResolver, UnboundResolver,
 };
 use hickory_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts};
 use hickory_resolver::{Name, TokioAsyncResolver};
@@ -191,5 +191,18 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
         })?,
     )?;
 
+    dns_mod.set(
+        "configure_test_resolver",
+        lua.create_function(move |lua, zones: mlua::Value| {
+            let mut resolver = TestResolver::default();
+            for zone in lua.from_value::<Vec<String>>(zones)? {
+                resolver = resolver.with_zone(&zone);
+            }
+
+            dns_resolver::reconfigure_resolver(resolver);
+            Ok(())
+        })?,
+    )?;
+
     Ok(())
 }