From 767397e0c1a761757e88c353c7850ea001dbcfdf Mon Sep 17 00:00:00 2001
From: Mohammad Nejati <ashtumashtum@gmail.com>
Date: Wed, 5 Jun 2024 12:59:40 +0000
Subject: [PATCH] read_size_hint does not exceed read_message_max

Fixes #2879
---
 .../beast/websocket/detail/impl_base.hpp      |  10 +-
 include/boost/beast/websocket/impl/read.hpp   |   2 +-
 include/boost/beast/websocket/impl/stream.hpp |   2 +-
 .../beast/websocket/impl/stream_impl.hpp      |   5 +-
 test/beast/websocket/read3.cpp                | 122 ++++++++++++++++++
 5 files changed, 137 insertions(+), 4 deletions(-)

diff --git a/include/boost/beast/websocket/detail/impl_base.hpp b/include/boost/beast/websocket/detail/impl_base.hpp
index 705ce72d..e5980bc0 100644
--- a/include/boost/beast/websocket/detail/impl_base.hpp
+++ b/include/boost/beast/websocket/detail/impl_base.hpp
@@ -311,8 +311,9 @@ struct impl_base<true>
     read_size_hint_pmd(
         std::size_t initial_size,
         bool rd_done,
+        std::uint64_t rd_msg_max,
         std::uint64_t rd_remain,
-        detail::frame_header const& rd_fh) const
+        frame_header const& rd_fh) const
     {
         using beast::detail::clamp;
         std::size_t result;
@@ -339,6 +340,9 @@ struct impl_base<true>
             initial_size, clamp(rd_remain));
     done:
         BOOST_ASSERT(result != 0);
+        // Ensure offered size does not exceed rd_msg_max
+        if(rd_msg_max)
+            result = clamp(result, rd_msg_max);
         return result;
     }
 };
@@ -461,6 +465,7 @@ struct impl_base<false>
     read_size_hint_pmd(
         std::size_t initial_size,
         bool rd_done,
+        std::uint64_t rd_msg_max,
         std::uint64_t rd_remain,
         frame_header const& rd_fh) const
     {
@@ -485,6 +490,9 @@ struct impl_base<false>
                 initial_size, clamp(rd_remain));
         }
         BOOST_ASSERT(result != 0);
+        // Ensure offered size does not exceed rd_msg_max
+        if(rd_msg_max)
+            result = clamp(result, rd_msg_max);
         return result;
     }
 };
diff --git a/include/boost/beast/websocket/impl/read.hpp b/include/boost/beast/websocket/impl/read.hpp
index 131ea6e2..a81b13b4 100644
--- a/include/boost/beast/websocket/impl/read.hpp
+++ b/include/boost/beast/websocket/impl/read.hpp
@@ -986,7 +986,7 @@ read_some(
     if(! limit)
         limit = (std::numeric_limits<std::size_t>::max)();
     auto const size =
-        clamp(read_size_hint(buffer), limit);
+        clamp(impl_->read_size_hint_db(buffer), limit);
     BOOST_ASSERT(size > 0);
     auto mb = beast::detail::dynamic_buffer_prepare(
         buffer, size, ec, error::buffer_overflow);
diff --git a/include/boost/beast/websocket/impl/stream.hpp b/include/boost/beast/websocket/impl/stream.hpp
index 52c23d31..7d43bbac 100644
--- a/include/boost/beast/websocket/impl/stream.hpp
+++ b/include/boost/beast/websocket/impl/stream.hpp
@@ -133,7 +133,7 @@ read_size_hint(
     std::size_t initial_size) const
 {
     return impl_->read_size_hint_pmd(
-        initial_size, impl_->rd_done,
+        initial_size, impl_->rd_done, impl_->rd_msg_max,
         impl_->rd_remain, impl_->rd_fh);
 }
 
diff --git a/include/boost/beast/websocket/impl/stream_impl.hpp b/include/boost/beast/websocket/impl/stream_impl.hpp
index 615b4beb..5113be91 100644
--- a/include/boost/beast/websocket/impl/stream_impl.hpp
+++ b/include/boost/beast/websocket/impl/stream_impl.hpp
@@ -300,7 +300,7 @@ struct stream<NextLayer, deflateSupported>::impl_type
         if(initial_size == 0)
             return 1; // buffer is full
         return this->read_size_hint_pmd(
-            initial_size, rd_done, rd_remain, rd_fh);
+            initial_size, rd_done, rd_msg_max, rd_remain, rd_fh);
     }
 
     template<class DynamicBuffer>
@@ -886,6 +886,9 @@ parse_fh(
                 return false;
             }
         }
+        // The final size of a deflated frame is unknown. In certain cases,
+        // post-inflation, it might shrink and become <= rd_msg_max.
+        // Therefore, we will verify the size during the inflation process.
         if(! this->rd_deflated())
         {
             if(rd_msg_max && beast::detail::sum_exceeds(
diff --git a/test/beast/websocket/read3.cpp b/test/beast/websocket/read3.cpp
index 718298c9..ba09a231 100644
--- a/test/beast/websocket/read3.cpp
+++ b/test/beast/websocket/read3.cpp
@@ -993,6 +993,127 @@ public:
         }
     }
 
+    /*
+     * Tests when a deflate frame doesn't prepare the supplied buffer
+     * more than rd_msg_max.
+     */
+    void
+    testIssue2879()
+    {
+        // contains a deflate block with a size of 288230380446614262 (truncated at the end)
+        auto frame = sbuf(
+            "\xc2\xff\x04\x00\x00\x00\xff\xff\x02\xf6"
+            "\x4c\x31\x0d\x0a\x66\x3a\x27\x00\x00\x69"
+            "\x20\x20\x20\xff\x46\x00\x00\x10\x00\x6f"
+            "\x29\x00\xd6\x2e\x31\x0d\x0a\x48\x6f\x73"
+            "\x74\x3a\x42\x42\x42\x52\x41\x41\x45\x42"
+            "\x42\x42\x42\x42\x42\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+            "\xff\xff\xff\xff\xff\xff\xff\xff\xff\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+            "\x00\x00\x00\x00\x00\x00\x00\x01\x01\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
+            "\x0f\x0f\x0f\x0f\x0f\x0f\x01\x01\x01\x01"
+            "\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+            "\x01\x43\x43\x43\x43\x43\x3f\x3f\x28\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x45\x54\x50\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x3f"
+            "\x3f\x3f\x3f\x3f\x3f\x3f\x3f\x43\x43\x43"
+            "\x43\x43\x43\x43\x43\x43\x43\x70\x67\x72"
+            "\x61\x64\x65\x00\x54\x00");
+
+        permessage_deflate pmd;
+        pmd.client_enable = true;
+        pmd.server_enable = true;
+
+        // read
+        {
+            net::io_context ioc;
+            stream<test::stream> wsc{ioc};
+            stream<test::stream> wss{ioc};
+            wsc.set_option(pmd);
+            wss.set_option(pmd);
+            wss.read_message_max(2048);
+            wsc.next_layer().connect(wss.next_layer());
+            wsc.async_handshake(
+                "localhost", "/", [](error_code){});
+            wss.async_accept([](error_code){});
+            ioc.run();
+            ioc.restart();
+            BEAST_EXPECT(wsc.is_open());
+            BEAST_EXPECT(wss.is_open());
+            net::write(wsc.next_layer(), frame);
+            error_code ec;
+            flat_buffer b;
+            wss.read(b, ec);
+            BEAST_EXPECTS(ec == error::message_too_big, ec.message());
+        }
+
+        // async read
+        {
+            net::io_context ioc;
+            stream<test::stream> wsc{ioc};
+            stream<test::stream> wss{ioc};
+            wsc.set_option(pmd);
+            wss.set_option(pmd);
+            wss.read_message_max(2048);
+            wsc.next_layer().connect(wss.next_layer());
+            wsc.async_handshake(
+                "localhost", "/", [](error_code){});
+            wss.async_accept([](error_code){});
+            ioc.run();
+            ioc.restart();
+            BEAST_EXPECT(wsc.is_open());
+            BEAST_EXPECT(wss.is_open());
+            net::write(wsc.next_layer(), frame);
+            error_code ec;
+            flat_buffer b;
+            wss.async_read(b,
+                [&ec](error_code ec_, std::size_t){ ec = ec_; });
+            ioc.run();
+            BEAST_EXPECTS(ec == error::message_too_big, ec.message());
+        }
+    }
+
     void
     testMoveOnly()
     {
@@ -1038,6 +1159,7 @@ public:
         testIssue807();
         testIssue954();
         testIssue1630();
+        testIssue2879();
         testIssueBF1();
         testIssueBF2();
         testMoveOnly();