From 8d76214193eceb0147896e226c1d0909b2717f0e Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Mon, 20 Mar 2023 21:13:38 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=A0=E9=80=92=E6=9C=AC=E5=9C=B0ip=E3=80=81?=
 =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=BB=B6=E8=BF=9F=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 switch/src/handle/heartbeat_handler.rs | 13 +++++++++++--
 switch/src/handle/recv_handler.rs      | 22 ++++++++++++++++------
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/switch/src/handle/heartbeat_handler.rs b/switch/src/handle/heartbeat_handler.rs
index 5a4df8a..b65f2d9 100644
--- a/switch/src/handle/heartbeat_handler.rs
+++ b/switch/src/handle/heartbeat_handler.rs
@@ -42,6 +42,13 @@ pub fn start_heartbeat(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Ve
     });
 }
 
+fn set_now_time(packet: &mut NetPacket<[u8; 16]>) -> io::Result<()> {
+    let current_time = Local::now().timestamp_millis() as u16;
+    let mut ping = PingPacket::new(packet.payload_mut())?;
+    ping.set_time(current_time);
+    Ok(())
+}
+
 fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>, current_device: Arc<AtomicCell<CurrentDeviceInfo>>) -> io::Result<()> {
     let mut net_packet = NetPacket::new([0u8; 16])?;
     net_packet.set_version(Version::V1);
@@ -53,9 +60,7 @@ fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<P
         let current_device = current_device.load();
         net_packet.set_source(current_device.virtual_ip());
         {
-            let current_time = Local::now().timestamp_millis() as u16;
             let mut ping = PingPacket::new(net_packet.payload_mut())?;
-            ping.set_time(current_time);
             let epoch = { device_list.lock().0 };
             ping.set_epoch(epoch);
         }
@@ -63,6 +68,7 @@ fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<P
             let mut route_list: Option<Vec<(Ipv4Addr, Route)>> = None;
             let peer_list = device_list.lock().1.clone();
             for peer in peer_list {
+                set_now_time(&mut net_packet)?;
                 net_packet.first_set_ttl(MAX_TTL);
                 net_packet.set_destination(peer.virtual_ip);
                 if sender.send_to_id(net_packet.buffer(), &peer.virtual_ip).is_err() {
@@ -78,6 +84,7 @@ fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<P
                     net_packet.first_set_ttl(2);
                     for (peer_ip, route) in route_list.iter() {
                         if peer_ip != &peer.virtual_ip && route.metric == 1 {
+                            set_now_time(&mut net_packet)?;
                             let _ = sender.send_to_route(net_packet.buffer(), &route.route_key());
                             num += 1;
                         }
@@ -88,12 +95,14 @@ fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<P
                 }
                 thread::sleep(Duration::from_millis(1));
             }
+            set_now_time(&mut net_packet)?;
             net_packet.set_destination(current_device.virtual_gateway());
             if let Err(e) = sender.send_to_addr(net_packet.buffer(), current_device.connect_server) {
                 log::warn!("connect_server:{:?},e:{:?}",current_device.connect_server,e);
             }
         } else {
             for (peer_ip, route) in sender.route_table().iter() {
+                set_now_time(&mut net_packet)?;
                 net_packet.set_destination(*peer_ip);
                 if let Err(e) = sender.send_to_route(net_packet.buffer(), &route.route_key()) {
                     log::warn!("peer_ip:{:?},route:{:?},e:{:?}",peer_ip,route,e);
diff --git a/switch/src/handle/recv_handler.rs b/switch/src/handle/recv_handler.rs
index 9388ce8..2741750 100644
--- a/switch/src/handle/recv_handler.rs
+++ b/switch/src/handle/recv_handler.rs
@@ -254,11 +254,20 @@ impl RecvHandler {
     fn control(&self, current_device: CurrentDeviceInfo, source: Ipv4Addr, mut net_packet: NetPacket<&mut [u8]>, route_key: &RouteKey) -> crate::Result<()> {
         match ControlPacket::new(net_packet.transport_protocol(), net_packet.payload())? {
             ControlPacket::PingPacket(_) => {
+                let metric = net_packet.source_ttl() - net_packet.ttl() + 1;
                 net_packet.set_transport_protocol(control_packet::Protocol::Pong.into());
                 net_packet.set_source(current_device.virtual_ip());
                 net_packet.set_destination(source);
                 net_packet.first_set_ttl(MAX_TTL);
                 self.channel.send_to_route(net_packet.buffer(), route_key)?;
+                if metric == 1 {
+                    if let Some(current_route) = self.channel.route(&source) {
+                        if current_route.metric > 1 {
+                            let route = Route::from(*route_key, 1, -1);
+                            self.channel.add_route(source, route);
+                        }
+                    }
+                }
             }
             ControlPacket::PongPacket(pong_packet) => {
                 let current_time = Local::now().timestamp_millis() as u16;
@@ -330,18 +339,19 @@ impl RecvHandler {
                     let nat_info = self.nat_test.nat_info();
                     punch_reply.public_ip_list = nat_info.public_ips.iter().map(|i| {
                         match i {
-                            IpAddr::V4(ip) => {
-                                u32::from_be_bytes(ip.octets())
-                            }
-                            IpAddr::V6(_) => {
-                                panic!()
-                            }
+                            IpAddr::V4(ip) => u32::from_be_bytes(ip.octets()),
+                            IpAddr::V6(_) => 0
                         }
                     }).collect();
                     punch_reply.public_port = nat_info.public_port as u32;
                     punch_reply.public_port_range = nat_info.public_port_range as u32;
                     punch_reply.nat_type =
                         protobuf::EnumOrUnknown::new(PunchNatType::from(nat_info.nat_type));
+                    punch_reply.local_ip = match nat_info.local_ip {
+                        IpAddr::V4(ip) => u32::from_be_bytes(ip.octets()),
+                        IpAddr::V6(_) => 0
+                    };
+                    punch_reply.local_port = nat_info.local_port as u32;
                     let bytes = punch_reply.write_to_bytes()?;
                     let mut net_packet =
                         NetPacket::new(vec![0u8; 12 + bytes.len()])?;