@samitouri / QOSAMI-WSL-2 / commits / 33dcac06

Fix potential channel corruption after cancelled ReceiveMessageOrClosed() (#40663)

* Fix potential channel corruption after cancelled ReceiveMessageOrClosed() * Apply PR feedback * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Blue committed May 29, 2026 at 11:08 UTC 33dcac06c93741ece6a7d1ff47b217a325d6e8c5

4 files changed +226 -31

src/shared/inc/SocketChannel.h

+3 -2

@@ -110,6 +110,7 @@ public:

110

111 #ifdef WIN32

112 m_exitEvents = std::move(other.m_exitEvents);

113 + m_pendingBytes = std::move(other.m_pendingBytes);

114 #endif

115 m_ignore_sequence = other.m_ignore_sequence;

116 m_sent_non_transaction_messages = other.m_sent_non_transaction_messages;

@@ -636,9 +637,8 @@ private:

637 auto io = CreateIO();

638

639 gsl::span<gsl::byte> message;

639 -

640 io.AddHandle(std::make_unique<windows::common::io::ReadSocketMessageHandle>(

641 - m_socket.get(), m_buffer, [&message](auto& received) { message = received; }));

641 + m_socket.get(), m_buffer, m_pendingBytes, [&message](auto& received) { message = received; }));

642

643 io.Run(TimeoutToMilliseconds(timeout));

644

@@ -723,6 +723,7 @@ private:

723 #ifdef WIN32

724

725 std::vector<HANDLE> m_exitEvents;

726 + std::vector<gsl::byte> m_pendingBytes;

727

728 #endif

729 uint32_t m_sent_non_transaction_messages = 0;

src/windows/common/HandleIO.cpp

+81 -26

@@ -31,14 +31,15 @@ LARGE_INTEGER InitializeFileOffset(HANDLE File)

31 return Offset;

32 }

34 -void CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)

34 +DWORD CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)

35 {

36 DWORD bytesTransferred{};

37 - if (CancelIoEx((HANDLE)Handle, &Overlapped))

37 + if (CancelIoEx((HANDLE)Handle, &Overlapped) || GetLastError() == ERROR_NOT_FOUND)

38 {

39 if constexpr (std::is_same_v<decltype(Handle), SOCKET>)

40 {

41 - if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, nullptr))

41 + DWORD flagsReturned{};

42 + if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, &flagsReturned))

43 {

44 auto error = WSAGetLastError();

45 LOG_LAST_ERROR_IF(error != WSAECONNABORTED && error != WSA_OPERATION_ABORTED && error != WSAECONNRESET);

@@ -56,9 +57,10 @@ void CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)

57 }

58 else

59 {

59 - // ERROR_NOT_FOUND is returned if there was no IO to cancel.

60 - LOG_LAST_ERROR_IF(GetLastError() != ERROR_NOT_FOUND);

60 + LOG_LAST_ERROR_MSG("Unexpected error while cancelling IO on handle: 0x%p", (void*)Handle);

61 }

62 +

63 + return bytesTransferred;

64 }

66 inline void UnregisterWait(HANDLE waitHandle) noexcept

@@ -528,8 +530,11 @@ void HTTPChunkBasedReadHandle::OnRead(const gsl::span<char>& Input)

530 // ReadSocketMessageHandle

531

532 ReadSocketMessageHandle::ReadSocketMessageHandle(

531 - HandleWrapper&& MovedSocket, std::vector<gsl::byte>& Buffer, std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage) :

532 - Socket(std::move(MovedSocket)), Buffer(Buffer), OnMessage(std::move(OnMessage))

533 + HandleWrapper&& MovedSocket,

534 + std::vector<gsl::byte>& Buffer,

535 + std::vector<gsl::byte>& PendingBytes,

536 + std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage) :

537 + Socket(std::move(MovedSocket)), Buffer(Buffer), PendingBytes(PendingBytes), OnMessage(std::move(OnMessage))

538 {

539 Overlapped.hEvent = Event.get();

540

@@ -537,13 +542,53 @@ ReadSocketMessageHandle::ReadSocketMessageHandle(

542 {

543 Buffer.resize(sizeof(MESSAGE_HEADER));

544 }

545 +

546 + if (PendingBytes.empty())

547 + {

548 + return;

549 + }

550 +

551 + // If bytes from a previously cancelled transaction are passed, process them now.

552 + if (Buffer.size() < PendingBytes.size())

553 + {

554 + Buffer.resize(PendingBytes.size());

555 + }

556 +

557 + std::copy(PendingBytes.begin(), PendingBytes.end(), Buffer.begin());

558 + CurrentOffset = PendingBytes.size();

559 + PendingBytes.clear();

560 +

561 + if (CurrentOffset < sizeof(MESSAGE_HEADER))

562 + {

563 + BytesRemaining = sizeof(MESSAGE_HEADER) - CurrentOffset;

564 + }

565 + else

566 + {

567 + BytesRemaining = 0;

568 + }

569 }

570

571 ReadSocketMessageHandle::~ReadSocketMessageHandle()

572 {

544 - if (State == IOHandleStatus::Pending)

573 + if (State != IOHandleStatus::Completed)

574 {

546 - CancelPendingIo((SOCKET)Socket.Get(), Overlapped);

575 + auto pendingSize = CurrentOffset;

576 +

577 + if (State == IOHandleStatus::Pending)

578 + {

579 + // Cancel the pending receive and move any bytes already buffered for the in-flight message into PendingBytes

580 + const auto socket = reinterpret_cast<SOCKET>(Socket.Get());

581 + pendingSize += CancelPendingIo(socket, Overlapped);

582 + }

583 +

584 + if (pendingSize > 0)

585 + {

586 + WI_ASSERT(pendingSize <= Buffer.size());

587 + PendingBytes.assign(Buffer.begin(), Buffer.begin() + pendingSize);

588 +

589 + WSL_LOG(

590 + "CanceledMessageRead", TraceLoggingValue(pendingSize, "TotalBytes"), TraceLoggingValue(Socket.Get(), "Socket"));

591 + }

592 }

593 }

594

@@ -601,40 +646,50 @@ void ReadSocketMessageHandle::ProcessRecvResult(DWORD BytesRead)

646 return;

647 }

648

649 + ProcessChunk();

650 +}

651 +

652 +bool ReadSocketMessageHandle::ProcessChunk()

653 +{

654 + const auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize;

655 +

656 if (ReadingHeader)

657 {

606 - auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize;

607 -

658 THROW_HR_IF_MSG(E_UNEXPECTED, messageSize < sizeof(MESSAGE_HEADER), "Unexpected message size: %u", messageSize);

659 THROW_HR_IF_MSG(E_UNEXPECTED, messageSize > 4 * 1024 * 1024, "Message size too large: %u", messageSize);

660

611 - if (messageSize == sizeof(MESSAGE_HEADER))

612 - {

613 - OnMessage(gsl::make_span(Buffer.data(), messageSize));

614 - State = IOHandleStatus::Completed;

615 - return;

616 - }

617 -

661 if (Buffer.size() < messageSize)

662 {

663 Buffer.resize(messageSize);

664 }

665

666 ReadingHeader = false;

624 - CurrentOffset = sizeof(MESSAGE_HEADER);

625 - BytesRemaining = messageSize - sizeof(MESSAGE_HEADER);

626 - }

627 - else

628 - {

629 - auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize;

630 - OnMessage(gsl::make_span(Buffer.data(), messageSize));

631 - State = IOHandleStatus::Completed;

667 + if (CurrentOffset < messageSize)

668 + {

669 + BytesRemaining = messageSize - CurrentOffset;

670 + }

671 +

672 + if (BytesRemaining > 0)

673 + {

674 + return true;

675 + }

676 }

677 +

678 + OnMessage(gsl::make_span(Buffer.data(), messageSize));

679 + State = IOHandleStatus::Completed;

680 + return false;

681 }

682

683 void ReadSocketMessageHandle::Schedule()

684 {

685 WI_ASSERT(State == IOHandleStatus::Standby);

686 +

687 + // Process previously received bytes, if any.

688 + if (BytesRemaining == 0 && !ProcessChunk())

689 + {

690 + return; // Message has been fully received, no need to schedule a receive.

691 + }

692 +

693 ScheduleRecv();

694 }

695

src/windows/common/HandleIO.h

+7 -1

@@ -178,7 +178,11 @@ public:

178 NON_COPYABLE(ReadSocketMessageHandle);

179 NON_MOVABLE(ReadSocketMessageHandle);

180

181 - ReadSocketMessageHandle(HandleWrapper&& Socket, std::vector<gsl::byte>& Buffer, std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage);

181 + ReadSocketMessageHandle(

182 + HandleWrapper&& Socket,

183 + std::vector<gsl::byte>& Buffer,

184 + std::vector<gsl::byte>& PendingBytes,

185 + std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage);

186 ~ReadSocketMessageHandle();

187

188 void Schedule() override;

@@ -188,9 +192,11 @@ public:

192 private:

193 void ScheduleRecv();

194 void ProcessRecvResult(DWORD BytesRead);

195 + bool ProcessChunk();

196

197 HandleWrapper Socket;

198 std::vector<gsl::byte>& Buffer;

199 + std::vector<gsl::byte>& PendingBytes;

200 std::function<void(const gsl::span<gsl::byte>& Message)> OnMessage;

201 wil::unique_event Event{wil::EventOptions::ManualReset};

202 OVERLAPPED Overlapped{};

test/windows/UnitTests.cpp

+135 -2

@@ -6842,14 +6842,17 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",

6842 // Drive a ReadSocketMessageHandle until completion and return the bytes delivered to its

6843 // OnMessage callback. If a non-success HRESULT is supplied, the call is expected to throw

6844 // that HRESULT instead, and the OnMessage callback must not be invoked.

6845 - auto readMessage = [](wil::unique_socket&& server, HRESULT expectedHr = S_OK) {

6845 + auto readMessage = [](wil::unique_socket&& server, HRESULT expectedHr = S_OK, std::vector<gsl::byte> pendingBytes = {}) {

6846 std::vector<gsl::byte> buffer;

6847 bool callbackInvoked = false;

6848 std::vector<gsl::byte> message;

6849

6850 wsl::windows::common::io::MultiHandleWait io;

6851 io.AddHandle(std::make_unique<wsl::windows::common::io::ReadSocketMessageHandle>(

6852 - wsl::windows::common::io::HandleWrapper{std::move(server)}, buffer, [&callbackInvoked, &message](const gsl::span<gsl::byte>& received) {

6852 + wsl::windows::common::io::HandleWrapper{std::move(server)},

6853 + buffer,

6854 + pendingBytes,

6855 + [&callbackInvoked, &message](const gsl::span<gsl::byte>& received) {

6856 callbackInvoked = true;

6857 message.assign(received.begin(), received.end());

6858 }));

@@ -6944,6 +6947,136 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",

6947

6948 readMessage(std::move(server), E_UNEXPECTED);

6949 }

6950 +

6951 + // Scenario 6: PendingBytes carries a complete header-only message left over from a

6952 + // previous aborted receive. The reader should deliver it without touching the socket.

6953 + {

6954 + auto [client, server] = MakeSocketPair();

6955 + client.reset(); // close the peer; we should still complete from PendingBytes alone.

6956 +

6957 + MESSAGE_HEADER header{};

6958 + header.MessageType = LxMiniInitMessageAny;

6959 + header.MessageSize = sizeof(header);

6960 + header.TransactionId = 77;

6961 + header.TransactionStep = 1;

6962 +

6963 + const auto* headerBytes = reinterpret_cast<const gsl::byte*>(&header);

6964 + std::vector<gsl::byte> pendingBytes(headerBytes, headerBytes + sizeof(header));

6965 +

6966 + const auto message = readMessage(std::move(server), S_OK, std::move(pendingBytes));

6967 + VERIFY_ARE_EQUAL(message.size(), sizeof(header));

6968 + VERIFY_IS_TRUE(std::memcmp(message.data(), &header, sizeof(header)) == 0);

6969 + }

6970 +

6971 + // Scenario 7: PendingBytes carries a complete message with a body left over from a

6972 + // previous aborted receive. The reader should deliver it without touching the socket.

6973 + {

6974 + auto [client, server] = MakeSocketPair();

6975 + client.reset();

6976 +

6977 + constexpr size_t bodySize = 32;

6978 + std::vector<gsl::byte> payload(sizeof(MESSAGE_HEADER) + bodySize);

6979 + auto* header = reinterpret_cast<MESSAGE_HEADER*>(payload.data());

6980 + header->MessageType = LxMiniInitMessageAny;

6981 + header->MessageSize = gsl::narrow_cast<unsigned int>(payload.size());

6982 + header->TransactionId = 81;

6983 + header->TransactionStep = 3;

6984 + for (size_t i = 0; i < bodySize; ++i)

6985 + {

6986 + payload[sizeof(MESSAGE_HEADER) + i] = static_cast<gsl::byte>(i ^ 0xA5);

6987 + }

6988 +

6989 + std::vector<gsl::byte> pendingBytes(payload.begin(), payload.end());

6990 +

6991 + const auto message = readMessage(std::move(server), S_OK, std::move(pendingBytes));

6992 + VERIFY_ARE_EQUAL(message.size(), payload.size());

6993 + VERIFY_IS_TRUE(std::memcmp(message.data(), payload.data(), payload.size()) == 0);

6994 + }

6995 +

6996 + // Scenario 8: PendingBytes carries only part of a header. The reader must fill in the

6997 + // rest of the header (and the body) from the socket and deliver the assembled message.

6998 + {

6999 + auto [client, server] = MakeSocketPair();

7000 +

7001 + constexpr size_t bodySize = 48;

7002 + constexpr size_t prebufferedBytes = 6; // less than sizeof(MESSAGE_HEADER) = 16

7003 + std::vector<gsl::byte> payload(sizeof(MESSAGE_HEADER) + bodySize);

7004 + auto* header = reinterpret_cast<MESSAGE_HEADER*>(payload.data());

7005 + header->MessageType = LxMiniInitMessageAny;

7006 + header->MessageSize = gsl::narrow_cast<unsigned int>(payload.size());

7007 + header->TransactionId = 91;

7008 + header->TransactionStep = 4;

7009 + for (size_t i = 0; i < bodySize; ++i)

7010 + {

7011 + payload[sizeof(MESSAGE_HEADER) + i] = static_cast<gsl::byte>(0xC3);

7012 + }

7013 +

7014 + std::vector<gsl::byte> pendingBytes(payload.begin(), payload.begin() + prebufferedBytes);

7015 + WriteSocket(client.get(), payload.data() + prebufferedBytes, payload.size() - prebufferedBytes);

7016 + client.reset();

7017 +

7018 + const auto message = readMessage(std::move(server), S_OK, std::move(pendingBytes));

7019 + VERIFY_ARE_EQUAL(message.size(), payload.size());

7020 + VERIFY_IS_TRUE(std::memcmp(message.data(), payload.data(), payload.size()) == 0);

7021 + }

7022 +

7023 + // Scenario 9: PendingBytes carries the full header plus part of the body. The reader

7024 + // must read the remaining body bytes from the socket and deliver the assembled message.

7025 + {

7026 + auto [client, server] = MakeSocketPair();

7027 +

7028 + constexpr size_t bodySize = 64;

7029 + constexpr size_t prebufferedBodyBytes = 12;

7030 + std::vector<gsl::byte> payload(sizeof(MESSAGE_HEADER) + bodySize);

7031 + auto* header = reinterpret_cast<MESSAGE_HEADER*>(payload.data());

7032 + header->MessageType = LxMiniInitMessageAny;

7033 + header->MessageSize = gsl::narrow_cast<unsigned int>(payload.size());

7034 + header->TransactionId = 92;

7035 + header->TransactionStep = 5;

7036 + for (size_t i = 0; i < bodySize; ++i)

7037 + {

7038 + payload[sizeof(MESSAGE_HEADER) + i] = static_cast<gsl::byte>(i & 0xFF);

7039 + }

7040 +

7041 + const size_t prebufferedBytes = sizeof(MESSAGE_HEADER) + prebufferedBodyBytes;

7042 + std::vector<gsl::byte> pendingBytes(payload.begin(), payload.begin() + prebufferedBytes);

7043 + WriteSocket(client.get(), payload.data() + prebufferedBytes, payload.size() - prebufferedBytes);

7044 + client.reset();

7045 +

7046 + const auto message = readMessage(std::move(server), S_OK, std::move(pendingBytes));

7047 + VERIFY_ARE_EQUAL(message.size(), payload.size());

7048 + VERIFY_IS_TRUE(std::memcmp(message.data(), payload.data(), payload.size()) == 0);

7049 + }

7050 +

7051 + // Scenario 10: PendingBytes contains an invalid (too-small) message size. The

7052 + // IO should detect this and throw E_UNEXPECTED without invoking OnMessage.

7053 + {

7054 + auto [client, server] = MakeSocketPair();

7055 + client.reset();

7056 +

7057 + MESSAGE_HEADER header{};

7058 + header.MessageType = LxMiniInitMessageAny;

7059 + header.MessageSize = sizeof(header) - 1; // invalid: smaller than the header itself

7060 + header.TransactionId = 99;

7061 + header.TransactionStep = 1;

7062 +

7063 + const auto* headerBytes = reinterpret_cast<const gsl::byte*>(&header);

7064 + std::vector<gsl::byte> pendingBytes{headerBytes, headerBytes + sizeof(header)};

7065 +

7066 + std::vector<gsl::byte> buffer;

7067 + bool callbackInvoked = false;

7068 + const auto hr = wil::ResultFromException([&]() {

7069 + wsl::windows::common::io::MultiHandleWait io;

7070 + io.AddHandle(std::make_unique<wsl::windows::common::io::ReadSocketMessageHandle>(

7071 + wsl::windows::common::io::HandleWrapper{std::move(server)},

7072 + buffer,

7073 + pendingBytes,

7074 + [&callbackInvoked](const gsl::span<gsl::byte>&) { callbackInvoked = true; }));

7075 + io.Run(std::chrono::seconds(60));

7076 + });

7077 + VERIFY_ARE_EQUAL(hr, E_UNEXPECTED);

7078 + VERIFY_IS_FALSE(callbackInvoked);

7079 + }

7080 }

7081

7082 TEST_METHOD(MultiHandleWaitAboveMaximumWaitObjects)