From f75691b4835a3cd1894ea10eafcf63647a5f7131 Mon Sep 17 00:00:00 2001
From: lbl8603 <49143209+lbl8603@users.noreply.github.com>
Date: Wed, 5 Jun 2024 22:38:51 +0800
Subject: [PATCH] =?UTF-8?q?=E8=BF=87=E6=BB=A4=E9=9D=9Evnt=E5=8C=85?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/handle/recv_data/mod.rs |  7 +++-
 vnt/src/nat/mod.rs              | 67 ++++++++++++++++++++-------------
 vnt/src/nat/stun.rs             |  3 --
 3 files changed, 45 insertions(+), 32 deletions(-)

diff --git a/vnt/src/handle/recv_data/mod.rs b/vnt/src/handle/recv_data/mod.rs
index bb14c20..88362a0 100644
--- a/vnt/src/handle/recv_data/mod.rs
+++ b/vnt/src/handle/recv_data/mod.rs
@@ -24,7 +24,7 @@ use crate::handle::{BaseConfigInfo, CurrentDeviceInfo, PeerDeviceInfo, SELF_IP};
 #[cfg(feature = "ip_proxy")]
 use crate::ip_proxy::IpProxyMap;
 use crate::nat::NatTest;
-use crate::protocol::NetPacket;
+use crate::protocol::{NetPacket, HEAD_LEN};
 use crate::tun_tap_device::tun_create_helper::DeviceAdapter;
 use crate::util::U64Adder;
 
@@ -50,6 +50,9 @@ impl<Call: VntCallback> RecvChannelHandler for RecvDataHandler<Call> {
         route_key: RouteKey,
         context: &ChannelContext,
     ) {
+        if buf.len() < HEAD_LEN {
+            return;
+        }
         //判断stun响应包
         if !route_key.is_tcp() {
             if let Ok(rs) = self
@@ -136,7 +139,7 @@ impl<Call: VntCallback> RecvDataHandler<Call> {
         let net_packet = NetPacket::new(buf)?;
         let extend = NetPacket::unchecked(extend);
         if net_packet.ttl() == 0 || net_packet.source_ttl() < net_packet.ttl() {
-            log::warn!("丢弃过时包:{:?}", net_packet.head());
+            log::warn!("丢弃过时包:{:?} {}", net_packet.head(), route_key.addr);
             return Ok(());
         }
         let current_device = self.current_device.load();
diff --git a/vnt/src/nat/mod.rs b/vnt/src/nat/mod.rs
index b442c30..6056716 100644
--- a/vnt/src/nat/mod.rs
+++ b/vnt/src/nat/mod.rs
@@ -224,34 +224,48 @@ impl NatTest {
         source_addr: SocketAddr,
         buf: &[u8],
     ) -> anyhow::Result<bool> {
-        if let Some(addr) = stun::recv_stun_response(buf) {
-            if let SocketAddr::V4(addr) = addr {
-                let mut check_fail = true;
-                let source_ip = match source_addr.ip() {
-                    IpAddr::V4(ip) => ip,
-                    IpAddr::V6(ip) => {
-                        if let Some(ip) = ip.to_ipv4_mapped() {
-                            ip
-                        } else {
-                            return Ok(false);
-                        }
-                    }
-                };
-                'a: for stun_server in &self.stun_server {
-                    for x in stun_server.to_socket_addrs()? {
-                        if source_addr.port() == x.port() {
-                            if let IpAddr::V4(ip) = x.ip() {
-                                if ip == source_ip {
-                                    check_fail = false;
-                                    break 'a;
-                                }
-                            };
-                        }
+        if buf[0] == 0x01 && buf[1] == 0x01 {
+            if let Some(addr) = stun::recv_stun_response(buf) {
+                if let Err(e) = self.recv_data_(index, source_addr, addr) {
+                    log::warn!("{:?}", e);
+                }
+            }
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+    fn recv_data_(
+        &self,
+        index: usize,
+        source_addr: SocketAddr,
+        addr: SocketAddr,
+    ) -> anyhow::Result<()> {
+        if let SocketAddr::V4(addr) = addr {
+            let mut check_fail = true;
+            let source_ip = match source_addr.ip() {
+                IpAddr::V4(ip) => ip,
+                IpAddr::V6(ip) => {
+                    if let Some(ip) = ip.to_ipv4_mapped() {
+                        ip
+                    } else {
+                        return Ok(());
                     }
                 }
-                if check_fail {
-                    return Ok(false);
+            };
+            'a: for stun_server in &self.stun_server {
+                for x in stun_server.to_socket_addrs()? {
+                    if source_addr.port() == x.port() {
+                        if let IpAddr::V4(ip) = x.ip() {
+                            if ip == source_ip {
+                                check_fail = false;
+                                break 'a;
+                            }
+                        };
+                    }
                 }
+            }
+            if !check_fail {
                 let ip = addr.ip();
                 if !ip.is_multicast()
                     && !ip.is_broadcast()
@@ -260,10 +274,9 @@ impl NatTest {
                     && !ip.is_private()
                 {
                     self.update_addr(index, *addr.ip(), addr.port());
-                    return Ok(true);
                 }
             }
         }
-        return Ok(false);
+        Ok(())
     }
 }
diff --git a/vnt/src/nat/stun.rs b/vnt/src/nat/stun.rs
index 05b9a58..c1af7ff 100644
--- a/vnt/src/nat/stun.rs
+++ b/vnt/src/nat/stun.rs
@@ -198,9 +198,6 @@ pub fn send_stun_request() -> Vec<u8> {
 }
 
 pub fn recv_stun_response(buf: &[u8]) -> Option<SocketAddr> {
-    if buf[0] != 0x01 && buf[1] != 0x01 {
-        return None;
-    }
     let msg = stun_format::Msg::from(buf);
     if let Some(tid) = msg.tid() {
         if tid & TAG != TAG {