From 98b5d9fd37da2e59c11b04d65fce985212071f36 Mon Sep 17 00:00:00 2001
From: taynpg <taynpg@163.com>
Date: Fri, 20 Dec 2024 13:22:28 +0800
Subject: [PATCH] =?UTF-8?q?safe=EF=BC=9A=E6=B7=BB=E5=8A=A0=E4=B8=80?=
 =?UTF-8?q?=E4=BA=9B=E5=9F=BA=E6=9C=AC=E7=9A=84=E9=98=B2=E5=BE=A1=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/server.cpp | 22 ++++++++++++++++------
 util/util.cpp     |  9 +++++++++
 util/util.h       |  3 ++-
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/server/server.cpp b/server/server.cpp
index 155c155..721397d 100644
--- a/server/server.cpp
+++ b/server/server.cpp
@@ -190,16 +190,26 @@ void CTcpServer::accept_client()
         if (!error) {
             auto endpoint = socket->remote_endpoint();
             std::string client_key = endpoint.address().to_string() + ":" + std::to_string(endpoint.port());
-            logger_->info("New connection from {}", client_key);
 
+            bool can = false;
             {
                 std::lock_guard<std::mutex> lock(cli_mut_);
-                auto cache = std::make_shared<ClientCache>();
-                cache->socket_ = socket;
-                client_map_[client_key] = cache;
+                if (client_map_.size() >= 100) {
+                    logger_->info("Max client connections reached. Closing connection from {}", client_key);
+                    socket->close();
+                } else {
+                    logger_->info("New connection from {}", client_key);
+                    auto cache = std::make_shared<ClientCache>();
+                    cache->socket_ = socket;
+                    client_map_[client_key] = cache;
+                    can = true;
+                }
+            }
+            if (can == false) {
+                std::this_thread::sleep_for(std::chrono::minutes(1));
+            } else {
+                client_threads_[client_key] = std::thread(&CTcpServer::th_client, this, socket, client_key);
             }
-
-            client_threads_[client_key] = std::thread(&CTcpServer::th_client, this, socket, client_key);
         }
         accept_client();
     });
diff --git a/util/util.cpp b/util/util.cpp
index 296dcba..b39bf83 100644
--- a/util/util.cpp
+++ b/util/util.cpp
@@ -1,5 +1,6 @@
 #include "util.h"
 #include <cstdint>
+#include <thread>
 
 std::shared_ptr<spdlog::logger> get_logger(const std::string& mark, const std::string& log_file)
 {
@@ -42,6 +43,14 @@ CFrameBuffer* CTransProtocal::parse(CMutBuffer& buffer)
     unsigned char header[] = {0xFF, 0xFE};
     unsigned char tail[] = {0xFF, 0xFF};
 
+    // 如果超出 1MB的内容都无法解析成功，则认为是有无效客户端参与链接。
+    if (buffer.get_len() > MAX_FRAME_SIZE) {
+        buffer.clear();
+        // 这里故意延迟。
+        std::this_thread::sleep_for(std::chrono::seconds(600));
+        return result;
+    }
+
     int find = buffer.index_of((const char*)header, sizeof(header));
     if (find < 0) {
         return result;
diff --git a/util/util.h b/util/util.h
index e24c38c..116ea8e 100644
--- a/util/util.h
+++ b/util/util.h
@@ -6,7 +6,8 @@
 #include <spdlog/sinks/stdout_color_sinks.h>
 #include <spdlog/spdlog.h>
 
-constexpr int g_BuffSize = 102400;
+constexpr size_t g_BuffSize = 102400;
+const size_t MAX_FRAME_SIZE = 10 * g_BuffSize;
 enum FrameType : int16_t {
     TYPE_DEFAULT = 0,
     TYPE_GET_LIST,