Apply PR feedback

Author: OneBlue

src/shared/inc/SocketChannel.h

@@ -110,6 +110,7 @@ class SocketChannel
 
 #ifdef WIN32
         m_exitEvents = std::move(other.m_exitEvents);
+        m_pendingBytes = std::move(other.m_pendingBytes);
 #endif
         m_ignore_sequence = other.m_ignore_sequence;
         m_sent_non_transaction_messages = other.m_sent_non_transaction_messages;
@@ -722,7 +723,7 @@ class SocketChannel
 #ifdef WIN32
 
     std::vector<HANDLE> m_exitEvents;
-    std::optional<std::vector<gsl::byte>> m_pendingBytes;
+    std::vector<gsl::byte> m_pendingBytes;
 
 #endif
     uint32_t m_sent_non_transaction_messages = 0;

src/windows/common/HandleIO.cpp

@@ -38,7 +38,8 @@ DWORD CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)
     {
         if constexpr (std::is_same_v<decltype(Handle), SOCKET>)
         {
-            if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, nullptr))
+            DWORD flagsReturned{};
+            if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, &flagsReturned))
             {
                 auto error = WSAGetLastError();
                 LOG_LAST_ERROR_IF(error != WSAECONNABORTED && error != WSA_OPERATION_ABORTED && error != WSAECONNRESET);
@@ -531,7 +532,7 @@ void HTTPChunkBasedReadHandle::OnRead(const gsl::span<char>& Input)
 ReadSocketMessageHandle::ReadSocketMessageHandle(
     HandleWrapper&& MovedSocket,
     std::vector<gsl::byte>& Buffer,
-    std::optional<std::vector<gsl::byte>>& PendingBytes,
+    std::vector<gsl::byte>& PendingBytes,
     std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage) :
     Socket(std::move(MovedSocket)), Buffer(Buffer), PendingBytes(PendingBytes), OnMessage(std::move(OnMessage))
 {
@@ -542,20 +543,20 @@ ReadSocketMessageHandle::ReadSocketMessageHandle(
         Buffer.resize(sizeof(MESSAGE_HEADER));
     }
 
-    if (!PendingBytes.has_value())
+    if (PendingBytes.empty())
     {
         return;
     }
 
     // If bytes from a previously cancelled transaction are passed, process them now.
-    if (Buffer.size() < PendingBytes->size())
+    if (Buffer.size() < PendingBytes.size())
     {
-        Buffer.resize(PendingBytes->size());
+        Buffer.resize(PendingBytes.size());
     }
 
-    std::copy(PendingBytes->begin(), PendingBytes->end(), Buffer.begin());
-
-    CurrentOffset = PendingBytes->size();
+    std::copy(PendingBytes.begin(), PendingBytes.end(), Buffer.begin());
+    CurrentOffset = PendingBytes.size();
+    PendingBytes.clear();
 
     if (CurrentOffset < sizeof(MESSAGE_HEADER))
     {
@@ -569,24 +570,24 @@ ReadSocketMessageHandle::ReadSocketMessageHandle(
 
 ReadSocketMessageHandle::~ReadSocketMessageHandle()
 {
-    if (State == IOHandleStatus::Completed)
-    {
-        PendingBytes.reset();
-    }
-    else if (State == IOHandleStatus::Pending)
+    if (State != IOHandleStatus::Completed)
     {
-        // Cancel the pending receive and move any bytes already buffered for the in-flight message into PendingBytes
-        const auto socket = reinterpret_cast<SOCKET>(Socket.Get());
-        auto receivedBytes = CancelPendingIo(socket, Overlapped);
+        auto pendingSize = CurrentOffset;
+
+        if (State == IOHandleStatus::Pending)
+        {
+            // Cancel the pending receive and move any bytes already buffered for the in-flight message into PendingBytes
+            const auto socket = reinterpret_cast<SOCKET>(Socket.Get());
+            pendingSize += CancelPendingIo(socket, Overlapped);
+        }
 
-        const auto totalBytes = CurrentOffset + receivedBytes;
-        if (totalBytes > 0)
+        if (pendingSize > 0)
         {
-            WI_ASSERT(totalBytes <= Buffer.size());
-            Buffer.resize(totalBytes);
-            PendingBytes.emplace(std::move(Buffer));
+            WI_ASSERT(pendingSize <= Buffer.size());
+            PendingBytes = {Buffer.begin(), Buffer.begin() + pendingSize};
 
-            WSL_LOG("CanceledMessageRead", TraceLoggingValue(totalBytes, "TotalBytes"), TraceLoggingValue(socket, "Socket"));
+            WSL_LOG(
+                "CanceledMessageRead", TraceLoggingValue(pendingSize, "TotalBytes"), TraceLoggingValue(Socket.Get(), "Socket"));
         }
     }
 }
@@ -663,7 +664,10 @@ bool ReadSocketMessageHandle::ProcessChunk()
         }
 
         ReadingHeader = false;
-        BytesRemaining = messageSize - CurrentOffset;
+        if (CurrentOffset < messageSize)
+        {
+            BytesRemaining = messageSize - CurrentOffset;
+        }
 
         if (BytesRemaining > 0)
         {
@@ -680,14 +684,10 @@ void ReadSocketMessageHandle::Schedule()
 {
     WI_ASSERT(State == IOHandleStatus::Standby);
 
-    // If a previous receive on this socket was aborted while bytes were already in our
-    // application buffer (see destructor), drain those chunks before issuing a new WSARecv.
-    while (State == IOHandleStatus::Standby && BytesRemaining == 0)
+    // Process previously received bytes, if any.
+    if (BytesRemaining == 0 && !ProcessChunk())
     {
-        if (!ProcessChunk())
-        {
-            return;
-        }
+        return; // Message has been fully received, not need to schedule a receive.
     }
 
     ScheduleRecv();

src/windows/common/HandleIO.h

@@ -181,7 +181,7 @@ class ReadSocketMessageHandle : public OverlappedIOHandle
     ReadSocketMessageHandle(
         HandleWrapper&& Socket,
         std::vector<gsl::byte>& Buffer,
-        std::optional<std::vector<gsl::byte>>& PendingBytes,
+        std::vector<gsl::byte>& PendingBytes,
         std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage);
     ~ReadSocketMessageHandle();
 
@@ -196,7 +196,7 @@ class ReadSocketMessageHandle : public OverlappedIOHandle
 
     HandleWrapper Socket;
     std::vector<gsl::byte>& Buffer;
-    std::optional<std::vector<gsl::byte>>& PendingBytes;
+    std::vector<gsl::byte>& PendingBytes;
     std::function<void(const gsl::span<gsl::byte>& Message)> OnMessage;
     wil::unique_event Event{wil::EventOptions::ManualReset};
     OVERLAPPED Overlapped{};

test/windows/UnitTests.cpp

@@ -6842,7 +6842,7 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",
         // Drive a ReadSocketMessageHandle until completion and return the bytes delivered to its
         // OnMessage callback. If a non-success HRESULT is supplied, the call is expected to throw
         // that HRESULT instead, and the OnMessage callback must not be invoked.
-        auto readMessage = [](wil::unique_socket&& server, HRESULT expectedHr = S_OK, std::optional<std::vector<gsl::byte>> pendingBytes = {}) {
+        auto readMessage = [](wil::unique_socket&& server, HRESULT expectedHr = S_OK, std::vector<gsl::byte> pendingBytes = {}) {
             std::vector<gsl::byte> buffer;
             bool callbackInvoked = false;
             std::vector<gsl::byte> message;
@@ -7049,7 +7049,7 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",
         }
 
         // Scenario 10: PendingBytes contains an invalid (too-small) message size. The
-        // constructor should detect this and throw E_UNEXPECTED without invoking OnMessage.
+        // IO should detect this and throw E_UNEXPECTED without invoking OnMessage.
         {
             auto [client, server] = MakeSocketPair();
             client.reset();
@@ -7061,7 +7061,7 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",
             header.TransactionStep = 1;
 
             const auto* headerBytes = reinterpret_cast<const gsl::byte*>(&header);
-            std::optional<std::vector<gsl::byte>> pendingBytes{std::vector<gsl::byte>{headerBytes, headerBytes + sizeof(header)}};
+            std::vector<gsl::byte> pendingBytes{headerBytes, headerBytes + sizeof(header)};
 
             std::vector<gsl::byte> buffer;
             bool callbackInvoked = false;