leyline/core/config/
dns.rs1use 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}