From 4ef71a5b4cfee5b7f092b6a561e803be36bc603e Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Wed, 9 Oct 2024 16:57:14 +0800
Subject: [PATCH] --max-sessions option

---
 src/args.rs |  5 +++++
 src/lib.rs  | 13 +++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/args.rs b/src/args.rs
index b991512..352cb57 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -107,6 +107,10 @@ pub struct Args {
     /// Exit immediately when fatal error occurs, useful for running as a service
     #[arg(long)]
     pub exit_on_fatal_error: bool,
+
+    /// Maximum number of sessions to be handled concurrently
+    #[arg(long, value_name = "number", default_value = "200")]
+    pub max_sessions: usize,
 }
 
 fn validate_tun(p: &str) -> Result<String> {
@@ -149,6 +153,7 @@ impl Default for Args {
             virtual_dns_pool: IpCidr::from_str("198.18.0.0/15").unwrap(),
             daemonize: false,
             exit_on_fatal_error: false,
+            max_sessions: 200,
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 4e500c4..554fe14 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -65,8 +65,6 @@ pub mod win_svc;
 
 const DNS_PORT: u16 = 53;
 
-const MAX_SESSIONS: u64 = 200;
-
 static TASK_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
 use std::sync::atomic::Ordering::Relaxed;
 
@@ -244,13 +242,15 @@ where
                 ip_stack_stream?
             }
         };
+        let max_sessions = args.max_sessions as u64;
         match ip_stack_stream {
             IpStackStream::Tcp(tcp) => {
-                if TASK_COUNT.load(Relaxed) > MAX_SESSIONS {
-                    log::warn!("Too many sessions that over {MAX_SESSIONS}, dropping new session");
+                if TASK_COUNT.load(Relaxed) > max_sessions {
                     if args.exit_on_fatal_error {
+                        log::info!("Too many sessions that over {max_sessions}, exiting...");
                         break;
                     }
+                    log::warn!("Too many sessions that over {max_sessions}, dropping new session");
                     continue;
                 }
                 log::trace!("Session count {}", TASK_COUNT.fetch_add(1, Relaxed) + 1);
@@ -272,11 +272,12 @@ where
                 });
             }
             IpStackStream::Udp(udp) => {
-                if TASK_COUNT.load(Relaxed) > MAX_SESSIONS {
-                    log::warn!("Too many sessions that over {MAX_SESSIONS}, dropping new session");
+                if TASK_COUNT.load(Relaxed) > max_sessions {
                     if args.exit_on_fatal_error {
+                        log::info!("Too many sessions that over {max_sessions}, exiting...");
                         break;
                     }
+                    log::warn!("Too many sessions that over {max_sessions}, dropping new session");
                     continue;
                 }
                 log::trace!("Session count {}", TASK_COUNT.fetch_add(1, Relaxed) + 1);