From 45468be3bff5bd6ff72e9f50f6239d6abe1e1805 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 18 Jun 2022 13:10:50 +0800
Subject: [PATCH] feat: add shutdown timeout

---
 src/server.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/server.rs b/src/server.rs
index 33746b5..55cec0b 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -662,6 +662,7 @@ async fn transfer<I, O>(inbound: I, outbound: O) -> Result<()>
     let (mut ro, mut wo) = tokio::io::split(outbound);
     // IO copy timeout 6 HOURS
     let tcp_io_copy_timeout = Duration::from_secs(6 * 3600);
+    let shutdown_tcp_timeout = Duration::from_secs(60);
     let client_to_server = async move {
         // let copy_result = time::timeout(tcp_io_copy_timeout, tokio::io::copy(&mut ri, &mut wo));
         let copy_result = time::timeout(
@@ -669,7 +670,7 @@ async fn transfer<I, O>(inbound: I, outbound: O) -> Result<()>
             copy_data(&mut ri, &mut wo),
         );
         let r = copy_result.await;
-        wo.shutdown().await.ok();
+        time::timeout(shutdown_tcp_timeout, wo.shutdown()).await.ok();
         match r {
             Err(e) => {
                 error!("TCP copy timeout: {}", e);
@@ -685,7 +686,7 @@ async fn transfer<I, O>(inbound: I, outbound: O) -> Result<()>
             copy_data(&mut ro, &mut wi),
         );
         let r = copy_result.await;
-        wi.shutdown().await.ok();
+        time::timeout(shutdown_tcp_timeout, wi.shutdown()).await.ok();
         match r {
             Err(e) => {
                 error!("TCP copy timeout: {}", e);