From 6b5d7beaa910a5d38ff4af5c5a670316f1a35d0d Mon Sep 17 00:00:00 2001
From: cstef <colin@petit-suisse.fr>
Date: Sat, 18 May 2024 13:22:23 +0200
Subject: [PATCH] fix: actually time out with max_time

---
 src/lib.rs            | 29 +++++++++++++++++++----------
 src/runner/filters.rs |  2 +-
 src/utils/logger.rs   | 27 ++++++++++++++++++++++++++-
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 92d8c75..6f4b98f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,13 +19,12 @@ use crate::{
         table::build_opts_table,
     },
 };
-use color_eyre::eyre::bail;
-use color_eyre::eyre::Result;
+use color_eyre::eyre::{bail, Result};
 use colored::Colorize;
 use futures::{future::abortable, FutureExt};
 use indicatif::HumanDuration;
 use itertools::Itertools;
-use log::{error, info, warn};
+use log::{debug, error, info, warn};
 use merge::Merge;
 use parking_lot::Mutex;
 use ptree::print_tree;
@@ -350,9 +349,12 @@ pub async fn _main(opts: Opts) -> Result<Tree<TreeData>> {
     };
     // Run the main function with a timeout if specified
     let (task, handle) = if let Some(max_time) = opts.max_time {
-        abortable(timeout(Duration::from_secs(max_time as u64), main_fun).into_inner())
+        log::debug!("Setting timeout to {}s", max_time);
+        abortable(timeout(Duration::from_secs(max_time as u64), main_fun))
     } else {
-        abortable(main_fun)
+        // NOTE: This is clearly not the best way to handle this because it adds a useless poll but whatever
+        // u64::MAX = 18_446_744_073_709_551_615 secs = 584'942'417'355 years (hopefully enough time to finish the scan)
+        abortable(timeout(Duration::from_secs(u64::MAX), main_fun))
     };
 
     let main_thread = tokio::spawn(task);
@@ -417,14 +419,21 @@ pub async fn _main(opts: Opts) -> Result<Tree<TreeData>> {
         tokio::spawn(ctrlc_task);
     let abort_res = main_thread.await?;
 
-    let thread_res = match abort_res {
+    let timeout_res = match abort_res {
         Ok(res) => Some(res),
-        Err(_) => None,
+        Err(e) => {
+            debug!("Aborted: {}", e);
+            None
+        }
     };
 
-    if thread_res.is_some() {
-        if let Err(e) = thread_res.unwrap() {
-            error!("{}", e);
+    if let Some(thread_res) = timeout_res {
+        if let Err(e) = thread_res {
+            debug!("Timeout reached: {}", e);
+            error!(
+                "Timeout reached after {}s",
+                opts.max_time.unwrap().to_string().bold()
+            );
         } else {
             if !opts.quiet {
                 println!(
diff --git a/src/runner/filters.rs b/src/runner/filters.rs
index 3e67f32..1e88059 100644
--- a/src/runner/filters.rs
+++ b/src/runner/filters.rs
@@ -405,7 +405,7 @@ pub fn is_html_directory(body: &str) -> bool {
 pub fn is_directory(response: &reqwest::Response, body: &str) -> bool {
     if let Some(content_type) = response.headers().get(reqwest::header::CONTENT_TYPE) {
         if content_type.to_str().unwrap().starts_with("text/html") {
-            log::debug!("{} is HTML", response.url());
+            // log::debug!("{} is HTML", response.url());
             if is_html_directory(body) {
                 log::debug!("{} is directory suitable for recursion", response.url());
                 return true;
diff --git a/src/utils/logger.rs b/src/utils/logger.rs
index 5bfa433..2fd4d4c 100644
--- a/src/utils/logger.rs
+++ b/src/utils/logger.rs
@@ -6,6 +6,16 @@ pub fn init_logger() {
     let env = Env::default().filter_or("RWALK_LOG", "info");
 
     Builder::from_env(env)
+        .filter_module("hyper_util::client::legacy::pool", log::LevelFilter::Warn)
+        .filter_module("reqwest::connect", log::LevelFilter::Warn)
+        .filter_module(
+            "hyper_util::client::legacy::connect::http",
+            log::LevelFilter::Warn,
+        )
+        .filter_module(
+            "hyper_util::client::legacy::connect::dns",
+            log::LevelFilter::Warn,
+        )
         .format(|buf, record| {
             let mut style = buf.style();
             match record.level() {
@@ -23,8 +33,23 @@ pub fn init_logger() {
                 log::Level::Debug => "⚙",
                 log::Level::Trace => "⚡",
             });
+            let module = match record.level() {
+                log::Level::Debug => Some(style.value(record.module_path().unwrap())),
+                log::Level::Trace => Some(style.value(record.module_path().unwrap())),
+                _ => None,
+            };
 
-            writeln!(buf, "{} {}", icon, record.args())
+            writeln!(
+                buf,
+                "{} {}{}",
+                icon,
+                if let Some(module) = module {
+                    format!("({}) ", module)
+                } else {
+                    "".to_string()
+                },
+                record.args()
+            )
         })
         .init();
 }