增加代理超时时间

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(())
 }