new: implemented adaptive timeout, now certain timeout sensitive plugins such as dns and port.scanner can adjust the workers timeout when getting positive results

Author: evilsocket

src/options.rs

@@ -67,10 +67,10 @@ pub(crate) struct Options {
     #[clap(long, value_enum, default_value_t = session::loot::OutputFormat::Text)]
     pub output_format: session::loot::OutputFormat,
     /// Connection timeout in milliseconds.
-    #[clap(long, default_value_t = 10000)]
+    #[clap(long, default_value_t = 1000)]
     pub timeout: u64,
     /// Number of attempts if a request fails.
-    #[clap(long, default_value_t = 3)]
+    #[clap(long, default_value_t = 1)]
     pub retries: usize,
     /// Delay in milliseconds to wait before a retry.
     #[clap(long, default_value_t = 1000)]

src/plugins/dns/mod.rs

@@ -25,22 +25,51 @@ super::manager::register_plugin! {
 
 #[derive(Clone)]
 pub(crate) struct DNS {
-    resolver: Option<TokioAsyncResolver>,
+    ips: Vec<IpAddr>,
     opts: options::Options,
+    core_options: Options,
     hits: Arc<DashMap<IpAddr, usize>>,
     domains: Arc<DashSet<String>>,
 }
 
 impl DNS {
     pub fn new() -> Self {
         DNS {
-            resolver: None,
+            ips: vec![],
             opts: options::Options::default(),
+            core_options: Options::default(),
             hits: Arc::new(DashMap::new()),
             domains: Arc::new(DashSet::new()),
         }
     }
 
+    async fn get_resolver(&self, timeout: Duration) -> Result<TokioAsyncResolver, Error> {
+        if !self.ips.is_empty() {
+            let nameserver_group =
+                NameServerConfigGroup::from_ips_clear(&self.ips, self.opts.dns_port, true);
+
+            let mut options = ResolverOpts::default();
+
+            options.num_concurrent_reqs = self.core_options.concurrency;
+            options.attempts = self.opts.dns_attempts;
+            options.timeout = timeout;
+            options.shuffle_dns_servers = true;
+
+            Ok(AsyncResolver::tokio(
+                ResolverConfig::from_parts(None, vec![], nameserver_group),
+                options,
+            ))
+        } else {
+            let (config, mut options) =
+                trust_dns_resolver::system_conf::read_system_conf().map_err(|e| e.to_string())?;
+
+            options.attempts = self.opts.dns_attempts;
+            options.timeout = timeout;
+
+            Ok(AsyncResolver::tokio(config, options))
+        }
+    }
+
     async fn filter(&self, addresses: Vec<IpAddr>) -> Vec<IpAddr> {
         // Some domains are configured to resolve any subdomain, whatever it is, to the same IP. We do
         // this filtering in order too many positives for an address and work around this behaviour.
@@ -170,8 +199,9 @@ impl Plugin for DNS {
     }
 
     async fn setup(&mut self, opts: &Options) -> Result<(), Error> {
+        self.core_options = opts.clone();
         self.opts = opts.dns.clone();
-        self.resolver = Some(if let Some(resolvers) = opts.dns.dns_resolvers.as_ref() {
+        self.ips = if let Some(resolvers) = opts.dns.dns_resolvers.as_ref() {
             let ips: Vec<IpAddr> = resolvers
                 .split(',')
                 .map(|s| s.trim())
@@ -181,25 +211,12 @@ impl Plugin for DNS {
 
             log::info!("using resolvers: {:?}", &ips);
 
-            let nameserver_group =
-                NameServerConfigGroup::from_ips_clear(&ips, opts.dns.dns_port, true);
-
-            let mut options = ResolverOpts::default();
-
-            options.num_concurrent_reqs = opts.concurrency;
-            options.attempts = opts.dns.dns_attempts;
-            options.timeout = Duration::from_millis(opts.timeout);
-            options.shuffle_dns_servers = true;
-
-            AsyncResolver::tokio(
-                ResolverConfig::from_parts(None, vec![], nameserver_group),
-                options,
-            )
+            ips
         } else {
             log::info!("using system resolver");
 
-            AsyncResolver::tokio_from_system_conf().map_err(|e| e.to_string())?
-        });
+            vec![]
+        };
 
         Ok(())
     }
@@ -224,10 +241,11 @@ impl Plugin for DNS {
         }
 
         // each worker will use its own resolver object instance
-        let resolver = self.resolver.as_ref().unwrap().clone();
-
+        let resolver = self.get_resolver(timeout).await?;
+        let started_at = std::time::Instant::now();
         // attempt resolving this subdomain to a one or more IP addresses
         if let Ok(response) = resolver.lookup_ip(&subdomain).await {
+            let elapsed = started_at.elapsed();
             // collect valid IPs
             let addresses: Vec<IpAddr> = response.iter().filter(|ip| !ip.is_loopback()).collect();
             // Some domains are configured to resolve any subdomain, whatever it is, to the same IP. We do
@@ -259,6 +277,7 @@ impl Plugin for DNS {
                 };
 
                 loot_data.push(("addresses".to_owned(), addr_data));
+                loot_data.push(("time_ms".to_owned(), elapsed.as_millis().to_string()));
 
                 let mut loot = vec![Loot::new("dns", &subdomain, loot_data)];
 

src/plugins/manager.rs

@@ -1,6 +1,6 @@
 use std::collections::BTreeMap;
 use std::sync::{LazyLock, Mutex};
-use std::time;
+use std::time::{self, Duration};
 
 use ansi_term::Style;
 use dashmap::DashSet;
@@ -134,7 +134,6 @@ pub(crate) async fn run(
 async fn worker(plugin: &dyn Plugin, unreachables: Arc<DashSet<Arc<str>>>, session: Arc<Session>) {
     log::debug!("worker started");
 
-    let timeout = time::Duration::from_millis(session.options.timeout);
     let retry_time: time::Duration = time::Duration::from_millis(session.options.retry_time);
 
     while let Ok(creds) = session.recv_credentials().await {
@@ -161,6 +160,7 @@ async fn worker(plugin: &dyn Plugin, unreachables: Arc<DashSet<Arc<str>>>, sessi
 
             // skip attempt if we had enough failures from this specific target
             if !unreachables.contains(creds.target.as_str()) {
+                let timeout = session.runtime.get_timeout();
                 match plugin.attempt(&creds, timeout).await {
                     Err(err) => {
                         errors += 1;
@@ -196,6 +196,21 @@ async fn worker(plugin: &dyn Plugin, unreachables: Arc<DashSet<Arc<str>>>, sessi
                         // do we have new loot?
                         if let Some(loots) = loot {
                             for loot in loots {
+                                // some plugins might return the elapsed time in the loot
+                                // if we have it, we can adjust the timeout to avoid waiting
+                                // for too long
+                                if let Some(mut elapsed) = loot.get_elapsed_time() {
+                                    if elapsed.as_millis() == 0 {
+                                        // make it at least 1ms
+                                        elapsed = Duration::from_millis(1);
+                                    }
+                                    // increase the elapsed time just to be safe
+                                    elapsed *= 10;
+                                    if elapsed < timeout {
+                                        session.runtime.set_timeout(elapsed.as_millis() as u64);
+                                    }
+                                }
+
                                 session.add_loot(loot).await.unwrap();
                             }
                         }

src/plugins/port_scanner/mod.rs

@@ -50,7 +50,10 @@ impl PortScanner {
             let mut data = vec![
                 ("transport".to_owned(), "tcp".to_owned()),
                 ("port".to_owned(), creds.username.to_owned()),
-                ("time".to_owned(), format!("{:?}", start.elapsed())),
+                (
+                    "time_ms".to_owned(),
+                    format!("{}", start.elapsed().as_millis()),
+                ),
             ];
 
             if !self.opts.port_scanner_no_banners {
@@ -143,7 +146,10 @@ impl PortScanner {
                 let mut data = vec![
                     ("transport".to_owned(), "udp".to_owned()),
                     ("port".to_owned(), creds.username.to_owned()),
-                    ("time".to_owned(), format!("{:?}", start.elapsed())),
+                    (
+                        "time_ms".to_owned(),
+                        format!("{}", start.elapsed().as_millis()),
+                    ),
                 ];
 
                 for (name, value) in grab_udp_banner(&buf[0..size]).await {

src/plugins/port_scanner/options.rs

@@ -1,5 +1,6 @@
 use std::fs::OpenOptions;
 use std::io::prelude::*;
+use std::time::Duration;
 use std::{fmt, path::Path};
 
 use ansi_term::Colour;
@@ -59,6 +60,15 @@ impl Loot {
         &self.data
     }
 
+    pub fn get_elapsed_time(&self) -> Option<Duration> {
+        if let Some(time) = self.data.get("time_ms")
+            && let Ok(time) = time.parse::<u64>()
+        {
+            return Some(Duration::from_millis(time));
+        }
+        None
+    }
+
     pub fn is_partial(&self) -> bool {
         self.partial
     }

src/session/loot.rs

@@ -109,7 +109,7 @@ pub(crate) struct Session {
     pub results: Mutex<Vec<Loot>>,
 
     #[serde(skip_serializing, skip_deserializing)]
-    runtime: Runtime,
+    pub runtime: Runtime,
 }
 
 impl Session {
@@ -129,7 +129,7 @@ impl Session {
             parse_target(target, 0)?;
         }
 
-        let runtime = Runtime::new(options.concurrency);
+        let runtime = Runtime::new(options.concurrency, options.timeout);
         let total = AtomicUsize::new(0);
         let done = AtomicUsize::new(0);
         let errors = AtomicUsize::new(0);
@@ -153,7 +153,7 @@ impl Session {
             let file = fs::File::open(path).map_err(|e| e.to_string())?;
             let mut session: Session = serde_json::from_reader(file).map_err(|e| e.to_string())?;
 
-            session.runtime = Runtime::new(session.options.concurrency);
+            session.runtime = Runtime::new(session.options.concurrency, session.options.timeout);
 
             Ok(Arc::new(session))
         } else {
@@ -341,8 +341,6 @@ impl Session {
     pub async fn report_runtime_statistics(&self) {
         let report_interval = time::Duration::from_millis(self.options.report_time);
         while !self.is_stop() {
-            tokio::time::sleep(report_interval).await;
-
             let total = self.get_total();
             let done = self.get_done();
             let perc = (done as f32 / total as f32) * 100.0;
@@ -371,6 +369,8 @@ impl Session {
             } else {
                 log::info!("{}", stats.to_text());
             }
+
+            tokio::time::sleep(report_interval).await;
         }
     }
 }

src/session/mod.rs

@@ -1,4 +1,7 @@
-use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::{
+    sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
+    time::Duration,
+};
 
 use super::Error;
 use crate::Credentials;
@@ -9,21 +12,23 @@ pub(crate) struct Runtime {
     creds_tx: async_channel::Sender<Credentials>,
     creds_rx: async_channel::Receiver<Credentials>,
     speed: AtomicUsize,
+    timeout_ms: AtomicU64,
 }
 
 impl Default for Runtime {
     fn default() -> Self {
-        Self::new(1)
+        Self::new(1, 10000)
     }
 }
 
 impl Runtime {
-    pub(crate) fn new(concurrency: usize) -> Self {
+    pub(crate) fn new(concurrency: usize, timeout_ms: u64) -> Self {
         // use a buffer of 100x the concurrency to avoid blocking on sending credentials
         let (creds_tx, creds_rx) = async_channel::bounded(concurrency * 100);
         Self {
             stop: AtomicBool::new(false),
             speed: AtomicUsize::new(0),
+            timeout_ms: AtomicU64::new(timeout_ms),
             creds_tx,
             creds_rx,
         }
@@ -54,4 +59,17 @@ impl Runtime {
     pub async fn recv_credentials(&self) -> Result<Credentials, Error> {
         self.creds_rx.recv().await.map_err(|e| e.to_string())
     }
+
+    pub fn get_timeout(&self) -> Duration {
+        Duration::from_millis(self.timeout_ms.load(Ordering::Relaxed))
+    }
+
+    pub fn set_timeout(&self, timeout_ms: u64) {
+        log::info!(
+            "adjusting timeout from {:?} to {}ms",
+            self.get_timeout(),
+            timeout_ms
+        );
+        self.timeout_ms.store(timeout_ms, Ordering::Relaxed);
+    }
 }