From 94d6caef7ec6fb377c8b3f9f3bf67711fd9f0cf7 Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Tue, 26 Dec 2023 22:27:16 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E8=BF=9E=E6=8E=A5tcp?=
 =?UTF-8?q?=E5=9C=B0=E5=9D=80=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vnt/src/core/mod.rs                    | 15 ++-------------
 vnt/src/handle/handshake_handler.rs    | 12 +++++++++++-
 vnt/src/handle/heartbeat_handler.rs    | 12 ++----------
 vnt/src/handle/registration_handler.rs | 11 ++++++++++-
 4 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs
index c3cb8e6..78abfe5 100644
--- a/vnt/src/core/mod.rs
+++ b/vnt/src/core/mod.rs
@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 use std::io;
+use std::net::TcpStream;
 use std::net::UdpSocket;
 use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
-use std::net::{SocketAddrV6, TcpStream};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -581,7 +581,7 @@ impl Config {
         token: String,
         device_id: String,
         name: String,
-        mut server_address: SocketAddr,
+        server_address: SocketAddr,
         server_address_str: String,
         mut stun_server: Vec<String>,
         in_ips: Vec<(u32, u32, Ipv4Addr)>,
@@ -606,17 +606,6 @@ impl Config {
                 x.push_str(":3478");
             }
         }
-        match server_address {
-            SocketAddr::V4(ipv4) => {
-                server_address = SocketAddr::V6(SocketAddrV6::new(
-                    ipv4.ip().to_ipv6_mapped(),
-                    ipv4.port(),
-                    0,
-                    0,
-                ))
-            }
-            SocketAddr::V6(_) => {}
-        }
         Self {
             tap,
             token,
diff --git a/vnt/src/handle/handshake_handler.rs b/vnt/src/handle/handshake_handler.rs
index 420a0da..8b68eb9 100644
--- a/vnt/src/handle/handshake_handler.rs
+++ b/vnt/src/handle/handshake_handler.rs
@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::net::{SocketAddr, SocketAddrV6};
 
 use crate::channel::channel::Context;
 use crate::channel::RouteKey;
@@ -168,6 +168,16 @@ fn send_recv(
         }
         Ok(len)
     } else {
+        let server_address = if let SocketAddr::V4(ipv4) = server_address {
+            SocketAddr::V6(SocketAddrV6::new(
+                ipv4.ip().to_ipv6_mapped(),
+                ipv4.port(),
+                0,
+                0,
+            ))
+        } else {
+            server_address
+        };
         if let Err(e) = main_channel.send_to(send_buf, server_address) {
             return Err(HandshakeEnum::Other(format!("send error:{}", e)));
         }
diff --git a/vnt/src/handle/heartbeat_handler.rs b/vnt/src/handle/heartbeat_handler.rs
index 97ebf69..bbc9386 100644
--- a/vnt/src/handle/heartbeat_handler.rs
+++ b/vnt/src/handle/heartbeat_handler.rs
@@ -1,5 +1,5 @@
 use std::io;
-use std::net::{Ipv4Addr, SocketAddr, SocketAddrV6, ToSocketAddrs};
+use std::net::{Ipv4Addr, ToSocketAddrs};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -138,17 +138,9 @@ async fn start_heartbeat_main_(
         let src = current_dev.virtual_ip();
         if count % 40 == 19 {
             if let Ok(mut addr) = server_address_str.to_socket_addrs() {
-                if let Some(mut addr) = addr.next() {
+                if let Some(addr) = addr.next() {
                     if addr != current_dev.connect_server {
                         let mut tmp = current_dev.clone();
-                        if let SocketAddr::V4(ipv4) = addr {
-                            addr = SocketAddr::V6(SocketAddrV6::new(
-                                ipv4.ip().to_ipv6_mapped(),
-                                ipv4.port(),
-                                0,
-                                0,
-                            ))
-                        }
                         tmp.connect_server = addr;
                         log::info!(
                             "服务端地址变化,旧地址:{}，新地址:{}",
diff --git a/vnt/src/handle/registration_handler.rs b/vnt/src/handle/registration_handler.rs
index 2970d95..a1c6f45 100644
--- a/vnt/src/handle/registration_handler.rs
+++ b/vnt/src/handle/registration_handler.rs
@@ -1,6 +1,6 @@
 use crossbeam_utils::atomic::AtomicCell;
 use std::io::{Read, Write};
-use std::net::{Ipv4Addr, SocketAddr};
+use std::net::{Ipv4Addr, SocketAddr, SocketAddrV6};
 use std::time::{Duration, Instant};
 
 use crate::channel::sender::ChannelSender;
@@ -82,6 +82,15 @@ pub fn registration(
         }
         &mut recv_buf[4..len]
     } else {
+        let server_address = match server_address {
+            SocketAddr::V4(ipv4) => SocketAddr::V6(SocketAddrV6::new(
+                ipv4.ip().to_ipv6_mapped(),
+                ipv4.port(),
+                0,
+                0,
+            )),
+            SocketAddr::V6(_) => server_address,
+        };
         if let Err(e) = main_channel.send_to(buf, server_address) {
             return Err(ReqEnum::Other(format!("send error:{}", e)));
         }