From 1f80e06d9dfd5aa30fea0940c948a078ccdfa4fd Mon Sep 17 00:00:00 2001
From: lbl8603 <49143209+lbl8603@users.noreply.github.com>
Date: Sat, 29 Jun 2024 17:28:30 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E7=AB=AF=E6=B5=81=E9=87=8F?=
 =?UTF-8?q?=E8=BF=87=E6=BB=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/handle/mod.rs              | 12 +++++++++++-
 vnt/src/handle/recv_data/server.rs | 11 ++++++++++-
 vnt/src/handle/recv_data/turn.rs   |  4 ++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/vnt/src/handle/mod.rs b/vnt/src/handle/mod.rs
index 5ba0312..e39305f 100644
--- a/vnt/src/handle/mod.rs
+++ b/vnt/src/handle/mod.rs
@@ -1,5 +1,5 @@
 use crossbeam_utils::atomic::AtomicCell;
-use std::net::{Ipv4Addr, SocketAddr};
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 pub mod callback;
 mod extension;
@@ -237,6 +237,16 @@ impl CurrentDeviceInfo {
     pub fn not_in_network(&self, ip: Ipv4Addr) -> bool {
         u32::from(ip) & u32::from(self.virtual_netmask) != u32::from(self.virtual_network)
     }
+    pub fn is_server_addr(&self, addr: SocketAddr) -> bool {
+        if self.connect_server == addr {
+            return true;
+        }
+        let f = |ip: IpAddr| match ip {
+            IpAddr::V4(v4) => Some(v4),
+            IpAddr::V6(v6) => v6.to_ipv4(),
+        };
+        addr.port() == self.connect_server.port() && f(addr.ip()) == f(self.connect_server.ip())
+    }
 }
 pub fn change_status(
     current_device: &AtomicCell<CurrentDeviceInfo>,
diff --git a/vnt/src/handle/recv_data/server.rs b/vnt/src/handle/recv_data/server.rs
index e746bb5..c3a018e 100644
--- a/vnt/src/handle/recv_data/server.rs
+++ b/vnt/src/handle/recv_data/server.rs
@@ -100,6 +100,15 @@ impl<Call: VntCallback, Device: DeviceWrite> PacketHandler for ServerPacketHandl
         context: &ChannelContext,
         current_device: &CurrentDeviceInfo,
     ) -> anyhow::Result<()> {
+        if !current_device.is_server_addr(route_key.addr) {
+            //拦截不是服务端的流量
+            log::info!(
+                "route_key={:?},不是来源于服务端地址{}",
+                route_key,
+                current_device.connect_server
+            );
+            return Ok(());
+        }
         context
             .route_table
             .update_read_time(&net_packet.source(), &route_key);
@@ -548,7 +557,7 @@ impl<Call: VntCallback, Device: DeviceWrite> ServerPacketHandler<Call, Device> {
                     //纪元不一致，可能有新客户端连接，向服务端拉取客户端列表
                     let mut poll_device = NetPacket::new_encrypt([0; 12 + ENCRYPTION_RESERVED])?;
                     poll_device.set_source(current_device.virtual_ip);
-                    poll_device.set_destination(GATEWAY_IP);
+                    poll_device.set_destination(current_device.virtual_gateway);
                     poll_device.set_default_version();
                     poll_device.set_gateway_flag(true);
                     poll_device.first_set_ttl(MAX_TTL);
diff --git a/vnt/src/handle/recv_data/turn.rs b/vnt/src/handle/recv_data/turn.rs
index f4d70ff..c356e68 100644
--- a/vnt/src/handle/recv_data/turn.rs
+++ b/vnt/src/handle/recv_data/turn.rs
@@ -27,6 +27,10 @@ impl PacketHandler for TurnPacketHandler {
         // ttl减一
         let ttl = net_packet.incr_ttl();
         if ttl > 0 {
+            if net_packet.is_gateway() {
+                // 暂时不转发服务端包
+                return Ok(());
+            }
             let destination = net_packet.destination();
             if let Some(route) = context.route_table.route_one(&destination) {
                 if route.addr == route_key.addr {