From 9495e3970017e7f7adb50080508524d9e423064e Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Mon, 18 Sep 2023 21:51:08 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96nat=E6=A0=A1=E9=AA=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/channel/punch.rs       |  7 +++++--
 vnt/src/handle/recv_handler.rs | 30 ++++++------------------------
 vnt/src/nat/mod.rs             | 26 +++++++++++++++-----------
 3 files changed, 26 insertions(+), 37 deletions(-)

diff --git a/vnt/src/channel/punch.rs b/vnt/src/channel/punch.rs
index 0b520a8..8a4367c 100644
--- a/vnt/src/channel/punch.rs
+++ b/vnt/src/channel/punch.rs
@@ -50,9 +50,12 @@ impl NatInfo {
         public_port_range: u16,
         local_ipv4_addr: SocketAddrV4,
         ipv6_addr: SocketAddrV6,
-        nat_type: NatType,
+        mut nat_type: NatType,
     ) -> Self {
-        public_ips.retain(|ip| !ip.is_loopback() && !ip.is_private());
+        public_ips.retain(|ip| !ip.is_loopback() && !ip.is_private() && !ip.is_unspecified());
+        if public_ips.len() > 1 {
+            nat_type = NatType::Symmetric;
+        }
         Self {
             public_ips,
             public_port,
diff --git a/vnt/src/handle/recv_handler.rs b/vnt/src/handle/recv_handler.rs
index 4423df5..f133a2a 100644
--- a/vnt/src/handle/recv_handler.rs
+++ b/vnt/src/handle/recv_handler.rs
@@ -431,18 +431,9 @@ impl ChannelDataHandler {
                 }
                 std::net::IpAddr::V6(_) => {}
             },
-            ControlPacket::AddrResponse(addr_packet) => {
-                if !addr_packet.ipv4().is_multicast()
-                    && !addr_packet.ipv4().is_broadcast()
-                    && !addr_packet.ipv4().is_unspecified()
-                    && !addr_packet.ipv4().is_loopback()
-                    && !addr_packet.ipv4().is_private()
-                    && addr_packet.port() != 0
-                {
-                    self.nat_test
-                        .update_addr(addr_packet.ipv4(), addr_packet.port())
-                }
-            }
+            ControlPacket::AddrResponse(addr_packet) => self
+                .nat_test
+                .update_addr(addr_packet.ipv4(), addr_packet.port()),
         }
         Ok(())
     }
@@ -630,18 +621,9 @@ impl ChannelDataHandler {
                 )
                 .await?;
             }
-            ControlPacket::AddrResponse(addr_packet) => {
-                if addr_packet.port() != 0
-                    && !addr_packet.ipv4().is_multicast()
-                    && !addr_packet.ipv4().is_broadcast()
-                    && !addr_packet.ipv4().is_unspecified()
-                    && !addr_packet.ipv4().is_loopback()
-                    && !addr_packet.ipv4().is_private()
-                {
-                    self.nat_test
-                        .update_addr(addr_packet.ipv4(), addr_packet.port())
-                }
-            }
+            ControlPacket::AddrResponse(addr_packet) => self
+                .nat_test
+                .update_addr(addr_packet.ipv4(), addr_packet.port()),
             _ => {}
         }
         Ok(())
diff --git a/vnt/src/nat/mod.rs b/vnt/src/nat/mod.rs
index 1b1b90e..8218f15 100644
--- a/vnt/src/nat/mod.rs
+++ b/vnt/src/nat/mod.rs
@@ -99,10 +99,18 @@ impl NatTest {
         self.info.lock().clone()
     }
     pub fn update_addr(&self, ip: Ipv4Addr, port: u16) {
-        let mut guard = self.info.lock();
-        guard.public_port = port;
-        if !guard.public_ips.contains(&ip) {
-            guard.public_ips.push(ip);
+        if !ip.is_multicast()
+            && !ip.is_broadcast()
+            && !ip.is_unspecified()
+            && !ip.is_loopback()
+            && !ip.is_private()
+            && port != 0
+        {
+            let mut guard = self.info.lock();
+            guard.public_port = port;
+            if !guard.public_ips.contains(&ip) {
+                guard.public_ips.push(ip);
+            }
         }
     }
     pub async fn re_test(
@@ -131,13 +139,9 @@ impl NatTest {
         ipv6_addr: SocketAddrV6,
     ) -> NatInfo {
         return match stun_test::stun_test_nat(stun_server.clone()).await {
-            Ok((nat_type, ips, port_range)) => {
-                let mut public_ips = Vec::new();
-                public_ips.push(Ipv4Addr::from(public_ip));
-                for ip in ips {
-                    if ip != public_ip {
-                        public_ips.push(ip);
-                    }
+            Ok((nat_type, mut public_ips, port_range)) => {
+                if !public_ips.contains(&public_ip) {
+                    public_ips.push(public_ip)
                 }
                 NatInfo::new(
                     public_ips,