From 5e018801dba3e9948c016ca00935801b556286f4 Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Sun, 17 Sep 2023 19:52:54 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=BB=A3=E7=90=86=E8=B6=85?=
 =?UTF-8?q?=E6=97=B6=E6=97=B6=E9=97=B4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/ip_proxy/tcp_proxy.rs | 55 +++++++++++++++++++++++++----------
 1 file changed, 40 insertions(+), 15 deletions(-)

diff --git a/vnt/src/ip_proxy/tcp_proxy.rs b/vnt/src/ip_proxy/tcp_proxy.rs
index c291780..bead280 100644
--- a/vnt/src/ip_proxy/tcp_proxy.rs
+++ b/vnt/src/ip_proxy/tcp_proxy.rs
@@ -2,6 +2,7 @@ use dashmap::DashMap;
 use std::io;
 use std::net::{SocketAddr, SocketAddrV4};
 use std::sync::Arc;
+use std::time::Duration;
 
 use tokio::net::{TcpListener, TcpStream};
 
@@ -30,19 +31,36 @@ impl TcpProxy {
                         if let Some(entry) = tcp_proxy_map.get(&sender_addr) {
                             let dest_addr = *entry.value();
                             drop(entry);
-                            let peer_tcp_stream = match TcpStream::connect(dest_addr).await {
-                                Ok(peer_tcp_stream) => peer_tcp_stream,
-                                Err(e) => {
-                                    log::warn!(
-                                        "tcp代理异常:{:?},来源:{},目标：{}",
-                                        e,
-                                        sender_addr,
-                                        dest_addr
-                                    );
-                                    continue;
-                                }
-                            };
+
                             tokio::spawn(async move {
+                                let peer_tcp_stream = match tokio::time::timeout(
+                                    Duration::from_secs(5),
+                                    TcpStream::connect(dest_addr),
+                                )
+                                .await
+                                {
+                                    Ok(peer_tcp_stream) => match peer_tcp_stream {
+                                        Ok(peer_tcp_stream) => peer_tcp_stream,
+                                        Err(e) => {
+                                            log::warn!(
+                                                "tcp代理异常:{:?},来源:{},目标：{}",
+                                                e,
+                                                sender_addr,
+                                                dest_addr
+                                            );
+                                            return;
+                                        }
+                                    },
+                                    Err(e) => {
+                                        log::warn!(
+                                            "tcp代理异常:{:?},来源:{},目标：{}",
+                                            e,
+                                            sender_addr,
+                                            dest_addr
+                                        );
+                                        return;
+                                    }
+                                };
                                 if let Err(e) = proxy(tcp_stream, peer_tcp_stream).await {
                                     log::warn!("{}->{},{}", sender_addr, dest_addr, e);
                                 }
@@ -68,8 +86,15 @@ async fn proxy(mut client: TcpStream, mut server: TcpStream) -> io::Result<()> {
     let client_to_server = tokio::io::copy(&mut client_reader, &mut server_writer);
     let server_to_client = tokio::io::copy(&mut server_reader, &mut client_writer);
 
-    let (r1, r2) = tokio::join!(client_to_server, server_to_client);
-    r1?;
-    r2?;
+    let (r1, r2) = tokio::join!(
+        tokio::time::timeout(Duration::from_secs(10), client_to_server),
+        tokio::time::timeout(Duration::from_secs(10), server_to_client)
+    );
+    if let Ok(r) = r1 {
+        r?;
+    }
+    if let Ok(r) = r2 {
+        r?;
+    }
     Ok(())
 }