Skip to main content

leyline/core/config/
dns.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::sync::Arc;
4
5use crate::tls::{Resolver, SystemResolver};
6
7use super::host::normalize_host;
8
9#[derive(Clone)]
10#[non_exhaustive]
11pub struct DnsConfig {
12    resolver: Arc<dyn Resolver>,
13    overrides: HashMap<String, Vec<SocketAddr>>,
14}
15
16impl std::fmt::Debug for DnsConfig {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        f.debug_struct("DnsConfig")
19            .field("overrides", &self.overrides)
20            .finish_non_exhaustive()
21    }
22}
23
24impl Default for DnsConfig {
25    fn default() -> Self {
26        Self {
27            resolver: Arc::new(SystemResolver),
28            overrides: HashMap::new(),
29        }
30    }
31}
32
33impl DnsConfig {
34    pub fn new() -> Self {
35        Self::default()
36    }
37
38    pub fn resolver(mut self, resolver: Arc<dyn Resolver>) -> Self {
39        self.resolver = resolver;
40        self
41    }
42
43    pub fn resolve_host(
44        mut self,
45        host: impl AsRef<str>,
46        addrs: impl IntoIterator<Item = SocketAddr>,
47    ) -> Self {
48        let host = normalize_host(host.as_ref());
49        let addrs: Vec<SocketAddr> = addrs.into_iter().collect();
50        if addrs.is_empty() {
51            self.overrides.remove(&host);
52        } else {
53            self.overrides.insert(host, addrs);
54        }
55        self
56    }
57
58    pub(crate) fn into_resolver(self) -> Arc<dyn Resolver> {
59        if self.overrides.is_empty() {
60            self.resolver
61        } else {
62            Arc::new(LayeredResolver {
63                resolver: self.resolver,
64                overrides: self.overrides,
65            })
66        }
67    }
68}
69
70impl From<Arc<dyn Resolver>> for DnsConfig {
71    fn from(resolver: Arc<dyn Resolver>) -> Self {
72        Self::new().resolver(resolver)
73    }
74}
75
76struct LayeredResolver {
77    resolver: Arc<dyn Resolver>,
78    overrides: HashMap<String, Vec<SocketAddr>>,
79}
80
81impl Resolver for LayeredResolver {
82    fn resolve<'a>(&'a self, host: &'a str, port: u16) -> crate::tls::ResolveFuture<'a> {
83        let key = normalize_host(host);
84        if let Some(addrs) = self.overrides.get(&key) {
85            let addrs = addrs
86                .iter()
87                .map(|addr| SocketAddr::new(addr.ip(), port))
88                .collect::<Vec<_>>();
89            return Box::pin(async move { Ok(addrs) });
90        }
91        self.resolver.resolve(host, port)
92    }
93}