From dcff7c0734b7ff5519e535307a2a6db6c104e7fc Mon Sep 17 00:00:00 2001
From: Jung-Sang Ahn <jungsang.ahn@gmail.com>
Date: Wed, 15 May 2024 14:27:18 -0700
Subject: [PATCH] Add more sanity checks on a new connection

* If there is an incoming random message, where `CRC_ON_ENTIRE_MESSAGE`
flag is randomly set, NuRaft tries to read the entire message for CRC
check. That requires allocating a memory blob with the given size
(also a random number) which most likely causes problem.

* Added more sanity checks before the memory allocation.
---
 include/libnuraft/msg_type.hxx |  2 +-
 src/asio_service.cxx           | 21 +++++++++++++++------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/include/libnuraft/msg_type.hxx b/include/libnuraft/msg_type.hxx
index dfdb2746..1f58ec9e 100644
--- a/include/libnuraft/msg_type.hxx
+++ b/include/libnuraft/msg_type.hxx
@@ -64,7 +64,7 @@ enum msg_type {
 
 inline bool ATTR_UNUSED is_valid_msg(msg_type type) {
     if ( type >= request_vote_request &&
-         type <= other_response ) {
+         type <= custom_notification_response ) {
         return true;
     }
     return false;
diff --git a/src/asio_service.cxx b/src/asio_service.cxx
index dd99f645..a1fd31fe 100644
--- a/src/asio_service.cxx
+++ b/src/asio_service.cxx
@@ -47,11 +47,8 @@ limitations under the License.
 #include <atomic>
 #include <ctime>
 #include <exception>
-#include <fstream>
 #include <list>
-#include <queue>
 #include <thread>
-#include <regex>
 
 #ifdef USE_BOOST_ASIO
     using namespace boost;
@@ -382,9 +379,21 @@ class rpc_session
             // byte marker = header_->get_byte();
             h_bs.pos(0);
             byte marker = h_bs.get_u8();
-            if (marker == 0x1) {
-                // Means that this is RPC_RESP, shouldn't happen.
-                p_er("Wrong packet: expected REQ, got RESP");
+            if (marker != 0x0) {
+                // Means that this is not RPC_REQ, shouldn't happen.
+                p_er("Wrong packet: expected REQ, got %u", marker);
+
+                if (impl_->get_options().corrupted_msg_handler_) {
+                    impl_->get_options().corrupted_msg_handler_(header_, nullptr);
+                }
+
+                this->stop();
+                return;
+            }
+
+            msg_type m_type = (msg_type)h_bs.get_u8();
+            if (!is_valid_msg(m_type)) {
+                p_er("Wrong message type: got %u", (uint8_t)m_type);
 
                 if (impl_->get_options().corrupted_msg_handler_) {
                     impl_->get_options().corrupted_msg_handler_(header_, nullptr);