diff --git a/boltconn/src/app.rs b/boltconn/src/app.rs index 4fd8641..1c1e76b 100644 --- a/boltconn/src/app.rs +++ b/boltconn/src/app.rs @@ -110,8 +110,13 @@ impl App { Err(e) => return Err(anyhow!("Parse dns config failed: {e}")), }; Arc::new( - Dns::with_config(outbound_iface.as_str(), config.dns.preference, group) - .map_err(|e| anyhow!("DNS failed to initialize: {e}"))?, + Dns::with_config( + outbound_iface.as_str(), + config.dns.preference, + &config.dns.hosts, + group, + ) + .map_err(|e| anyhow!("DNS failed to initialize: {e}"))?, ) }; diff --git a/boltconn/src/common/host_matcher.rs b/boltconn/src/common/host_matcher.rs index e9e0a43..278cef9 100644 --- a/boltconn/src/common/host_matcher.rs +++ b/boltconn/src/common/host_matcher.rs @@ -35,12 +35,16 @@ impl HostMatcher { } false } + + pub fn builder() -> HostMatcherBuilder { + HostMatcherBuilder::new() + } } pub struct HostMatcherBuilder(Vec<(String, HostType)>); impl HostMatcherBuilder { - pub fn new() -> Self { + fn new() -> Self { Self(Vec::new()) } @@ -64,7 +68,7 @@ impl HostMatcherBuilder { #[test] fn test_matcher() { - let mut builder = HostMatcherBuilder::new(); + let mut builder = HostMatcher::builder(); builder.add_suffix("telemetry.google.com"); builder.add_suffix("analytics.google.com"); builder.add_exact("test.google.com"); @@ -77,7 +81,7 @@ fn test_matcher() { assert!(!matcher.matches("me.notgoogle.com")); assert!(!matcher.matches("ogle.com")); assert!(!matcher.matches("t-02.test.google.com")); - let mut builder = HostMatcherBuilder::new(); + let mut builder = HostMatcher::builder(); builder.add_suffix("ogle.com"); let matcher = builder.build(); assert!(matcher.matches("hi.ogle.com")); diff --git a/boltconn/src/config/config.rs b/boltconn/src/config/config.rs index 7fb712f..7a3a6d3 100644 --- a/boltconn/src/config/config.rs +++ b/boltconn/src/config/config.rs @@ -71,6 +71,8 @@ pub struct RawDnsConfig { pub preference: DnsPreference, pub bootstrap: Vec, pub nameserver: Vec, + #[serde(default = "default_hosts")] + pub hosts: HashMap, } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -183,6 +185,10 @@ fn default_module() -> Vec { Default::default() } +fn default_hosts() -> HashMap { + Default::default() +} + fn default_dns_pref() -> DnsPreference { DnsPreference::PreferIpv4 } diff --git a/boltconn/src/dispatch/ruleset.rs b/boltconn/src/dispatch/ruleset.rs index f4bcf52..eb4a11d 100644 --- a/boltconn/src/dispatch/ruleset.rs +++ b/boltconn/src/dispatch/ruleset.rs @@ -124,7 +124,7 @@ impl RuleSetBuilder { pub fn new(name: &str, payload: &RuleSchema) -> Option { let mut retval = Self { name: name.to_string(), - domain: HostMatcherBuilder::new(), + domain: HostMatcher::builder(), domain_keyword: vec![], ip_cidr: Default::default(), local_ip_cidr: Default::default(), @@ -286,7 +286,7 @@ impl RuleSetBuilder { }); Self { name: name.to_string(), - domain: HostMatcherBuilder::new(), + domain: HostMatcher::builder(), domain_keyword: vec![], ip_cidr: table, local_ip_cidr: Default::default(), diff --git a/boltconn/src/network/dns/dns.rs b/boltconn/src/network/dns/dns.rs index 3049cd3..8413bf1 100644 --- a/boltconn/src/network/dns/dns.rs +++ b/boltconn/src/network/dns/dns.rs @@ -1,5 +1,6 @@ use crate::config::DnsPreference; use crate::network::dns::dns_table::DnsTable; +use crate::network::dns::hosts::HostsResolver; use crate::network::dns::provider::IfaceProvider; use arc_swap::ArcSwap; use hickory_proto::op::{Message, MessageType, ResponseCode}; @@ -7,6 +8,7 @@ use hickory_proto::rr::{DNSClass, RData, Record, RecordType}; use hickory_resolver::config::*; use hickory_resolver::name_server::{GenericConnector, RuntimeProvider}; use hickory_resolver::AsyncResolver; +use std::collections::HashMap; use std::io; use std::io::Result; use std::net::IpAddr; @@ -16,6 +18,7 @@ use std::time::Duration; pub struct GenericDns { table: DnsTable, preference: DnsPreference, + host_resolver: ArcSwap, resolvers: ArcSwap>>>, } @@ -25,6 +28,7 @@ impl Dns { pub fn with_config( iface_name: &str, preference: DnsPreference, + hosts: &HashMap, configs: Vec, ) -> anyhow::Result { let mut resolvers = Vec::new(); @@ -36,9 +40,11 @@ impl Dns { GenericConnector::new(IfaceProvider::new(iface_name)), )); } + let host_resolver = HostsResolver::new(hosts); Ok(Dns { table: DnsTable::new(), preference, + host_resolver: ArcSwap::new(Arc::new(host_resolver)), resolvers: ArcSwap::new(Arc::new(resolvers)), }) } @@ -70,6 +76,7 @@ impl GenericDns

{ Self { table: DnsTable::new(), preference, + host_resolver: ArcSwap::new(Arc::new(HostsResolver::empty())), resolvers: ArcSwap::new(Arc::new(vec![resolver])), } } @@ -126,6 +133,9 @@ impl GenericDns

{ } pub async fn genuine_lookup(&self, domain_name: &str) -> Option { + if let Some(ip) = self.host_resolver.load().resolve(domain_name) { + return Some(ip); + } match self.preference { DnsPreference::Ipv4Only => self.genuine_lookup_v4(domain_name).await, DnsPreference::Ipv6Only => self.genuine_lookup_v6(domain_name).await, diff --git a/boltconn/src/network/dns/hosts.rs b/boltconn/src/network/dns/hosts.rs new file mode 100644 index 0000000..bab926c --- /dev/null +++ b/boltconn/src/network/dns/hosts.rs @@ -0,0 +1,53 @@ +use crate::common::host_matcher::HostMatcher; +use std::collections::HashMap; +use std::net::IpAddr; + +pub(super) struct HostsResolver { + matcher: HostMatcher, + exact_resolver: HashMap, + suffix_resolver: Vec<(String, IpAddr)>, +} + +impl HostsResolver { + pub fn new(hosts: &HashMap) -> Self { + let mut exact_resolver = HashMap::new(); + let mut suffix_resolver = Vec::new(); + let mut builder = HostMatcher::builder(); + for (host, ip) in hosts { + if let Some(stripped_host) = host.strip_prefix("*.") { + suffix_resolver.push((stripped_host.to_string(), *ip)); + builder.add_suffix(stripped_host); + } else { + exact_resolver.insert(host.to_string(), *ip); + builder.add_exact(host); + } + } + Self { + matcher: builder.build(), + exact_resolver, + suffix_resolver, + } + } + + pub fn empty() -> Self { + Self { + matcher: HostMatcher::builder().build(), + exact_resolver: HashMap::new(), + suffix_resolver: Vec::new(), + } + } + + pub fn resolve(&self, host: &str) -> Option { + if !self.matcher.matches(host) { + return None; + } else if let Some(ip) = self.exact_resolver.get(host) { + return Some(*ip); + } + for (suffix, ip) in &self.suffix_resolver { + if host.ends_with(suffix) { + return Some(*ip); + } + } + None + } +} diff --git a/boltconn/src/network/dns/mod.rs b/boltconn/src/network/dns/mod.rs index e899405..704f110 100644 --- a/boltconn/src/network/dns/mod.rs +++ b/boltconn/src/network/dns/mod.rs @@ -1,6 +1,7 @@ #[allow(clippy::module_inception)] mod dns; mod dns_table; +mod hosts; mod provider; use crate::network::dns::provider::IfaceProvider;