From 7d62d1e20e39b2840ef95aefc415c2c98db9baaf Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Fri, 17 Jan 2025 16:54:57 +0100
Subject: [PATCH] Support size limit

---
 src/common/network/tcp_client_socket.cpp | 5 +++--
 src/common/network/tcp_client_socket.hpp | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/common/network/tcp_client_socket.cpp b/src/common/network/tcp_client_socket.cpp
index efb4ad34..fc5921f1 100644
--- a/src/common/network/tcp_client_socket.cpp
+++ b/src/common/network/tcp_client_socket.cpp
@@ -56,11 +56,12 @@ namespace network
         return data.empty();
     }
 
-    std::optional<std::string> tcp_client_socket::receive()
+    std::optional<std::string> tcp_client_socket::receive(const std::optional<size_t> max_size)
     {
         char buffer[0x2000];
+        const auto size = std::min(sizeof(buffer), max_size.value_or(sizeof(buffer)));
 
-        const auto result = recv(this->get_socket(), buffer, static_cast<int>(sizeof(buffer)), 0);
+        const auto result = recv(this->get_socket(), buffer, static_cast<int>(size), 0);
         if (result > 0)
         {
             return std::string(buffer, result);
diff --git a/src/common/network/tcp_client_socket.hpp b/src/common/network/tcp_client_socket.hpp
index ec032601..187829be 100644
--- a/src/common/network/tcp_client_socket.hpp
+++ b/src/common/network/tcp_client_socket.hpp
@@ -21,7 +21,7 @@ namespace network
 
         [[maybe_unused]] bool send(const void* data, size_t size) const;
         [[maybe_unused]] bool send(std::string_view data) const;
-        std::optional<std::string> receive();
+        std::optional<std::string> receive(std::optional<size_t> max_size = std::nullopt);
 
         std::optional<address> get_target() const;