tokio fixes

Cargo.toml

@@ -5,11 +5,13 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.89"
-reqwest = { version = "0.12.7", default-features = false, features = ["json", "blocking", "rustls-tls"] }
+reqwest = { version = "0.12.7", default-features = false, features = ["json", "rustls-tls"] }
 serde = { version = "1.0.210", features = ["derive"] }
 thiserror = "1.0.63"
 trust-dns-resolver = "0.21"
 rayon = "1.10.0"
 tokio = { version = "1.40.0", features = ["full"] }
 futures = "0.3.30"
 tokio-stream = "0.1.16"
+tracing = "0.1.40"
+tracing-subscriber = "0.3.18"

src/main.rs

@@ -4,6 +4,7 @@ use std::{
     env,
     time::{Duration, Instant},
 };
+use tracing::info;
 
 mod error;
 pub use error::Error;
@@ -13,8 +14,15 @@ mod subdomains;
 use model::Subdomain;
 mod common_ports;
 
+const PORTS_CONCURRENCY: usize = 200;
+const SUBDOMAINS_CONCURRENCY: usize = 100;
+const CRT_SH_TIMEOUT: Duration = Duration::from_secs(10);
+const PORT_TIMEOUT: Duration = Duration::from_secs(1);
+
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
+    tracing_subscriber::fmt::init();
+
     let args: Vec<String> = env::args().collect();
 
     if args.len() != 2 {
@@ -23,32 +31,41 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let target = args[1].as_str();
 
-    let http_timeout = Duration::from_secs(10);
-    let http_client = Client::builder().timeout(http_timeout).build()?;
+    let http_client = Client::builder().timeout(CRT_SH_TIMEOUT).build()?;
 
-    let ports_concurrency = 200;
-    let subdomains_concurrency = 100;
     let scan_start = Instant::now();
 
     let subdomains = subdomains::enumerate(&http_client, target).await?;
 
+    let domain_search = scan_start.elapsed();
+    info!(
+        "Found {} subdomains in {:?}",
+        subdomains.len(),
+        domain_search
+    );
+
     // Concurrent stream method 1: Using buffer_unordered + collect
     let scan_result: Vec<Subdomain> = stream::iter(subdomains.into_iter())
-        .map(|subdomain| ports::scan_ports(ports_concurrency, subdomain))
-        .buffer_unordered(subdomains_concurrency)
+        .map(|subdomain: Subdomain| tokio::spawn(ports::scan_ports(PORTS_CONCURRENCY, subdomain)))
+        .buffer_unordered(SUBDOMAINS_CONCURRENCY)
+        .map(|result| result.unwrap())
         .collect()
         .await;
-
-    let scan_duration = scan_start.elapsed();
-    println!("Scan completed in {:?}", scan_duration);
+
+    let complete_scan = scan_start.elapsed();
+    info!("Scan completed in {:?}", complete_scan);
+    info!(
+        "Excess time: {:?}",
+        complete_scan - domain_search - PORT_TIMEOUT
+    );
 
     for subdomain in scan_result {
         println!("{}:", &subdomain.domain);
         for port in &subdomain.open_ports {
             println!("    {}: open", port.port);
         }
 
-        println!("");
+        println!();
     }
 
     Ok(())

src/model.rs

@@ -3,6 +3,7 @@ use serde::Deserialize;
 #[derive(Debug, Clone)]
 pub struct Subdomain {
     pub domain: String,
+    #[allow(dead_code)]
     pub has_address: bool,
     pub open_ports: Vec<Port>,
 }

src/ports.rs

@@ -1,68 +1,58 @@
 use crate::{
     common_ports::MOST_COMMON_PORTS_100,
     model::{Port, Subdomain},
+    PORT_TIMEOUT,
+};
+use futures::{stream, StreamExt};
+use std::{
+    net::{SocketAddr, ToSocketAddrs},
+    time::Instant,
 };
-use futures::StreamExt;
-use std::net::{SocketAddr, ToSocketAddrs};
-use std::time::Duration;
 use tokio::net::TcpStream;
-use tokio::sync::mpsc;
 
-pub async fn scan_ports(concurrency: usize, subdomain: Subdomain) -> Subdomain {
-    let mut ret = subdomain.clone();
-    let socket_addresses: Vec<SocketAddr> = format!("{}:1024", subdomain.domain)
+#[tracing::instrument(skip(subdomain), fields(subdomain = subdomain.domain))]
+pub async fn scan_ports(concurrency: usize, mut subdomain: Subdomain) -> Subdomain {
+    tracing::debug!("Scanning");
+    let now = Instant::now();
+
+    let socket_addresses: Vec<SocketAddr> = format!("{}:1024", &subdomain.domain)
         .to_socket_addrs()
         .expect("port scanner: Creating socket address")
         .collect();
 
-    if socket_addresses.len() == 0 {
+    if socket_addresses.is_empty() {
         return subdomain;
     }
 
     let socket_address = socket_addresses[0];
 
-    // Concurrent stream method 3: using channels
-    let (input_tx, input_rx) = mpsc::channel(concurrency);
-    let (output_tx, output_rx) = mpsc::channel(concurrency);
-
-    tokio::spawn(async move {
-        for port in MOST_COMMON_PORTS_100 {
-            let _ = input_tx.send(*port).await;
-        }
-    });
+    subdomain.open_ports = stream::iter(
+        // we clone to avoid some borrowing issues
+        MOST_COMMON_PORTS_100.to_owned(),
+    )
+    .map(|port| scan_port(socket_address, port))
+    .buffer_unordered(concurrency)
+    .filter(|p| futures::future::ready(p.is_open))
+    .collect::<Vec<Port>>()
+    .await;
 
-    let input_rx_stream = tokio_stream::wrappers::ReceiverStream::new(input_rx);
-    input_rx_stream
-        .for_each_concurrent(concurrency, |port| {
-            let output_tx = output_tx.clone();
-            async move {
-                let port = scan_port(socket_address, port).await;
-                if port.is_open {
-                    let _ = output_tx.send(port).await;
-                }
-            }
-        })
-        .await;
-    // close channel
-    drop(output_tx);
+    tracing::debug!("Scanned in {:?}", now.elapsed());
 
-    let output_rx_stream = tokio_stream::wrappers::ReceiverStream::new(output_rx);
-    ret.open_ports = output_rx_stream.collect().await;
-
-    ret
+    subdomain
 }
 
+#[tracing::instrument(skip(socket_address))]
 async fn scan_port(mut socket_address: SocketAddr, port: u16) -> Port {
-    let timeout = Duration::from_secs(3);
+    tracing::trace!("Scanning");
     socket_address.set_port(port);
+    let start = Instant::now();
 
     let is_open = matches!(
-        tokio::time::timeout(timeout, TcpStream::connect(&socket_address)).await,
+        tokio::time::timeout(PORT_TIMEOUT, TcpStream::connect(&socket_address)).await,
         Ok(Ok(_)),
     );
 
-    Port {
-        port: port,
-        is_open,
-    }
+    let time_taken = start.elapsed();
+    tracing::trace!("Open {} in {:?}", port, time_taken);
+    Port { port, is_open }
 }

src/subdomains.rs

@@ -16,7 +16,7 @@ type DnsResolver = AsyncResolver<GenericConnection, GenericConnectionProvider<To
 
 pub async fn enumerate(http_client: &Client, target: &str) -> Result<Vec<Subdomain>, Error> {
     let entries: Vec<CrtShEntry> = http_client
-        .get(&format!("https://crt.sh/?q=%25.{}&output=json", target))
+        .get(format!("https://crt.sh/?q=%25.{}&output=json", target))
         .send()
         .await?
         .json()
@@ -31,14 +31,13 @@ pub async fn enumerate(http_client: &Client, target: &str) -> Result<Vec<Subdoma
     // clean and dedup results
     let mut subdomains: HashSet<String> = entries
         .into_iter()
-        .map(|entry| {
+        .flat_map(|entry| {
             entry
                 .name_value
                 .split("\n")
                 .map(|subdomain| subdomain.trim().to_string())
                 .collect::<Vec<String>>()
         })
-        .flatten()
         .filter(|subdomain: &String| subdomain != target)
         .filter(|subdomain: &String| !subdomain.contains("*"))
         .collect();