✨ Implement HTTP ping

mokeyish · Mar 11, 2023 · 11e6d50 · 11e6d50
1 parent 36c0454
commit 11e6d50
Show file tree

Hide file tree

Showing 7 changed files with 333 additions and 50 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -51,12 +51,13 @@ tracing-subscriber = { version = "0.3", features = [
 ] }
 # tracing-appender = "0.2"
 tokio = { version = "1.21", features = ["time", "rt", "signal", "macros"] }
+tokio-rustls = "0.23.4"
 url = "2.3.1"
 trust-dns-proto = { version = "0.22.0", git = "https://github.com/bluejekyll/trust-dns.git", rev = "5492bde"}
 trust-dns-resolver = { version = "0.22.0", features = ["serde-config"] , git = "https://github.com/bluejekyll/trust-dns.git", rev = "5492bde"}
 trust-dns-server = { version = "0.22.0", features = ["resolver"] , git = "https://github.com/bluejekyll/trust-dns.git", rev = "5492bde"}
 webpki-roots = "0.22.1"
-rustls = "0.20.0"
+rustls = { version = "0.20.0", features = ["dangerous_configuration"]}
 rustls-native-certs = { version = "0.6.2", git = "https://github.com/mokeyish/rustls-native-certs.git" }
 lru = { version = "0.10", default-features = false}
 once_cell = "1.16.0"

diff --git a/src/dns.rs b/src/dns.rs
@@ -26,6 +26,7 @@ pub struct DnsContext {
     pub fastest_speed: Duration,
     pub source: LookupFrom,
     pub no_cache: bool,
+    pub background: bool
 }
 
 impl DnsContext {
@@ -38,6 +39,7 @@ impl DnsContext {
             fastest_speed: Default::default(),
             source: Default::default(),
             no_cache: false,
+            background: false
         }
     }
 

diff --git a/src/dns_conf.rs b/src/dns_conf.rs
@@ -1376,6 +1376,8 @@ pub enum SpeedCheckMode {
     None,
     Ping,
     Tcp(u16),
+    Http(u16),
+    Https(u16),
 }
 
 impl Default for SpeedCheckMode {
@@ -1393,6 +1395,14 @@ impl FromStr for SpeedCheckMode {
             Ok(SpeedCheckMode::Ping)
         } else if let Some(port) = s.strip_prefix("tcp:") {
             u16::from_str(port).map(SpeedCheckMode::Tcp).map_err(|_| ())
+        } else if let Some(port) = s.strip_prefix("http:") {
+            u16::from_str(port)
+                .map(SpeedCheckMode::Http)
+                .map_err(|_| ())
+        } else if let Some(port) = s.strip_prefix("https:") {
+            u16::from_str(port)
+                .map(SpeedCheckMode::Https)
+                .map_err(|_| ())
         } else {
             Err(())
         }

diff --git a/src/dns_mw_cache.rs b/src/dns_mw_cache.rs
@@ -237,6 +237,8 @@ impl DomainPrefetcher {
                                 let now = Instant::now();
                                 let mut ctx =
                                     DnsContext::new(query.name(), cfg.clone(), Default::default());
+
+                                ctx.background = true;
 
                                 if let Ok(lookup) =
                                     client.execute(&mut ctx, &query.clone().into()).await

diff --git a/src/dns_mw_ns.rs b/src/dns_mw_ns.rs
@@ -105,9 +105,9 @@ impl Middleware<DnsContext, DnsRequest, DnsResponse, DnsError> for NameServerMid
         if rtype.is_ip_addr() {
             let cfg = ctx.cfg();
 
-            let opts = match ctx.domain_rule.as_ref() {
+            let mut opts = match ctx.domain_rule.as_ref() {
                 Some(rule) => LookupIpOptions {
-                    response_strategy: rule.response_mode.unwrap_or_else(|| cfg.response_mode()),
+                    response_strategy: rule.get(|n|n.response_mode).unwrap_or_else(|| cfg.response_mode()),
                     speed_check_mode: if rule.speed_check_mode.is_empty() {
                         cfg.speed_check_mode().clone()
                     } else {
@@ -128,6 +128,10 @@ impl Middleware<DnsContext, DnsRequest, DnsResponse, DnsError> for NameServerMid
                 },
             };
 
+            if ctx.background {
+                opts.response_strategy = ResponseMode::FastestIp;
+            }
+
             match lookup_ip(name_server_group.deref(), name.clone(), rtype, &opts).await {
                 Ok(lookup_ip) => Ok(lookup_ip.into()),
                 Err(err) => Err(err),
@@ -197,6 +201,28 @@ async fn lookup_ip(
                                 .map(|ip| PingAddr::Tcp(SocketAddr::new(*ip, *port)))
                                 .collect::<Vec<_>>()
                         }
+                        SpeedCheckMode::Http(port) => {
+                            debug!(
+                                "Speed test {} http ping {:?} port {}",
+                                lookup_ip.query().name(),
+                                ips,
+                                port
+                            );
+                            ips.iter()
+                                .map(|ip| PingAddr::Http(SocketAddr::new(*ip, *port)))
+                                .collect::<Vec<_>>()
+                        },
+                        SpeedCheckMode::Https(port) => {
+                            debug!(
+                                "Speed test {} https ping {:?} port {}",
+                                lookup_ip.query().name(),
+                                ips,
+                                port
+                            );
+                            ips.iter()
+                                .map(|ip| PingAddr::Https(SocketAddr::new(*ip, *port)))
+                                .collect::<Vec<_>>()
+                        },
                     };
 
                     ping_tasks.push(ping_fastest(ping_dests, ping_ops).boxed());

diff --git a/src/dns_rule.rs b/src/dns_rule.rs
@@ -212,6 +212,14 @@ impl DomainRuleTreeNode {
     pub fn zone(&self) -> Option<&Arc<DomainRuleTreeNode>> {
         self.zone.as_ref()
     }
+
+    pub fn get<T>(&self, f: impl Fn(&Self) -> Option<T>) -> Option<T> {
+        f(self).or_else(|| {
+            self.zone()
+                .map(|z| f(z))
+                .unwrap_or_default()
+        })
+    }
 }
 
 impl Deref for DomainRuleTreeNode {