From 9b9b2023d7ceccd590944137afc133781b138a6e Mon Sep 17 00:00:00 2001
From: helintongh <agh6399@gmail.com>
Date: Sun, 12 Nov 2023 13:59:54 +0800
Subject: [PATCH] feat: coro server limit websocket body max size

---
 include/cinatra/coro_http_connection.hpp | 15 ++++-
 tests/test_coro_http_server.cpp          | 78 ++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/include/cinatra/coro_http_connection.hpp b/include/cinatra/coro_http_connection.hpp
index 52652c7c..38b2651b 100644
--- a/include/cinatra/coro_http_connection.hpp
+++ b/include/cinatra/coro_http_connection.hpp
@@ -333,6 +333,16 @@ class coro_http_connection
           }
           payload = body_;
         }
+
+        if (max_part_size_ != 0 && payload_length > max_part_size_) {
+          std::string close_reason = "message_too_big";
+          std::string close_msg = ws_.format_close_payload(
+              close_code::too_big, close_reason.data(), close_reason.size());
+          co_await write_websocket(close_msg, opcode::close);
+          close();
+          break;
+        }
+
         ws_frame_type type = ws_.parse_payload(payload);
 
         switch (type) {
@@ -404,6 +414,8 @@ class coro_http_connection
     conn_id_ = conn_id;
   }
 
+  void set_ws_max_size(uint64_t max_size) { max_part_size_ = max_size; }
+
   template <typename AsioBuffer>
   async_simple::coro::Lazy<std::pair<std::error_code, size_t>> async_read(
       AsioBuffer &&buffer, size_t size_to_read) noexcept {
@@ -535,6 +547,7 @@ class coro_http_connection
   std::function<void(const uint64_t &conn_id)> quit_cb_ = nullptr;
   bool checkout_timeout_ = false;
   std::atomic<std::chrono::system_clock::time_point> last_rwtime_;
+  uint64_t max_part_size_ = 8 * 1024 * 1024;
 
   websocket ws_;
 #ifdef CINATRA_ENABLE_SSL
@@ -543,4 +556,4 @@ class coro_http_connection
   bool use_ssl_ = false;
 #endif
 };
-}  // namespace cinatra
\ No newline at end of file
+}  // namespace cinatra
diff --git a/tests/test_coro_http_server.cpp b/tests/test_coro_http_server.cpp
index 57ab04c5..80188049 100644
--- a/tests/test_coro_http_server.cpp
+++ b/tests/test_coro_http_server.cpp
@@ -642,6 +642,84 @@ TEST_CASE("check connecton timeout") {
   CHECK(server.connection_count() == 0);
 }
 
+TEST_CASE("test websocket with message max_size limit") {
+  cinatra::coro_http_server server(1, 9001);
+  server.set_http_handler<cinatra::GET>(
+      "/ws_echo1",
+      [](cinatra::coro_http_request &req,
+         cinatra::coro_http_response &resp) -> async_simple::coro::Lazy<void> {
+        REQUIRE(req.get_content_type() == cinatra::content_type::websocket);
+        cinatra::websocket_result result{};
+
+        while (true) {
+          req.get_conn()->set_ws_max_size(65536);
+          result = co_await req.get_conn()->read_websocket();
+          if (result.ec) {
+            break;
+          }
+
+          if (result.type == cinatra::ws_frame_type::WS_CLOSE_FRAME) {
+            REQUIRE(result.data.empty());
+            break;
+          }
+
+          auto ec = co_await req.get_conn()->write_websocket(result.data);
+          if (ec) {
+            break;
+          }
+        }
+      });
+  server.async_start();
+
+  auto client = std::make_shared<cinatra::coro_http_client>();
+
+  SUBCASE("medium message - 16 bit length") {
+    std::string medium_message(
+        65535, 'x');  // 65,535 'x' characters for the medium message test.
+
+    client->on_ws_msg([medium_message](cinatra::resp_data data) {
+      if (data.net_err) {
+        std::cout << "ws_msg net error " << data.net_err.message() << "\n";
+        return;
+      }
+
+      size_t size = data.resp_body.size();
+      std::cout << "ws msg len: " << data.resp_body.size() << std::endl;
+      REQUIRE(data.resp_body == medium_message);
+    });
+
+    async_simple::coro::syncAwait(
+        client->async_ws_connect("ws://127.0.0.1:9001/ws_echo1"));
+    async_simple::coro::syncAwait(client->async_send_ws(medium_message));
+    async_simple::coro::syncAwait(client->async_send_ws_close());
+  }
+
+  SUBCASE("large message - 64 bit length") {
+    std::string large_message(
+        70000, 'x');  // 70,000 'x' characters for the large message test.
+
+    client->on_ws_msg([large_message](cinatra::resp_data data) {
+      if (data.net_err) {
+        std::cout << "ws_msg net error " << data.net_err.message() << "\n";
+        return;
+      }
+
+      std::cout << "ws msg len: " << data.resp_body.size() << std::endl;
+    });
+
+    client->on_ws_close([](std::string_view reason) {
+      REQUIRE(reason.size() > 0);
+    });
+
+    async_simple::coro::syncAwait(
+        client->async_ws_connect("ws://127.0.0.1:9001/ws_echo1"));
+    async_simple::coro::syncAwait(client->async_send_ws(large_message));
+    async_simple::coro::syncAwait(client->async_send_ws_close());
+  }
+
+  server.stop();
+}
+
 #ifdef CINATRA_ENABLE_SSL
 TEST_CASE("test ssl server") {
   cinatra::coro_http_server server(1, 9001);