From 63b77aa6171b3bfde35c45352575e219a07ad469 Mon Sep 17 00:00:00 2001
From: lanthora <lanthora@outlook.com>
Date: Fri, 28 Jun 2024 14:33:08 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E8=B7=AF=E7=94=B1=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=E5=AF=BC=E8=87=B4=E7=9A=84=E5=B1=80=E5=9F=9F=E7=BD=91?=
 =?UTF-8?q?=E5=AF=B9=E7=AD=89=E8=BF=9E=E6=8E=A5=E5=A4=B1=E6=95=88=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CMakeLists.txt     |  2 +-
 src/core/client.cc | 19 +++++++++++++++----
 src/core/client.h  |  2 ++
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 54e0560..a4e32fb 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.18.4)
 
-project(candy LANGUAGES C CXX VERSION 5.9.1)
+project(candy LANGUAGES C CXX VERSION 5.9.2)
 
 option(CANDY_NOEXE  "Don't build executable")
 option(CANDY_DEVEL  "Build development library")
diff --git a/src/core/client.cc b/src/core/client.cc
index 839781e..98283b6 100644
--- a/src/core/client.cc
+++ b/src/core/client.cc
@@ -124,6 +124,7 @@ int Client::setLocalhost(std::string ip) {
 int Client::run() {
     std::lock_guard lock(this->runningMutex);
     this->running = true;
+    this->localP2PDisabled = false;
 
     if (startWsThread()) {
         spdlog::critical("start websocket client thread failed");
@@ -432,6 +433,10 @@ void Client::sendDiscoveryMessage(uint32_t dst) {
 }
 
 void Client::sendLocalPeerConnMessage(const PeerInfo &peer) {
+    if (this->localP2PDisabled) {
+        return;
+    }
+
     LocalPeerConnMessage header;
     header.ge.subtype = GeSubType::LOCAL_PEER_CONN;
     header.ge.extra = 0;
@@ -572,6 +577,7 @@ void Client::handleSysRtMessage(WebSocketMessage &message) {
         spdlog::warn("invalid system route message: {:n}", spdlog::to_hex(message.buffer));
         return;
     }
+    this->localP2PDisabled = true;
     SysRouteMessage *header = (SysRouteMessage *)message.buffer.c_str();
     SysRouteItem *rt = header->rtTable;
     std::unique_lock lock(this->sysRtTableMutex);
@@ -610,6 +616,9 @@ void Client::handleGeneralMessage(WebSocketMessage &message) {
 }
 
 void Client::handleLocalPeerConnMessage(WebSocketMessage &message) {
+    if (this->localP2PDisabled) {
+        return;
+    }
     if (message.buffer.size() < sizeof(LocalPeerConnMessage)) {
         spdlog::warn("invalid local peer conn message: {:n}", spdlog::to_hex(message.buffer));
         return;
@@ -1067,7 +1076,7 @@ int Client::sendHeartbeatMessage(const PeerInfo &peer) {
 
     if ((peer.getState() == PeerState::PREPARING || peer.getState() == PeerState::SYNCHRONIZING ||
          peer.getState() == PeerState::CONNECTING) &&
-        (peer.local.ip && peer.local.port)) {
+        (!this->localP2PDisabled && peer.local.ip && peer.local.port)) {
         heartbeat.ip = Address::hostToNet(this->selfInfo.local.ip);
         heartbeat.port = Address::hostToNet(this->selfInfo.local.port);
         message.ip = peer.local.ip;
@@ -1211,12 +1220,14 @@ int Client::handleHeartbeatMessage(const UdpMessage &message) {
         return -1;
     }
 
-    if (isLocalIp(message.ip)) {
+    if (!isLocalIp(message.ip)) {
+        peer.wide.ip = message.ip;
+        peer.wide.port = message.port;
+    } else if (!this->localP2PDisabled) {
         peer.local.ip = message.ip;
         peer.local.port = message.port;
     } else {
-        peer.wide.ip = message.ip;
-        peer.wide.port = message.port;
+        return 0;
     }
 
     if (isLocalIp(message.ip) || !isLocalIp(peer.real.ip)) {
diff --git a/src/core/client.h b/src/core/client.h
index 46ac928..157698d 100644
--- a/src/core/client.h
+++ b/src/core/client.h
@@ -7,6 +7,7 @@
 #include "tun/tun.h"
 #include "utility/random.h"
 #include "websocket/client.h"
+#include <atomic>
 #include <functional>
 #include <list>
 #include <map>
@@ -153,6 +154,7 @@ class Client {
     uint64_t tickTick = randomUint32();
     uint32_t discoveryInterval;
     std::mutex cryptMutex;
+    std::atomic<bool> localP2PDisabled;
 
     // Route
     void showCandyRtChange(const CandyRouteEntry &entry);