fix: when releasing web socket stream, more data must result in DataReceived() (#1045)

Author: richardapeters

services/network/WebSocketServerConnectionObserver.cpp

@@ -40,13 +40,14 @@ namespace services
 
     void WebSocketServerConnectionObserver::DataReceived()
     {
-        auto startSize = receiveBuffer.size();
-
         receivingState->DataReceived();
         sendingState->CheckForSomethingToDo();
 
-        if (startSize != receiveBuffer.size())
+        if (moreDataReceived)
+        {
+            moreDataReceived = false;
             services::Connection::Observer().DataReceived();
+        }
     }
 
     void WebSocketServerConnectionObserver::Detaching()
@@ -97,6 +98,12 @@ namespace services
     void WebSocketServerConnectionObserver::ReceiveStreamAllocatable()
     {
         receivingState->DataReceived();
+
+        if (moreDataReceived)
+        {
+            moreDataReceived = false;
+            services::Connection::Observer().DataReceived();
+        }
     }
 
     void WebSocketServerConnectionObserver::SendStreamAllocatable()
@@ -238,6 +245,7 @@ namespace services
 
             while (growSize != 0)
             {
+                connection.moreDataReceived = true;
                 auto range = connection.receiveBuffer.contiguous_range(connection.receiveBuffer.end() - growSize);
                 stream >> range;
                 growSize -= range.size();

services/network/WebSocketServerConnectionObserver.hpp

@@ -199,6 +199,7 @@ namespace services
         infra::SharedPtr<void> keepAliveWhileWriting;
         bool sendBufferReadyForSending = false;
         infra::BoundedVector<uint8_t>::WithMaxSize<8> pongBuffer;
+        bool moreDataReceived = false;
     };
 }
 

services/network/test/TestWebSocketServerConnectionObserver.cpp

@@ -203,8 +203,11 @@ TEST_F(WebSocketServerConnectionObserverTest, receive_frame_with_larger_payload_
     testing::InSequence s;
     EXPECT_CALL(connectionObserver, DataReceived()).WillOnce(testing::Invoke([this]()
         {
+            EXPECT_CALL(connectionObserver, DataReceived()).WillOnce(testing::Invoke([this]()
+                {
+                    CheckDataReceived(std::vector<uint8_t>(8, 0x91));
+                }));
             CheckDataReceived(std::vector<uint8_t>(512, 0x91));
-            CheckDataReceived(std::vector<uint8_t>(8, 0x91));
         }));
 
     connection.SimulateDataReceived(receiveDataPayload);