Update the VM termination logic to enforce timeouts and avoid hanging if init is stuck during termination (#40431)

* Update the VM termination logic to enforce timeouts and avoid hang if init is stuck during session termination * Save state * Save state * Rethink IO logic * Fix build * Apply PR feedback * Fix error check * Apply PR feedback * Apply PR feedback

Blue committed May 11, 2026 at 20:32 UTC 818ae873a1e6cf65e0c928f250a54d96a030fb3c
19 files changed +791 -124
src/linux/init/WSLCInit.cpp
+1 -1
@@ -1115,4 +1115,4 @@ int WSLCEntryPoint(int Argc, char* Argv[])
1115 reboot(RB_POWER_OFF);
1116
1117 return 0;
1118 -}
\ No newline at end of file
1118 +}
src/shared/inc/SocketChannel.h
+132 -32
@@ -63,21 +63,31 @@ public:
63 void SendResultMessage(TResult value);
64
65 template <typename TMessage>
66 - std::pair<TMessage*, gsl::span<gsl::byte>> ReceiveOrClosed(TTimeout timeout = DefaultSocketTimeout);
66 + std::pair<TMessage*, gsl::span<gsl::byte>> ReceiveOrClosed();
67
68 template <typename TMessage>
69 - TMessage& Receive(gsl::span<gsl::byte>* responseSpan = nullptr, TTimeout timeout = DefaultSocketTimeout);
69 + TMessage& Receive(gsl::span<gsl::byte>* responseSpan = nullptr);
70
71 private:
72 - Transaction(SocketChannel& channel, uint32_t id) :
73 - m_channel(channel), m_id(id), m_step(static_cast<uint32_t>(TRANSACTION_STEP::REQUEST))
72 + Transaction(SocketChannel& channel, uint32_t id, TTimeout timeout) :
73 + m_channel(channel), m_id(id), m_step(static_cast<uint32_t>(TRANSACTION_STEP::REQUEST)), m_deadline(ComputeDeadline(timeout))
74 {
75 }
76
77 + TTimeout RemainingTimeout() const;
78 +
79 + static std::optional<std::chrono::steady_clock::time_point> ComputeDeadline(TTimeout timeout);
80 +
81 SocketChannel& m_channel;
82 uint32_t m_id;
83 /** Use uint32_t as step can go beyond FIRST_REPLY */
84 uint32_t m_step;
85 + std::optional<std::chrono::steady_clock::time_point> m_deadline;
86 +
87 +#ifndef WIN32
88 + // This is required because the Linux timeout logic requires a pointer. Returning a pointer is OK because only thread can use a transaction at a given time.
89 + mutable timeval m_timeoutStorage{};
90 +#endif
91 };
92
93 class SocketChannel
@@ -99,7 +109,7 @@ public:
109 m_socket = std::move(other.m_socket);
110
111 #ifdef WIN32
102 - m_exitEvent = std::move(other.m_exitEvent);
112 + m_exitEvents = std::move(other.m_exitEvents);
113 #endif
114 m_ignore_sequence = other.m_ignore_sequence;
115 m_sent_non_transaction_messages = other.m_sent_non_transaction_messages;
@@ -115,15 +125,27 @@ public:
125
126 #ifdef WIN32
127
118 - SocketChannel(TSocket&& socket, std::string&& name, HANDLE exitEvent) :
119 - m_socket(std::move(socket)), m_exitEvent(exitEvent), m_name(std::move(name))
128 + SocketChannel(TSocket&& socket, std::string&& name, std::vector<HANDLE>&& exitEvents) :
129 + m_socket(std::move(socket)), m_exitEvents(std::move(exitEvents)), m_name(std::move(name))
130 + {
131 + }
132 +
133 + std::vector<HANDLE> SetExitEvents(std::vector<HANDLE>&& exitEvents)
134 + {
135 + std::vector<HANDLE> oldEvents;
136 + std::swap(oldEvents, m_exitEvents);
137 + return oldEvents;
138 + }
139 +
140 + const std::vector<HANDLE>& GetExitEvents() const
141 {
142 + return m_exitEvents;
143 }
144
145 #endif
146
147 template <typename TMessage>
126 - void SendMessage(gsl::span<gsl::byte> span, uint32_t transactionStep = static_cast<uint32_t>(TRANSACTION_STEP::NONE), uint32_t transactionId = 0)
148 + void SendMessage(gsl::span<gsl::byte> span, uint32_t transactionStep = static_cast<uint32_t>(TRANSACTION_STEP::NONE), uint32_t transactionId = 0, TTimeout timeout = DefaultSocketTimeout)
149 {
150 // Ensure that no other thread is using this channel.
151 const std::unique_lock<std::mutex> lock{m_sendMutex, std::try_to_lock};
@@ -161,13 +183,15 @@ public:
183
184 #ifdef WIN32
185
164 - auto sentBytes = wsl::windows::common::socket::Send(m_socket.get(), span, m_exitEvent);
186 + auto io = CreateIO();
187 + io.AddHandle(std::make_unique<windows::common::relay::WriteHandle>(m_socket.get(), span));
188 +
189 + io.Run(TimeoutToMilliseconds(timeout));
190
191 WSL_LOG(
192 "SentMessage",
193 TraceLoggingValue(m_name.c_str(), "Name"),
169 - TraceLoggingValue(reinterpret_cast<const TMessage*>(span.data())->PrettyPrint().c_str(), "Content"),
170 - TraceLoggingValue(sentBytes, "SentBytes"));
194 + TraceLoggingValue(reinterpret_cast<const TMessage*>(span.data())->PrettyPrint().c_str(), "Content"));
195
196 #else
197
@@ -206,7 +230,7 @@ public:
230 }
231
232 template <typename TMessage>
209 - void SendMessage(TMessage& message, uint32_t transactionStep = static_cast<uint32_t>(TRANSACTION_STEP::NONE), uint32_t transactionId = 0)
233 + void SendMessage(TMessage& message, uint32_t transactionStep = static_cast<uint32_t>(TRANSACTION_STEP::NONE), uint32_t transactionId = 0, TTimeout timeout = DefaultSocketTimeout)
234 {
235 // Catch situations where the other SendMessage() method should be used
236 const auto& header = GetMessageHeader(message);
@@ -220,7 +244,7 @@ public:
244 #endif
245 }
246
223 - SendMessage<TMessage>(gslhelpers::struct_as_writeable_bytes(message), transactionStep, transactionId);
247 + SendMessage<TMessage>(gslhelpers::struct_as_writeable_bytes(message), transactionStep, transactionId, timeout);
248 }
249
250 template <typename TResult>
@@ -267,7 +291,7 @@ public:
291 m_received_non_transaction_messages++;
292 }
293
270 - receivedSpan = ReceiveImpl(TMessage::Type, timeout);
294 + receivedSpan = ReceiveImpl(timeout);
295 if (receivedSpan.empty())
296 {
297
@@ -508,24 +532,24 @@ public:
532 return *message;
533 }
534
511 - Transaction StartTransaction()
535 + Transaction StartTransaction(TTimeout timeout = DefaultSocketTimeout)
536 {
537 uint32_t transactionId = m_transaction_id_seed++;
514 - return wsl::shared::Transaction(*this, transactionId);
538 + return wsl::shared::Transaction(*this, transactionId, timeout);
539 }
540
517 - Transaction ReceiveTransaction()
541 + Transaction ReceiveTransaction(TTimeout timeout = DefaultSocketTimeout)
542 {
543 // Transaction id should follow the received one on the receive end.
520 - return wsl::shared::Transaction(*this, 0);
544 + return wsl::shared::Transaction(*this, 0, timeout);
545 }
546
547 template <typename TSentMessage>
548 typename TSentMessage::TResponse& Transaction(gsl::span<gsl::byte> message, gsl::span<gsl::byte>* responseSpan = nullptr, TTimeout timeout = DefaultSocketTimeout)
549 {
526 - auto transaction = StartTransaction();
527 - transaction.Send<TSentMessage>(message);
528 - return transaction.Receive<typename TSentMessage::TResponse>(responseSpan, timeout);
550 + auto transaction = StartTransaction(timeout);
551 + transaction.template Send<TSentMessage>(message);
552 + return transaction.template Receive<typename TSentMessage::TResponse>(responseSpan);
553 }
554
555 template <typename TSentMessage>
@@ -537,12 +561,12 @@ public:
561 }
562
563 template <typename TSentMessage>
540 - TSentMessage::TResponse& Transaction()
564 + TSentMessage::TResponse& Transaction(TTimeout timeout = DefaultSocketTimeout)
565 {
566 TSentMessage message{};
567 message.Header.MessageSize = sizeof(message);
568 message.Header.MessageType = TSentMessage::Type;
545 - return Transaction<TSentMessage>(message);
569 + return Transaction<TSentMessage>(message, nullptr, timeout);
570 }
571
572 void Close()
@@ -581,15 +605,49 @@ public:
605
606 private:
607 #ifdef WIN32
608 + windows::common::relay::MultiHandleWait CreateIO() const
609 + {
610 + wsl::windows::common::relay::MultiHandleWait io;
611
585 - gsl::span<gsl::byte> ReceiveImpl(auto expectedMessage, TTimeout timeout)
612 + for (const auto event : m_exitEvents)
613 + {
614 + io.AddHandle(
615 + std::make_unique<windows::common::relay::EventHandle>(
616 + event,
617 + [this, event]() { THROW_HR_MSG(E_ABORT, "Exit event 0x%p signaled on channel: %hs", event, m_name.c_str()); }),
618 + windows::common::relay::MultiHandleWait::CancelOnCompleted | windows::common::relay::MultiHandleWait::NeedNotComplete);
619 + }
620 +
621 + return io;
622 + }
623 +
624 + static std::optional<std::chrono::milliseconds> TimeoutToMilliseconds(TTimeout timeout)
625 {
587 - return wsl::shared::socket::RecvMessage(m_socket.get(), m_buffer, m_exitEvent, timeout);
626 + if (timeout == INFINITE)
627 + {
628 + return std::nullopt;
629 + }
630 +
631 + return std::chrono::milliseconds{timeout};
632 + }
633 +
634 + gsl::span<gsl::byte> ReceiveImpl(TTimeout timeout)
635 + {
636 + auto io = CreateIO();
637 +
638 + gsl::span<gsl::byte> message;
639 +
640 + io.AddHandle(std::make_unique<windows::common::relay::ReadSocketMessageHandle>(
641 + m_socket.get(), m_buffer, [&message](auto& received) { message = received; }));
642 +
643 + io.Run(TimeoutToMilliseconds(timeout));
644 +
645 + return message;
646 }
647
648 #else
649
592 - gsl::span<gsl::byte> ReceiveImpl(auto expectedMessage, TTimeout timeout)
650 + gsl::span<gsl::byte> ReceiveImpl(TTimeout timeout)
651 {
652 return wsl::shared::socket::RecvMessage(m_socket.get(), m_buffer, timeout);
653 }
@@ -664,7 +722,7 @@ private:
722
723 #ifdef WIN32
724
667 - HANDLE m_exitEvent{};
725 + std::vector<HANDLE> m_exitEvents;
726
727 #endif
728 uint32_t m_sent_non_transaction_messages = 0;
@@ -676,10 +734,52 @@ private:
734 std::mutex m_receiveMutex;
735 };
736
737 +inline std::optional<std::chrono::steady_clock::time_point> Transaction::ComputeDeadline(TTimeout timeout)
738 +{
739 +#ifdef WIN32
740 + if (timeout == INFINITE)
741 + {
742 + return std::nullopt;
743 + }
744 +
745 + return std::chrono::steady_clock::now() + std::chrono::milliseconds{timeout};
746 +#else
747 + if (timeout == nullptr)
748 + {
749 + return std::nullopt;
750 + }
751 +
752 + return std::chrono::steady_clock::now() + std::chrono::seconds{timeout->tv_sec} + std::chrono::microseconds{timeout->tv_usec};
753 +#endif
754 +}
755 +
756 +inline TTimeout Transaction::RemainingTimeout() const
757 +{
758 + if (!m_deadline.has_value())
759 + {
760 + return DefaultSocketTimeout;
761 + }
762 +
763 +#ifdef WIN32
764 +
765 + auto remaining = std::chrono::duration_cast<std::chrono::milliseconds>(*m_deadline - std::chrono::steady_clock::now());
766 + return remaining.count() > 0 ? static_cast<DWORD>(remaining.count()) : 0;
767 +
768 +#else
769 +
770 + auto remaining = std::chrono::duration_cast<std::chrono::milliseconds>(*m_deadline - std::chrono::steady_clock::now());
771 +
772 + m_timeoutStorage.tv_sec = static_cast<time_t>(remaining.count() / 1000);
773 + m_timeoutStorage.tv_usec = static_cast<suseconds_t>((remaining.count() % 1000) * 1000);
774 + return &m_timeoutStorage;
775 +
776 +#endif
777 +}
778 +
779 template <typename TMessage>
780 void Transaction::Send(gsl::span<gsl::byte> span)
781 {
682 - m_channel.SendMessage<TMessage>(span, m_step, m_id);
782 + m_channel.SendMessage<TMessage>(span, m_step, m_id, RemainingTimeout());
783 m_step++;
784 }
785
@@ -701,9 +801,9 @@ void Transaction::SendResultMessage(TResult value)
801 }
802
803 template <typename TMessage>
704 -std::pair<TMessage*, gsl::span<gsl::byte>> Transaction::ReceiveOrClosed(TTimeout timeout)
804 +std::pair<TMessage*, gsl::span<gsl::byte>> Transaction::ReceiveOrClosed()
805 {
706 - auto result = m_channel.ReceiveMessageOrClosed<TMessage>(timeout, m_step, m_id);
806 + auto result = m_channel.ReceiveMessageOrClosed<TMessage>(RemainingTimeout(), m_step, m_id);
807 if (m_step == static_cast<uint32_t>(TRANSACTION_STEP::REQUEST) && result.first != nullptr)
808 {
809 // Use the request's id for the reply side transaction.
@@ -715,9 +815,9 @@ std::pair<TMessage*, gsl::span<gsl::byte>> Transaction::ReceiveOrClosed(TTimeout
815 }
816
817 template <typename TMessage>
718 -TMessage& Transaction::Receive(gsl::span<gsl::byte>* responseSpan, TTimeout timeout)
818 +TMessage& Transaction::Receive(gsl::span<gsl::byte>* responseSpan)
819 {
720 - auto& message = m_channel.ReceiveMessage<TMessage>(responseSpan, timeout, m_step, m_id);
820 + auto& message = m_channel.ReceiveMessage<TMessage>(responseSpan, RemainingTimeout(), m_step, m_id);
821 if (m_step == static_cast<uint32_t>(TRANSACTION_STEP::REQUEST))
822 {
823 // Use the request's id for the reply side transaction.
src/windows/common/DnsTunnelingChannel.cpp
+1 -1
@@ -6,7 +6,7 @@
6 using wsl::core::networking::DnsTunnelingChannel;
7
8 DnsTunnelingChannel::DnsTunnelingChannel(wil::unique_socket&& socket, DnsTunnelingCallback&& reportDnsRequest) :
9 - m_channel{std::move(socket), "DnsTunneling", m_stopEvent.get()}, m_reportDnsRequest(std::move(reportDnsRequest))
9 + m_channel{std::move(socket), "DnsTunneling", {m_stopEvent.get()}}, m_reportDnsRequest(std::move(reportDnsRequest))
10 {
11 WSL_LOG("DnsTunnelingChannel::DnsTunnelingChannel [Windows]", TraceLoggingValue(m_channel.Socket(), "socket"));
12
src/windows/common/GnsChannel.cpp
+1 -1
@@ -7,7 +7,7 @@
7 using namespace wsl::shared;
8 using wsl::core::GnsChannel;
9
10 -GnsChannel::GnsChannel(wil::unique_socket&& socket) : m_channel(std::move(socket), "GNS", m_stopEvent.get())
10 +GnsChannel::GnsChannel(wil::unique_socket&& socket) : m_channel(std::move(socket), "GNS", {m_stopEvent.get()})
11 {
12 WSL_LOG("GnsChannel::GnsChannel", TraceLoggingValue(m_channel.Socket(), "socket"));
13 }
src/windows/common/GnsPortTrackerChannel.cpp
+1 -1
@@ -12,7 +12,7 @@ GnsPortTrackerChannel::GnsPortTrackerChannel(
12 const std::function<void(const std::string&, bool)>& InterfaceStateCallback) :
13 m_callback(Callback),
14 m_interfaceStateCallback(InterfaceStateCallback),
15 - m_channel(std::move(Socket), "GNSPortTracker", m_stopEvent.get())
15 + m_channel(std::move(Socket), "GNSPortTracker", {m_stopEvent.get()})
16 {
17 m_thread = std::thread{std::bind(&GnsPortTrackerChannel::Run, this)};
18 }
src/windows/common/relay.cpp
+187 -37
@@ -25,6 +25,7 @@ using wsl::windows::common::relay::LineBasedReadHandle;
25 using wsl::windows::common::relay::MultiHandleWait;
26 using wsl::windows::common::relay::OverlappedIOHandle;
27 using wsl::windows::common::relay::ReadHandle;
28 +using wsl::windows::common::relay::ReadSocketMessageHandle;
29 using wsl::windows::common::relay::RelayHandle;
30 using wsl::windows::common::relay::ScopedMultiRelay;
31 using wsl::windows::common::relay::ScopedRelay;
@@ -44,6 +45,36 @@ LARGE_INTEGER InitializeFileOffset(HANDLE File)
45 return Offset;
46 }
47
48 +void CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)
49 +{
50 + DWORD bytesTransferred{};
51 + if (CancelIoEx((HANDLE)Handle, &Overlapped))
52 + {
53 + if constexpr (std::is_same_v<decltype(Handle), SOCKET>)
54 + {
55 + if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, nullptr))
56 + {
57 + auto error = WSAGetLastError();
58 + LOG_LAST_ERROR_IF(error != WSAECONNABORTED && error != WSA_OPERATION_ABORTED && error != WSAECONNRESET);
59 + }
60 + }
61 + else
62 + {
63 + static_assert(std::is_same_v<decltype(Handle), HANDLE>);
64 + if (!GetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true))
65 + {
66 + auto error = GetLastError();
67 + LOG_LAST_ERROR_IF(error != ERROR_CONNECTION_ABORTED && error != ERROR_OPERATION_ABORTED);
68 + }
69 + }
70 + }
71 + else
72 + {
73 + // ERROR_NOT_FOUND is returned if there was no IO to cancel.
74 + LOG_LAST_ERROR_IF(GetLastError() != ERROR_NOT_FOUND);
75 + }
76 +}
77 +
78 } // namespace
79
80 std::thread wsl::windows::common::relay::CreateThread(_In_ HANDLE InputHandle, _In_ HANDLE OutputHandle, _In_opt_ HANDLE ExitHandle, _In_ size_t BufferSize)
@@ -1128,20 +1159,7 @@ ReadHandle::~ReadHandle()
1159 {
1160 if (State == IOHandleStatus::Pending)
1161 {
1131 - DWORD bytesRead{};
1132 - if (CancelIoEx(Handle.Get(), &Overlapped))
1133 - {
1134 - if (!GetOverlappedResult(Handle.Get(), &Overlapped, &bytesRead, true))
1135 - {
1136 - auto error = GetLastError();
1137 - LOG_LAST_ERROR_IF(error != ERROR_CONNECTION_ABORTED && error != ERROR_OPERATION_ABORTED);
1138 - }
1139 - }
1140 - else
1141 - {
1142 - // ERROR_NOT_FOUND is returned if there was no IO to cancel.
1143 - LOG_LAST_ERROR_IF(GetLastError() != ERROR_NOT_FOUND);
1144 - }
1162 + CancelPendingIo(Handle.Get(), Overlapped);
1163 }
1164 }
1165
@@ -1155,12 +1173,13 @@ void ReadHandle::Schedule()
1173 DWORD bytesRead{};
1174 Overlapped.Offset = Offset.LowPart;
1175 Overlapped.OffsetHigh = Offset.HighPart;
1158 - if (ReadFile(Handle.Get(), Buffer.data(), static_cast<DWORD>(Buffer.size()), &bytesRead, &Overlapped))
1176 + auto* bufferData = reinterpret_cast<char*>(Buffer.Span().data());
1177 + if (ReadFile(Handle.Get(), bufferData, static_cast<DWORD>(Buffer.Size()), &bytesRead, &Overlapped))
1178 {
1179 Offset.QuadPart += bytesRead;
1180
1181 // Signal the read.
1163 - OnRead(gsl::make_span<char>(Buffer.data(), static_cast<size_t>(bytesRead)));
1182 + OnRead(gsl::make_span<char>(bufferData, static_cast<size_t>(bytesRead)));
1183
1184 // ReadFile completed immediately, process the result right away.
1185 if (bytesRead == 0)
@@ -1211,7 +1230,7 @@ void ReadHandle::Collect()
1230 Offset.QuadPart += bytesRead;
1231
1232 // Signal the read.
1214 - OnRead(gsl::make_span<char>(Buffer.data(), static_cast<size_t>(bytesRead)));
1233 + OnRead(gsl::make_span<char>(reinterpret_cast<char*>(Buffer.Span().data()), static_cast<size_t>(bytesRead)));
1234
1235 // Transition to Complete if this was a zero byte read.
1236 if (bytesRead == 0)
@@ -1442,30 +1461,159 @@ void HTTPChunkBasedReadHandle::OnRead(const gsl::span<char>& Input)
1461 }
1462 }
1463
1445 -WriteHandle::WriteHandle(HandleWrapper&& MovedHandle, const std::vector<char>& Buffer) :
1446 - Handle(std::move(MovedHandle)), Buffer(Buffer), Offset(InitializeFileOffset(Handle.Get()))
1464 +ReadSocketMessageHandle::ReadSocketMessageHandle(
1465 + HandleWrapper&& MovedSocket, std::vector<gsl::byte>& Buffer, std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage) :
1466 + Socket(std::move(MovedSocket)), Buffer(Buffer), OnMessage(std::move(OnMessage))
1467 {
1468 Overlapped.hEvent = Event.get();
1469 +
1470 + if (Buffer.size() < sizeof(MESSAGE_HEADER))
1471 + {
1472 + Buffer.resize(sizeof(MESSAGE_HEADER));
1473 + }
1474 }
1475
1451 -WriteHandle::~WriteHandle()
1476 +ReadSocketMessageHandle::~ReadSocketMessageHandle()
1477 {
1478 if (State == IOHandleStatus::Pending)
1479 {
1455 - DWORD bytesRead{};
1456 - if (CancelIoEx(Handle.Get(), &Overlapped))
1480 + CancelPendingIo((SOCKET)Socket.Get(), Overlapped);
1481 + }
1482 +}
1483 +
1484 +void ReadSocketMessageHandle::ScheduleRecv()
1485 +{
1486 + Event.ResetEvent();
1487 +
1488 + auto target = gsl::make_span(Buffer).subspan(CurrentOffset, BytesRemaining);
1489 + WSABUF wsaBuf = {gsl::narrow_cast<ULONG>(target.size()), reinterpret_cast<CHAR*>(target.data())};
1490 + DWORD bytesRead{};
1491 + DWORD flags = 0;
1492 + if (WSARecv(reinterpret_cast<SOCKET>(Socket.Get()), &wsaBuf, 1, &bytesRead, &flags, &Overlapped, nullptr) == 0)
1493 + {
1494 + ProcessRecvResult(bytesRead);
1495 + }
1496 + else
1497 + {
1498 + auto error = WSAGetLastError();
1499 + if (error == WSAECONNABORTED || error == WSAECONNRESET)
1500 {
1458 - if (!GetOverlappedResult(Handle.Get(), &Overlapped, &bytesRead, true))
1459 - {
1460 - auto error = GetLastError();
1461 - LOG_LAST_ERROR_IF(error != ERROR_CONNECTION_ABORTED && error != ERROR_OPERATION_ABORTED);
1462 - }
1501 + ProcessRecvResult(0);
1502 + return;
1503 }
1464 - else
1504 +
1505 + THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != WSA_IO_PENDING, "Socket: 0x%p", (void*)Socket.Get());
1506 +
1507 + State = IOHandleStatus::Pending;
1508 + }
1509 +}
1510 +
1511 +void ReadSocketMessageHandle::ProcessRecvResult(DWORD BytesRead)
1512 +{
1513 + if (BytesRead == 0)
1514 + {
1515 + // If the socket was closed before any bytes of the next message were read, signal a clean end-of-stream.
1516 + // If some bytes were already buffered, the peer closed mid-message which is a protocol error.
1517 + THROW_HR_IF_MSG(
1518 + E_UNEXPECTED,
1519 + CurrentOffset > 0,
1520 + "Socket closed before a complete message could be read. ReadingHeader: %d, CurrentOffset: %zu, BytesRemaining: %zu",
1521 + ReadingHeader,
1522 + CurrentOffset,
1523 + BytesRemaining);
1524 +
1525 + OnMessage({});
1526 + State = IOHandleStatus::Completed;
1527 + return;
1528 + }
1529 +
1530 + CurrentOffset += BytesRead;
1531 + BytesRemaining -= BytesRead;
1532 +
1533 + if (BytesRemaining > 0)
1534 + {
1535 + return;
1536 + }
1537 +
1538 + if (ReadingHeader)
1539 + {
1540 + auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize;
1541 +
1542 + THROW_HR_IF_MSG(E_UNEXPECTED, messageSize < sizeof(MESSAGE_HEADER), "Unexpected message size: %u", messageSize);
1543 + THROW_HR_IF_MSG(E_UNEXPECTED, messageSize > 4 * 1024 * 1024, "Message size too large: %u", messageSize);
1544 +
1545 + if (messageSize == sizeof(MESSAGE_HEADER))
1546 {
1466 - // ERROR_NOT_FOUND is returned if there was no IO to cancel.
1467 - LOG_LAST_ERROR_IF(GetLastError() != ERROR_NOT_FOUND);
1547 + OnMessage(gsl::make_span(Buffer.data(), messageSize));
1548 + State = IOHandleStatus::Completed;
1549 + return;
1550 }
1551 +
1552 + if (Buffer.size() < messageSize)
1553 + {
1554 + Buffer.resize(messageSize);
1555 + }
1556 +
1557 + ReadingHeader = false;
1558 + CurrentOffset = sizeof(MESSAGE_HEADER);
1559 + BytesRemaining = messageSize - sizeof(MESSAGE_HEADER);
1560 + }
1561 + else
1562 + {
1563 + auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize;
1564 + OnMessage(gsl::make_span(Buffer.data(), messageSize));
1565 + State = IOHandleStatus::Completed;
1566 + }
1567 +}
1568 +
1569 +void ReadSocketMessageHandle::Schedule()
1570 +{
1571 + WI_ASSERT(State == IOHandleStatus::Standby);
1572 + ScheduleRecv();
1573 +}
1574 +
1575 +void ReadSocketMessageHandle::Collect()
1576 +{
1577 + WI_ASSERT(State == IOHandleStatus::Pending);
1578 +
1579 + State = IOHandleStatus::Standby;
1580 +
1581 + DWORD bytesRead{};
1582 + DWORD flags{};
1583 + if (!WSAGetOverlappedResult(reinterpret_cast<SOCKET>(Socket.Get()), &Overlapped, &bytesRead, FALSE, &flags))
1584 + {
1585 + long error = WSAGetLastError();
1586 + THROW_WIN32_IF(error, error != WSAECONNABORTED && error != WSAECONNRESET);
1587 +
1588 + WI_ASSERT(bytesRead == 0);
1589 + }
1590 +
1591 + ProcessRecvResult(bytesRead);
1592 +}
1593 +
1594 +HANDLE ReadSocketMessageHandle::GetHandle() const
1595 +{
1596 + return Event.get();
1597 +}
1598 +
1599 +WriteHandle::WriteHandle(HandleWrapper&& MovedHandle, const std::vector<char>& Source) :
1600 + Handle(std::move(MovedHandle)), Buffer(Source.size()), Offset(InitializeFileOffset(Handle.Get()))
1601 +{
1602 + std::memcpy(Buffer.Span().data(), Source.data(), Source.size());
1603 + Overlapped.hEvent = Event.get();
1604 +}
1605 +
1606 +WriteHandle::WriteHandle(HandleWrapper&& MovedHandle, gsl::span<gsl::byte> Source) :
1607 + Handle(std::move(MovedHandle)), Buffer(Source), Offset(InitializeFileOffset(Handle.Get()))
1608 +{
1609 + Overlapped.hEvent = Event.get();
1610 +}
1611 +
1612 +WriteHandle::~WriteHandle()
1613 +{
1614 + if (State == IOHandleStatus::Pending)
1615 + {
1616 + CancelPendingIo(Handle.Get(), Overlapped);
1617 }
1618 }
1619
@@ -1479,13 +1627,14 @@ void WriteHandle::Schedule()
1627 Overlapped.OffsetHigh = Offset.HighPart;
1628
1629 // Schedule the write.
1630 + const auto buffer = Buffer.Span();
1631 DWORD bytesWritten{};
1483 - if (WriteFile(Handle.Get(), Buffer.data(), static_cast<DWORD>(Buffer.size()), &bytesWritten, &Overlapped))
1632 + if (WriteFile(Handle.Get(), buffer.data(), static_cast<DWORD>(buffer.size()), &bytesWritten, &Overlapped))
1633 {
1634 Offset.QuadPart += bytesWritten;
1635
1487 - Buffer.erase(Buffer.begin(), Buffer.begin() + bytesWritten);
1488 - if (Buffer.empty())
1636 + Buffer.Consume(bytesWritten);
1637 + if (Buffer.Size() == 0)
1638 {
1639 State = IOHandleStatus::Completed;
1640 }
@@ -1512,8 +1661,8 @@ void WriteHandle::Collect()
1661 THROW_IF_WIN32_BOOL_FALSE(GetOverlappedResult(Handle.Get(), &Overlapped, &bytesWritten, false));
1662 Offset.QuadPart += bytesWritten;
1663
1515 - Buffer.erase(Buffer.begin(), Buffer.begin() + bytesWritten);
1516 - if (Buffer.empty())
1664 + Buffer.Consume(bytesWritten);
1665 + if (Buffer.Size() == 0)
1666 {
1667 State = IOHandleStatus::Completed;
1668 }
@@ -1525,7 +1674,8 @@ void WriteHandle::Push(const gsl::span<char>& Content)
1674 WI_ASSERT(State == IOHandleStatus::Standby || State == IOHandleStatus::Completed);
1675 WI_ASSERT(!Content.empty());
1676
1528 - Buffer.insert(Buffer.end(), Content.begin(), Content.end());
1677 + // Resize() throws E_UNEXPECTED if Buffer does not own its storage.
1678 + Buffer.Append(Content);
1679
1680 State = IOHandleStatus::Standby;
1681 }
@@ -1698,4 +1848,4 @@ void DockerIORelayHandle::OnRead(const gsl::span<char>& Buffer)
1848 // If no handle is active, expect a header.
1849 ProcessNextHeader();
1850 }
1701 -}
\ No newline at end of file
1851 +}
src/windows/common/relay.hpp
+92 -2
@@ -231,6 +231,68 @@ private:
231 std::function<void()> OnClose;
232 };
233
234 +// A buffer that may either own its underlying storage (constructed from a size, allocating an
235 +// internal std::vector<char>) or borrow it from a caller-provided gsl::span<gsl::byte>.
236 +class BufferWrapper
237 +{
238 +public:
239 + DEFAULT_MOVABLE(BufferWrapper);
240 + NON_COPYABLE(BufferWrapper);
241 +
242 + explicit BufferWrapper(size_t size) : m_owned(std::in_place, size)
243 + {
244 + }
245 +
246 + explicit BufferWrapper(gsl::span<gsl::byte> span) : m_unowned(span)
247 + {
248 + }
249 +
250 + bool Owned() const noexcept
251 + {
252 + return m_owned.has_value();
253 + }
254 +
255 + void Resize(size_t size)
256 + {
257 + THROW_HR_IF_MSG(E_UNEXPECTED, !Owned(), "BufferWrapper::Resize called on a non-owned buffer");
258 + m_owned->resize(size);
259 + }
260 +
261 + void Append(gsl::span<char> Span)
262 + {
263 + THROW_HR_IF_MSG(E_UNEXPECTED, !Owned(), "BufferWrapper::Append called on a non-owned buffer");
264 +
265 + m_owned->insert(m_owned->end(), Span.begin(), Span.end());
266 + }
267 +
268 + void Consume(size_t bytes) noexcept
269 + {
270 + WI_ASSERT(bytes <= Size());
271 + if (Owned())
272 + {
273 + m_owned->erase(m_owned->begin(), m_owned->begin() + bytes);
274 + }
275 + else
276 + {
277 + m_unowned = m_unowned.subspan(bytes);
278 + }
279 + }
280 +
281 + gsl::span<gsl::byte> Span() noexcept
282 + {
283 + return Owned() ? gsl::make_span(reinterpret_cast<gsl::byte*>(m_owned->data()), m_owned->size()) : m_unowned;
284 + }
285 +
286 + size_t Size() const noexcept
287 + {
288 + return Owned() ? m_owned->size() : m_unowned.size();
289 + }
290 +
291 +private:
292 + std::optional<std::vector<char>> m_owned;
293 + gsl::span<gsl::byte> m_unowned;
294 +};
295 +
296 class OverlappedIOHandle
297 {
298 public:
@@ -282,7 +344,7 @@ private:
344 std::function<void(const gsl::span<char>& Buffer)> OnRead;
345 wil::unique_event Event{wil::EventOptions::ManualReset};
346 OVERLAPPED Overlapped{};
285 - std::vector<char> Buffer = std::vector<char>(LX_RELAY_BUFFER_SIZE);
347 + BufferWrapper Buffer{LX_RELAY_BUFFER_SIZE};
348 LARGE_INTEGER Offset{};
349 };
350
@@ -343,6 +405,33 @@ private:
405 bool ExpectHeader = true;
406 };
407
408 +class ReadSocketMessageHandle : public OverlappedIOHandle
409 +{
410 +public:
411 + NON_COPYABLE(ReadSocketMessageHandle);
412 + NON_MOVABLE(ReadSocketMessageHandle);
413 +
414 + ReadSocketMessageHandle(HandleWrapper&& Socket, std::vector<gsl::byte>& Buffer, std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage);
415 + ~ReadSocketMessageHandle();
416 +
417 + void Schedule() override;
418 + void Collect() override;
419 + HANDLE GetHandle() const override;
420 +
421 +private:
422 + void ScheduleRecv();
423 + void ProcessRecvResult(DWORD BytesRead);
424 +
425 + HandleWrapper Socket;
426 + std::vector<gsl::byte>& Buffer;
427 + std::function<void(const gsl::span<gsl::byte>& Message)> OnMessage;
428 + wil::unique_event Event{wil::EventOptions::ManualReset};
429 + OVERLAPPED Overlapped{};
430 + bool ReadingHeader = true;
431 + size_t BytesRemaining = sizeof(MESSAGE_HEADER);
432 + size_t CurrentOffset = 0;
433 +};
434 +
435 class WriteHandle : public OverlappedIOHandle
436 {
437 public:
@@ -350,6 +439,7 @@ public:
439 NON_MOVABLE(WriteHandle);
440
441 WriteHandle(HandleWrapper&& Handle, const std::vector<char>& Buffer = {});
442 + WriteHandle(HandleWrapper&& Handle, gsl::span<gsl::byte> Span);
443 ~WriteHandle();
444 void Schedule() override;
445 void Collect() override;
@@ -360,7 +450,7 @@ private:
450 HandleWrapper Handle;
451 wil::unique_event Event{wil::EventOptions::ManualReset};
452 OVERLAPPED Overlapped{};
363 - std::vector<char> Buffer;
453 + BufferWrapper Buffer;
454 LARGE_INTEGER Offset{};
455 };
456
src/windows/service/exe/LxssCreateProcess.h
+2 -2
@@ -85,11 +85,11 @@ public:
85 wsl::shared::MessageWriter<CREATE_PROCESS_MESSAGE> message(LxInitCreateProcess);
86 message.WriteString(message->PathIndex, Path);
87 gsl::copy(as_bytes(gsl::span(ArgumentsData)), message.InsertBuffer(message->CommandLineIndex, ArgumentsData.size()));
88 - auto transaction = channel.StartTransaction();
88 + auto transaction = channel.StartTransaction(Timeout);
89 transaction.Send<CREATE_PROCESS_MESSAGE>(message.Span());
90
91 auto readResult = [&]() {
92 - const auto& message = transaction.Receive<RESULT_MESSAGE<int32_t>>(nullptr, Timeout);
92 + const auto& message = transaction.Receive<RESULT_MESSAGE<int32_t>>();
93 return message.Result;
94 };
95
src/windows/service/exe/LxssUserSession.cpp
+1 -1
@@ -2208,7 +2208,7 @@ try
2208 {
2209 wsl::windows::common::wslutil::SetThreadDescription(L"Telemetry");
2210
2211 - wsl::shared::SocketChannel channel(std::move(socket), "Telemetry", m_vmTerminating.get());
2211 + wsl::shared::SocketChannel channel(std::move(socket), "Telemetry", {m_vmTerminating.get()});
2212
2213 // Check if drvfs notifications are enabled for the user.
2214 bool drvFsNotifications{};
src/windows/service/exe/WslCoreInstance.cpp
+3 -3
@@ -271,7 +271,7 @@ void WslCoreInstance::CreateLxProcess(
271
272 void WslCoreInstance::ReadOOBEResult(wil::unique_socket&& Socket, wsl::windows::service::DistributionRegistration&& registration)
273 {
274 - wsl::shared::SocketChannel channel(std::move(Socket), "OOBE", m_destroyingEvent.get());
274 + wsl::shared::SocketChannel channel(std::move(Socket), "OOBE", {m_destroyingEvent.get()});
275
276 const auto* oobeResult = channel.ReceiveMessageOrClosed<LX_INIT_OOBE_RESULT>().first;
277
@@ -475,9 +475,9 @@ bool WslCoreInstance::RequestStop(_In_ bool Force)
475 terminateMessage.Header.MessageSize = sizeof(terminateMessage);
476 terminateMessage.Force = Force;
477
478 - auto transaction = m_initChannel->GetChannel().StartTransaction();
478 + auto transaction = m_initChannel->GetChannel().StartTransaction(m_socketTimeout);
479 transaction.Send(terminateMessage);
480 - auto [message, span] = transaction.ReceiveOrClosed<RESULT_MESSAGE<bool>>(m_socketTimeout);
480 + auto [message, span] = transaction.ReceiveOrClosed<RESULT_MESSAGE<bool>>();
481 if (message)
482 {
483 shutdown = message->Result;
src/windows/service/exe/WslCoreVm.cpp
+9 -8
@@ -429,7 +429,8 @@ void WslCoreVm::Initialize(const GUID& VmId, const wil::shared_handle& UserToken
429 }
430
431 // Accept a connection from mini_init with a receive timeout so the service does not get stuck waiting for a response from the VM.
432 - m_miniInitChannel = wsl::shared::SocketChannel{AcceptConnection(m_vmConfig.KernelBootTimeout), "mini_init", m_terminatingEvent.get()};
432 + m_miniInitChannel =
433 + wsl::shared::SocketChannel{AcceptConnection(m_vmConfig.KernelBootTimeout), "mini_init", {m_terminatingEvent.get()}};
434
435 // Accept the connection from the Linux guest for notifications.
436 m_notifyChannel = AcceptConnection(m_vmConfig.KernelBootTimeout);
@@ -1059,7 +1060,7 @@ void WslCoreVm::CollectCrashDumps(wil::unique_socket&& listenSocket) const
1060 DWORD receiveTimeout = m_vmConfig.KernelBootTimeout;
1061 THROW_LAST_ERROR_IF(setsockopt(socket->get(), SOL_SOCKET, SO_RCVTIMEO, (const char*)&receiveTimeout, sizeof(receiveTimeout)) == SOCKET_ERROR);
1062
1062 - auto channel = wsl::shared::SocketChannel{std::move(socket.value()), "crash_dump", m_terminatingEvent.get()};
1063 + auto channel = wsl::shared::SocketChannel{std::move(socket.value()), "crash_dump", {m_terminatingEvent.get()}};
1064
1065 auto transaction = channel.ReceiveTransaction();
1066 gsl::span<gsl::byte> responseSpan;
@@ -1960,7 +1961,7 @@ WslCoreVm::DiskMountResult WslCoreVm::MountDiskLockHeld(
1961 transaction.Send<LX_MINI_INIT_MOUNT_MESSAGE>(message.Span());
1962
1963 // Accept a connection from mini_init
1963 - wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", m_terminatingEvent.get()};
1964 + wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", {m_terminatingEvent.get()}};
1965
1966 // Get the mount result from mini_init
1967 auto [mountResult, step] = GetMountResult(channel);
@@ -2090,7 +2091,7 @@ void WslCoreVm::WaitForPmemDeviceInVm(_In_ ULONG PmemId)
2091 channel = {
2092 AcceptConnection(m_vmConfig.KernelBootTimeout),
2093 "WaitForPmem",
2093 - m_terminatingEvent.get(),
2094 + {m_terminatingEvent.get()},
2095 };
2096 }
2097
@@ -2399,7 +2400,7 @@ void WslCoreVm::ResizeDistribution(_In_ ULONG Lun, _In_ HANDLE OutputHandle, _In
2400 auto transaction = m_miniInitChannel.StartTransaction();
2401 transaction.Send(message);
2402
2402 - wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "ResizeDistribution", m_terminatingEvent.get()};
2403 + wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "ResizeDistribution", {m_terminatingEvent.get()}};
2404 auto outputChannel = AcceptConnection(m_vmConfig.KernelBootTimeout);
2405
2406 wsl::windows::common::relay::ScopedRelay outputRelay(std::move(outputChannel), OutputHandle);
@@ -2478,7 +2479,7 @@ std::pair<int, LX_MINI_MOUNT_STEP> WslCoreVm::UnmountDisk(_In_ const AttachedDis
2479 transaction.Send(message);
2480
2481 // Accept a connection from mini_init.
2481 - wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", m_terminatingEvent.get()};
2482 + wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", {m_terminatingEvent.get()}};
2483
2484 // Get the unmount result from mini_init
2485 return GetMountResult(channel);
@@ -2494,7 +2495,7 @@ std::pair<int, LX_MINI_MOUNT_STEP> WslCoreVm::UnmountVolume(_In_ const AttachedD
2495 transaction.Send<LX_MINI_INIT_UNMOUNT_MESSAGE>(message.Span());
2496
2497 // Accept a connection from mini_init.
2497 - wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", m_terminatingEvent.get()};
2498 + wsl::shared::SocketChannel channel{AcceptConnection(m_vmConfig.KernelBootTimeout), "MountResult", {m_terminatingEvent.get()}};
2499
2500 // Get the unmount result from mini_init.
2501 return GetMountResult(channel);
@@ -2551,7 +2552,7 @@ try
2552 break;
2553 }
2554
2554 - wsl::shared::SocketChannel channel{std::move(socket.value()), "VirtioFs", m_terminatingEvent.get()};
2555 + wsl::shared::SocketChannel channel{std::move(socket.value()), "VirtioFs", {m_terminatingEvent.get()}};
2556 std::thread([this, channel = std::move(channel)]() mutable {
2557 try
2558 {
src/windows/wslcsession/DockerHTTPClient.cpp
+1 -1
@@ -571,7 +571,7 @@ wil::unique_socket DockerHTTPClient::ConnectSocket()
571
572 // Connect the new hvsocket.
573 wsl::shared::SocketChannel newChannel{
574 - wsl::windows::common::hvsocket::Connect(m_vmId, response.Port, m_exitingEvent, m_connectTimeoutMs), "DockerClient", m_exitingEvent};
574 + wsl::windows::common::hvsocket::Connect(m_vmId, response.Port, m_exitingEvent, m_connectTimeoutMs), "DockerClient", {m_exitingEvent}};
575 lock.reset();
576
577 // Connect that socket to the docker unix socket.
src/windows/wslcsession/WSLCProcessControl.cpp
+1 -1
@@ -201,7 +201,7 @@ void DockerExecProcessControl::OnContainerReleased() noexcept
201 }
202
203 VMProcessControl::VMProcessControl(WSLCVirtualMachine& VirtualMachine, int Pid, wil::unique_socket&& TtyControl) :
204 - m_pid(Pid), m_ttyControlChannel(std::move(TtyControl), "TtyControl", VirtualMachine.TerminatingEvent()), m_vm(&VirtualMachine)
204 + m_pid(Pid), m_ttyControlChannel(std::move(TtyControl), "TtyControl", {VirtualMachine.TerminatingEvent()}), m_vm(&VirtualMachine)
205 {
206 }
207
src/windows/wslcsession/WSLCSession.cpp
+22 -15
@@ -258,7 +258,7 @@ try
258 TraceLoggingValue(Settings->CreatorPid, "CreatorPid"));
259
260 // Create the VM.
261 - m_virtualMachine.emplace(Vm, Settings);
261 + m_virtualMachine.emplace(Vm, Settings, m_sessionTerminatingEvent.get());
262
263 // Make sure that everything is destroyed correctly if an exception is thrown.
264 auto errorCleanup = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [&]() { LOG_IF_FAILED(Terminate()); });
@@ -2078,7 +2078,12 @@ HRESULT WSLCSession::PruneVolumes(const WSLCPruneVolumesOptions* /*Options*/, WS
2078
2079 int WSLCSession::StopProcess(ServiceRunningProcess& Process, DWORD TerminateTimeoutMs, DWORD KillTimeoutMs)
2080 {
2081 - LOG_IF_FAILED(Process.Get().Signal(WSLCSignalSIGTERM));
2081 + auto signalResult = Process.Get().Signal(WSLCSignalSIGTERM);
2082 + if (FAILED(signalResult))
2083 + {
2084 + LOG_HR_MSG(signalResult, "Failed to terminate process %i", Process.Get().GetPid());
2085 + return -1;
2086 + }
2087
2088 try
2089 {
@@ -2437,22 +2442,24 @@ try
2442 }
2443 else
2444 {
2440 - // Stop dockerd first, then containerd (dockerd is a client of containerd).
2441 - // N.B. dockerd waits a couple seconds if there are any outstanding HTTP request sockets opened.
2442 - if (m_dockerdProcess.has_value())
2445 + if (m_virtualMachine)
2446 {
2444 - auto dockerdExitCode = StopProcess(m_dockerdProcess.value(), c_processTerminateTimeoutMs, c_processKillTimeoutMs);
2445 - WSL_LOG("DockerdExit", TraceLoggingValue(dockerdExitCode, "code"));
2446 - }
2447 + m_virtualMachine->OnSessionTerminated();
2448
2448 - if (m_containerdProcess.has_value())
2449 - {
2450 - auto containerdExitCode = StopProcess(m_containerdProcess.value(), c_processTerminateTimeoutMs, c_processKillTimeoutMs);
2451 - WSL_LOG("ContainerdExit", TraceLoggingValue(containerdExitCode, "code"));
2452 - }
2449 + // Stop dockerd first, then containerd (dockerd is a client of containerd).
2450 + // N.B. dockerd waits a couple seconds if there are any outstanding HTTP request sockets opened.
2451 + if (m_dockerdProcess.has_value())
2452 + {
2453 + auto dockerdExitCode = StopProcess(m_dockerdProcess.value(), c_processTerminateTimeoutMs, c_processKillTimeoutMs);
2454 + WSL_LOG("DockerdExit", TraceLoggingValue(dockerdExitCode, "code"));
2455 + }
2456 +
2457 + if (m_containerdProcess.has_value())
2458 + {
2459 + auto containerdExitCode = StopProcess(m_containerdProcess.value(), c_processTerminateTimeoutMs, c_processKillTimeoutMs);
2460 + WSL_LOG("ContainerdExit", TraceLoggingValue(containerdExitCode, "code"));
2461 + }
2462
2454 - if (m_virtualMachine)
2455 - {
2463 // N.B. dockerd has exited by this point, so unmounting the VHD is safe since no container can be running.
2464 try
2465 {
src/windows/wslcsession/WSLCVirtualMachine.cpp
+32 -17
@@ -248,12 +248,13 @@ VMPortMapping& VMPortMapping::operator=(VMPortMapping&& Other)
248 return *this;
249 }
250
251 -WSLCVirtualMachine::WSLCVirtualMachine(_In_ IWSLCVirtualMachine* Vm, _In_ const WSLCSessionInitSettings* Settings) :
251 +WSLCVirtualMachine::WSLCVirtualMachine(_In_ IWSLCVirtualMachine* Vm, _In_ const WSLCSessionInitSettings* Settings, _In_ HANDLE SessionTerminatingEvent) :
252 m_vm(Vm),
253 m_featureFlags(static_cast<WSLCFeatureFlags>(Settings->FeatureFlags)),
254 m_networkingMode(Settings->NetworkingMode),
255 m_bootTimeoutMs(Settings->BootTimeoutMs),
256 - m_rootVhdType(Settings->RootVhdTypeOverride ? Settings->RootVhdTypeOverride : "ext4")
256 + m_rootVhdType(Settings->RootVhdTypeOverride ? Settings->RootVhdTypeOverride : "ext4"),
257 + m_sessionTerminatingEvent(SessionTerminatingEvent)
258 {
259 // N.B. The constructor should not run any operation that could throw, so the destructor runs even if the VM fails to boot.
260 }
@@ -272,13 +273,14 @@ void WSLCVirtualMachine::Initialize()
273 wil::unique_socket socket;
274 THROW_IF_FAILED(m_vm->AcceptConnection(reinterpret_cast<HANDLE*>(&socket)));
275
275 - m_initChannel = wsl::shared::SocketChannel{std::move(socket), "mini_init", m_vmTerminatingEvent.get()};
276 + m_initChannel = wsl::shared::SocketChannel{std::move(socket), "mini_init", {m_vmTerminatingEvent.get(), m_sessionTerminatingEvent}};
277
278 // Create a thread to watch for exited processes.
279 auto [__, ___, childChannel] = Fork(WSLC_FORK::Thread);
280 + childChannel.SetExitEvents({m_vmTerminatingEvent.get()});
281
282 WSLC_WATCH_PROCESSES watchMessage{};
281 - auto watchTransaction = childChannel.StartTransaction();
283 + auto watchTransaction = childChannel.StartTransaction(m_initChannelTimeout);
284 watchTransaction.Send(watchMessage);
285
286 THROW_HR_IF(E_FAIL, watchTransaction.Receive<RESULT_MESSAGE<uint32_t>>().Result != 0);
@@ -493,7 +495,7 @@ void WSLCVirtualMachine::Unmount(_In_ const char* Path)
495 wsl::shared::MessageWriter<WSLC_UNMOUNT> message;
496 message.WriteString(Path);
497
496 - const auto& response = subChannel.Transaction<WSLC_UNMOUNT>(message.Span());
498 + const auto& response = subChannel.Transaction<WSLC_UNMOUNT>(message.Span(), nullptr, m_initChannelTimeout);
499
500 // TODO: Return errno to caller
501 THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_NOT_FOUND), response.Result == EINVAL);
@@ -511,7 +513,7 @@ void WSLCVirtualMachine::DetachDisk(_In_ ULONG Lun)
513 // Detach it from the guest
514 WSLC_DETACH message;
515 message.Lun = Lun;
514 - const auto& response = m_initChannel.Transaction(message);
516 + const auto& response = m_initChannel.Transaction(message, nullptr, m_initChannelTimeout);
517
518 // TODO: Return errno to caller
519 THROW_HR_IF(E_FAIL, response.Result != 0);
@@ -539,7 +541,7 @@ std::tuple<int32_t, int32_t, wsl::shared::SocketChannel> WSLCVirtualMachine::For
541 message.ForkType = Type;
542 message.TtyColumns = static_cast<uint16_t>(TtyColumns);
543 message.TtyRows = static_cast<uint16_t>(TtyRows);
542 - const auto& response = Channel.Transaction(message);
544 + const auto& response = Channel.Transaction(message, nullptr, m_initChannelTimeout);
545 port = response.Port;
546 pid = response.Pid;
547 ptyMaster = response.PtyMasterFd;
@@ -547,9 +549,10 @@ std::tuple<int32_t, int32_t, wsl::shared::SocketChannel> WSLCVirtualMachine::For
549
550 THROW_HR_IF_MSG(E_FAIL, pid <= 0, "fork() returned %i", pid);
551
550 - auto socket = wsl::windows::common::hvsocket::Connect(m_vmId, port, m_vmTerminatingEvent.get(), m_bootTimeoutMs);
552 + auto socket = wsl::windows::common::hvsocket::Connect(m_vmId, port, m_vmTerminatingEvent.get(), m_initChannelTimeout);
553
552 - return std::make_tuple(pid, ptyMaster, wsl::shared::SocketChannel{std::move(socket), std::to_string(pid), m_vmTerminatingEvent.get()});
554 + return std::make_tuple(
555 + pid, ptyMaster, wsl::shared::SocketChannel{std::move(socket), std::to_string(pid), std::vector<HANDLE>(Channel.GetExitEvents())});
556 }
557
558 WSLCVirtualMachine::ConnectedSocket WSLCVirtualMachine::ConnectSocket(wsl::shared::SocketChannel& Channel, int32_t Fd)
@@ -557,12 +560,12 @@ WSLCVirtualMachine::ConnectedSocket WSLCVirtualMachine::ConnectSocket(wsl::share
560 WSLC_ACCEPT message{};
561 message.Fd = Fd;
562
560 - auto transaction = Channel.StartTransaction();
563 + auto transaction = Channel.StartTransaction(m_initChannelTimeout);
564 transaction.Send(message);
565 const auto& response = transaction.Receive<WSLC_ACCEPT::TResponse>();
566
567 ConnectedSocket socket;
565 - socket.Socket = wsl::windows::common::hvsocket::Connect(m_vmId, response.Result);
568 + socket.Socket = wsl::windows::common::hvsocket::Connect(m_vmId, response.Result, m_vmTerminatingEvent.get(), m_initChannelTimeout);
569
570 // If the FD was unspecified, read the Linux file descriptor from the guest.
571 if (Fd == -1)
@@ -583,7 +586,7 @@ std::string WSLCVirtualMachine::GetVhdDevicePath(ULONG Lun)
586 message.Header.MessageSize = sizeof(message);
587 message.Header.MessageType = WSLC_GET_DISK::Type;
588 message.ScsiLun = Lun;
586 - const auto& response = m_initChannel.Transaction(message);
589 + const auto& response = m_initChannel.Transaction(message, nullptr, m_initChannelTimeout);
590 THROW_HR_IF_MSG(E_FAIL, response.Result != 0, "Failed to get disk path, init returned: %lu", response.Result);
591
592 return response.Buffer;
@@ -686,7 +689,7 @@ Microsoft::WRL::ComPtr<WSLCProcess> WSLCVirtualMachine::CreateLinuxProcessImpl(
689 relayMessage.Socket = tty->Fd;
690 relayMessage.TtyControl = ttyControlhandle.Fd; // N.B. Fd is set to -1 if unset.
691 {
689 - auto relayTransaction = childChannel.StartTransaction();
692 + auto relayTransaction = childChannel.StartTransaction(m_initChannelTimeout);
693 relayTransaction.Send(relayMessage);
694 }
695
@@ -700,7 +703,7 @@ Microsoft::WRL::ComPtr<WSLCProcess> WSLCVirtualMachine::CreateLinuxProcessImpl(
703 registerProcess(grandChildPid);
704
705 {
703 - auto execTransaction = grandChildChannel.StartTransaction();
706 + auto execTransaction = grandChildChannel.StartTransaction(m_initChannelTimeout);
707 execTransaction.Send<WSLC_EXEC>(Message.Span());
708 auto [execResponse, execSpan] = execTransaction.ReceiveOrClosed<RESULT_MESSAGE<int32_t>>();
709 result = execResponse != nullptr ? execResponse->Result : 0;
@@ -717,7 +720,7 @@ Microsoft::WRL::ComPtr<WSLCProcess> WSLCVirtualMachine::CreateLinuxProcessImpl(
720 {
721 registerProcess(pid);
722
720 - auto execTransaction = childChannel.StartTransaction();
723 + auto execTransaction = childChannel.StartTransaction(m_initChannelTimeout);
724 execTransaction.Send<WSLC_EXEC>(Message.Span());
725 auto [execResponse, execSpan] = execTransaction.ReceiveOrClosed<RESULT_MESSAGE<int32_t>>();
726 auto result = execResponse != nullptr ? execResponse->Result : 0;
@@ -806,7 +809,7 @@ void WSLCVirtualMachine::Signal(_In_ LONG Pid, _In_ int Signal)
809 WSLC_SIGNAL message;
810 message.Pid = Pid;
811 message.Signal = Signal;
809 - const auto& response = m_initChannel.Transaction(message);
812 + const auto& response = m_initChannel.Transaction(message, nullptr, m_initChannelTimeout);
813
814 THROW_HR_IF(E_FAIL, response.Result != 0);
815 }
@@ -1123,6 +1126,17 @@ void WSLCVirtualMachine::OnProcessReleased(int Pid)
1126 });
1127 }
1128
1129 +void WSLCVirtualMachine::OnSessionTerminated()
1130 +{
1131 + std::lock_guard lock{m_lock};
1132 +
1133 + // Don't cancel init transactions on the session termination event, since that event is set.
1134 + m_initChannel.SetExitEvents({m_vmTerminatingEvent.get()});
1135 +
1136 + // Set a lower timeout for init transactions since we're terminating.
1137 + m_initChannelTimeout = 15 * 1000;
1138 +}
1139 +
1140 std::shared_ptr<VmPortAllocation> WSLCVirtualMachine::TryAllocatePort(uint16_t Port, int Family, int Protocol)
1141 {
1142 std::lock_guard lock{m_lock};
@@ -1202,7 +1216,8 @@ void WSLCVirtualMachine::CollectCrashDumps(wil::unique_socket&& listenSocket)
1216 constexpr DWORD timeout = 30 * 1000;
1217 THROW_LAST_ERROR_IF(setsockopt(socket->get(), SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof(timeout)) == SOCKET_ERROR);
1218
1205 - auto channel = wsl::shared::SocketChannel{std::move(socket.value()), "crash_dump", m_vmTerminatingEvent.get()};
1219 + auto channel = wsl::shared::SocketChannel{
1220 + std::move(socket.value()), "crash_dump", {m_vmTerminatingEvent.get(), m_sessionTerminatingEvent}};
1221
1222 auto transaction = channel.ReceiveTransaction();
1223 gsl::span<gsl::byte> responseSpan;
src/windows/wslcsession/WSLCVirtualMachine.h
+5 -1
@@ -120,7 +120,7 @@ public:
120
121 using TPrepareCommandLine = std::function<void(const std::vector<ConnectedSocket>&)>;
122
123 - WSLCVirtualMachine(_In_ IWSLCVirtualMachine* Vm, _In_ const WSLCSessionInitSettings* Settings);
123 + WSLCVirtualMachine(_In_ IWSLCVirtualMachine* Vm, _In_ const WSLCSessionInitSettings* Settings, _In_ HANDLE SessionTerminatingEvent);
124 ~WSLCVirtualMachine();
125
126 void Initialize();
@@ -134,6 +134,7 @@ public:
134 void Signal(_In_ LONG Pid, _In_ int Signal);
135
136 void OnProcessReleased(int Pid);
137 + void OnSessionTerminated();
138
139 std::shared_ptr<VmPortAllocation> TryAllocatePort(uint16_t Port, int Family, int Protocol);
140 std::shared_ptr<VmPortAllocation> AllocatePort(int Family, int Protocol);
@@ -223,8 +224,11 @@ private:
224 std::vector<std::weak_ptr<VMProcessControl>> m_trackedProcesses;
225
226 wil::unique_event m_vmTerminatingEvent{wil::EventOptions::ManualReset};
227 + HANDLE m_sessionTerminatingEvent{};
228
229 wsl::shared::SocketChannel m_initChannel;
230 + DWORD m_initChannelTimeout = 30 * 1000;
231 +
232 wil::unique_handle m_portRelayChannelRead;
233 wil::unique_handle m_portRelayChannelWrite;
234
test/windows/Common.cpp
+38
@@ -2615,6 +2615,32 @@ std::string ReadToString(SOCKET Handle)
2615 return output;
2616 }
2617
2618 +std::pair<wil::unique_socket, wil::unique_socket> MakeSocketPair()
2619 +{
2620 + wil::unique_socket listenSocket{WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, nullptr, 0, WSA_FLAG_OVERLAPPED)};
2621 + THROW_LAST_ERROR_IF(!listenSocket);
2622 +
2623 + sockaddr_in bindAddr{};
2624 + bindAddr.sin_family = AF_INET;
2625 + bindAddr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
2626 + bindAddr.sin_port = 0;
2627 + THROW_LAST_ERROR_IF(bind(listenSocket.get(), reinterpret_cast<sockaddr*>(&bindAddr), sizeof(bindAddr)) == SOCKET_ERROR);
2628 + THROW_LAST_ERROR_IF(listen(listenSocket.get(), 1) == SOCKET_ERROR);
2629 +
2630 + sockaddr_in boundAddr{};
2631 + int boundAddrLen = sizeof(boundAddr);
2632 + THROW_LAST_ERROR_IF(getsockname(listenSocket.get(), reinterpret_cast<sockaddr*>(&boundAddr), &boundAddrLen) == SOCKET_ERROR);
2633 +
2634 + wil::unique_socket clientSocket{WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, nullptr, 0, WSA_FLAG_OVERLAPPED)};
2635 + THROW_LAST_ERROR_IF(!clientSocket);
2636 + THROW_LAST_ERROR_IF(connect(clientSocket.get(), reinterpret_cast<sockaddr*>(&boundAddr), sizeof(boundAddr)) == SOCKET_ERROR);
2637 +
2638 + wil::unique_socket serverSocket{accept(listenSocket.get(), nullptr, nullptr)};
2639 + THROW_LAST_ERROR_IF(!serverSocket);
2640 +
2641 + return {std::move(clientSocket), std::move(serverSocket)};
2642 +}
2643 +
2644 std::string ReadToString(HANDLE Handle)
2645 {
2646 std::string output;
@@ -2948,3 +2974,15 @@ void SetPathAccess(const std::filesystem::path& path, DWORD Permissions, ACCESS_
2974 THROW_IF_WIN32_ERROR(SetNamedSecurityInfoW(
2975 const_cast<LPWSTR>(path.c_str()), SE_FILE_OBJECT, DACL_SECURITY_INFORMATION, nullptr, nullptr, newAcl.get(), nullptr));
2976 }
2977 +
2978 +void WriteSocket(SOCKET Socket, const void* data, size_t size)
2979 +{
2980 + while (size > 0)
2981 + {
2982 + auto result = send(Socket, static_cast<const char*>(data), gsl::narrow_cast<int>(size), 0);
2983 + VERIFY_IS_TRUE(result > 0);
2984 +
2985 + size -= result;
2986 + data = static_cast<const char*>(data) + result;
2987 + }
2988 +}
test/windows/Common.h
+6
@@ -601,6 +601,10 @@ void ValidateOutput(LPCWSTR CommandLine, const std::wstring& ExpectedOutput, con
601 std::string ReadToString(SOCKET Handle);
602 std::string ReadToString(HANDLE Handle);
603
604 +// Connects a pair of overlapped TCP sockets via an anonymous bind on the loopback interface.
605 +// Returns {client, server}.
606 +std::pair<wil::unique_socket, wil::unique_socket> MakeSocketPair();
607 +
608 std::wstring ReadFileContent(const std::string& Path);
609 std::wstring ReadFileContent(const std::wstring& Path);
610
@@ -671,3 +675,5 @@ void VerifyAreEqualUnordered(const std::vector<T>& expected, const std::vector<T
675 }
676
677 void SetPathAccess(const std::filesystem::path& path, DWORD Permissions, ACCESS_MODE Mode);
678 +
679 +void WriteSocket(SOCKET Socket, const void* data, size_t size);
test/windows/UnitTests.cpp
+256
@@ -6654,5 +6654,261 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",
6654 -1);
6655 }
6656
6657 + TEST_METHOD(ReadSocketMessageHandle)
6658 + {
6659 + // Drive a ReadSocketMessageHandle until completion and return the bytes delivered to its
6660 + // OnMessage callback. If a non-success HRESULT is supplied, the call is expected to throw
6661 + // that HRESULT instead, and the OnMessage callback must not be invoked.
6662 + auto readMessage = [](wil::unique_socket&& server, HRESULT expectedHr = S_OK) {
6663 + std::vector<gsl::byte> buffer;
6664 + bool callbackInvoked = false;
6665 + std::vector<gsl::byte> message;
6666 +
6667 + wsl::windows::common::relay::MultiHandleWait io;
6668 + io.AddHandle(std::make_unique<wsl::windows::common::relay::ReadSocketMessageHandle>(
6669 + wsl::windows::common::relay::HandleWrapper{std::move(server)}, buffer, [&callbackInvoked, &message](const gsl::span<gsl::byte>& received) {
6670 + callbackInvoked = true;
6671 + message.assign(received.begin(), received.end());
6672 + }));
6673 +
6674 + const auto hr = wil::ResultFromException([&]() { io.Run(std::chrono::seconds(60)); });
6675 + VERIFY_ARE_EQUAL(hr, expectedHr);
6676 + VERIFY_ARE_EQUAL(callbackInvoked, SUCCEEDED(expectedHr));
6677 + return message;
6678 + };
6679 +
6680 + // Scenario 1: A complete header-only message is delivered intact.
6681 + {
6682 + auto [client, server] = MakeSocketPair();
6683 +
6684 + MESSAGE_HEADER header{};
6685 + header.MessageType = LxMiniInitMessageAny;
6686 + header.MessageSize = sizeof(header);
6687 + header.TransactionId = 7;
6688 + header.TransactionStep = 1;
6689 + WriteSocket(client.get(), &header, sizeof(header));
6690 + client.reset();
6691 +
6692 + const auto message = readMessage(std::move(server));
6693 + VERIFY_ARE_EQUAL(message.size(), sizeof(header));
6694 + VERIFY_IS_TRUE(std::memcmp(message.data(), &header, sizeof(header)) == 0);
6695 + }
6696 +
6697 + // Scenario 2: A complete message with a payload body is delivered intact.
6698 + {
6699 + auto [client, server] = MakeSocketPair();
6700 +
6701 + constexpr size_t bodySize = 128;
6702 + std::vector<gsl::byte> payload(sizeof(MESSAGE_HEADER) + bodySize);
6703 + auto* header = reinterpret_cast<MESSAGE_HEADER*>(payload.data());
6704 + header->MessageType = LxMiniInitMessageAny;
6705 + header->MessageSize = gsl::narrow_cast<unsigned int>(payload.size());
6706 + header->TransactionId = 42;
6707 + header->TransactionStep = 2;
6708 + for (size_t i = 0; i < bodySize; ++i)
6709 + {
6710 + payload[sizeof(MESSAGE_HEADER) + i] = static_cast<gsl::byte>(i & 0xFF);
6711 + }
6712 + WriteSocket(client.get(), payload.data(), payload.size());
6713 + client.reset();
6714 +
6715 + const auto message = readMessage(std::move(server));
6716 + VERIFY_ARE_EQUAL(message.size(), payload.size());
6717 + VERIFY_IS_TRUE(std::memcmp(message.data(), payload.data(), payload.size()) == 0);
6718 + }
6719 +
6720 + // Scenario 3: Sender closes without writing any bytes. The reader should observe a clean
6721 + // end-of-stream and signal completion with an empty span (no exception).
6722 + {
6723 + auto [client, server] = MakeSocketPair();
6724 + client.reset();
6725 +
6726 + const auto message = readMessage(std::move(server));
6727 + VERIFY_ARE_EQUAL(message.size(), static_cast<size_t>(0));
6728 + }
6729 +
6730 + // Scenario 4: Sender closes after sending fewer bytes than a full header.
6731 + // The reader should treat this as a protocol error and throw E_UNEXPECTED.
6732 + {
6733 + auto [client, server] = MakeSocketPair();
6734 +
6735 + std::array<gsl::byte, sizeof(MESSAGE_HEADER) - 1> partialHeader{};
6736 + std::memset(partialHeader.data(), 0xCC, partialHeader.size());
6737 + WriteSocket(client.get(), partialHeader.data(), partialHeader.size());
6738 + client.reset();
6739 +
6740 + readMessage(std::move(server), E_UNEXPECTED);
6741 + }
6742 +
6743 + // Scenario 5: Sender provides a complete header but closes after sending only part of the body.
6744 + // The reader should treat this as a protocol error and throw E_UNEXPECTED.
6745 + {
6746 + auto [client, server] = MakeSocketPair();
6747 +
6748 + constexpr size_t fullBodySize = 64;
6749 + constexpr size_t partialBodySize = 16;
6750 + MESSAGE_HEADER header{};
6751 + header.MessageType = LxMiniInitMessageAny;
6752 + header.MessageSize = gsl::narrow_cast<unsigned int>(sizeof(header) + fullBodySize);
6753 + header.TransactionId = 11;
6754 + header.TransactionStep = 1;
6755 + WriteSocket(client.get(), &header, sizeof(header));
6756 +
6757 + std::array<gsl::byte, partialBodySize> partialBody{};
6758 + std::memset(partialBody.data(), 0x55, partialBody.size());
6759 + WriteSocket(client.get(), partialBody.data(), partialBody.size());
6760 + client.reset();
6761 +
6762 + readMessage(std::move(server), E_UNEXPECTED);
6763 + }
6764 + }
6765 +
6766 + TEST_METHOD(SocketChannel)
6767 + {
6768 + // Read exactly `size` bytes from a raw socket into the destination buffer.
6769 + auto recvAll = [](SOCKET socket, void* destination, size_t size) {
6770 + auto* cursor = static_cast<char*>(destination);
6771 + size_t total = 0;
6772 + while (total < size)
6773 + {
6774 + const auto received = recv(socket, cursor + total, gsl::narrow_cast<int>(size - total), 0);
6775 + VERIFY_IS_TRUE(received > 0);
6776 + total += static_cast<size_t>(received);
6777 + }
6778 + };
6779 +
6780 + // Scenario 1: SendMessage produces the expected wire format on the peer socket.
6781 + // The header should carry the auto-stamped TransactionId (1 for the first non-transaction
6782 + // message) and a NONE transaction step, and the payload bytes should match exactly.
6783 + {
6784 + auto [client, server] = MakeSocketPair();
6785 + wsl::shared::SocketChannel channel{std::move(client), "client"};
6786 +
6787 + RESULT_MESSAGE<int32_t> message{};
6788 + message.Header.MessageType = RESULT_MESSAGE<int32_t>::Type;
6789 + message.Header.MessageSize = sizeof(message);
6790 + message.Result = static_cast<int32_t>(0xCAFEBABE);
6791 + channel.SendMessage(message);
6792 +
6793 + std::array<gsl::byte, sizeof(message)> received{};
6794 + recvAll(server.get(), received.data(), received.size());
6795 +
6796 + const auto* header = reinterpret_cast<const MESSAGE_HEADER*>(received.data());
6797 + VERIFY_ARE_EQUAL(header->MessageType, RESULT_MESSAGE<int32_t>::Type);
6798 + VERIFY_ARE_EQUAL(header->MessageSize, gsl::narrow_cast<unsigned int>(sizeof(message)));
6799 + VERIFY_ARE_EQUAL(header->TransactionId, 1u);
6800 + VERIFY_ARE_EQUAL(header->TransactionStep, static_cast<unsigned int>(TRANSACTION_STEP::NONE));
6801 +
6802 + const auto* payload = reinterpret_cast<const RESULT_MESSAGE<int32_t>*>(received.data());
6803 + VERIFY_ARE_EQUAL(payload->Result, static_cast<int32_t>(0xCAFEBABE));
6804 + }
6805 +
6806 + // Scenario 2: Two channels can round-trip a typed message end-to-end.
6807 + {
6808 + auto [a, b] = MakeSocketPair();
6809 + wsl::shared::SocketChannel sender{std::move(a), "sender"};
6810 + wsl::shared::SocketChannel receiver{std::move(b), "receiver"};
6811 +
6812 + RESULT_MESSAGE<int32_t> message{};
6813 + message.Header.MessageType = RESULT_MESSAGE<int32_t>::Type;
6814 + message.Header.MessageSize = sizeof(message);
6815 + message.Result = 1234;
6816 + sender.SendMessage(message);
6817 +
6818 + auto& received = receiver.ReceiveMessage<RESULT_MESSAGE<int32_t>>();
6819 + VERIFY_ARE_EQUAL(received.Header.MessageType, RESULT_MESSAGE<int32_t>::Type);
6820 + VERIFY_ARE_EQUAL(received.Header.MessageSize, gsl::narrow_cast<unsigned int>(sizeof(message)));
6821 + VERIFY_ARE_EQUAL(received.Header.TransactionId, 1u);
6822 + VERIFY_ARE_EQUAL(received.Result, 1234);
6823 + }
6824 +
6825 + // Scenario 3: An exit event signaled while ReceiveMessage is waiting causes the call to
6826 + // throw E_ABORT. Pre-signaling avoids racing a worker thread with the call.
6827 + {
6828 + auto [client, server] = MakeSocketPair();
6829 + wil::unique_event exitEvent{wil::EventOptions::ManualReset};
6830 + exitEvent.SetEvent();
6831 + wsl::shared::SocketChannel channel{std::move(server), "server", std::vector<HANDLE>{exitEvent.get()}};
6832 +
6833 + const auto hr = wil::ResultFromException([&]() { channel.ReceiveMessage<RESULT_MESSAGE<int32_t>>(); });
6834 + VERIFY_ARE_EQUAL(hr, E_ABORT);
6835 + }
6836 +
6837 + // Scenario 4: ReceiveMessageOrClosed on a peer that closed the socket without sending
6838 + // any data returns {nullptr, empty span} and does not throw.
6839 + {
6840 + auto [client, server] = MakeSocketPair();
6841 + wsl::shared::SocketChannel channel{std::move(server), "server"};
6842 + client.reset();
6843 +
6844 + auto [message, span] = channel.ReceiveMessageOrClosed<RESULT_MESSAGE<int32_t>>();
6845 + VERIFY_IS_NULL(message);
6846 + VERIFY_ARE_EQUAL(span.size(), static_cast<size_t>(0));
6847 + }
6848 +
6849 + // Scenario 5: A message arriving with a TransactionId other than the next expected
6850 + // sequence number (the first message must have id 1) is rejected with E_UNEXPECTED.
6851 + {
6852 + auto [client, server] = MakeSocketPair();
6853 + wsl::shared::SocketChannel channel{std::move(server), "server"};
6854 +
6855 + RESULT_MESSAGE<int32_t> message{};
6856 + message.Header.MessageType = RESULT_MESSAGE<int32_t>::Type;
6857 + message.Header.MessageSize = sizeof(message);
6858 + message.Header.TransactionId = 99;
6859 + message.Header.TransactionStep = static_cast<unsigned int>(TRANSACTION_STEP::NONE);
6860 + message.Result = 0;
6861 + WriteSocket(client.get(), &message, sizeof(message));
6862 +
6863 + const auto hr = wil::ResultFromException([&]() { channel.ReceiveMessage<RESULT_MESSAGE<int32_t>>(); });
6864 + VERIFY_ARE_EQUAL(hr, E_UNEXPECTED);
6865 + }
6866 +
6867 + // Scenario 6: A transaction-tagged message arriving on a channel waiting for a
6868 + // non-transaction message is rejected with E_UNEXPECTED.
6869 + {
6870 + auto [client, server] = MakeSocketPair();
6871 + wsl::shared::SocketChannel channel{std::move(server), "server"};
6872 +
6873 + RESULT_MESSAGE<int32_t> message{};
6874 + message.Header.MessageType = RESULT_MESSAGE<int32_t>::Type;
6875 + message.Header.MessageSize = sizeof(message);
6876 + message.Header.TransactionId = 1;
6877 + message.Header.TransactionStep = static_cast<unsigned int>(TRANSACTION_STEP::REQUEST);
6878 + message.Result = 0;
6879 + WriteSocket(client.get(), &message, sizeof(message));
6880 +
6881 + const auto hr = wil::ResultFromException([&]() { channel.ReceiveMessage<RESULT_MESSAGE<int32_t>>(); });
6882 + VERIFY_ARE_EQUAL(hr, E_UNEXPECTED);
6883 + }
6884 +
6885 + // Scenario 7: A header-only message arriving when a larger message type is expected
6886 + // is rejected with E_UNEXPECTED because the received span is too small for the type.
6887 + {
6888 + auto [client, server] = MakeSocketPair();
6889 + wsl::shared::SocketChannel channel{std::move(server), "server"};
6890 +
6891 + MESSAGE_HEADER header{};
6892 + header.MessageType = RESULT_MESSAGE<int32_t>::Type;
6893 + header.MessageSize = sizeof(header);
6894 + header.TransactionId = 1;
6895 + header.TransactionStep = static_cast<unsigned int>(TRANSACTION_STEP::NONE);
6896 + WriteSocket(client.get(), &header, sizeof(header));
6897 +
6898 + const auto hr = wil::ResultFromException([&]() { channel.ReceiveMessage<RESULT_MESSAGE<int32_t>>(); });
6899 + VERIFY_ARE_EQUAL(hr, E_UNEXPECTED);
6900 + }
6901 +
6902 + // Scenario 8: ReceiveMessage with a finite timeout on an idle socket throws
6903 + // HRESULT_FROM_WIN32(ERROR_TIMEOUT) once the timeout elapses.
6904 + {
6905 + auto [client, server] = MakeSocketPair();
6906 + wsl::shared::SocketChannel channel{std::move(server), "server"};
6907 +
6908 + const auto hr = wil::ResultFromException([&]() { channel.ReceiveMessage<RESULT_MESSAGE<int32_t>>(nullptr, 100); });
6909 + VERIFY_ARE_EQUAL(hr, HRESULT_FROM_WIN32(ERROR_TIMEOUT));
6910 + }
6911 + }
6912 +
6913 }; // namespace UnitTests
6914 } // namespace UnitTests