From 0411248a93fc662a9c1090e1135b5791056340b8 Mon Sep 17 00:00:00 2001
From: "E. Almqvist" <elalmqvist@gmail.com>
Date: Fri, 12 Aug 2022 19:44:45 +0200
Subject: [PATCH] Worker refactor

---
 src/scanner.rs | 41 +++++++++++++++++++++++++++--------------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/scanner.rs b/src/scanner.rs
index 7fe98ba..e5b89a8 100644
--- a/src/scanner.rs
+++ b/src/scanner.rs
@@ -2,8 +2,8 @@ use crate::ipv4::{IPv4, IPv4Range};
 use anyhow::Result;
 use log::info;
 use std::net::{SocketAddr, TcpStream};
-use std::thread;
 use std::thread::JoinHandle;
+use std::{panic, thread};
 
 fn tcp_scan(mut target: IPv4, target_port: u16) -> bool {
     let _dest = target
@@ -23,16 +23,16 @@ fn create_scan_thread(
     thread_id: u32,
     ip_range: IPv4Range,
     target_port: u16,
-) -> JoinHandle<Vec<bool>> {
+) -> JoinHandle<Vec<(u32, bool)>> {
     thread::spawn(move || {
         info!("Starting thread worker #{}", thread_id);
-        let mut results: Vec<bool> = Vec::new();
+        let mut results: Vec<(u32, bool)> = Vec::new();
 
         // do the scan thing
         ip_range.into_iter().for_each(|id| {
             let target = IPv4::new(id as u64);
             let result = tcp_scan(target, target_port);
-            results.push(result);
+            results.push((id, result));
         });
 
         results
@@ -44,13 +44,19 @@ fn create_scan_worker(
     ips_per_thread: u32,
     target_port: u16,
     ignorelist: Vec<u32>,
-) -> JoinHandle<Vec<bool>> {
+) -> thread::Result<Vec<(u32, bool)>> {
     let (f, t) = (
         (thread_id * ips_per_thread),
         ((thread_id + 1) * ips_per_thread),
     );
     let range = IPv4Range::new(f, t, ignorelist);
-    create_scan_thread(thread_id, range, target_port)
+    let thread = create_scan_thread(thread_id, range, target_port);
+    let result = thread.join();
+
+    match result {
+        Ok(_) => return result,
+        Err(e) => panic!("** Worker panic! {:?}", e),
+    }
 }
 
 fn create_scan_workers(
@@ -59,35 +65,42 @@ fn create_scan_workers(
     target_port: u16,
     num_threads: u32,
     ignorelist: Option<Vec<u32>>,
-) -> Vec<JoinHandle<Vec<bool>>> {
+) -> Vec<(u32, bool)> {
     println!("Starting wwmap...");
 
     let ips_per_thread = (((to - from) as f32) / num_threads as f32) as u32;
     let ips_left = (to - from) - (num_threads * ips_per_thread); // how many ips we have left after the first threads
 
     // container for all of our threads
-    let mut threads: Vec<JoinHandle<Vec<bool>>> = Vec::new();
+    // let mut threads: Vec<JoinHandle<Vec<bool>>> = Vec::new();
+
+    // container for all of our results
+    let mut results: Vec<(u32, bool)> = Vec::new();
 
     for thread_id in 0..num_threads {
         let id_ignorelist = ignorelist.clone().unwrap_or_default();
 
         // Create a worker
-        let worker = create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist);
-        threads.push(worker);
+        let mut worker_results =
+            create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist).unwrap();
+
+        results.append(&mut worker_results);
     }
 
     // Clean up the rest
     if ips_left > 0 {
         let id_ignorelist = ignorelist.clone().unwrap_or_default();
-        threads.push(create_scan_worker(
-            threads.len() as u32 + 1,
+        let mut worker_results = create_scan_worker(
+            results.len() as u32 + 1,
             ips_per_thread,
             target_port,
             id_ignorelist,
-        ));
+        )
+        .unwrap();
+        results.append(&mut worker_results);
     }
 
-    threads
+    results
 }
 
 pub fn start_scan(