cleanup: refactor wslrelay and other minor cleanup (#14099)

* cleanup: refactor wslrelay and other minor cleanup

* pr feedback

---------

Co-authored-by: Ben Hillis <benhill@ntdev.microsoft.com>

src/windows/common/CMakeLists.txt

@@ -1,6 +1,7 @@
 set(SOURCES
     ConsoleProgressBar.cpp
     ConsoleProgressIndicator.cpp
+    ConsoleState.cpp
     DeviceHostProxy.cpp
     disk.cpp
     Distribution.cpp
@@ -32,7 +33,6 @@ set(SOURCES
     string.cpp
     SubProcess.cpp
     svccomm.cpp
-    ConsoleState.cpp
     WslClient.cpp
     WslCoreConfig.cpp
     WslCoreFilesystem.cpp
@@ -74,6 +74,7 @@ set(HEADERS
     ../inc/wslrelay.h
     ConsoleProgressBar.h
     ConsoleProgressIndicator.h
+    ConsoleState.h
     DeviceHostProxy.h
     disk.hpp
     Distribution.h
@@ -107,7 +108,6 @@ set(HEADERS
     Stringify.h
     SubProcess.h
     svccomm.hpp
-    ConsoleState.h
     WslClient.h
     WslCoreConfig.h
     WslCoreFilesystem.h

src/windows/common/ConsoleState.cpp

@@ -9,6 +9,7 @@ Module Name:
 Abstract:
 
     This file contains function definitions for the ConsoleState helper class.
+
 --*/
 
 #include "precomp.h"

src/windows/common/helpers.cpp

@@ -178,6 +178,11 @@ private:
         launcher.AddOption(wslrelay::disable_telemetry_option);
     }
 
+    if (WI_IsFlagSet(Flags, LaunchWslRelayFlags::ConnectPipe))
+    {
+        launcher.AddOption(wslrelay::connect_pipe_option);
+    }
+
     return launcher.Launch(UserToken, WI_IsFlagSet(Flags, LaunchWslRelayFlags::HideWindow));
 }
 } // namespace
@@ -501,11 +506,10 @@ bool wsl::windows::common::helpers::IsWslSupportInterfacePresent()
 void wsl::windows::common::helpers::LaunchDebugConsole(
     _In_ LPCWSTR PipeName, _In_ bool ConnectExistingPipe, _In_ HANDLE UserToken, _In_opt_ HANDLE LogFile, _In_ bool DisableTelemetry)
 {
-    wslrelay::RelayMode relayMode;
+    LaunchWslRelayFlags flags{};
     wil::unique_hfile pipe;
     if (ConnectExistingPipe)
     {
-        relayMode = wslrelay::RelayMode::DebugConsoleRelay;
         // Connect to an existing pipe. The connection should be:
         //     Asynchronous (FILE_FLAG_OVERLAPPED)
         //     Anonymous (SECURITY_SQOS_PRESENT | SECURITY_ANONYMOUS)
@@ -515,21 +519,20 @@ void wsl::windows::common::helpers::LaunchDebugConsole(
     }
     else
     {
-        relayMode = wslrelay::RelayMode::DebugConsole;
-        // Create a new pipe server. The pipe should be:
+        // Create a new pipe server the child process will connect to. The pipe should be:
         //     Bi-directional: PIPE_ACCESS_DUPLEX
         //     Asynchronous: FILE_FLAG_OVERLAPPED
         //     Raw: PIPE_TYPE_BYTE | PIPE_READMODE_BYTE
         //     Blocking: PIPE_WAIT
+        WI_SetFlag(flags, LaunchWslRelayFlags::ConnectPipe);
         pipe.reset(CreateNamedPipeW(
             PipeName, (PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED), (PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT), 1, LX_RELAY_BUFFER_SIZE, LX_RELAY_BUFFER_SIZE, 0, nullptr));
     }
 
     THROW_LAST_ERROR_IF(!pipe);
 
-    LaunchWslRelayFlags flags{};
     WI_SetFlagIf(flags, LaunchWslRelayFlags::DisableTelemetry, DisableTelemetry);
-    wil::unique_handle info{LaunchWslRelay(relayMode, LogFile, nullptr, pipe.get(), {}, nullptr, UserToken, flags)};
+    wil::unique_handle info{LaunchWslRelay(wslrelay::RelayMode::DebugConsole, LogFile, nullptr, pipe.get(), {}, nullptr, UserToken, flags)};
 }
 
 [[nodiscard]] wil::unique_handle wsl::windows::common::helpers::LaunchInteropServer(
@@ -550,7 +553,7 @@ void wsl::windows::common::helpers::LaunchKdRelay(_In_ LPCWSTR PipeName, _In_ HA
 
     THROW_LAST_ERROR_IF(!pipe);
 
-    LaunchWslRelayFlags flags{};
+    LaunchWslRelayFlags flags = LaunchWslRelayFlags::ConnectPipe;
     WI_SetFlagIf(flags, LaunchWslRelayFlags::DisableTelemetry, DisableTelemetry);
     wil::unique_handle info{LaunchWslRelay(wslrelay::RelayMode::KdRelay, nullptr, nullptr, pipe.get(), Port, ExitEvent, UserToken, flags)};
 }

src/windows/common/helpers.hpp

@@ -62,7 +62,8 @@ enum class LaunchWslRelayFlags
 {
     None = 0,
     DisableTelemetry = 1,
-    HideWindow = 2
+    HideWindow = 2,
+    ConnectPipe = 4
 };
 
 DEFINE_ENUM_FLAG_OPERATORS(LaunchWslRelayFlags);

src/windows/inc/wslrelay.h

@@ -19,7 +19,6 @@ enum RelayMode
 {
     Invalid = -1,
     DebugConsole,
-    DebugConsoleRelay,
     PortRelay,
     KdRelay
 };
@@ -32,4 +31,5 @@ LPCWSTR const pipe_option = L"--pipe";
 LPCWSTR const exit_event_option = L"--exit-event";
 LPCWSTR const port_option = L"--port";
 LPCWSTR const disable_telemetry_option = L"--disable-telemetry";
+LPCWSTR const connect_pipe_option = L"--connect-pipe";
 } // namespace wslrelay

src/windows/wslrelay/main.cpp

@@ -38,9 +38,10 @@ try
     wslrelay::RelayMode mode{wslrelay::RelayMode::Invalid};
     wil::unique_handle pipe{};
     wil::unique_handle exitEvent{};
-    int port{};
+    uint32_t port{};
     GUID vmId{};
     bool disableTelemetry = !wsl::shared::OfficialBuild;
+    bool connectPipe = false;
 
     ArgumentParser parser(GetCommandLineW(), wslrelay::binary_name);
     parser.AddArgument(Integer(reinterpret_cast<int&>(mode)), wslrelay::mode_option);
@@ -50,17 +51,29 @@ try
     parser.AddArgument(Handle{exitEvent}, wslrelay::exit_event_option);
     parser.AddArgument(Integer{port}, wslrelay::port_option);
     parser.AddArgument(disableTelemetry, wslrelay::disable_telemetry_option);
+    parser.AddArgument(connectPipe, wslrelay::connect_pipe_option);
     parser.Parse();
 
     // Initialize logging.
     WslTraceLoggingInitialize(LxssTelemetryProvider, disableTelemetry);
     auto cleanup = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [] { WslTraceLoggingUninitialize(); });
 
+    // Ensure that the other end of the pipe has connected if required.
+    if (connectPipe)
+    {
+        std::vector<HANDLE> exitEvents;
+        if (exitEvent)
+        {
+            exitEvents.push_back(exitEvent.get());
+        }
+
+        wsl::windows::common::helpers::ConnectPipe(pipe.get(), (15 * 1000), exitEvents);
+    }
+
     // Perform the requested operation.
     switch (mode)
     {
     case wslrelay::RelayMode::DebugConsole:
-    case wslrelay::RelayMode::DebugConsoleRelay:
     {
         // If not relaying to a file, create a console window.
         if (!handle)
@@ -68,12 +81,6 @@ try
             wsl::windows::common::helpers::CreateConsole(L"WSL Debug Console");
         }
 
-        if (mode == wslrelay::RelayMode::DebugConsole)
-        {
-            // Ensure that the other end of the pipe has connected.
-            wsl::windows::common::helpers::ConnectPipe(pipe.get(), (15 * 1000));
-        }
-
         // Relay the contents of the pipe to the output handle.
         wsl::windows::common::relay::InterruptableRelay(pipe.get(), handle ? handle.get() : GetStdHandle(STD_OUTPUT_HANDLE));
 
@@ -98,9 +105,6 @@ try
     {
         THROW_HR_IF(E_INVALIDARG, port == 0);
 
-        // Ensure that the other end of the pipe has connected.
-        wsl::windows::common::helpers::ConnectPipe(pipe.get(), (15 * 1000), {exitEvent.get()});
-
         // Bind, listen, and accept a connection on the specified port.
         const wil::unique_socket listenSocket(WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, nullptr, 0, WSA_FLAG_OVERLAPPED));
         THROW_LAST_ERROR_IF(!listenSocket);
@@ -126,7 +130,7 @@ try
     }
 
     default:
-        THROW_HR(E_INVALIDARG);
+        THROW_HR_MSG(E_INVALIDARG, "Invalid relay mode %d specified.", static_cast<int>(mode));
     }
 
     return 0;