From 57ed454c938d5e457037be9a86deab2afb44c21a Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Wed, 20 Sep 2023 11:11:31 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=86=85=E7=BD=91ip=E6=96=AD?=
 =?UTF-8?q?=E7=BA=BF=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/ip_proxy/tcp_proxy.rs | 37 ++++++++++++++++++++++++++---------
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/vnt/src/ip_proxy/tcp_proxy.rs b/vnt/src/ip_proxy/tcp_proxy.rs
index bad11f7..f9dc750 100644
--- a/vnt/src/ip_proxy/tcp_proxy.rs
+++ b/vnt/src/ip_proxy/tcp_proxy.rs
@@ -3,8 +3,10 @@ use std::io;
 use std::net::{SocketAddr, SocketAddrV4};
 use std::sync::Arc;
 use std::time::Duration;
-
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
 use tokio::net::{TcpListener, TcpStream};
+use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 
 pub struct TcpProxy {
     tcp_listener: TcpListener,
@@ -79,15 +81,32 @@ impl TcpProxy {
     }
 }
 
-async fn proxy(mut client: TcpStream, mut server: TcpStream) -> io::Result<()> {
-    let (mut client_reader, mut client_writer) = client.split();
-    let (mut server_reader, mut server_writer) = server.split();
+async fn proxy(client: TcpStream, server: TcpStream) -> io::Result<()> {
+    let (client_read, client_write) = client.into_split();
+    let (server_read, server_write) = server.into_split();
+    tokio::spawn(async move {
+        if let Err(e) = copy(client_read, server_write).await {
+            log::warn!("{:?}",e);
+        }
+    });
+    copy(server_read, client_write).await
+}
 
-    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);
-    tokio::select! {
-        _ = tokio::time::timeout(Duration::from_secs(10), client_to_server) =>{},
-        _ = tokio::time::timeout(Duration::from_secs(10), server_to_client) =>{},
+async fn copy(mut read: OwnedReadHalf, mut write: OwnedWriteHalf) -> io::Result<()> {
+    let mut buf = [0; 10240];
+    loop {
+        tokio::select! {
+            result = read.read(&mut buf) =>{
+                let len = result?;
+                if len==0{
+                    break;
+                }
+                write.write_all(&buf[..len]).await?;
+            }
+            _ = tokio::time::sleep(Duration::from_secs(300)) =>{
+                break;
+            }
+        }
     }
     Ok(())
 }