From 043ea80cbda9c7591ce7cfee599b6401a6f0b506 Mon Sep 17 00:00:00 2001
From: Zihaoxu2008 <zihaoxu2008@outlook.com>
Date: Sat, 25 Jan 2025 20:28:46 +0800
Subject: [PATCH] =?UTF-8?q?=E5=8E=BB=E4=BD=A0=E5=A6=88=E7=9A=84=E7=BA=BF?=
 =?UTF-8?q?=E7=A8=8B=E6=B1=A0=EF=BC=8C=E6=96=B0=E5=A2=9Euse=5Fthread=5Fpoo?=
 =?UTF-8?q?l=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config.json          |  1 +
 src/bot_socks5.rs    |  5 ++++-
 src/configuration.rs |  6 +++--
 src/main.rs          | 53 +++++++++++++++++++++++++++++---------------
 4 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/config.json b/config.json
index b571202..bfc5193 100644
--- a/config.json
+++ b/config.json
@@ -4,6 +4,7 @@
   "proxy_url": "",
   "username_prefix": "test",
   "worker_threads": 100,
+  "use_thread_pool": false,
   "count": 300,
   "reco_words": ["验证程序已启用","AntiAttack"],
   "login_commands": ["reg helloa helloa"],
diff --git a/src/bot_socks5.rs b/src/bot_socks5.rs
index 7512c54..aba4ccf 100644
--- a/src/bot_socks5.rs
+++ b/src/bot_socks5.rs
@@ -269,8 +269,9 @@ impl<T: StreamType> Bot<T> {
         let mut login_command_flag = false;
         let mut status = 0;
         let mut iter = 0;
+        let mut flag = true;
         println!("[{}/WORKER] Worker started!", self.username);
-        loop {
+        while flag {
             if !status_rx.is_empty() {
                 status = status_rx.recv().await.unwrap();
             }
@@ -320,7 +321,9 @@ impl<T: StreamType> Bot<T> {
                 iter += 1;
             }
             sleep(Duration::from_millis(1));
+            flag = self.conn.lock().await.is_alive();
         }
+        println!("[{}/WORKER] Worker stopped.", self.username);
     }
 
     // 运行bot
diff --git a/src/configuration.rs b/src/configuration.rs
index 1a76988..f2dc23a 100644
--- a/src/configuration.rs
+++ b/src/configuration.rs
@@ -14,7 +14,8 @@ pub struct Configuration {
     spam_file: String,
     spam_delay_ms: i32,
     random_prefix: bool,
-    pub worker_threads: i32
+    pub worker_threads: i32,
+    pub use_thread_pool: bool
 }
 
 impl Configuration {
@@ -57,7 +58,8 @@ impl Configuration {
             spam_file: "".to_string(),
             spam_delay_ms: 0,
             random_prefix: false,
-            worker_threads: 100
+            worker_threads: 100,
+            use_thread_pool: true
         }
     }
 
diff --git a/src/main.rs b/src/main.rs
index 77cb15a..e33c604 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,7 +9,7 @@ use reqwest;
 use serde_json::Value;
 // use tokio::task;
 use std::sync::Arc;
-// use std::thread;
+use std::thread;
 use tokio::sync::Mutex;
 use crate::utils::generate_username;
 use crate::configuration::Configuration;
@@ -31,7 +31,7 @@ async fn main() {
     // 一个互斥锁，防止有多个线程同时请求uuproxy
     let thread_lock = Arc::new(Mutex::new(()));
     
-    let mut handles = Vec::new();
+    // let mut handles = Vec::new();
     let configuration = Configuration::new("config.json");
     let configuration_rc = Arc::new(configuration);
     let mut spam_file = None;
@@ -60,19 +60,23 @@ async fn main() {
         // });
         // 生成一个线程
         // let handle = tokio::task::spawn(attack_thread(lock_clone,configuration_rc_clone,variables_clone));
-        let handle = rt.spawn(attack_thread(lock_clone,configuration_rc_clone,variables_clone));
-        // let handle = thread::spawn(move || {
-        //     tokio::runtime::Builder::new_current_thread()
-        //         .enable_all()
-        //         .thread_stack_size(8 * 1024 * 1024)
-        //         .build()
-        //         .unwrap()
-        //         .block_on(async move {
-        //             attack_thread(lock_clone,configuration_rc_clone,variables_clone).await;
-        //             // println!("[MAIN] Thread {} created", i);
-        //         });
-        //});
-        handles.push(handle);
+        // let handle;
+        if configuration_rc.use_thread_pool {
+            rt.spawn(attack_thread(lock_clone,configuration_rc_clone,variables_clone));
+        }else{
+            thread::spawn(move || {
+                tokio::runtime::Builder::new_current_thread()
+                    .enable_all()
+                    .thread_stack_size(8 * 1024 * 1024)
+                    .build()
+                    .unwrap()
+                    .block_on(async move {
+                        attack_thread(lock_clone,configuration_rc_clone,variables_clone).await;
+                        // println!("[MAIN] Thread {} created", i);
+                    });
+            });
+        }
+        // handles.push(handle);
         // println!("[MAIN] Thread {} created", i) // 这个太吵了（
     }
     
@@ -141,9 +145,22 @@ async fn attack_thread(lock: Arc<Mutex<()>>, config: Arc<Configuration>, variabl
                 // 插入生成update_worker线程的代码
                 let mut bot_clone = bot.clone();
                 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u8>();
-                tokio::spawn(async move {
-                    bot_clone.update_worker(rx).await;
-                });
+                if config.use_thread_pool {
+                    tokio::spawn(async move {
+                        bot_clone.update_worker(rx).await;
+                    });
+                }else {
+                    thread::spawn(|| {
+                        tokio::runtime::Builder::new_current_thread()
+                            .enable_all()
+                            .thread_stack_size(8 * 1024 * 1024)
+                            .build()
+                            .unwrap()
+                            .block_on(async move {
+                                bot_clone.update_worker(rx).await
+                            });
+                    });
+                }
                 match bot.run(tx).await {
                     Ok(_) => {
                         if bot.should_restart() {