Fix broken WSLCorePort channel after receive timeout

src/linux/init/GnsEngine.cpp

@@ -25,12 +25,13 @@ constexpr auto c_ipStrings = {"ip", "ip6"};
 const char* c_loopbackInterfaceName = "lo";
 
 GnsEngine::GnsEngine(
+    wsl::shared::SocketChannel& channel,
     const NotificationRoutine& notificationRoutine,
     const StatusRoutine& statusRoutine,
     NetworkManager& manager,
     std::optional<int> dnsTunnelingFd,
     const std::string& dnsTunnelingIpAddress) :
-    notificationRoutine(notificationRoutine), statusRoutine(statusRoutine), manager(manager)
+    channel(channel), notificationRoutine(notificationRoutine), statusRoutine(statusRoutine), manager(manager)
 {
     if (dnsTunnelingFd.has_value())
     {
@@ -363,11 +364,11 @@ void GnsEngine::ProcessLinkChange(Interface& interface, const wsl::shared::hns::
     }
 }
 
-std::tuple<bool, int> GnsEngine::ProcessNextMessage()
+std::tuple<bool, int> GnsEngine::ProcessNextMessage(wsl::shared::Transaction& transaction)
 {
     int return_value = 0;
 
-    auto payload = notificationRoutine();
+    auto payload = notificationRoutine(transaction);
     if (!payload.has_value())
     {
         GNS_LOG_ERROR("Received empty message, exiting");
@@ -723,22 +724,23 @@ void GnsEngine::run()
 
     while (true)
     {
+        auto transaction = channel.ReceiveTransaction();
         try
         {
             GNS_LOG_INFO("Processing Next Message");
-            auto [should_continue, return_value] = ProcessNextMessage();
+            auto [should_continue, return_value] = ProcessNextMessage(transaction);
             if (!should_continue)
             {
                 break;
             }
 
             GNS_LOG_INFO("Processing Next Message Successful ({:#x})", return_value);
-            statusRoutine(return_value, "");
+            statusRoutine(return_value, "", transaction);
         }
         catch (const std::exception& e)
         {
             GNS_LOG_ERROR("Error while processing message: {}", e.what());
-            statusRoutine(-1, e.what());
+            statusRoutine(-1, e.what(), transaction);
         }
     }
 

src/linux/init/GnsEngine.h

@@ -21,10 +21,11 @@ class GnsEngine
         std::optional<GUID> AdapterId;
     };
 
-    using NotificationRoutine = std::function<std::optional<Message>()>;
-    using StatusRoutine = std::function<void(int, const std::string&)>;
+    using NotificationRoutine = std::function<std::optional<Message>(wsl::shared::Transaction&)>;
+    using StatusRoutine = std::function<void(int, const std::string&, wsl::shared::Transaction&)>;
 
     GnsEngine(
+        wsl::shared::SocketChannel& channel,
         const NotificationRoutine& notificationRoutine,
         const StatusRoutine& statusRoutine,
         NetworkManager& manager,
@@ -36,7 +37,7 @@ class GnsEngine
     void run();
 
 private:
-    std::tuple<bool, int> ProcessNextMessage();
+    std::tuple<bool, int> ProcessNextMessage(wsl::shared::Transaction& transaction);
 
     void ProcessNotification(const nlohmann::json& payload, Interface& interface);
 
@@ -69,6 +70,7 @@ class GnsEngine
     void ProcessNotificationImpl(
         Interface& interface, const nlohmann::json& payload, void (GnsEngine::*routine)(Interface&, const T&, wsl::shared::hns::ModifyRequestType));
 
+    wsl::shared::SocketChannel& channel;
     const NotificationRoutine& notificationRoutine;
     const StatusRoutine& statusRoutine;
     NetworkManager& manager;

src/linux/init/binfmt.cpp

@@ -174,7 +174,8 @@ try
     // Send the create process message to the interop server.
     //
 
-    channel.SendMessage<LX_INIT_CREATE_NT_PROCESS_UTILITY_VM>(Span);
+    auto transaction = channel.StartTransaction();
+    transaction.Send<LX_INIT_CREATE_NT_PROCESS_UTILITY_VM>(Span);
 
     //
     // Accept connections from the interop server.

src/linux/init/config.cpp

@@ -334,7 +334,7 @@ try
 CATCH_LOG()
 
 void ConfigHandleInteropMessage(
-    wsl::shared::SocketChannel& ResponseChannel,
+    wsl::shared::Transaction& Transaction,
     wsl::shared::SocketChannel& InteropChannel,
     bool Elevated,
     gsl::span<gsl::byte> Message,
@@ -350,7 +350,7 @@ Routine Description:
 
 Arguments:
 
-    ResponseChannel - Supplies channel used to send responses.
+    Transaction - Supplies transaction used to send responses.
 
     InteropChannel - Supplies a channel to the host to be used for create
         process requests.
@@ -381,7 +381,7 @@ try
 
     case LxInitMessageQueryDrvfsElevated:
     {
-        ResponseChannel.SendResultMessage<bool>(Elevated);
+        Transaction.SendResultMessage<bool>(Elevated);
         break;
     }
 
@@ -397,15 +397,15 @@ try
         auto Value = UtilGetEnvironmentVariable(Query->Buffer);
         wsl::shared::MessageWriter<LX_INIT_QUERY_ENVIRONMENT_VARIABLE> Response(LxInitMessageQueryEnvironmentVariable);
         Response.WriteString(Value);
-        ResponseChannel.SendMessage<LX_INIT_QUERY_ENVIRONMENT_VARIABLE>(Response.Span());
+        Transaction.Send<LX_INIT_QUERY_ENVIRONMENT_VARIABLE>(Response.Span());
     }
 
     break;
 
     case LxInitMessageQueryFeatureFlags:
     {
         assert(Config.FeatureFlags.has_value());
-        ResponseChannel.SendResultMessage<int32_t>(Config.FeatureFlags.value());
+        Transaction.SendResultMessage<int32_t>(Config.FeatureFlags.value());
         break;
     }
 
@@ -419,7 +419,7 @@ try
         }
 
         bool success = false;
-        auto sendResponse = wil::scope_exit([&]() { ResponseChannel.SendResultMessage<bool>(success); });
+        auto sendResponse = wil::scope_exit([&]() { Transaction.SendResultMessage<bool>(success); });
 
         if (!Config.BootInit || Config.InitPid.value_or(0) != getpid())
         {
@@ -435,7 +435,7 @@ try
 
     case LxInitMessageQueryNetworkingMode:
         assert(Config.NetworkingMode.has_value());
-        ResponseChannel.SendResultMessage<uint8_t>(static_cast<uint8_t>(Config.NetworkingMode.value()));
+        Transaction.SendResultMessage<uint8_t>(static_cast<uint8_t>(Config.NetworkingMode.value()));
         break;
 
     case LxInitMessageQueryVmId:
@@ -446,7 +446,7 @@ try
             Response.WriteString(Config.VmId.value());
         }
 
-        ResponseChannel.SendMessage<LX_INIT_QUERY_VM_ID>(Response.Span());
+        Transaction.Send<LX_INIT_QUERY_VM_ID>(Response.Span());
         break;
     }
 
@@ -618,7 +618,7 @@ try
 }
 CATCH_LOG()
 
-int ConfigInitializeInstance(wsl::shared::SocketChannel& Channel, gsl::span<gsl::byte> Buffer, wsl::linux::WslDistributionConfig& Config)
+int ConfigInitializeInstance(const std::function<void(const gsl::span<gsl::byte>&)>& SendResponse, gsl::span<gsl::byte> Buffer, wsl::linux::WslDistributionConfig& Config)
 
 /*++
 
@@ -632,7 +632,7 @@ Routine Description:
 
 Arguments:
 
-    MessageFd - Supplies a file descriptor to send the response message.
+    SendResponse - Supplies a function to send the response message.
 
     Buffer - Supplies the message buffer.
 
@@ -923,7 +923,7 @@ try
         Response.WriteString(Response->VersionIndex, Version->c_str());
     }
 
-    Channel.SendMessage<LX_INIT_CONFIGURATION_INFORMATION_RESPONSE>(Response.Span());
+    SendResponse(Response.Span());
 
     //
     // Accept the interop connection.
@@ -973,13 +973,14 @@ try
                     continue;
                 }
 
-                auto [Message, Span] = ClientChannel.ReceiveMessageOrClosed<MESSAGE_HEADER>();
+                auto transaction = ClientChannel.ReceiveTransaction();
+                auto [Message, Span] = transaction.ReceiveOrClosed<MESSAGE_HEADER>();
                 if (Message == nullptr)
                 {
                     continue;
                 }
 
-                ConfigHandleInteropMessage(ClientChannel, InteropChannel, Elevated, Span, Message, Config);
+                ConfigHandleInteropMessage(transaction, InteropChannel, Elevated, Span, Message, Config);
             }
         });
 
@@ -2186,7 +2187,7 @@ Return Value:
     return Result;
 }
 
-int ConfigRemountDrvFs(gsl::span<gsl::byte> Buffer, wsl::shared::SocketChannel& Channel, const wsl::linux::WslDistributionConfig& Config)
+int ConfigRemountDrvFs(gsl::span<gsl::byte> Buffer, wsl::shared::Transaction& Transaction, const wsl::linux::WslDistributionConfig& Config)
 
 /*++
 
@@ -2207,7 +2208,7 @@ Return Value:
 
 --*/
 {
-    Channel.SendResultMessage<int32_t>(ConfigRemountDrvFsImpl(Buffer, Config));
+    Transaction.SendResultMessage<int32_t>(ConfigRemountDrvFsImpl(Buffer, Config));
 
     return 0;
 }

src/linux/init/config.h

@@ -20,6 +20,7 @@ Module Name:
 #include <set>
 #include <string_view>
 #include <optional>
+#include <functional>
 #include "SocketChannel.h"
 #include "WslDistributionConfig.h"
 
@@ -399,7 +400,7 @@ std::set<std::pair<unsigned int, std::string>> ConfigGetMountedDrvFsVolumes(void
 std::vector<std::pair<std::string, std::string>> ConfigGetWslgEnvironmentVariables(const wsl::linux::WslDistributionConfig& Config);
 
 void ConfigHandleInteropMessage(
-    wsl::shared::SocketChannel& ResponseChannel,
+    wsl::shared::Transaction& Transaction,
     wsl::shared::SocketChannel& InteropChannel,
     bool Elevated,
     gsl::span<gsl::byte> Message,
@@ -408,7 +409,7 @@ void ConfigHandleInteropMessage(
 
 void ConfigInitializeCgroups(wsl::linux::WslDistributionConfig& Config);
 
-int ConfigInitializeInstance(wsl::shared::SocketChannel& Channel, gsl::span<gsl::byte> Buffer, wsl::linux::WslDistributionConfig& Config);
+int ConfigInitializeInstance(const std::function<void(const gsl::span<gsl::byte>&)>& SendResponse, gsl::span<gsl::byte> Buffer, wsl::linux::WslDistributionConfig& Config);
 
 void ConfigMountDrvFsVolumes(unsigned int DrvFsVolumes, uid_t OwnerUid, std::optional<bool> Admin, const wsl::linux::WslDistributionConfig& Config);
 
@@ -420,7 +421,7 @@ int ConfigRegisterBinfmtInterpreter(void);
 
 int ConfigSetMountNamespace(bool Elevated);
 
-int ConfigRemountDrvFs(gsl::span<gsl::byte> Buffer, wsl::shared::SocketChannel& Channel, const wsl::linux::WslDistributionConfig& Config);
+int ConfigRemountDrvFs(gsl::span<gsl::byte> Buffer, wsl::shared::Transaction& Transaction, const wsl::linux::WslDistributionConfig& Config);
 
 int ConfigRemountDrvFsImpl(gsl::span<gsl::byte> Buffer, const wsl::linux::WslDistributionConfig& Config);
 

src/linux/init/drvfs.cpp

@@ -209,8 +209,9 @@ Return Value:
     QueryPortMessage.MessageType = LxInitMessageQueryDrvfsElevated;
     QueryPortMessage.MessageSize = sizeof(QueryPortMessage);
 
-    channel.SendMessage(QueryPortMessage);
-    return channel.ReceiveMessage<RESULT_MESSAGE<bool>>().Result;
+    auto transaction = channel.StartTransaction();
+    transaction.Send(QueryPortMessage);
+    return transaction.Receive<RESULT_MESSAGE<bool>>().Result;
 }
 
 int MountFilesystem(const char* FsType, const char* Source, const char* Target, const char* Options, int* ExitCode)