feat: nameserver policy

boltconn/src/app.rs

@@ -8,7 +8,7 @@ use crate::external::{
 };
 use crate::intercept::{InterceptModifier, InterceptionManager};
 use crate::network::configure::TunConfigure;
-use crate::network::dns::{new_bootstrap_resolver, parse_dns_config, Dns};
+use crate::network::dns::{new_bootstrap_resolver, parse_dns_config, Dns, NameserverPolicies};
 use crate::network::tun_device::TunDevice;
 use crate::platform::get_default_v4_route;
 use crate::proxy::{
@@ -113,10 +113,18 @@ impl App {
                 }
                 Err(e) => return Err(anyhow!("Parse dns config failed: {e}")),
             };
+            let ns_policy = NameserverPolicies::new(
+                &config.dns.nameserver_policy,
+                Some(&bootstrap),
+                outbound_iface.as_str(),
+            )
+            .await
+            .map_err(|e| anyhow!("Parse nameserver policy failed: {e}"))?;
             Arc::new(Dns::with_config(
                 outbound_iface.as_str(),
                 config.dns.preference,
                 &config.dns.hosts,
+                ns_policy,
                 group,
             ))
         };
@@ -360,6 +368,12 @@ impl App {
         let bootstrap =
             new_bootstrap_resolver(&self.outbound_iface, config.dns.bootstrap.as_slice())?;
         let group = parse_dns_config(config.dns.nameserver.iter(), Some(&bootstrap)).await?;
+        let ns_policy = NameserverPolicies::new(
+            &config.dns.nameserver_policy,
+            Some(&bootstrap),
+            self.outbound_iface.as_str(),
+        )
+        .await?;
         let dispatching = {
             let builder =
                 DispatchingBuilder::new(self.dns.clone(), mmdb.clone(), &loaded_config, &ruleset)?;
@@ -377,6 +391,8 @@ impl App {
         );
 
         self.dns.replace_resolvers(&self.outbound_iface, group);
+        self.dns.replace_ns_policy(ns_policy);
+        self.dns.replace_hosts(&config.dns.hosts);
 
         // start atomic replacing
         self.api_dispatching_handler.store(dispatching.clone());

boltconn/src/common/host_matcher.rs

@@ -57,6 +57,15 @@ impl HostMatcherBuilder {
             .push((host.chars().rev().collect(), HostType::Suffix))
     }
 
+    /// Automatically add a host to the matcher, determining the type based on wildcards.
+    pub fn add_auto(&mut self, host: &str) {
+        if let Some(stripped_host) = host.strip_prefix("*.") {
+            self.add_suffix(stripped_host);
+        } else {
+            self.add_exact(host);
+        }
+    }
+
     pub fn build(self) -> HostMatcher {
         HostMatcher(Trie::from_iter(self.0))
     }

boltconn/src/network/dns/dns.rs

@@ -1,6 +1,7 @@
 use crate::config::DnsPreference;
 use crate::network::dns::dns_table::DnsTable;
 use crate::network::dns::hosts::HostsResolver;
+use crate::network::dns::ns_policy::NameserverPolicies;
 use crate::network::dns::provider::IfaceProvider;
 use arc_swap::ArcSwap;
 use hickory_proto::op::{Message, MessageType, ResponseCode};
@@ -19,6 +20,7 @@ pub struct GenericDns<P: RuntimeProvider> {
     table: DnsTable,
     preference: DnsPreference,
     host_resolver: ArcSwap<HostsResolver>,
+    ns_policy: ArcSwap<NameserverPolicies>,
     resolvers: ArcSwap<Vec<AsyncResolver<GenericConnector<P>>>>,
 }
 
@@ -29,6 +31,7 @@ impl Dns {
         iface_name: &str,
         preference: DnsPreference,
         hosts: &HashMap<String, IpAddr>,
+        ns_policy: NameserverPolicies,
         configs: Vec<NameServerConfigGroup>,
     ) -> Dns {
         let mut resolvers = Vec::new();
@@ -45,10 +48,20 @@ impl Dns {
             table: DnsTable::new(),
             preference,
             host_resolver: ArcSwap::new(Arc::new(host_resolver)),
+            ns_policy: ArcSwap::new(Arc::new(ns_policy)),
             resolvers: ArcSwap::new(Arc::new(resolvers)),
         }
     }
 
+    pub fn replace_hosts(&self, hosts: &HashMap<String, IpAddr>) {
+        self.host_resolver
+            .store(Arc::new(HostsResolver::new(hosts)));
+    }
+
+    pub fn replace_ns_policy(&self, ns_policy: NameserverPolicies) {
+        self.ns_policy.store(Arc::new(ns_policy));
+    }
+
     pub fn replace_resolvers(&self, iface_name: &str, configs: Vec<NameServerConfigGroup>) {
         let mut resolvers = Vec::new();
         for config in configs {
@@ -72,6 +85,7 @@ impl<P: RuntimeProvider> GenericDns<P> {
             table: DnsTable::new(),
             preference,
             host_resolver: ArcSwap::new(Arc::new(HostsResolver::empty())),
+            ns_policy: ArcSwap::new(Arc::new(NameserverPolicies::empty())),
             resolvers: ArcSwap::new(Arc::new(vec![resolver])),
         }
     }
@@ -95,34 +109,54 @@ impl<P: RuntimeProvider> GenericDns<P> {
 
     async fn genuine_lookup_v4(&self, domain_name: &str) -> Option<IpAddr> {
         for r in self.resolvers.load().iter() {
-            if let Ok(r) =
-                tokio::time::timeout(Duration::from_secs(5), r.ipv4_lookup(domain_name)).await
-            {
-                if let Ok(result) = r {
-                    if let Some(i) = result.iter().next() {
-                        return Some(i.0.into());
-                    }
+            if let Some(ip) = Self::genuine_lookup_one_v4(domain_name, r).await {
+                return Some(ip);
+            }
+        }
+        None
+    }
+
+    async fn genuine_lookup_one_v4<R: RuntimeProvider>(
+        domain_name: &str,
+        resolver: &AsyncResolver<GenericConnector<R>>,
+    ) -> Option<IpAddr> {
+        if let Ok(r) =
+            tokio::time::timeout(Duration::from_secs(5), resolver.ipv4_lookup(domain_name)).await
+        {
+            if let Ok(result) = r {
+                if let Some(i) = result.iter().next() {
+                    return Some(i.0.into());
                 }
-            } else {
-                tracing::debug!("DNS lookup for {domain_name} timeout");
             }
+        } else {
+            tracing::debug!("DNS lookup for {domain_name} timeout");
         }
         None
     }
 
     async fn genuine_lookup_v6(&self, domain_name: &str) -> Option<IpAddr> {
         for r in self.resolvers.load().iter() {
-            if let Ok(r) =
-                tokio::time::timeout(Duration::from_secs(5), r.ipv6_lookup(domain_name)).await
-            {
-                if let Ok(result) = r {
-                    if let Some(i) = result.iter().next() {
-                        return Some(i.0.into());
-                    }
+            if let Some(ip) = Self::genuine_lookup_one_v6(domain_name, r).await {
+                return Some(ip);
+            }
+        }
+        None
+    }
+
+    async fn genuine_lookup_one_v6<R: RuntimeProvider>(
+        domain_name: &str,
+        resolver: &AsyncResolver<GenericConnector<R>>,
+    ) -> Option<IpAddr> {
+        if let Ok(r) =
+            tokio::time::timeout(Duration::from_secs(5), resolver.ipv6_lookup(domain_name)).await
+        {
+            if let Ok(result) = r {
+                if let Some(i) = result.iter().next() {
+                    return Some(i.0.into());
                 }
-            } else {
-                tracing::debug!("DNS lookup for {domain_name} timeout");
             }
+        } else {
+            tracing::debug!("DNS lookup for {domain_name} timeout");
         }
         None
     }
@@ -131,6 +165,26 @@ impl<P: RuntimeProvider> GenericDns<P> {
         if let Some(ip) = self.host_resolver.load().resolve(domain_name) {
             return Some(ip);
         }
+        if let Some(resolver) = self.ns_policy.load().resolve(domain_name) {
+            return match self.preference {
+                DnsPreference::Ipv4Only => Self::genuine_lookup_one_v4(domain_name, resolver).await,
+                DnsPreference::Ipv6Only => Self::genuine_lookup_one_v6(domain_name, resolver).await,
+                DnsPreference::PreferIpv4 => {
+                    if let Some(a) = Self::genuine_lookup_one_v4(domain_name, resolver).await {
+                        Some(a)
+                    } else {
+                        Self::genuine_lookup_one_v6(domain_name, resolver).await
+                    }
+                }
+                DnsPreference::PreferIpv6 => {
+                    if let Some(a) = Self::genuine_lookup_one_v6(domain_name, resolver).await {
+                        Some(a)
+                    } else {
+                        Self::genuine_lookup_one_v4(domain_name, resolver).await
+                    }
+                }
+            };
+        }
         match self.preference {
             DnsPreference::Ipv4Only => self.genuine_lookup_v4(domain_name).await,
             DnsPreference::Ipv6Only => self.genuine_lookup_v6(domain_name).await,

boltconn/src/network/dns/mod.rs

@@ -12,6 +12,7 @@ use hickory_resolver::config::{
 };
 use hickory_resolver::name_server::GenericConnector;
 use hickory_resolver::AsyncResolver;
+pub use ns_policy::NameserverPolicies;
 use std::net::{IpAddr, SocketAddr};
 
 fn add_tls_server(

boltconn/src/network/dns/ns_policy.rs

@@ -7,7 +7,7 @@ use hickory_resolver::AsyncResolver;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 
-pub(super) struct NameserverPolicies {
+pub struct NameserverPolicies {
     matchers: Vec<(HostMatcher, AsyncResolver<GenericConnector<IfaceProvider>>)>,
 }
 
@@ -31,12 +31,12 @@ impl NameserverPolicies {
 
             // clustering
             match builder.entry(key) {
-                Entry::Occupied(mut e) => e.get_mut().0.add_suffix(host),
+                Entry::Occupied(mut e) => e.get_mut().0.add_auto(host),
                 Entry::Vacant(e) => {
                     let ns_config =
                         parse_single_dns(e.key().0.as_str(), e.key().1.as_str(), bootstrap).await?;
                     let mut matcher = HostMatcher::builder();
-                    matcher.add_suffix(host);
+                    matcher.add_auto(host);
                     e.insert((matcher, ns_config));
                 }
             }
@@ -55,4 +55,19 @@ impl NameserverPolicies {
             .collect();
         Ok(Self { matchers: res })
     }
+
+    pub fn empty() -> Self {
+        Self {
+            matchers: Vec::new(),
+        }
+    }
+
+    pub fn resolve(&self, host: &str) -> Option<&AsyncResolver<GenericConnector<IfaceProvider>>> {
+        for (matcher, resolver) in &self.matchers {
+            if matcher.matches(host) {
+                return Some(resolver);
+            }
+        }
+        None
+    }
 }