feat(scanner): add cancel token from task manager

Author: Cyberdolfi

src/scanning/crawler.rs

@@ -12,18 +12,32 @@ use crate::scanning::scanner::{scan, ScanConfig};
 use crate::scanning::utils::{format_time, prettier_ping_result, save_server};
 
 pub async fn crawl(gen_config: IpGenerator) {
-    let _ = TaskManager::spawn("Crawler", move |_cancel_token| async move {
+    let _ = TaskManager::spawn("Crawler", move |cancel_token| async move {
         loop {
+            if cancel_token.is_cancelled() {
+                logger::warning("Crawler shutting down...".to_string())
+                    .prefix("Crawler").send().await;
+                break;
+            }
+
             let start_time = Instant::now();
 
             let ports = vec![25565]; // Example for later implementation
             let targets: Vec<(Ipv4Addr, u16)> = gen_config.generate()
+                .take_while(|_| {
+                    let cancelled = cancel_token.is_cancelled();
+                    async move { !cancelled }
+                })
                 .flat_map(|ip| {
                     stream::iter(ports.clone().into_iter().map(move |port| (ip, port)))
                 })
                 .collect()
                 .await;
 
+            if cancel_token.is_cancelled() && targets.is_empty() {
+                break;
+            }
+
             let total_targets = targets.len();
             let mut found_batch: Vec<(ServerInfo, ServerHistory)> = Vec::new();
 
@@ -54,6 +68,12 @@ pub async fn crawl(gen_config: IpGenerator) {
             while let Some(maybe_result) = scan_stream.next().await {
                 processed_count += 1;
 
+                if cancel_token.is_cancelled() {
+                    logger::warning("Scan interrupted. Saving results...".to_string())
+                        .prefix("Crawler").send().await;
+                    break;
+                }
+
                 // Success
                 if let Some(result) = maybe_result {
                     let parsed = parse_server(result.ip, result.port, result.ping.clone(), result.query, result.join);
@@ -70,12 +90,11 @@ pub async fn crawl(gen_config: IpGenerator) {
                     );
                     output.push_str(&prettier_ping_result(result.ping).await);
                     logger::success(output).prefix("Crawler").send().await;
-                }
 
-                // Database insert
-                if processed_count % 10000 == 0 {
-                    let batch_to_insert = std::mem::take(&mut found_batch);
-                    save_server(&batch_to_insert).await;
+                    if found_batch.len() >= 30 {
+                        let batch_to_insert = std::mem::take(&mut found_batch);
+                        save_server(&batch_to_insert).await;
+                    }
                 }
 
                 // Progress calc
@@ -103,6 +122,10 @@ pub async fn crawl(gen_config: IpGenerator) {
                 save_server(&found_batch).await;
             }
 
+            if cancel_token.is_cancelled() {
+                break;
+            }
+
             // Finished
             let elapsed_time = start_time.elapsed();
 

src/scanning/file_scanner.rs

@@ -15,7 +15,7 @@ use crate::scanning::utils::{count_lines_fast, format_time, prettier_ping_result
 
 // TODO: Make a live stream from the file to scanner, because its more RAM efficient
 pub async fn scan_file(path: String) {
-    let _ = TaskManager::spawn("File Scanner", move |_cancel_token| async move {
+    let _ = TaskManager::spawn("File Scanner", move |cancel_token| async move {
         let file = match File::open(&path).await {
             Ok(f) => f,
             Err(e) => {
@@ -43,6 +43,12 @@ pub async fn scan_file(path: String) {
 
         // Read file
         while let Ok(Some(line)) = lines.next_line().await {
+            if cancel_token.is_cancelled() {
+                logger::warning("File reading interrupted.".to_string())
+                    .prefix("File Scanner").send().await;
+                return;
+            }
+
             let line = line.trim();
             // Skip comments
             if line.is_empty() || line.starts_with('#') { continue; }
@@ -86,6 +92,12 @@ pub async fn scan_file(path: String) {
 
         // Scan stream
         while let Some(maybe_result) = scan_stream.next().await {
+            if cancel_token.is_cancelled() {
+                logger::warning("Scan interrupted. Saving results...".to_string())
+                    .prefix("File Scanner").send().await;
+                break;
+            }
+
             processed_count += 1;
 
             // Success
@@ -105,7 +117,7 @@ pub async fn scan_file(path: String) {
                 output.push_str(&prettier_ping_result(result.ping).await);
                 logger::success(output).prefix("File Scanner").send().await;
 
-                if found_batch.len() >= 50 {
+                if found_batch.len() >= 30 {
                     let batch_to_insert = std::mem::take(&mut found_batch);
                     save_server(&batch_to_insert).await;
                 }