Merge branch 'pr17' into merge-pr17

# Conflicts:
#	src/config.rs
#	src/listener.rs
#	src/main.rs

Author: therealaleph

README.md

@@ -44,6 +44,7 @@ You should get a Cloudflare IP (usually starts with `104.`, `172.67.`, `141.101.
 
 ```json
 {
+  "graceful_shutdown_sec": 0,
   "listeners": [
     {
       "listen": "0.0.0.0:40443",
@@ -66,6 +67,10 @@ Replace `CLOUDFLARE_IP` with the IP from step 1. The `fake_sni` can be any domai
 | `keepalive_time_sec` | Seconds of idle before TCP keepalive probes begin (default: `11`) |
 | `keepalive_interval_sec` | Seconds between individual TCP keepalive probes (default: `2`) |
 
+| Top-level Field | Description |
+|---|---|
+| `graceful_shutdown_sec` | Seconds to wait for active connections to finish after receiving a shutdown signal. `0` exits immediately (default: `0`) |
+
 Multiple listeners are supported -- each maps to one upstream.
 
 ### Step 3: Edit your v2ray/xray config
@@ -145,6 +150,7 @@ nslookup myserver.example.com
 
 ```json
 {
+  "graceful_shutdown_sec": 0,
   "listeners": [
     {
       "listen": "0.0.0.0:40443",
@@ -167,6 +173,10 @@ nslookup myserver.example.com
 | `keepalive_time_sec` | ثانیه‌های بی‌فعالیتی قبل از شروع پروب‌های TCP keepalive (پیش‌فرض: `11`) |
 | `keepalive_interval_sec` | فاصله زمانی بین پروب‌های TCP keepalive به ثانیه (پیش‌فرض: `2`) |
 
+| فیلد سطح بالا | توضیح |
+|---|---|
+| `graceful_shutdown_sec` | مدت انتظار (به ثانیه) برای اتمام اتصالات فعال پس از دریافت سیگنال خاموشی. مقدار `0` بلافاصله خارج می‌شود (پیش‌فرض: `0`) |
+
 ### مرحله ۳: تغییر کانفیگ v2ray/xray
 
 در کانفیگ VLESS/VMess خود این تغییرات را بدهید:

src/config.rs

@@ -22,12 +22,18 @@ fn default_buffer_size() -> usize {
     8
 }
 
+fn default_graceful_shutdown_sec() -> u64 {
+    0
+}
+
 #[derive(Debug, Deserialize)]
 pub struct Config {
     pub idle_timeout: Option<u64>,
     #[serde(default = "default_buffer_size")]
     pub buffer_size: usize,
     pub listeners: Vec<ListenerConfig>,
+    #[serde(default = "default_graceful_shutdown_sec")]
+    pub graceful_shutdown_sec: u64,
 }
 
 #[derive(Debug, Deserialize)]

src/listener.rs

@@ -2,7 +2,9 @@ use std::net::IpAddr;
 use std::time::Duration;
 
 use tokio::net::TcpListener;
+use tokio::task::JoinSet;
 use tracing::{error, info, warn};
+use tokio_util::sync::CancellationToken;
 
 use crate::config::ListenerConfig;
 use crate::handler;
@@ -26,6 +28,7 @@ pub async fn run_listener(
     cmd_tx: std::sync::mpsc::Sender<SnifferCommand>,
     idle_timeout: Option<u64>,
     buffer_size: usize,
+    token: CancellationToken,
 ) {
     let listener = match TcpListener::bind(lc.listen).await {
         Ok(l) => {
@@ -38,8 +41,15 @@ pub async fn run_listener(
         }
     };
 
+    let mut tasks = JoinSet::new();
+
     loop {
-        match listener.accept().await {
+        let accepted = tokio::select! {
+            result = listener.accept() => result,
+            _ = token.cancelled() => break,
+        };
+
+        match accepted {
             Ok((stream, peer)) => {
                 let upstream = lc.connect;
                 let sni = lc.fake_sni.clone();
@@ -49,7 +59,7 @@ pub async fn run_listener(
                 let handshake_timeout = lc.handshake_timeout_sec;
                 let keepalive_time = lc.keepalive_time_sec;
                 let keepalive_interval = lc.keepalive_interval_sec;
-                tokio::spawn(async move {
+                tasks.spawn(async move {
                     tracing::debug!(peer = %peer, "accepted connection");
                     handler::handle_connection(stream, upstream, sni, lip, tx, conn_timeout, handshake_timeout, keepalive_time, keepalive_interval, idle_timeout, buffer_size).await;
                 });
@@ -65,4 +75,8 @@ pub async fn run_listener(
             }
         }
     }
+
+    info!(listen = %lc.listen, "stopped accepting, draining active connections");
+    while tasks.join_next().await.is_some() {}
+    info!(listen = %lc.listen, "all connections drained");
 }

src/main.rs

@@ -13,9 +13,11 @@ mod sniffer;
 use std::net::IpAddr;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
+use std::time::Duration;
 
 use tracing::{error, info};
 use tracing_subscriber::EnvFilter;
+use tokio_util::sync::CancellationToken;
 
 fn main() {
     tracing_subscriber::fmt()
@@ -91,6 +93,7 @@ fn main() {
     let (cmd_tx, cmd_rx) = std::sync::mpsc::channel::<proto::SnifferCommand>();
 
     let stop = Arc::new(AtomicBool::new(false));
+    let token = CancellationToken::new();
 
     let sniffer_stop = stop.clone();
     let sniffer_local_ips = local_ips.clone();
@@ -103,24 +106,35 @@ fn main() {
         .expect("failed to spawn sniffer thread");
 
     let rt = tokio::runtime::Runtime::new().expect("failed to create tokio runtime");
+    let graceful_shutdown_sec = cfg.graceful_shutdown_sec;
     rt.block_on(async {
-        let signal_stop = stop.clone();
-        tokio::spawn(async move {
-            shutdown::wait_for_signal(signal_stop).await;
-            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-            std::process::exit(0);
-        });
-
         let mut handles = Vec::new();
         for lc in cfg.listeners {
             let tx = cmd_tx.clone();
             let lip = resolve_local_ip(lc.connect.ip()).unwrap_or(local_ips[0]);
-            handles.push(tokio::spawn(listener::run_listener(lc, lip, tx, cfg.idle_timeout, cfg.buffer_size)));
+            let tk = token.clone();
+            handles.push(tokio::spawn(listener::run_listener(lc, lip, tx, cfg.idle_timeout, cfg.buffer_size, tk)));
         }
 
-        for h in handles {
-            let _ = h.await;
+        shutdown::wait_for_signal(stop, token).await;
+
+        if graceful_shutdown_sec == 0 {
+            info!("graceful_shutdown_sec=0, exiting immediately");
+        } else {
+            info!("waiting up to {}s for active connections to drain", graceful_shutdown_sec);
+
+            let drain_all = async {
+                for h in handles {
+                    let _ = h.await;
+                }
+            };
+
+            if tokio::time::timeout(Duration::from_secs(graceful_shutdown_sec), drain_all).await.is_err() {
+                info!("drain timeout ({}s), forcing exit", graceful_shutdown_sec);
+            }
         }
+
+        info!("shutdown complete");
     });
 }
 

src/shutdown.rs

@@ -2,15 +2,20 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use tracing::info;
+use tokio_util::sync::CancellationToken;
 
-pub async fn wait_for_signal(stop: Arc<AtomicBool>) {
-    let ctrl_c = tokio::signal::ctrl_c();
+pub async fn wait_for_signal(stop: Arc<AtomicBool>, token: CancellationToken) {
+    wait_for_os_signal().await;
+    stop.store(true, Ordering::Relaxed);
+    token.cancel();
+}
 
+async fn wait_for_os_signal() {
+    let ctrl_c = tokio::signal::ctrl_c();
     #[cfg(unix)]
     {
         use tokio::signal::unix::{signal, SignalKind};
         let mut sigterm = signal(SignalKind::terminate()).expect("failed to register SIGTERM");
-
         tokio::select! {
             _ = ctrl_c => {
                 info!("received SIGINT, shutting down");
@@ -20,12 +25,9 @@ pub async fn wait_for_signal(stop: Arc<AtomicBool>) {
             }
         }
     }
-
     #[cfg(not(unix))]
     {
         ctrl_c.await.expect("failed to listen for Ctrl+C");
         info!("received Ctrl+C, shutting down");
     }
-
-    stop.store(true, Ordering::Relaxed);
 }