feat: v0.2.0-rc, optimize code

src/app.rs

@@ -19,13 +19,14 @@ use super::service::HostConfig;
 pub struct ProxyApp {
     tls: bool,
     lookup_dns: bool,
+    dns_domains: Option<Vec<String>>,
     host_config_map: HashMap<String, HostConfig>,
     dns_resolver: TokioAsyncResolver,
     dns_resolver_cache: RwLock<HashMap<String, String>>,
 }
 
 impl ProxyApp {
-    pub fn new(tls: bool, lookup_dns: bool, host_config_map: HashMap<String, HostConfig>) -> Self {
+    pub fn new(tls: bool, lookup_dns: bool, dns_domains: Option<Vec<String>>, host_config_map: HashMap<String, HostConfig>) -> 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<String> {
         {
@@ -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(|| "<none>".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::<Vec<_>>().len();
+        let sub_domain = domain_pattern_chars.skip(1).take(domain_pattern_len - 2).collect::<String>();
+        return domain.contains(&sub_domain);
+    }
+    if domain_pattern.starts_with('*') {
+        let sub_domain = domain_pattern_chars.skip(1).collect::<String>();
+        return domain.ends_with(&sub_domain);
+    }
+    if domain_pattern.ends_with('*') {
+        let domain_pattern_len = domain_pattern.chars().collect::<Vec<_>>().len();
+        let sub_domain = domain_pattern_chars.take(domain_pattern_len - 1).collect::<String>();
+        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*"));
 }

src/config.rs

@@ -26,6 +26,7 @@ impl ProxyConfig {
 pub struct ProxyGroup {
     pub port: u16,
     pub lookup_dns: Option<bool>,
+    pub dns_domains: Option<Vec<String>>,
     pub tls: Option<ProxyTls>,
     pub proxy_map: Option<HashMap<String, ProxyItem>>,
 }

src/main.rs

@@ -56,8 +56,10 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
         let listen_address = format!("0.0.0.0:{}", group.port);
         let host_config_map = build_host_config_map(group);
         let lookup_dns = group.lookup_dns.unwrap_or(false);
+        let dns_domains = group.dns_domains.clone();
 
         log::info!("Listen at: {}, tls: {}, lookup_dns: {}", listen_address, group.tls.is_some(), lookup_dns);
+        log::debug!("Dns domains: {:#?}", dns_domains);
 
         match &group.tls {
             None => {
@@ -65,6 +67,7 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
                     &server.configuration,
                     &listen_address,
                     lookup_dns,
+                    dns_domains,
                     host_config_map,
                 )));
             }
@@ -73,6 +76,7 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
                     &server.configuration,
                     &listen_address,
                     lookup_dns,
+                    dns_domains,
                     proxy_tls,
                     host_config_map,
                 )));

src/service.rs

@@ -83,9 +83,10 @@ pub fn proxy_service_tcp(
     server_conf: &Arc<ServerConf>,
     listen_addr: &str,
     lookup_dns: bool,
+    dns_domains: Option<Vec<String>>,
     host_config_map: HashMap<String, HostConfig>,
 ) -> 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<ServerConf>,
     listen_addr: &str,
     lookup_dns: bool,
+    dns_domains: Option<Vec<String>>,
     proxy_tls: &ProxyTls,
     host_config_map: HashMap<String, HostConfig>,
 ) -> 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| {