Add test wsla rootfs initialization in wsl client (#13723)

CMakeLists.txt

@@ -65,6 +65,7 @@ find_nuget_package(Microsoft.WSL.Kernel KERNEL /build/native)
 find_nuget_package(Microsoft.WSL.bsdtar BSDTARD /build/native/bin)
 find_nuget_package(Microsoft.WSL.LinuxSdk LINUXSDK /)
 find_nuget_package(Microsoft.WSL.TestDistro TEST_DISTRO /)
+find_nuget_package(Microsoft.WSL.WSLATestDistro WSLA_TEST_DISTRO /)
 find_nuget_package(Microsoft.WSLg WSLG /build/native/bin)
 find_nuget_package(StrawberryPerl PERL /)
 find_nuget_package(vswhere VSWHERE /tools)
@@ -380,6 +381,10 @@ if (DEFINED WSL_DEV_BINARY_PATH) # Development shortcut to make the package smal
                             WSL_GPU_LIB_PATH="${WSL_DEV_BINARY_PATH}/lib")
 endif()
 
+if (NOT DEFINED OFFICIAL_BUILD AND ${TARGET_PLATFORM} STREQUAL "x64")
+    add_compile_definitions(WSLA_TEST_DISTRO_PATH="${WSLA_TEST_DISTRO_SOURCE_DIR}/wslatestrootfs.vhd")
+endif()
+
 # Common include paths
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/wil/include)
 include_directories(${WSLDEPS_SOURCE_DIR}/include)

packages.config

@@ -24,6 +24,7 @@
   <package id="Microsoft.WSL.LxUtil.amd64fre" version="10.0.26100.1-240331-1435.ge-release" />
   <package id="Microsoft.WSL.LxUtil.arm64fre" version="10.0.26100.1-240331-1435.ge-release" />
   <package id="Microsoft.WSL.TestDistro" version="2.5.7-47" />
+  <package id="Microsoft.WSL.WSLATestDistro" version="0.1.0" />
   <package id="Microsoft.WSLg" version="1.0.71" />
   <package id="Microsoft.Xaml.Behaviors.WinUI.Managed" version="3.0.0" />
   <package id="StrawberryPerl" version="5.28.0.1" />

src/windows/common/WslClient.cpp

@@ -21,6 +21,7 @@ Abstract:
 #include "wslaservice.h"
 #include "WSLAApi.h"
 #include "WSLAProcessLauncher.h"
+#include "WslCoreFilesystem.h"
 
 #define BASH_PATH L"/bin/bash"
 
@@ -1527,13 +1528,17 @@ int RunDebugShell()
 // Temporary debugging tool for WSLA
 int WslaShell(_In_ std::wstring_view commandLine)
 {
-#ifdef WSL_SYSTEM_DISTRO_PATH
+#ifdef WSLA_TEST_DISTRO_PATH
 
-    std::wstring vhd = TEXT(WSL_SYSTEM_DISTRO_PATH);
+    std::wstring vhd = TEXT(WSLA_TEST_DISTRO_PATH);
+    std::string shell = "/bin/sh";
+    std::string fsType = "squashfs";
 
 #else
 
     std::wstring vhd = wsl::windows::common::wslutil::GetMsiPackagePath().value() + L"/system.vhd";
+    std::string shell = "/bin/bash";
+    std::string fsType = "ext4";
 
 #endif
 
@@ -1543,8 +1548,7 @@ int WslaShell(_In_ std::wstring_view commandLine)
     settings.MemoryMb = 1024;
     settings.BootTimeoutMs = 30000;
     settings.NetworkingMode = WSLANetworkingModeNAT;
-    std::string shell = "/bin/bash";
+    std::wstring containerRootVhd;
-    std::string fsType = "ext4";
     bool help = false;
 
     ArgumentParser parser(std::wstring{commandLine}, WSL_BINARY_NAME);
@@ -1554,25 +1558,39 @@ int WslaShell(_In_ std::wstring_view commandLine)
     parser.AddArgument(Integer(settings.MemoryMb), L"--memory");
     parser.AddArgument(Integer(settings.CpuCount), L"--cpu");
     parser.AddArgument(Utf8String(fsType), L"--fstype");
+    parser.AddArgument(containerRootVhd, L"--container-vhd");
     parser.AddArgument(help, L"--help");
     parser.Parse();
 
     if (help)
     {
         const auto usage = std::format(
-            LR"({} --wsla [--vhd </path/to/vhd>] [--shell </path/to/shell>] [--memory <memory-mb>] [--cpu <cpus>] [--dns-tunneling] [--fstype <fstype>] [--new-api] [--help])",
+            LR"({} --wsla [--vhd </path/to/vhd>] [--shell </path/to/shell>] [--memory <memory-mb>] [--cpu <cpus>] [--dns-tunneling] [--fstype <fstype>] [--container-vhd </path/to/vhd>] [--help])",
             WSL_BINARY_NAME);
 
         wprintf(L"%ls\n", usage.c_str());
         return 1;
     }
 
+    if (!containerRootVhd.empty())
+    {
+        settings.ContainerRootVhd = containerRootVhd.c_str();
+
+        if (!std::filesystem::exists(containerRootVhd))
+        {
+            auto token = wil::open_current_access_token();
+            auto tokenInfo = wil::get_token_information<TOKEN_USER>(token.get());
+            wsl::core::filesystem::CreateVhd(containerRootVhd.c_str(), 5368709120 /* 5 GB */, tokenInfo->User.Sid, FALSE, FALSE);
+            settings.FormatContainerRootVhd = TRUE;
+        }
+    }
+
     wil::com_ptr<IWSLAUserSession> userSession;
     THROW_IF_FAILED(CoCreateInstance(__uuidof(WSLAUserSession), nullptr, CLSCTX_LOCAL_SERVER, IID_PPV_ARGS(&userSession)));
     wsl::windows::common::security::ConfigureForCOMImpersonation(userSession.get());
 
     wil::com_ptr<IWSLAVirtualMachine> virtualMachine;
-    WSLA_SESSION_SETTINGS sessionSettings{L"my-display-name"};
+    WSLA_SESSION_SETTINGS sessionSettings{L"WSLA Test Session"};
     wil::com_ptr<IWSLASession> session;
     settings.RootVhd = vhd.c_str();
     settings.RootVhdType = fsType.c_str();
@@ -1581,6 +1599,13 @@ int WslaShell(_In_ std::wstring_view commandLine)
 
     wsl::windows::common::security::ConfigureForCOMImpersonation(userSession.get());
 
+    if (!containerRootVhd.empty())
+    {
+        wsl::windows::common::WSLAProcessLauncher initProcessLauncher{shell, {shell, "/etc/lsw-init.sh"}};
+        auto initProcess = initProcessLauncher.Launch(*session);
+        THROW_HR_IF(E_FAIL, initProcess.WaitAndCaptureOutput().Code != 0);
+    }
+
     wsl::windows::common::WSLAProcessLauncher launcher{shell, {shell}, {"TERM=xterm-256color"}, ProcessFlags::None};
     launcher.AddFd(WSLA_PROCESS_FD{.Fd = 0, .Type = WSLAFdTypeTerminalInput});
     launcher.AddFd(WSLA_PROCESS_FD{.Fd = 1, .Type = WSLAFdTypeTerminalOutput});

src/windows/wsl/CMakeLists.txt

@@ -10,7 +10,8 @@ add_executable(wsl ${SOURCES} ${HEADERS})
 target_link_libraries(wsl
                       ${COMMON_LINK_LIBRARIES}
                       common
-                      runtimeobject.lib)
+                      runtimeobject.lib
+                      VirtDisk.lib)
 
 target_precompile_headers(wsl REUSE_FROM common)
 set_target_properties(wsl PROPERTIES FOLDER windows)

src/windows/wslaservice/exe/WSLAVirtualMachine.cpp

@@ -358,8 +358,6 @@ void WSLAVirtualMachine::Start()
         wsl::windows::common::hcs::ModifyComputeSystem(m_computeSystem.get(), wsl::shared::ToJsonW(gpuRequest).c_str());
     }
 
-    ConfigureMounts();
-
     auto [__, ___, childChannel] = Fork(WSLA_FORK::Thread);
 
     WSLA_WATCH_PROCESSES watchMessage{};
@@ -368,6 +366,8 @@ void WSLAVirtualMachine::Start()
     THROW_HR_IF(E_FAIL, childChannel.ReceiveMessage<RESULT_MESSAGE<uint32_t>>().Result != 0);
 
     m_processExitThread = std::thread(std::bind(&WSLAVirtualMachine::WatchForExitedProcesses, this, std::move(childChannel)));
+
+    ConfigureMounts();
 }
 
 void WSLAVirtualMachine::ConfigureMounts()
@@ -385,7 +385,19 @@ void WSLAVirtualMachine::ConfigureMounts()
         MountGpuLibraries("/usr/lib/wsl/lib", "/usr/lib/wsl/drivers", WSLAMountFlagsNone);
     }
 
-    // TODO: Mount storage VHD here.
+    if (m_settings.ContainerRootVhd) // TODO: re-think how container root settings should work at the session level API.
+    {
+        auto [_, containerRootDevice] = AttachDisk(m_settings.ContainerRootVhd, false);
+
+        if (m_settings.FormatContainerRootVhd)
+        {
+            ServiceProcessLauncher formatProcessLauncher{"/usr/sbin/mkfs.ext4", {"/usr/sbin/mkfs.ext4", containerRootDevice}};
+            auto formatProcess = formatProcessLauncher.Launch(*this);
+            THROW_HR_IF(E_FAIL, formatProcess.WaitAndCaptureOutput().Code != 0);
+        }
+
+        Mount(m_initChannel, containerRootDevice.c_str(), "/root", "ext4", "rw", 0);
+    }
 }
 
 void WSLAVirtualMachine::WatchForExitedProcesses(wsl::shared::SocketChannel& Channel)

src/windows/wslaservice/inc/wslaservice.idl

@@ -225,6 +225,8 @@ struct _VIRTUAL_MACHINE_SETTINGS {  // TODO: Delete once the new API is wired.
     BOOL EnableGPU;
     LPCWSTR RootVhd; // Temporary option to provide the root VHD. TODO: Remove once runtime VHD is available.
     LPCSTR RootVhdType; // Temporary option to provide the root VHD. TODO: Remove once runtime VHD is available.
+    LPCWSTR ContainerRootVhd;
+    BOOL FormatContainerRootVhd;
 } VIRTUAL_MACHINE_SETTINGS;
 
 

src/windows/wslg/CMakeLists.txt

@@ -12,7 +12,8 @@ add_dependencies(wslg
 
 target_link_libraries(wslg
                       ${COMMON_LINK_LIBRARIES}
-                      common)
+                      common
+                      VirtDisk.lib)
 
 target_precompile_headers(wslg REUSE_FROM common)
 set_target_properties(wslg PROPERTIES FOLDER windows)