收到带确认的心跳再进入 CONNECTED 状态

lanthora · Sep 8, 2023 · 26720be · 26720be
1 parent 13c1458
commit 26720be
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 deletions.
diff --git a/src/core/client.cc b/src/core/client.cc
@@ -760,8 +760,13 @@ int Client::handleHeartbeatMessage(const UdpMessage &message) {
         spdlog::debug("peer port does not match, update: old {:08x} new {:08x}", peer.ip, message.ip);
         peer.port = message.port;
     }
-    peer.count = 0;
-    peer.updateState(PeerState::CONNECTED);
+    if (!peer.ack) {
+        peer.ack = 1;
+    }
+    if (heartbeat->ack) {
+        peer.count = 0;
+        peer.updateState(PeerState::CONNECTED);
+    }
     return 0;
 }
 
@@ -771,6 +776,7 @@ int Client::sendHeartbeat(const PeerInfo &peer) {
     heartbeat.tun = Address::hostToNet(this->tun.getIP());
     heartbeat.ip = Address::hostToNet(this->selfInfo.ip);
     heartbeat.port = Address::hostToNet(this->selfInfo.port);
+    heartbeat.ack = peer.ack;
 
     UdpMessage message;
     message.ip = peer.ip;

diff --git a/src/core/message.h b/src/core/message.h
@@ -91,6 +91,7 @@ struct PeerHeartbeatMessage {
     uint32_t tun;
     uint32_t ip;
     uint16_t port;
+    uint8_t ack;
 } __attribute__((packed));
 
 struct PeerRawIPv4Message {

diff --git a/src/peer/common.cc b/src/peer/common.cc
@@ -12,6 +12,7 @@ void PeerInfo::reset() {
     this->ip = 0;
     this->port = 0;
     this->count = 0;
+    this->ack = 0;
     this->key.clear();
 }
 

diff --git a/src/peer/peer.h b/src/peer/peer.h
@@ -23,6 +23,7 @@ class PeerInfo {
     uint32_t ip;
     uint16_t port;
     uint32_t count;
+    uint8_t ack;
 
     void reset();
     int updateKey(const std::string &password);