From d735aa949e4311b0455dd9e408cc5edf0061bbab Mon Sep 17 00:00:00 2001
From: Vinnie Falco <vinnie.falco@gmail.com>
Date: Fri, 10 Feb 2017 11:15:14 -0500
Subject: [PATCH] Fix race in writes during reads:

fix #261

This fixes a rare condition when responding to a ping or
close frame where the wr_block_ stream variable is not
correctly set for a short window of time.
---
 CHANGELOG.md                          |  6 +++++
 include/beast/websocket/impl/read.ipp | 33 ++++++++++++++++++++-------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0b913b49..377d7677 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,9 @@
+1.0.0-b29
+
+* Fix race in writes during WebSocket reads
+
+--------------------------------------------------------------------------------
+
 1.0.0-b28
 
 * Split out and rename test stream classes
diff --git a/include/beast/websocket/impl/read.ipp b/include/beast/websocket/impl/read.ipp
index 6e65beea..dfecdbae 100644
--- a/include/beast/websocket/impl/read.ipp
+++ b/include/beast/websocket/impl/read.ipp
@@ -499,6 +499,8 @@ operator()(error_code ec,
             //------------------------------------------------------------------
 
             case do_pong_resume:
+                BOOST_ASSERT(! d.ws.wr_block_);
+                d.ws.wr_block_ = &d;
                 d.state = do_pong_resume + 1;
                 d.ws.get_io_service().post(bind_handler(
                     std::move(*this), ec, bytes_transferred));
@@ -511,8 +513,7 @@ operator()(error_code ec,
                     ec = boost::asio::error::operation_aborted;
                     goto upcall;
                 }
-                d.state = do_pong;
-                break; // VFALCO fall through?
+                // [[fallthrough]]
 
             //------------------------------------------------------------------
 
@@ -520,14 +521,21 @@ operator()(error_code ec,
                 if(d.ws.wr_close_)
                 {
                     // ignore ping when closing
+                    if(d.ws.wr_block_)
+                    {
+                        BOOST_ASSERT(d.ws.wr_block_ == &d);
+                        d.ws.wr_block_ = nullptr;
+                    }
                     d.fb.reset();
                     d.state = do_read_fh;
                     break;
                 }
                 // send pong
+                if(! d.ws.wr_block_)
+                    d.ws.wr_block_ = &d;
+                else
+                    BOOST_ASSERT(d.ws.wr_block_ == &d);
                 d.state = do_pong + 1;
-                BOOST_ASSERT(! d.ws.wr_block_);
-                d.ws.wr_block_ = &d;
                 boost::asio::async_write(d.ws.stream_,
                     d.fb.data(), std::move(*this));
                 return;
@@ -541,18 +549,25 @@ operator()(error_code ec,
             //------------------------------------------------------------------
 
             case do_close_resume:
+                BOOST_ASSERT(! d.ws.wr_block_);
+                d.ws.wr_block_ = &d;
                 d.state = do_close_resume + 1;
+                // The current context is safe but might not be
+                // the same as the one for this operation (since
+                // we are being called from a write operation).
+                // Call post to make sure we are invoked the same
+                // way as the final handler for this operation.
                 d.ws.get_io_service().post(bind_handler(
                     std::move(*this), ec, bytes_transferred));
                 return;
 
             case do_close_resume + 1:
+                BOOST_ASSERT(d.ws.wr_block_ == &d);
                 if(d.ws.failed_)
                 {
                     // call handler
-                    d.state = do_call_handler;
                     ec = boost::asio::error::operation_aborted;
-                    break;
+                    goto upcall;
                 }
                 if(d.ws.wr_close_)
                 {
@@ -566,10 +581,12 @@ operator()(error_code ec,
             //------------------------------------------------------------------
 
             case do_close:
+                if(! d.ws.wr_block_)
+                    d.ws.wr_block_ = &d;
+                else
+                    BOOST_ASSERT(d.ws.wr_block_ == &d);
                 d.state = do_teardown;
                 d.ws.wr_close_ = true;
-                BOOST_ASSERT(! d.ws.wr_block_);
-                d.ws.wr_block_ = &d;
                 boost::asio::async_write(d.ws.stream_,
                     d.fb.data(), std::move(*this));
                 return;