From 48b41c4b22997bad9dcfc73788b6b2d190d29352 Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Mon, 4 Mar 2024 21:11:41 +0800
Subject: [PATCH] =?UTF-8?q?[mio]=20=E4=BF=AE=E5=A4=8Dip=E4=BB=A3=E7=90=86?=
 =?UTF-8?q?=E7=9B=B8=E5=85=B3=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/icmp_proxy.rs |  7 +++----
 vnt/src/ip_proxy/tcp_proxy.rs  | 17 ++++++++---------
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/vnt/src/ip_proxy/icmp_proxy.rs b/vnt/src/ip_proxy/icmp_proxy.rs
index 3509966..7472927 100644
--- a/vnt/src/ip_proxy/icmp_proxy.rs
+++ b/vnt/src/ip_proxy/icmp_proxy.rs
@@ -85,14 +85,13 @@ fn icmp_proxy(
     poll.registry()
         .register(&mut icmp_socket, SERVER, Interest::READABLE)?;
     let mut events = Events::with_capacity(32);
-    let stop = Waker::new(poll.registry(), NOTIFY)?;
-    let worker = stop_manager.add_listener("icmp_proxy".into(), move || {
+    let stop = Arc::new(Waker::new(poll.registry(), NOTIFY)?);
+    let _stop = stop.clone();
+    let _worker = stop_manager.add_listener("icmp_proxy".into(), move || {
         if let Err(e) = stop.wake() {
             log::warn!("stop icmp_proxy:{:?}", e);
         }
     })?;
-    // linux上收不到通知，原因未知
-    drop(worker);
     let mut buf = [0u8; 65535 - 20 - 8];
     loop {
         poll.poll(&mut events, None)?;
diff --git a/vnt/src/ip_proxy/tcp_proxy.rs b/vnt/src/ip_proxy/tcp_proxy.rs
index db1b312..aedabde 100644
--- a/vnt/src/ip_proxy/tcp_proxy.rs
+++ b/vnt/src/ip_proxy/tcp_proxy.rs
@@ -101,14 +101,13 @@ fn tcp_proxy(
     let mut events = Events::with_capacity(32);
     let mut tcp_map: HashMap<usize, ProxyValue> = HashMap::with_capacity(16);
     let mut mapping: HashMap<usize, usize> = HashMap::with_capacity(16);
-    let stop = Waker::new(poll.registry(), NOTIFY)?;
-    let worker = stop_manager.add_listener("tcp_proxy".into(), move || {
+    let stop = Arc::new(Waker::new(poll.registry(), NOTIFY)?);
+    let _stop = stop.clone();
+    let _worker = stop_manager.add_listener("tcp_proxy".into(), move || {
         if let Err(e) = stop.wake() {
             log::warn!("stop tcp_proxy:{:?}", e);
         }
     })?;
-    // linux上收不到通知，原因未知
-    drop(worker);
     loop {
         poll.poll(&mut events, None)?;
         if stop_manager.is_stop() {
@@ -376,6 +375,11 @@ fn readable_handle(
     let mut buf = [0; BUF_LEN];
 
     loop {
+        if mid_buf.len() >= BUF_LEN {
+            // 达到上限不再继续读取
+            log::warn!("达到上限不再继续读取 {:?}->{:?}",stream1,stream2);
+            return Ok(());
+        }
         match stream1.read(&mut buf) {
             Ok(len) => {
                 if len == 0 {
@@ -410,11 +414,6 @@ fn readable_handle(
                 }
                 mid_buf.reserve(buf.len());
                 mid_buf.put_slice(buf);
-                if mid_buf.len() >= BUF_LEN {
-                    // 达到上限不再继续读取
-                    log::warn!("达到上限不再继续读取");
-                    return Ok(());
-                }
             }
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {