mod_dns_resolver/
lib.rs

1use ::config::{any_err, get_or_create_sub_module, serialize_options, SerdeWrappedValue};
2use dns_resolver::{
3    get_resolver, ptr_host, resolve_a_or_aaaa, reverse_ip, set_mx_concurrency_limit,
4    set_mx_negative_cache_ttl, set_mx_timeout, AggregateResolver, HickoryResolver,
5    IpLookupStrategy, MailExchanger, Resolver, TestResolver,
6};
7use kumo_address::host_or_socket::HostOrSocketAddress;
8use mlua::{Lua, LuaSerdeExt, Value};
9use parking_lot::Mutex;
10use std::collections::HashMap;
11use std::net::IpAddr;
12use std::str::FromStr;
13use std::sync::{Arc, LazyLock};
14use std::time::Duration;
15
16mod config;
17mod hickory_backend;
18mod resolv_conf_loader;
19mod unbound_backend;
20
21use crate::config::DnsResolverConfig;
22
23static RESOLVERS: LazyLock<Mutex<HashMap<String, Arc<Box<dyn Resolver>>>>> =
24    LazyLock::new(|| Mutex::new(HashMap::new()));
25
26pub fn get_resolver_instance(
27    opt_resolver_name: &Option<String>,
28) -> anyhow::Result<Arc<Box<dyn Resolver>>> {
29    if let Some(name) = opt_resolver_name {
30        return RESOLVERS
31            .lock()
32            .get(name)
33            .cloned()
34            .ok_or_else(|| anyhow::anyhow!("resolver {name} is not defined"));
35    }
36
37    Ok(get_resolver())
38}
39
40pub fn get_opt_resolver(
41    opt_resolver_name: &Option<String>,
42) -> anyhow::Result<Option<Arc<Box<dyn Resolver>>>> {
43    if let Some(name) = opt_resolver_name {
44        let r = RESOLVERS
45            .lock()
46            .get(name)
47            .cloned()
48            .ok_or_else(|| anyhow::anyhow!("resolver {name} is not defined"))?;
49        Ok(Some(r))
50    } else {
51        Ok(None)
52    }
53}
54
55pub fn register(lua: &Lua) -> anyhow::Result<()> {
56    let dns_mod = get_or_create_sub_module(lua, "dns")?;
57
58    dns_mod.set(
59        "lookup_mx",
60        lua.create_async_function(|lua, domain: String| async move {
61            let mx = MailExchanger::resolve(&domain).await.map_err(any_err)?;
62            Ok(lua.to_value_with(&*mx, serialize_options()))
63        })?,
64    )?;
65
66    dns_mod.set(
67        "set_mx_concurrency_limit",
68        lua.create_function(move |_lua, limit: usize| {
69            set_mx_concurrency_limit(limit);
70            Ok(())
71        })?,
72    )?;
73
74    dns_mod.set(
75        "set_mx_timeout",
76        lua.create_function(move |lua, duration: Value| {
77            let duration: duration_serde::Wrap<Duration> = lua.from_value(duration)?;
78            set_mx_timeout(duration.into_inner()).map_err(any_err)
79        })?,
80    )?;
81
82    dns_mod.set(
83        "set_mx_negative_cache_ttl",
84        lua.create_function(move |lua, duration: Value| {
85            let duration: duration_serde::Wrap<Duration> = lua.from_value(duration)?;
86            set_mx_negative_cache_ttl(duration.into_inner()).map_err(any_err)
87        })?,
88    )?;
89
90    dns_mod.set(
91        "ptr_host",
92        lua.create_function(move |_lua, ip: String| {
93            let ip: IpAddr = ip.parse().map_err(any_err)?;
94            Ok(ptr_host(ip))
95        })?,
96    )?;
97
98    dns_mod.set(
99        "reverse_ip",
100        lua.create_function(move |_lua, ip: String| {
101            let ip: IpAddr = ip.parse().map_err(any_err)?;
102            Ok(reverse_ip(ip))
103        })?,
104    )?;
105
106    dns_mod.set(
107        "rbl_lookup",
108        lua.create_async_function(
109            |_lua, (ip_str, bl_domain, opt_resolver_name): (String, String, Option<String>)| async move {
110                let resolver = get_resolver_instance(&opt_resolver_name).map_err(any_err)?;
111
112                let address: HostOrSocketAddress = ip_str.parse().map_err(any_err)?;
113                let reversed_ip = reverse_ip(address.ip().ok_or_else(||mlua::Error::external(format!("{ip_str} is not a valid IpAddr or SocketAddr")))?);
114                let name = format!("{reversed_ip}.{bl_domain}.");
115
116                let answers = resolver.resolve_ip(&name).await.map_err(any_err)?;
117                match answers.first() {
118                    Some(ip) => {
119                        let txt = resolver.resolve_txt(&name).await.map(|a| a.as_txt().join("")).ok();
120                        Ok((Some(ip.to_string()), txt))
121                    }
122                    None => {
123                        Ok((None, None))
124                    }
125                }
126            },
127        )?,
128    )?;
129
130    dns_mod.set(
131        "lookup_ptr",
132        lua.create_async_function(
133            |lua, (ip_str, opt_resolver_name): (String, Option<String>)| async move {
134                let resolver = get_resolver_instance(&opt_resolver_name).map_err(any_err)?;
135                let addr = std::net::IpAddr::from_str(&ip_str).map_err(any_err)?;
136                let answer = resolver.resolve_ptr(addr).await.map_err(any_err)?;
137                Ok(lua.to_value_with(&*answer, serialize_options()))
138            },
139        )?,
140    )?;
141
142    dns_mod.set(
143        "lookup_txt",
144        lua.create_async_function(
145            |_lua, (domain, opt_resolver_name): (String, Option<String>)| async move {
146                let resolver = get_resolver_instance(&opt_resolver_name).map_err(any_err)?;
147                let answer = resolver.resolve_txt(&domain).await.map_err(any_err)?;
148                Ok(answer.as_txt())
149            },
150        )?,
151    )?;
152
153    dns_mod.set(
154        "lookup_addr",
155        lua.create_async_function(
156            |_lua,
157             (domain, opt_resolver_name, strategy): (
158                String,
159                Option<String>,
160                Option<SerdeWrappedValue<IpLookupStrategy>>,
161            )| async move {
162                let opt_resolver = get_opt_resolver(&opt_resolver_name).map_err(any_err)?;
163                let result = resolve_a_or_aaaa(
164                    &domain,
165                    opt_resolver.as_ref().map(|r| &***r),
166                    strategy.map(|v| v.0).unwrap_or_default(),
167                )
168                .await
169                .map_err(any_err)?;
170                let result: Vec<String> = result
171                    .into_iter()
172                    .map(|item| item.addr.to_string())
173                    .collect();
174                Ok(result)
175            },
176        )?,
177    )?;
178
179    #[derive(serde::Deserialize, Debug)]
180    #[serde(untagged)]
181    enum ZoneSpec {
182        /// An insecure (non-DNSSEC) zone.
183        Insecure(String),
184        /// A zone with an explicit DNSSEC secure flag.
185        Detailed {
186            zone: String,
187            #[serde(default)]
188            secure: bool,
189        },
190    }
191
192    #[derive(serde::Deserialize, Debug)]
193    #[serde(untagged)]
194    enum TestResolverConfig {
195        /// A bare list of zones (each an insecure string or `{zone, secure}`).
196        Zones(Vec<ZoneSpec>),
197        /// The full form, which additionally supports forcing SERVFAIL for a
198        /// set of owner names.
199        Detailed {
200            zones: Vec<ZoneSpec>,
201            #[serde(default)]
202            servfail: Vec<String>,
203        },
204    }
205
206    impl TestResolverConfig {
207        fn make_resolver(&self) -> anyhow::Result<TestResolver> {
208            let mut resolver = TestResolver::default();
209
210            let (zones, servfail): (&[ZoneSpec], &[String]) = match self {
211                Self::Zones(zones) => (zones, &[]),
212                Self::Detailed { zones, servfail } => (zones, servfail),
213            };
214
215            for zone in zones {
216                resolver = match zone {
217                    ZoneSpec::Insecure(zone) => resolver.with_zone(zone),
218                    ZoneSpec::Detailed { zone, secure: true } => resolver.with_secure_zone(zone),
219                    ZoneSpec::Detailed {
220                        zone,
221                        secure: false,
222                    } => resolver.with_zone(zone),
223                }
224                .map_err(|err| anyhow::anyhow!("{err}"))?;
225            }
226
227            for name in servfail {
228                resolver = resolver.with_servfail(name);
229            }
230
231            Ok(resolver)
232        }
233    }
234
235    #[derive(serde::Deserialize, Debug)]
236    enum KumoResolverConfig {
237        Hickory(DnsResolverConfig),
238        HickorySystemConfig,
239        Unbound(DnsResolverConfig),
240        Test(TestResolverConfig),
241        Aggregate(Vec<KumoResolverConfig>),
242    }
243
244    impl KumoResolverConfig {
245        fn make_resolver(&self, path: &str) -> anyhow::Result<Box<dyn Resolver>> {
246            match self {
247                Self::Hickory(config) => Ok(Box::new(
248                    hickory_backend::build_hickory_resolver(config)
249                        .map_err(|e| anyhow::anyhow!("{path}: {e}"))?,
250                )),
251                Self::HickorySystemConfig => Ok(Box::new(HickoryResolver::new()?)),
252                Self::Unbound(config) => Ok(Box::new(
253                    unbound_backend::build_unbound_resolver(config)
254                        .map_err(|e| anyhow::anyhow!("{path}: {e}"))?,
255                )),
256                Self::Test(config) => Ok(Box::new(config.make_resolver()?)),
257                Self::Aggregate(children) => {
258                    let mut resolver = AggregateResolver::new();
259                    for (idx, child) in children.iter().enumerate() {
260                        let child_path = format!("{path}.Aggregate[{idx}]");
261                        resolver.push_resolver(child.make_resolver(&child_path)?);
262                    }
263                    Ok(Box::new(resolver))
264                }
265            }
266        }
267    }
268
269    dns_mod.set(
270        "configure_resolver",
271        lua.create_function(move |lua, config: mlua::Value| {
272            match lua.from_value::<KumoResolverConfig>(config.clone()) {
273                Ok(config) => {
274                    let resolver = config
275                        .make_resolver("configure_resolver")
276                        .map_err(any_err)?;
277                    dns_resolver::reconfigure_resolver(resolver);
278                    Ok(())
279                }
280                Err(err1) => match lua.from_value::<DnsResolverConfig>(config) {
281                    Ok(config) => {
282                        let resolver = hickory_backend::build_hickory_resolver(&config)
283                            .map_err(any_err)?;
284                        dns_resolver::reconfigure_resolver(resolver);
285                        Ok(())
286                    }
287                    Err(err2) => Err(mlua::Error::external(format!(
288                        "failed to parse config as either KumoResolverConfig ({err1:#}) or DnsResolverConfig ({err2:#})"
289                    ))),
290                },
291            }
292        })?,
293    )?;
294
295    dns_mod.set(
296        "define_resolver",
297        lua.create_function(move |lua, (name, config): (String, mlua::Value)| {
298            let config = lua
299                .from_value::<KumoResolverConfig>(config.clone())
300                .map_err(any_err)?;
301            let path = format!("define_resolver({name:?})");
302            let resolver = config.make_resolver(&path).map_err(any_err)?;
303
304            RESOLVERS.lock().insert(name, resolver.into());
305
306            Ok(())
307        })?,
308    )?;
309
310    dns_mod.set(
311        "configure_unbound_resolver",
312        lua.create_function(move |lua, config: mlua::Value| {
313            let config: DnsResolverConfig = lua.from_value(config)?;
314            let resolver = unbound_backend::build_unbound_resolver(&config).map_err(any_err)?;
315            dns_resolver::reconfigure_resolver(resolver);
316            Ok(())
317        })?,
318    )?;
319
320    dns_mod.set(
321        "configure_test_resolver",
322        lua.create_function(move |lua, config: mlua::Value| {
323            let config = lua
324                .from_value::<TestResolverConfig>(config)
325                .map_err(any_err)?;
326            let resolver = config.make_resolver().map_err(any_err)?;
327            dns_resolver::reconfigure_resolver(resolver);
328            Ok(())
329        })?,
330    )?;
331
332    dns_mod.set(
333        "load_resolv_conf",
334        lua.create_function(move |lua, path: Option<String>| {
335            let config = resolv_conf_loader::load_resolv_conf(path.as_deref()).map_err(any_err)?;
336            lua.to_value_with(&config, serialize_options())
337        })?,
338    )?;
339
340    Ok(())
341}
mod_dns_resolver/lib.rs

mod_dns_resolver/
lib.rs
