From eddadacba7fe4cd00cbc58ebc72e3820cdd654ef Mon Sep 17 00:00:00 2001
From: Vinnie Falco <vinnie.falco@gmail.com>
Date: Wed, 27 Dec 2017 12:30:44 -0800
Subject: [PATCH] Control callback is invoked on the execution context:

fix #954

Fix a defect where the control callback could execute
on the same stack as the initiating function.
---
 CHANGELOG.md                                |  4 ++
 include/boost/beast/websocket/impl/read.ipp | 37 ++++++++++++++++-
 test/beast/websocket/read2.cpp              | 45 +++++++++++++++++++++
 3 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 508783b0..93a9d591 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@ Version 151:
 
 * Sanitizer failures are errors
 
+WebSocket:
+
+* Control callback is invoked on the execution context
+
 --------------------------------------------------------------------------------
 
 Version 150:
diff --git a/include/boost/beast/websocket/impl/read.ipp b/include/boost/beast/websocket/impl/read.ipp
index 8ea8204a..a93888eb 100644
--- a/include/boost/beast/websocket/impl/read.ipp
+++ b/include/boost/beast/websocket/impl/read.ipp
@@ -236,6 +236,17 @@ operator()(
                 // Handle ping frame
                 if(ws_.rd_fh_.op == detail::opcode::ping)
                 {
+                    if(ws_.ctrl_cb_)
+                    {
+                        if(! cont_)
+                        {
+                            BOOST_ASIO_CORO_YIELD
+                            boost::asio::post(
+                                ws_.get_executor(),
+                                    std::move(*this));
+                            BOOST_ASSERT(cont_);
+                        }
+                    }
                     {
                         auto const b = buffers_prefix(
                             clamp(ws_.rd_fh_.len),
@@ -249,7 +260,8 @@ operator()(
                         if(ws_.status_ == status::closing)
                             goto loop;
                         if(ws_.ctrl_cb_)
-                            ws_.ctrl_cb_(frame_type::ping, payload);
+                            ws_.ctrl_cb_(
+                                frame_type::ping, payload);
                         ws_.rd_fb_.reset();
                         ws_.template write_ping<
                             flat_static_buffer_base>(ws_.rd_fb_,
@@ -309,6 +321,18 @@ operator()(
                 // Handle pong frame
                 if(ws_.rd_fh_.op == detail::opcode::pong)
                 {
+                    // Ignore pong when closing
+                    if(! ws_.wr_close_ && ws_.ctrl_cb_)
+                    {
+                        if(! cont_)
+                        {
+                            BOOST_ASIO_CORO_YIELD
+                            boost::asio::post(
+                                ws_.get_executor(),
+                                    std::move(*this));
+                            BOOST_ASSERT(cont_);
+                        }
+                    }
                     auto const cb = buffers_prefix(clamp(
                         ws_.rd_fh_.len), ws_.rd_buf_.data());
                     auto const len = buffer_size(cb);
@@ -325,6 +349,17 @@ operator()(
                 // Handle close frame
                 BOOST_ASSERT(ws_.rd_fh_.op == detail::opcode::close);
                 {
+                    if(ws_.ctrl_cb_)
+                    {
+                        if(! cont_)
+                        {
+                            BOOST_ASIO_CORO_YIELD
+                            boost::asio::post(
+                                ws_.get_executor(),
+                                    std::move(*this));
+                            BOOST_ASSERT(cont_);
+                        }
+                    }
                     auto const cb = buffers_prefix(clamp(
                         ws_.rd_fh_.len), ws_.rd_buf_.data());
                     auto const len = buffer_size(cb);
diff --git a/test/beast/websocket/read2.cpp b/test/beast/websocket/read2.cpp
index 3cc1f232..726fee59 100644
--- a/test/beast/websocket/read2.cpp
+++ b/test/beast/websocket/read2.cpp
@@ -59,6 +59,50 @@ public:
         BEAST_EXPECT(n == 0);
     }
 
+    /*
+        When the internal read buffer contains a control frame and
+        stream::async_read_some is called, it is possible for the control
+        callback to be invoked on the caller's stack instead of through
+        the executor associated with the final completion handler.
+    */
+    void
+    testIssue954()
+    {
+        echo_server es{log};
+        boost::asio::io_context ioc;
+        stream<test::stream> ws{ioc};
+        ws.next_layer().connect(es.stream());
+        ws.handshake("localhost", "/");
+        // message followed by ping
+        ws.next_layer().append({
+            "\x81\x00"
+            "\x89\x00",
+            4});
+        bool called_cb = false;
+        bool called_handler = false;
+        ws.control_callback(
+            [&called_cb](frame_type, string_view)
+            {
+                called_cb = true;
+            });
+        multi_buffer b;
+        ws.async_read(b,
+            [&](error_code, std::size_t)
+            {
+                called_handler = true;
+            });
+        BEAST_EXPECT(! called_cb);
+        BEAST_EXPECT(! called_handler);
+        ioc.run();
+        BEAST_EXPECT(! called_cb);
+        BEAST_EXPECT(called_handler);
+        ws.async_read(b,
+            [&](error_code, std::size_t)
+            {
+            });
+        BEAST_EXPECT(! called_cb);
+    }
+
     /*  Bishop Fox Hybrid Assessment issue 1
 
         Happens with permessage-deflate enabled and a
@@ -242,6 +286,7 @@ public:
     {
         testIssue802();
         testIssue807();
+        testIssue954();
         testIssueBF1();
         testIssueBF2();
     }