diff --git a/proxy_config.json b/proxy_config.json index a088b43..b95aea1 100644 --- a/proxy_config.json +++ b/proxy_config.json @@ -2,7 +2,10 @@ "groups": [ { "port": 443, - "lookup_dns": true, + "lookup_dns": false, + "dns_domains": [ + "*baidu*" + ], "tls": { "issuer_cert": "cert.pem", "issuer_key": "cert.key" diff --git a/src/app.rs b/src/app.rs index 2899177..62cbf8d 100644 --- a/src/app.rs +++ b/src/app.rs @@ -19,13 +19,14 @@ use super::service::HostConfig; pub struct ProxyApp { tls: bool, lookup_dns: bool, + dns_domains: Option>, host_config_map: HashMap, dns_resolver: TokioAsyncResolver, dns_resolver_cache: RwLock>, } impl ProxyApp { - pub fn new(tls: bool, lookup_dns: bool, host_config_map: HashMap) -> Self { + pub fn new(tls: bool, lookup_dns: bool, dns_domains: Option>, host_config_map: HashMap) -> Self { let dns_resolver = TokioAsyncResolver::tokio( ResolverConfig::default(), ResolverOpts::default(), @@ -33,12 +34,32 @@ impl ProxyApp { ProxyApp { tls, lookup_dns, + dns_domains, host_config_map, dns_resolver, dns_resolver_cache: Default::default(), } } + fn lookup_dns(&self, hostname: &str) -> bool { + if self.lookup_dns { + return true; + } + if let Some(dns_names) = &self.dns_domains { + for domain_pattern in dns_names { + // Only support below patterns: + // - example.com + // - google.* + // - *.google.com + // - *google* + if domain_match(hostname, domain_pattern) { + return true; + } + } + } + false + } + // just only support IPv4 async fn lookup_ipv4(&self, hostname: &str) -> Option { { @@ -119,7 +140,7 @@ impl ProxyHttp for ProxyApp { ))); } - if self.lookup_dns { + if self.lookup_dns(&hostname) { if let Some(address) = self.lookup_ipv4(&hostname).await { let peer_addr = format!("{}:{}", address, if self.tls { 443 } else { 80 }); log::info!("DNS peer: {} --> {}", hostname, peer_addr); @@ -188,4 +209,38 @@ impl ProxyHttp for ProxyApp { ).unwrap_or_else(|| "".into()) ); } +} + +#[inline] +fn domain_match(domain: &str, domain_pattern: &str) -> bool { + if domain == domain_pattern { + return true; + } + let domain_pattern_chars = domain_pattern.chars(); + if domain_pattern.starts_with('*') && domain_pattern.ends_with('*') { + let domain_pattern_len = domain_pattern.chars().collect::>().len(); + let sub_domain = domain_pattern_chars.skip(1).take(domain_pattern_len - 2).collect::(); + return domain.contains(&sub_domain); + } + if domain_pattern.starts_with('*') { + let sub_domain = domain_pattern_chars.skip(1).collect::(); + return domain.ends_with(&sub_domain); + } + if domain_pattern.ends_with('*') { + let domain_pattern_len = domain_pattern.chars().collect::>().len(); + let sub_domain = domain_pattern_chars.take(domain_pattern_len - 1).collect::(); + return domain.starts_with(&sub_domain); + } + false +} + +#[test] +fn test_domain_match() { + assert!(domain_match("example.com", "example.com")); + assert!(!domain_match("www.example.com", "example.com")); + assert!(domain_match("www.example.com", "*example.com")); + assert!(domain_match("www.example.com", "*example.*")); + assert!(domain_match("example.com", "example.*")); + assert!(domain_match("example.com", "example*")); + assert!(domain_match("example.com", "*ample*")); } \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index 2f3f929..349c8ed 100644 --- a/src/config.rs +++ b/src/config.rs @@ -26,6 +26,7 @@ impl ProxyConfig { pub struct ProxyGroup { pub port: u16, pub lookup_dns: Option, + pub dns_domains: Option>, pub tls: Option, pub proxy_map: Option>, } diff --git a/src/main.rs b/src/main.rs index 5cbe803..34e8ce8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -56,8 +56,10 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec { @@ -65,6 +67,7 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec Vec, listen_addr: &str, lookup_dns: bool, + dns_domains: Option>, host_config_map: HashMap, ) -> impl pingora::services::Service { - let proxy_app = ProxyApp::new(false, lookup_dns, host_config_map); + let proxy_app = ProxyApp::new(false, lookup_dns, dns_domains, host_config_map); let mut service = http_proxy_service(server_conf, proxy_app); service.add_tcp(listen_addr); @@ -97,10 +98,11 @@ pub fn proxy_service_tls( server_conf: &Arc, listen_addr: &str, lookup_dns: bool, + dns_domains: Option>, proxy_tls: &ProxyTls, host_config_map: HashMap, ) -> impl pingora::services::Service { - let proxy_app = ProxyApp::new(true, lookup_dns, host_config_map); + let proxy_app = ProxyApp::new(true, lookup_dns, dns_domains, host_config_map); let mut service = http_proxy_service(server_conf, proxy_app); let callback = Box::new(Callback::new(proxy_tls).unwrap_or_else(|e| {