Merge pull request #844 from eynhaender/fix/chunking

Fix websocket body writes to frame as one multi-part message.

Author: evoskuil

include/bitcoin/network/net/socket.hpp

@@ -272,7 +272,7 @@ class BCT_API socket
     void async_read_some(const asio::mutable_buffer& buffer,
         const count_handler& handler) NOEXCEPT;
     void async_write(const asio::const_buffer& buffer, bool binary,
-        const count_handler& handler) NOEXCEPT;
+        bool finish, const count_handler& handler) NOEXCEPT;
     void async_read_http(http::flat_buffer& buffer, http::request& request,
         const count_handler& handler) NOEXCEPT;
     void async_write_http(http::response&& response,

src/net/socket.cpp

@@ -263,9 +263,12 @@ asio::ssl::socket& socket::get_ssl() NOEXCEPT
 // ----------------------------------------------------------------------------
 // protected
 
-// write message in a single frame (ws) or unframed fixed size bytes (tcp).
+// write message frame (ws, finish closes the message) or unframed fixed
+// size bytes (tcp). Beast applies the binary/text setting only at the start
+// of a message, so it is harmless to pass it on every frame of a
+// multi-frame message (see boost::beast::websocket::stream::binary).
 void socket::async_write(const asio::const_buffer& buffer, bool binary,
-    const count_handler& handler) NOEXCEPT
+    bool finish, const count_handler& handler) NOEXCEPT
 {
     BC_ASSERT(stranded());
 
@@ -275,7 +278,7 @@ void socket::async_write(const asio::const_buffer& buffer, bool binary,
         {
             VARIANT_DISPATCH_METHOD(get_ws(), binary(binary));
             VARIANT_DISPATCH_METHOD(get_ws(),
-                async_write(buffer, std::bind(&socket::handle_async,
+                async_write_some(finish, buffer, std::bind(&socket::handle_async,
                     shared_from_this(), _1, _2, handler, "async_write")));
         }
         else

src/net/socket_body.cpp

@@ -200,7 +200,10 @@ void socket::do_body_write(boost_code ec, size_t total,
     out->more = buffer.value().second;
     const auto& data = buffer.value().first;
 
-    async_write(data, out->writer.binary(),
+    // A ws message closes on the final chunk (finish = !more), so a body
+    // that spans multiple get() passes is delivered as one multi-frame
+    // message rather than one whole message per chunk.
+    async_write(data, out->writer.binary(), !out->more,
         std::bind(&socket::handle_body_write,
             shared_from_this(), _1, _2, total, out, handler));
 }
@@ -275,7 +278,7 @@ void socket::do_body_notify(boost_code ec, size_t total,
     const auto& data = buffer.value().first;
 
     // TODO: derive websocket binary/text from body type mapping.
-    async_write(data, false,
+    async_write(data, false, !out->more,
         std::bind(&socket::handle_body_notify,
             shared_from_this(), _1, _2, total, out, handler));
 }

src/net/socket_ws.cpp

@@ -70,7 +70,10 @@ void socket::do_ws_write(const asio::const_buffer& in, bool binary,
     const count_handler& handler) NOEXCEPT
 {
     BC_ASSERT(stranded());
-    async_write(in, binary, handler);
+
+    // Single-shot write of an already-whole buffer, always closes the
+    // message (finish = true).
+    async_write(in, binary, true, handler);
 }
 
 // WS (event).

test/messages/json_body_writer.cpp

@@ -97,6 +97,63 @@ BOOST_AUTO_TEST_CASE(json_body_writer__get__simple_object__success_expected_no_m
     BOOST_REQUIRE(!buffer.get().second);
 }
 
+// A size_hint smaller than the serialized payload forces writer.get() to be
+// called multiple times (more == true on all but the last pass). This is the
+// chunking that drives socket::async_write's finish flag (finish = !more)
+// for websocket body writes: the fix relies on every non-final chunk
+// reporting more == true and only the final chunk reporting more == false.
+// boost::json::serializer::read() fills the given buffer to capacity on
+// every pass but the last, so for this size_hint and payload the split is
+// deterministic: four 8-byte-or-fewer chunks.
+BOOST_AUTO_TEST_CASE(json_body_writer__get__size_hint_smaller_than_payload__four_chunks_reassemble_with_terminal_no_more)
+{
+    const std::string_view expected1{ R"({"key":")" };
+    const std::string_view expected2{ "xxxxxxxx" };
+    const std::string_view expected3{ "xxxxxxxx" };
+    const std::string_view expected4{ R"(xxxx"})" };
+    const asio::const_buffer out1{ expected1.data(), expected1.size() };
+    const asio::const_buffer out2{ expected2.data(), expected2.size() };
+    const asio::const_buffer out3{ expected3.data(), expected3.size() };
+    const asio::const_buffer out4{ expected4.data(), expected4.size() };
+
+    json::body<>::value_type body{};
+    body.model = object{ { "key", std::string(20, 'x') } };
+    body.size_hint = 8;
+    json::body<>::writer writer(body);
+    boost_code ec{};
+    writer.init(ec);
+    BOOST_REQUIRE(!ec);
+    BOOST_REQUIRE(!writer.done());
+
+    const auto buffer1 = writer.get(ec);
+    BOOST_REQUIRE(!ec);
+    BOOST_REQUIRE(buffer1.has_value());
+    BOOST_REQUIRE(buffer1.get().first == out1);
+    BOOST_REQUIRE(buffer1.get().second);
+    BOOST_REQUIRE(!writer.done());
+
+    const auto buffer2 = writer.get(ec);
+    BOOST_REQUIRE(!ec);
+    BOOST_REQUIRE(buffer2.has_value());
+    BOOST_REQUIRE(buffer2.get().first == out2);
+    BOOST_REQUIRE(buffer2.get().second);
+    BOOST_REQUIRE(!writer.done());
+
+    const auto buffer3 = writer.get(ec);
+    BOOST_REQUIRE(!ec);
+    BOOST_REQUIRE(buffer3.has_value());
+    BOOST_REQUIRE(buffer3.get().first == out3);
+    BOOST_REQUIRE(buffer3.get().second);
+    BOOST_REQUIRE(!writer.done());
+
+    const auto buffer4 = writer.get(ec);
+    BOOST_REQUIRE(!ec);
+    BOOST_REQUIRE(buffer4.has_value());
+    BOOST_REQUIRE(buffer4.get().first == out4);
+    BOOST_REQUIRE(!buffer4.get().second);
+    BOOST_REQUIRE(writer.done());
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 
 ////#endif // HAVE_SLOW_TESTS

test/net/socket.cpp

@@ -18,6 +18,8 @@
  */
 #include "../test.hpp"
 
+#include <future>
+
 BOOST_AUTO_TEST_SUITE(socket_tests)
 
 class socket_accessor
@@ -231,4 +233,115 @@ BOOST_AUTO_TEST_CASE(socket__write__disconnected__bad_stream)
     BOOST_REQUIRE(pool.join());
 }
 
+// Regression test for a fix to socket::async_write: it used to write every
+// body_write() chunk as its own whole websocket message (finish always
+// true), splitting one logical response into N independent messages. The
+// fix passes finish = !out->more, so only the final chunk closes the
+// message. A real (non-library) websocket client is used as ground truth:
+// beast's blocking read() returns only once a full message (finish) is
+// received, so a single read() call assembling the entire multi-chunk body
+// proves the chunks were sent as one multi-frame message rather than as
+// several.
+BOOST_AUTO_TEST_CASE(socket__body_write__websocket_multiple_chunks__single_message)
+{
+    using namespace std::chrono_literals;
+
+    const logger log{};
+    threadpool pool(2);
+    connector::parameters params{ .maximum_request = 1'000'000u };
+
+    // Bind a loopback acceptor on an ephemeral port.
+    asio::strand accept_strand(pool.service().get_executor());
+    asio::acceptor acceptor(accept_strand);
+    boost_code ec{};
+    const asio::endpoint bind_endpoint(asio::ipv4::loopback(), 0);
+
+    acceptor.open(bind_endpoint.protocol(), ec);
+    BOOST_REQUIRE(!ec);
+    acceptor.set_option(asio::reuse_address(true), ec);
+    BOOST_REQUIRE(!ec);
+    acceptor.bind(bind_endpoint, ec);
+    BOOST_REQUIRE(!ec);
+    acceptor.listen(1, ec);
+    BOOST_REQUIRE(!ec);
+    const auto port = acceptor.local_endpoint().port();
+
+    const auto server = std::make_shared<socket_accessor>(log, pool.service(), std::move(params));
+    const auto buffer = std::make_shared<http::flat_buffer>();
+    const auto request = std::make_shared<http::request>();
+
+    // Small size_hint relative to the payload forces writer.get() across
+    // many passes, i.e. many socket::async_write calls for one logical body.
+    json::body<>::value_type content{};
+    content.model = boost::json::object{ { "key", std::string(4000, 'x') } };
+    content.size_hint = 32;
+    const auto expected = boost::json::serialize(content.model);
+
+    const auto response = std::make_shared<http::response>();
+    response->result(boost::beast::http::status::ok);
+    response->body() = std::move(content);
+
+    const auto upgrade_result = std::make_shared<std::promise<code>>();
+    const auto write_result = std::make_shared<std::promise<code>>();
+    auto upgrade_future = upgrade_result->get_future();
+    auto write_future = write_result->get_future();
+
+    server->accept(acceptor,
+        [=](const code& accept_ec) mutable
+        {
+            if (accept_ec)
+            {
+                upgrade_result->set_value(accept_ec);
+                write_result->set_value(accept_ec);
+                return;
+            }
+
+            server->http_read(*buffer, *request,
+                [=](const code& read_ec, size_t) mutable
+                {
+                    upgrade_result->set_value(read_ec);
+                    if (read_ec != error::upgraded)
+                    {
+                        write_result->set_value(read_ec);
+                        return;
+                    }
+
+                    server->body_write(std::move(*response),
+                        [=](const code& write_ec, size_t) NOEXCEPT
+                        {
+                            write_result->set_value(write_ec);
+                        });
+                });
+        });
+
+    // Real (blocking) websocket client, independent of the code under test.
+    asio::context client_service;
+    asio::socket raw(client_service);
+    boost_code client_ec{};
+    raw.connect({ asio::ipv4::loopback(), port }, client_ec);
+    BOOST_REQUIRE(!client_ec);
+
+    ws::socket client(std::move(raw));
+    client.handshake("127.0.0.1", "/", client_ec);
+    BOOST_REQUIRE(!client_ec);
+
+    // One blocking read of a "complete message" must assemble every chunk.
+    http::flat_buffer read_buffer{};
+    client.read(read_buffer, client_ec);
+    BOOST_REQUIRE(!client_ec);
+    BOOST_REQUIRE(client.got_text());
+
+    const auto received = boost::beast::buffers_to_string(read_buffer.data());
+    BOOST_REQUIRE_EQUAL(received, expected);
+
+    BOOST_REQUIRE(upgrade_future.wait_for(2s) == std::future_status::ready);
+    BOOST_REQUIRE_EQUAL(upgrade_future.get(), error::upgraded);
+    BOOST_REQUIRE(write_future.wait_for(2s) == std::future_status::ready);
+    BOOST_REQUIRE_EQUAL(write_future.get(), error::success);
+
+    server->stop();
+    pool.stop();
+    BOOST_REQUIRE(pool.join());
+}
+
 BOOST_AUTO_TEST_SUITE_END()