diff --git a/contrib/epee/include/net/levin_protocol_handler_async.h b/contrib/epee/include/net/levin_protocol_handler_async.h index 86cee9a11..756560c2a 100644 --- a/contrib/epee/include/net/levin_protocol_handler_async.h +++ b/contrib/epee/include/net/levin_protocol_handler_async.h @@ -481,18 +481,26 @@ public: return false; } - temp = std::move(m_fragment_buffer); - m_fragment_buffer.clear(); + temp.swap(m_fragment_buffer); std::memcpy(std::addressof(m_current_head), std::addressof(temp[0]), sizeof(bucket_head2)); + const std::uint64_t inner_size = SWAP64LE(m_current_head.m_cb); + buff_to_invoke = {reinterpret_cast(temp.data()) + sizeof(bucket_head2), temp.size() - sizeof(bucket_head2)}; + if (buff_to_invoke.size() < inner_size) + { + MERROR(m_connection_context << "Invalid fragmented buffer size: " << buff_to_invoke.size() << " vs " << inner_size); + return false; + } + + buff_to_invoke = {buff_to_invoke.data(), std::size_t(inner_size)}; + const size_t max_bytes = m_connection_context.get_max_bytes(m_current_head.m_command); - if(m_current_head.m_cb > std::min(max_packet_size, max_bytes)) + if(buff_to_invoke.size() > std::min(max_packet_size, max_bytes)) { MERROR(m_connection_context << "Maximum packet size exceed!, m_max_packet_size = " << std::min(max_packet_size, max_bytes) << ", packet header received " << m_current_head.m_cb << ", command " << m_current_head.m_command << ", connection will be closed."); return false; } - buff_to_invoke = {reinterpret_cast(temp.data()) + sizeof(bucket_head2), temp.size() - sizeof(bucket_head2)}; } bool is_response = (m_oponent_protocol_ver == LEVIN_PROTOCOL_VER_1 && m_current_head.m_flags&LEVIN_PACKET_RESPONSE); diff --git a/tests/unit_tests/epee_levin_protocol_handler_async.cpp b/tests/unit_tests/epee_levin_protocol_handler_async.cpp index 9dbb2c19b..4ea6fde90 100644 --- a/tests/unit_tests/epee_levin_protocol_handler_async.cpp +++ b/tests/unit_tests/epee_levin_protocol_handler_async.cpp @@ -506,7 +506,6 @@ TEST_F(positive_test_connection_to_levin_protocol_handler_calls, handler_process } std::string compare_buffer(1024 * 4, 'c'); - compare_buffer.resize(((1024 - sizeof(epee::levin::bucket_head2)) * 5) - sizeof(epee::levin::bucket_head2)); // add padding zeroes ASSERT_EQ(4u, m_commands_handler.notify_counter()); ASSERT_EQ(0u, m_commands_handler.invoke_counter()); @@ -652,3 +651,24 @@ TEST_F(test_levin_protocol_handler__hanle_recv_with_invalid_data, handles_short_ ASSERT_FALSE(m_conn->m_protocol_handler.handle_recv(m_buf.data(), m_buf.size())); } + +TEST_F(test_levin_protocol_handler__hanle_recv_with_invalid_data, handles_bad_cb) +{ + m_req_head.m_cb = sizeof(epee::levin::bucket_head2); + m_req_head.m_flags = LEVIN_PACKET_BEGIN; + m_req_head.m_command = 0; + + epee::levin::bucket_head2 inner{}; + inner.m_cb = 2; + m_in_data.resize(sizeof(epee::levin::bucket_head2)); + prepare_buf(); + + ASSERT_TRUE(m_conn->m_protocol_handler.handle_recv(m_buf.data(), m_buf.size())); + + m_req_head.m_cb = 1; + m_req_head.m_flags = LEVIN_PACKET_END; + m_in_data.resize(1); + prepare_buf(); + + ASSERT_FALSE(m_conn->m_protocol_handler.handle_recv(m_buf.data(), m_buf.size())); +}