Rethink the Accept() logic to differentiate between errors and cancellation (#14156)

* Bring relay changes

* Redesign Accept() logic to differenciate between cancellation and errors

* Prepare for PR

* Apply PR feedback

Author: OneBlue

src/shared/inc/defs.h

@@ -22,6 +22,18 @@ Module Name:
 #define _wcsicmp wcscasecmp
 #endif
 
+#define NON_COPYABLE(Type) \
+    Type(const Type&) = delete; \
+    Type& operator=(const Type&) = delete;
+
+#define NON_MOVABLE(Type) \
+    Type(Type&&) = delete; \
+    Type& operator=(Type&&) = delete;
+
+#define DEFAULT_MOVABLE(Type) \
+    Type(Type&&) = default; \
+    Type& operator=(Type&&) = default;
+
 namespace wsl::shared {
 
 inline constexpr std::uint32_t VersionMajor = WSL_PACKAGE_VERSION_MAJOR;

src/windows/common/hvsocket.cpp

@@ -39,11 +39,14 @@ void InitializeWildcardSocketAddress(_Out_ PSOCKADDR_HV Address)
 }
 } // namespace
 
-wil::unique_socket wsl::windows::common::hvsocket::Accept(
-    _In_ SOCKET ListenSocket, _In_ int Timeout, _In_opt_ HANDLE ExitHandle, _In_ const std::source_location& Location)
+std::optional<wil::unique_socket> wsl::windows::common::hvsocket::CancellableAccept(
+    _In_ SOCKET ListenSocket, _In_ DWORD Timeout, _In_opt_ HANDLE ExitHandle, _In_ const std::source_location& Location)
 {
     wil::unique_socket Socket = Create();
-    wsl::windows::common::socket::Accept(ListenSocket, Socket.get(), Timeout, ExitHandle, Location);
+    if (!socket::CancellableAccept(ListenSocket, Socket.get(), Timeout, ExitHandle, Location))
+    {
+        return {};
+    }
 
     return Socket;
 }

src/windows/common/hvsocket.hpp

@@ -19,9 +19,9 @@ Module Name:
 
 namespace wsl::windows::common::hvsocket {
 
-wil::unique_socket Accept(
+std::optional<wil::unique_socket> CancellableAccept(
     _In_ SOCKET ListenSocket,
-    _In_ int Timeout,
+    _In_ DWORD Timeout,
     _In_opt_ HANDLE ExitHandle = nullptr,
     const std::source_location& Location = std::source_location::current());
 

src/windows/common/relay.cpp

@@ -16,8 +16,14 @@ Module Name:
 #include "relay.hpp"
 #pragma hdrstop
 
+using wsl::windows::common::relay::EventHandle;
+using wsl::windows::common::relay::HandleWrapper;
+using wsl::windows::common::relay::IOHandleStatus;
+using wsl::windows::common::relay::MultiHandleWait;
+using wsl::windows::common::relay::OverlappedIOHandle;
 using wsl::windows::common::relay::ScopedMultiRelay;
 using wsl::windows::common::relay::ScopedRelay;
+using wsl::windows::common::relay::SingleAcceptHandle;
 
 namespace {
 
@@ -108,7 +114,7 @@ wsl::windows::common::relay::InterruptableRead(
             return 0;
         }
 
-        THROW_LAST_ERROR_IF(lastError != ERROR_IO_PENDING);
+        THROW_LAST_ERROR_IF_MSG(lastError != ERROR_IO_PENDING, "Handle: 0x%p", (void*)InputHandle);
 
         auto cancelRead = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [&] {
             CancelIoEx(InputHandle, Overlapped);
@@ -569,4 +575,215 @@ try
         }
     }
 }
-CATCH_LOG()
+CATCH_LOG()
+
+void MultiHandleWait::AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags)
+{
+    m_handles.emplace_back(flags, std::move(handle));
+}
+
+void MultiHandleWait::Cancel()
+{
+    m_cancel = true;
+}
+bool MultiHandleWait::Run(std::optional<std::chrono::milliseconds> Timeout)
+{
+    m_cancel = false; // Run may be called multiple times.
+
+    std::optional<std::chrono::steady_clock::time_point> deadline;
+
+    if (Timeout.has_value())
+    {
+        deadline = std::chrono::steady_clock::now() + Timeout.value();
+    }
+
+    // Run until all handles are completed.
+
+    while (!m_handles.empty() && !m_cancel)
+    {
+        // Schedule IO on each handle until all are either pending, or completed.
+        for (size_t i = 0; i < m_handles.size(); i++)
+        {
+            while (m_handles[i].second->GetState() == IOHandleStatus::Standby)
+            {
+                try
+                {
+                    m_handles[i].second->Schedule();
+                }
+                catch (...)
+                {
+                    if (WI_IsFlagSet(m_handles[i].first, Flags::IgnoreErrors))
+                    {
+                        m_handles[i].second.reset(); // Reset the handle so it can be deleted.
+                    }
+                    else
+                    {
+                        throw;
+                    }
+                }
+            }
+        }
+
+        // Remove completed handles from m_handles.
+        for (auto it = m_handles.begin(); it != m_handles.end();)
+        {
+            if (!it->second)
+            {
+                it = m_handles.erase(it);
+            }
+            else if (it->second->GetState() == IOHandleStatus::Completed)
+            {
+                if (WI_IsFlagSet(it->first, Flags::CancelOnCompleted))
+                {
+                    m_cancel = true; // Cancel the IO if a handle with CancelOnCompleted is in the completed state.
+                }
+
+                it = m_handles.erase(it);
+            }
+            else
+            {
+                ++it;
+            }
+        }
+
+        if (m_handles.empty() || m_cancel)
+        {
+            break;
+        }
+
+        // Wait for the next operation to complete.
+        std::vector<HANDLE> waitHandles;
+        for (const auto& e : m_handles)
+        {
+            waitHandles.emplace_back(e.second->GetHandle());
+        }
+
+        DWORD waitTimeout = INFINITE;
+        if (deadline.has_value())
+        {
+            auto miliseconds =
+                std::chrono::duration_cast<std::chrono::milliseconds>(deadline.value() - std::chrono::steady_clock::now()).count();
+
+            waitTimeout = static_cast<DWORD>(std::max(0LL, miliseconds));
+        }
+
+        auto result = WaitForMultipleObjects(static_cast<DWORD>(waitHandles.size()), waitHandles.data(), false, waitTimeout);
+        if (result == WAIT_TIMEOUT)
+        {
+            THROW_WIN32(ERROR_TIMEOUT);
+        }
+        else if (result >= WAIT_OBJECT_0 && result < WAIT_OBJECT_0 + m_handles.size())
+        {
+            auto index = result - WAIT_OBJECT_0;
+
+            try
+            {
+                m_handles[index].second->Collect();
+            }
+            catch (...)
+            {
+                if (WI_IsFlagSet(m_handles[index].first, Flags::IgnoreErrors))
+                {
+                    m_handles.erase(m_handles.begin() + index);
+                }
+                else
+                {
+                    throw;
+                }
+            }
+        }
+        else
+        {
+            THROW_LAST_ERROR_MSG("Timeout: %lu, Count: %llu", waitTimeout, waitHandles.size());
+        }
+    }
+
+    return !m_cancel;
+}
+
+IOHandleStatus OverlappedIOHandle::GetState() const
+{
+    return State;
+}
+
+EventHandle::EventHandle(HandleWrapper&& Handle, std::function<void()>&& OnSignalled) :
+    Handle(std::move(Handle)), OnSignalled(std::move(OnSignalled))
+{
+}
+
+void EventHandle::Schedule()
+{
+    State = IOHandleStatus::Pending;
+}
+
+void EventHandle::Collect()
+{
+    State = IOHandleStatus::Completed;
+    OnSignalled();
+}
+
+HANDLE EventHandle::GetHandle() const
+{
+    return Handle.Get();
+}
+
+SingleAcceptHandle::SingleAcceptHandle(HandleWrapper&& ListenSocket, HandleWrapper&& AcceptedSocket, std::function<void()>&& OnAccepted) :
+    ListenSocket(std::move(ListenSocket)), AcceptedSocket(std::move(AcceptedSocket)), OnAccepted(std::move(OnAccepted))
+{
+    Overlapped.hEvent = Event.get();
+}
+
+SingleAcceptHandle::~SingleAcceptHandle()
+{
+    if (State == IOHandleStatus::Pending)
+    {
+        LOG_IF_WIN32_BOOL_FALSE(CancelIoEx(ListenSocket.Get(), &Overlapped));
+
+        DWORD bytesProcessed{};
+        DWORD flagsReturned{};
+        if (!WSAGetOverlappedResult((SOCKET)ListenSocket.Get(), &Overlapped, &bytesProcessed, TRUE, &flagsReturned))
+        {
+            auto error = GetLastError();
+            LOG_LAST_ERROR_IF(error != ERROR_CONNECTION_ABORTED && error != ERROR_OPERATION_ABORTED);
+        }
+    }
+}
+
+void SingleAcceptHandle::Schedule()
+{
+    WI_ASSERT(State == IOHandleStatus::Standby);
+
+    // Schedule the accept.
+    DWORD bytesReturned{};
+    if (AcceptEx((SOCKET)ListenSocket.Get(), (SOCKET)AcceptedSocket.Get(), &AcceptBuffer, 0, sizeof(SOCKADDR_STORAGE), sizeof(SOCKADDR_STORAGE), &bytesReturned, &Overlapped))
+    {
+        // Accept completed immediately.
+        State = IOHandleStatus::Completed;
+        OnAccepted();
+    }
+    else
+    {
+        auto error = WSAGetLastError();
+        THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_IO_PENDING, "Handle: 0x%p", (void*)ListenSocket.Get());
+
+        State = IOHandleStatus::Pending;
+    }
+}
+
+void SingleAcceptHandle::Collect()
+{
+    WI_ASSERT(State == IOHandleStatus::Pending);
+
+    DWORD bytesReceived{};
+    DWORD flagsReturned{};
+
+    THROW_IF_WIN32_BOOL_FALSE(WSAGetOverlappedResult((SOCKET)ListenSocket.Get(), &Overlapped, &bytesReceived, false, &flagsReturned));
+
+    State = IOHandleStatus::Completed;
+    OnAccepted();
+}
+
+HANDLE SingleAcceptHandle::GetHandle() const
+{
+    return Event.get();
+}