From 3b2e85a27ea8dc64d858aeb6c3441bdb75c3e170 Mon Sep 17 00:00:00 2001
From: Kibae Shin <kibae.shin@gmail.com>
Date: Sun, 17 Sep 2023 13:54:24 +0900
Subject: [PATCH] fix: tcp recv buffer overflow

---
 src/transport/tcp/tcp_server.hpp  |  4 +++-
 src/transport/tcp/tcp_session.cpp | 16 +---------------
 2 files changed, 4 insertions(+), 16 deletions(-)

diff --git a/src/transport/tcp/tcp_server.hpp b/src/transport/tcp/tcp_server.hpp
index a798286..1dc9006 100644
--- a/src/transport/tcp/tcp_server.hpp
+++ b/src/transport/tcp/tcp_server.hpp
@@ -15,6 +15,8 @@
 #define NTOHLL(x) ((((uint64_t)ntohl(x)) << 32) + ntohl(x >> 32))
 #endif
 
+#define MAX_RECV_BUF_LENGTH 1024 * 1024 * 4
+
 namespace onnxruntime_server::transport::tcp {
 	struct protocol_header {
 		int16_t type;
@@ -29,7 +31,7 @@ namespace onnxruntime_server::transport::tcp {
 	  private:
 		asio::socket socket;
 		tcp_server *server;
-		char chunk[1024];
+		char chunk[MAX_RECV_BUF_LENGTH];
 		std::string buffer;
 
 		onnxruntime_server::task::benchmark request_time;
diff --git a/src/transport/tcp/tcp_session.cpp b/src/transport/tcp/tcp_session.cpp
index d184a2f..6773638 100644
--- a/src/transport/tcp/tcp_session.cpp
+++ b/src/transport/tcp/tcp_session.cpp
@@ -19,15 +19,13 @@ void onnxruntime_server::transport::tcp::tcp_session::close() {
 	server->remove_session(shared_from_this());
 }
 
-#define MAX_LENGTH 1024 * 16
-
 void Orts::transport::tcp::tcp_session::do_read() {
 	if (_remote_endpoint.empty())
 		_remote_endpoint =
 			socket.remote_endpoint().address().to_string() + ":" + std::to_string(socket.remote_endpoint().port());
 
 	socket.async_read_some(
-		boost::asio::buffer(chunk, MAX_LENGTH),
+		boost::asio::buffer(chunk, MAX_RECV_BUF_LENGTH),
 		[self = shared_from_this()](boost::system::error_code ec, std::size_t length) {
 			if (!ec) {
 				self->buffer.append(self->chunk, length);
@@ -42,18 +40,6 @@ void Orts::transport::tcp::tcp_session::do_read() {
 				header.length = NTOHLL(header.length);
 				header.json_length = NTOHLL(header.json_length);
 				header.post_length = NTOHLL(header.post_length);
-				// assert(header.length == header.json_length + header.post_length);
-
-				/*
-				// check buffer size
-				if (header.length > MAX_BUFFER_LIMIT) {
-					PLOG(L_WARNING) << self->get_remote_endpoint()
-									<< " transport::session::do_read: Buffer size is too large: " << header.length
-									<< std::endl;
-					self->close();
-					return;
-				}
-				 */
 
 				// continue to read
 				if (self->buffer.size() < sizeof(protocol_header) + header.length) {