diff --git a/src/core/client.cc b/src/core/client.cc index 86727a8c..fbe07f96 100644 --- 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 index de3535f5..10d166c7 100644 --- 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 index e7abfcda..c53baded 100644 --- 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 index 39004448..24d336ec 100644 --- 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);