优化代码

switch/src/core/mod.rs

@@ -38,14 +38,10 @@ impl Switch {
         let virtual_gateway = Ipv4Addr::from(response.virtual_gateway);
         let virtual_netmask = Ipv4Addr::from(response.virtual_netmask);
         let current_device = Arc::new(AtomicCell::new(CurrentDeviceInfo::new(virtual_ip, virtual_gateway, virtual_netmask, config.server_address)));
-        let local_addr = channel.local_addr()?;
-        let local_ip = if local_addr.ip().is_unspecified() {
-            local_ip_address::local_ip().unwrap_or(local_addr.ip())
-        } else {
-            local_addr.ip()
-        };
+        let local_ip = crate::nat::local_ip()?;
+        let local_port = channel.local_addr()?.port();
         // NAT检测
-        let nat_test = NatTest::new(config.nat_test_server.clone(), Ipv4Addr::from(response.public_ip), response.public_port as u16, local_ip, local_addr.port());
+        let nat_test = NatTest::new(config.nat_test_server.clone(), Ipv4Addr::from(response.public_ip), response.public_port as u16, local_ip, local_port);
         // tun通道
         let (tun_writer, tun_reader) = tun_device::create_tun(virtual_ip, virtual_netmask, virtual_gateway)?;
 

switch/src/handle/heartbeat_handler.rs

@@ -59,7 +59,7 @@ fn start_heartbeat_(sender: Sender<Ipv4Addr>, device_list: Arc<Mutex<(u16, Vec<P
             let epoch = { device_list.lock().0 };
             ping.set_epoch(epoch);
         }
-        if count % 7 == 0 {
+        if count < 7 || count % 7 == 0 {
             let mut route_list: Option<Vec<(Ipv4Addr, Route)>> = None;
             let peer_list = device_list.lock().1.clone();
             for peer in peer_list {

switch/src/handle/recv_handler.rs

@@ -17,6 +17,7 @@ use packet::ip::ipv4::packet::IpV4Packet;
 use crate::error::Error;
 use crate::handle::{check_dest, ConnectStatus, CurrentDeviceInfo, PeerDeviceInfo};
 use crate::handle::registration_handler::Register;
+use crate::nat;
 use crate::nat::NatTest;
 use crate::proto::message::{DeviceList, PunchInfo, PunchNatType, RegistrationResponse};
 use crate::protocol::{control_packet, MAX_TTL, NetPacket, Protocol, service_packet, turn_packet, Version};
@@ -179,13 +180,9 @@ impl RecvHandler {
             service_packet::Protocol::RegistrationRequest => {}
             service_packet::Protocol::RegistrationResponse => {
                 let response = RegistrationResponse::parse_from_bytes(net_packet.payload())?;
-                let local_addr = self.channel.local_addr()?;
-                let local_ip = if local_addr.ip().is_unspecified() {
-                    local_ip_address::local_ip().unwrap_or(local_addr.ip())
-                } else {
-                    local_addr.ip()
-                };
-                let nat_info = self.nat_test.re_test(Ipv4Addr::from(response.public_ip), response.public_port as u16, local_ip, local_addr.port());
+                let local_port = self.channel.local_addr()?.port();
+                let local_ip = nat::local_ip()?;
+                let nat_info = self.nat_test.re_test(Ipv4Addr::from(response.public_ip), response.public_port as u16, local_ip, local_port);
                 self.channel.set_nat_type(nat_info.nat_type)?;
                 let new_ip = Ipv4Addr::from(response.virtual_ip);
                 let current_ip = current_device.virtual_ip();
@@ -296,7 +293,7 @@ impl RecvHandler {
                 }
             }
             ControlPacket::PunchRequest => {
-                log::info!("PunchRequest route_key:{:?}",route_key);
+                // log::info!("PunchRequest route_key:{:?}",route_key);
                 //回应
                 net_packet.set_transport_protocol(control_packet::Protocol::PunchResponse.into());
                 net_packet.set_source(current_device.virtual_ip());
@@ -307,7 +304,7 @@ impl RecvHandler {
                 self.channel.add_route(source, route);
             }
             ControlPacket::PunchResponse => {
-                log::info!("PunchResponse route_key:{:?}",route_key);
+                // log::info!("PunchResponse route_key:{:?}",route_key);
                 let route = Route::from(*route_key, 1, -1);
                 self.channel.add_route(net_packet.source(), route);
             }

switch/src/nat/mod.rs

@@ -1,3 +1,4 @@
+use std::io;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::sync::Arc;
 use parking_lot::Mutex;
@@ -6,6 +7,14 @@ use crate::proto::message::PunchNatType;
 
 pub mod check;
 
+use std::net::UdpSocket;
+
+pub fn local_ip() -> io::Result<IpAddr> {
+    let socket = UdpSocket::bind("0.0.0.0:0")?;
+    socket.connect("8.8.8.8:80")?;
+    let addr = socket.local_addr()?;
+    Ok(addr.ip())
+}
 
 #[derive(Clone)]
 pub struct NatTest {