处理客户端 Websocket 写失败

lanthora · Mar 14, 2024 · 2d0d5e4 · 2d0d5e4
1 parent 1bc92c4
commit 2d0d5e4
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 28 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.18.4)
 
-project(candy LANGUAGES C CXX VERSION 5.2)
+project(candy LANGUAGES C CXX VERSION 5.3)
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wfatal-errors")
 set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fsanitize=address")

diff --git a/src/core/client.cc b/src/core/client.cc
@@ -215,8 +215,8 @@ void Client::handleWebSocketMessage() {
                 break;
 
             // 动态地址响应包,启动 TUN 设备并发送 Auth 包
-            case MessageType::DHCP:
-                handleDynamicAddressMessage(message);
+            case MessageType::EXPECTED:
+                handleExpectedAddressMessage(message);
                 break;
 
             // 对端连接请求包
@@ -308,7 +308,10 @@ void Client::sendForwardMessage(const std::string &buffer) {
     WebSocketMessage message;
     message.buffer.push_back(MessageType::FORWARD);
     message.buffer.append(buffer);
-    ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send forward message failed");
+        Candy::shutdown();
+    }
 }
 
 void Client::sendVirtualMacMessage() {
@@ -317,7 +320,10 @@ void Client::sendVirtualMacMessage() {
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&buffer), sizeof(buffer));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send virtual mac message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -329,12 +335,15 @@ void Client::sendDynamicAddressMessage() {
         return;
     }
 
-    DynamicAddressMessage header(address.getCidr());
+    ExpectedAddressMessage header(address.getCidr());
     header.updateHash(this->password);
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&header), sizeof(header));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send expected address message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -351,7 +360,10 @@ void Client::sendAuthMessage() {
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&header), sizeof(AuthHeader));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send auth message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -364,7 +376,10 @@ void Client::sendPeerConnMessage(const PeerInfo &peer, uint32_t ip, uint16_t por
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&header), sizeof(PeerConnMessage));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send peer conn message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -376,7 +391,10 @@ void Client::sendDiscoveryMessage(uint32_t dst) {
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&header), sizeof(DiscoveryMessage));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send discovery conn message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -391,7 +409,10 @@ void Client::sendLocalPeerConnMessage(const PeerInfo &peer, uint32_t ip, uint16_
 
     WebSocketMessage message;
     message.buffer.assign((char *)(&header), sizeof(LocalPeerConnMessage));
-    this->ws.write(message);
+    if (this->ws.write(message)) {
+        spdlog::critical("send peer conn message failed");
+        Candy::shutdown();
+    }
     return;
 }
 
@@ -410,14 +431,14 @@ void Client::handleForwardMessage(WebSocketMessage &message) {
     tryDirectConnection(Address::netToHost(header->saddr));
 }
 
-void Client::handleDynamicAddressMessage(WebSocketMessage &message) {
-    if (message.buffer.size() < sizeof(DynamicAddressMessage)) {
+void Client::handleExpectedAddressMessage(WebSocketMessage &message) {
+    if (message.buffer.size() < sizeof(ExpectedAddressMessage)) {
         spdlog::warn("invalid expected address message: len {}", message.buffer.length());
         spdlog::debug("expected address buffer: {:n}", spdlog::to_hex(message.buffer));
         return;
     }
 
-    DynamicAddressMessage *header = (DynamicAddressMessage *)message.buffer.c_str();
+    ExpectedAddressMessage *header = (ExpectedAddressMessage *)message.buffer.c_str();
 
     Address address;
     if (address.cidrUpdate(header->cidr)) {

diff --git a/src/core/client.h b/src/core/client.h
@@ -87,7 +87,7 @@ class Client {
     void sendDiscoveryMessage(uint32_t dst);
     void sendLocalPeerConnMessage(const PeerInfo &peer, uint32_t ip, uint16_t port);
     void handleForwardMessage(WebSocketMessage &message);
-    void handleDynamicAddressMessage(WebSocketMessage &message);
+    void handleExpectedAddressMessage(WebSocketMessage &message);
     void handlePeerConnMessage(WebSocketMessage &message);
     void handleDiscoveryMessage(WebSocketMessage &message);
     void handleGeneralMessage(WebSocketMessage &message);

diff --git a/src/core/message.cc b/src/core/message.cc
@@ -47,20 +47,20 @@ ForwardHeader::ForwardHeader() {
     this->type = MessageType::FORWARD;
 }
 
-DynamicAddressMessage::DynamicAddressMessage(const std::string &cidr) {
-    this->type = MessageType::DHCP;
+ExpectedAddressMessage::ExpectedAddressMessage(const std::string &cidr) {
+    this->type = MessageType::EXPECTED;
     this->timestamp = Time::hostToNet(Time::unixTime());
     std::strcpy(this->cidr, cidr.c_str());
 }
 
-void DynamicAddressMessage::updateHash(const std::string &password) {
+void ExpectedAddressMessage::updateHash(const std::string &password) {
     std::string data;
     data.append(password);
     data.append((char *)&this->timestamp, sizeof(this->timestamp));
     SHA256((unsigned char *)data.data(), data.size(), this->hash);
 }
 
-bool DynamicAddressMessage::check(const std::string &password) {
+bool ExpectedAddressMessage::check(const std::string &password) {
     int64_t localTime = Time::unixTime();
     int64_t remoteTime = Time::netToHost(this->timestamp);
     if (std::abs(localTime - remoteTime) > 30) {

diff --git a/src/core/message.h b/src/core/message.h
@@ -15,7 +15,7 @@ namespace MessageType {
 
 constexpr uint8_t AUTH = 0;
 constexpr uint8_t FORWARD = 1;
-constexpr uint8_t DHCP = 2;
+constexpr uint8_t EXPECTED = 2;
 constexpr uint8_t PEER = 3;
 constexpr uint8_t VMAC = 4;
 constexpr uint8_t DISCOVERY = 5;
@@ -41,13 +41,13 @@ struct ForwardHeader {
     ForwardHeader();
 } __attribute__((packed));
 
-struct DynamicAddressMessage {
+struct ExpectedAddressMessage {
     uint8_t type;
     int64_t timestamp;
     char cidr[32];
     uint8_t hash[SHA256_DIGEST_LENGTH];
 
-    DynamicAddressMessage(const std::string &cidr);
+    ExpectedAddressMessage(const std::string &cidr);
     void updateHash(const std::string &password);
     bool check(const std::string &password);
 } __attribute__((packed));

diff --git a/src/core/server.cc b/src/core/server.cc
@@ -107,8 +107,8 @@ void Server::handleWebSocketMessage() {
         if (message.type == WebSocketMessageType::Message) {
             uint8_t msgType = message.buffer.front();
             switch (msgType) {
-            case MessageType::DHCP:
-                handleDynamicAddressMessage(message);
+            case MessageType::EXPECTED:
+                handleExpectedAddressMessage(message);
                 break;
             case MessageType::VMAC:
                 handleVirtualMacMessage(message);
@@ -226,14 +226,14 @@ void Server::handleForwardMessage(WebSocketMessage &message) {
     return;
 }
 
-void Server::handleDynamicAddressMessage(WebSocketMessage &message) {
-    if (message.buffer.length() < sizeof(DynamicAddressMessage)) {
+void Server::handleExpectedAddressMessage(WebSocketMessage &message) {
+    if (message.buffer.length() < sizeof(ExpectedAddressMessage)) {
         spdlog::warn("invalid dynamic address message: len {}", message.buffer.length());
         this->ws.close(message.conn);
         return;
     }
 
-    DynamicAddressMessage *header = (DynamicAddressMessage *)message.buffer.data();
+    ExpectedAddressMessage *header = (ExpectedAddressMessage *)message.buffer.data();
     if (!header->check(this->password)) {
         spdlog::warn("dynamic address header check failed: buffer {:n}", spdlog::to_hex(message.buffer));
         this->ws.close(message.conn);

diff --git a/src/core/server.h b/src/core/server.h
@@ -25,7 +25,7 @@ class Server {
 
     void handleAuthMessage(WebSocketMessage &message);
     void handleForwardMessage(WebSocketMessage &message);
-    void handleDynamicAddressMessage(WebSocketMessage &message);
+    void handleExpectedAddressMessage(WebSocketMessage &message);
     void handlePeerConnMessage(WebSocketMessage &message);
     void handleVirtualMacMessage(WebSocketMessage &message);
     void handleDiscoveryMessage(WebSocketMessage &message);