⚡Optimize cache prefetch module

Cargo.toml

@@ -3,7 +3,7 @@ name = "smartdns"
 version = "0.8.4"
 authors = ["YISH <[email protected]>"]
 edition = "2021"
-rust-version = "1.70.0"
+rust-version = "1.75.0"
 
 keywords = ["DNS", "BIND", "dig", "named", "dnssec", "SmartDNS", "Dnsmasq"]
 categories = ["network-programming"]

src/app.rs

@@ -15,11 +15,12 @@ use crate::{
     config::ServerOpts,
     dns::{DnsRequest, DnsResponse, SerialMessage},
     dns_conf::RuntimeConfig,
+    dns_error::LookupError,
     dns_mw::{DnsMiddlewareBuilder, DnsMiddlewareHandler},
     dns_mw_cache::DnsCache,
     log,
     server::{DnsHandle, IncomingDnsRequest, ServerHandle},
-    third_ext::FutureJoinAllExt as _,
+    third_ext::{FutureJoinAllExt as _, FutureTimeoutExt},
 };
 
 pub struct App {
@@ -412,7 +413,28 @@ async fn process(
                             response_header.set_authoritative(false);
 
                             let response = {
-                                match handler.search(&request, &server_opts).await {
+                                let res =
+                                    handler
+                                        .search(&request, &server_opts)
+                                        .timeout(Duration::from_secs(
+                                            if server_opts.is_background { 60 } else { 5 },
+                                        ))
+                                        .await
+                                        .unwrap_or_else(|_| {
+                                            let query = request.query().original().to_owned();
+                                            log::warn!(
+                                                "Query {} {} {} timeout.",
+                                                query.name(),
+                                                query.query_type(),
+                                                if server_opts.is_background {
+                                                    "in background"
+                                                } else {
+                                                    ""
+                                                }
+                                            );
+                                            Err(LookupError::no_records_found(query, 10))
+                                        });
+                                match res {
                                     Ok(lookup) => lookup,
                                     Err(e) => {
                                         if e.is_nx_domain() {

src/cli.rs

@@ -138,6 +138,10 @@ struct CompatibleCli {
     /// Verbose screen.
     #[arg(short = 'x', long)]
     verbose: bool,
+
+    /// ignore segment fault signal
+    #[arg(short = 'S')]
+    segment_fault_signal: bool,
 }
 
 impl From<CompatibleCli> for Cli {
@@ -147,6 +151,7 @@ impl From<CompatibleCli> for Cli {
             pid,
             verbose,
             foreground,
+            segment_fault_signal: _,
         }: CompatibleCli,
     ) -> Self {
         if !foreground {
@@ -299,4 +304,18 @@ mod tests {
 
         assert_eq!(cli.log_level(), Some(log::Level::INFO));
     }
+
+    #[test]
+    fn test_cli_args_parse_compatible_run_4() {
+        let cli = Cli::parse_from(["smartdns", "-f", "-c", "/etc/smartdns.conf", "-S"]);
+        assert!(matches!(
+            cli.command,
+            Commands::Run {
+                conf: Some(_),
+                pid: None,
+            }
+        ));
+
+        assert_eq!(cli.log_level(), Some(log::Level::INFO));
+    }
 }

src/config/speed_mode.rs

@@ -48,7 +48,7 @@ impl std::fmt::Debug for SpeedCheckMode {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Clone, PartialEq, Eq, Hash)]
 pub struct SpeedCheckModeList(pub Vec<SpeedCheckMode>);
 
 impl SpeedCheckModeList {
@@ -72,6 +72,16 @@ impl From<Vec<SpeedCheckMode>> for SpeedCheckModeList {
     }
 }
 
+impl std::fmt::Debug for SpeedCheckModeList {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for (i, m) in self.0.iter().enumerate() {
+            let last = i == self.len() - 1;
+            write!(f, "{:?}{}", m, if !last { ", " } else { "" })?;
+        }
+        Ok(())
+    }
+}
+
 impl std::ops::Deref for SpeedCheckModeList {
     type Target = Vec<SpeedCheckMode>;
 

src/dns.rs

@@ -193,6 +193,17 @@ mod serial_message {
                 .into())
         }
     }
+
+    impl TryFrom<SerialMessage> for Message {
+        type Error = ProtoError;
+
+        fn try_from(value: SerialMessage) -> Result<Self, Self::Error> {
+            match value {
+                SerialMessage::Raw(message, _, _) => Ok(message),
+                SerialMessage::Bytes(bytes, _, _) => Message::from_vec(&bytes),
+            }
+        }
+    }
 }
 
 mod request {

src/dns_client.rs

@@ -684,25 +684,30 @@ impl GenericResolver for NameServer {
         let name = name.into_name()?;
         let options: LookupOptions = options.into();
 
+        let query = Query::query(name, options.record_type);
+
+        let client_subnet = options.client_subnet.or(self.opts.client_subnet);
+
+        if options.client_subnet.is_none() {
+            if let Some(subnet) = client_subnet.as_ref() {
+                log::debug!(
+                    "query name: {} type: {} subnet: {}/{}",
+                    query.name(),
+                    query.query_type(),
+                    subnet.addr(),
+                    subnet.scope_prefix(),
+                );
+            }
+        }
+
         let request_options = {
             let opts = &self.options();
             let mut request_opts = DnsRequestOptions::default();
             request_opts.recursion_desired = opts.recursion_desired;
-            request_opts.use_edns = opts.edns0;
+            request_opts.use_edns = opts.edns0 || client_subnet.is_some();
             request_opts
         };
 
-        let query = Query::query(name, options.record_type);
-
-        let client_subnet = options.client_subnet.or(self.opts.client_subnet);
-
-        log::debug!(
-            "query name: {} type: {}, {:?}",
-            query.name(),
-            query.query_type(),
-            client_subnet
-        );
-
         let req = DnsRequest::new(
             build_message(query, request_options, client_subnet, options.is_dnssec),
             request_options,

src/dns_conf.rs

@@ -134,6 +134,35 @@ impl RuntimeConfig {
                 DEFAULT_GROUP
             );
         }
+
+        info!(
+            "cache: {}",
+            if self.cache_size() > 0 {
+                format!("size({})", self.cache_size())
+            } else {
+                "OFF".to_string()
+            }
+        );
+
+        if self.cache_size() > 0 {
+            info!(
+                "cache persist: {}",
+                if self.cache_persist() { "YES" } else { "NO" }
+            );
+
+            info!(
+                "domain prefetch: {}",
+                if self.prefetch_domain() { "ON" } else { "OFF" }
+            );
+        }
+
+        info!(
+            "speed check mode: {}",
+            match self.speed_check_mode() {
+                Some(mode) => format!("{:?}", mode),
+                None => "OFF".to_string(),
+            }
+        );
     }
 
     pub fn server_name(&self) -> Name {

src/dns_error.rs

@@ -36,6 +36,22 @@ pub enum LookupError {
     Io(Arc<io::Error>),
 }
 
+impl PartialEq for LookupError {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::ResponseCode(l0), Self::ResponseCode(r0)) => l0 == r0,
+            (Self::Proto(l0), Self::Proto(r0)) => l0.to_string() == r0.to_string(),
+            (Self::ResolveError(l0), Self::ResolveError(r0)) => l0.to_string() == r0.to_string(),
+            #[cfg(feature = "hickory-recursor")]
+            (Self::RecursiveError(l0), Self::RecursiveError(r0)) => {
+                l0.to_string() == r0.to_string()
+            }
+            (Self::Io(l0), Self::Io(r0)) => l0.to_string() == r0.to_string(),
+            _ => core::mem::discriminant(self) == core::mem::discriminant(other),
+        }
+    }
+}
+
 impl LookupError {
     pub fn is_nx_domain(&self) -> bool {
         matches!(self, Self::ResponseCode(resc) if resc.eq(&ResponseCode::NXDomain))