Refactor the session creation logic to match the new API

CMakeLists.txt

@@ -379,10 +379,10 @@ if (DEFINED WSL_DEV_BINARY_PATH) # Development shortcut to make the package smal
                             WSL_KERNEL_MODULES_PATH="${WSL_DEV_BINARY_PATH}/modules.vhd"
                             WSL_DEV_INSTALL_PATH="${WSL_DEV_BINARY_PATH}"
                             WSL_GPU_LIB_PATH="${WSL_DEV_BINARY_PATH}/lib")
-endif()
 
-if (NOT OFFICIAL_BUILD AND ${TARGET_PLATFORM} STREQUAL "x64")
-    add_compile_definitions(WSLA_TEST_DISTRO_PATH="${WSLA_TEST_DISTRO_SOURCE_DIR}/wslatestrootfs.vhd")
+    if (NOT OFFICIAL_BUILD AND ${TARGET_PLATFORM} STREQUAL "x64")
+        add_compile_definitions(WSLA_TEST_DISTRO_PATH="${WSLA_TEST_DISTRO_SOURCE_DIR}/wslatestrootfs.vhd")
+    endif()
 endif()
 
 # Common include paths

src/windows/common/WslClient.cpp

@@ -1525,43 +1525,38 @@ int RunDebugShell()
     THROW_HR(HCS_E_CONNECTION_CLOSED);
 }
 
+DEFINE_ENUM_FLAG_OPERATORS(WSLAFeatureFlags);
+
 // Temporary debugging tool for WSLA
 int WslaShell(_In_ std::wstring_view commandLine)
 {
-#ifdef WSLA_TEST_DISTRO_PATH
+    WSLA_SESSION_SETTINGS sessionSettings{};
+    sessionSettings.DisplayName = L"WSLAShell";
+    sessionSettings.CpuCount = 4;
+    sessionSettings.MemoryMb = 4096;
+    sessionSettings.NetworkingMode = WSLANetworkingModeNAT;
+    sessionSettings.BootTimeoutMs = 30 * 1000;
+    sessionSettings.MaximumStorageSizeMb = 4096;
 
-    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
-
-    VIRTUAL_MACHINE_SETTINGS settings{};
-    settings.CpuCount = 4;
-    settings.DisplayName = L"WSLA";
-    settings.MemoryMb = 1024;
-    settings.BootTimeoutMs = 30000;
-    settings.NetworkingMode = WSLANetworkingModeNAT;
-    std::wstring containerRootVhd;
     std::string containerImage;
     bool help = false;
     std::wstring debugShell;
 
+    std::wstring storagePath;
+    std::wstring rootVhdOverride;
+    std::string rootVhdTypeOverride;
     ArgumentParser parser(std::wstring{commandLine}, WSL_BINARY_NAME);
-    parser.AddArgument(vhd, L"--vhd");
+    parser.AddArgument(rootVhdOverride, L"--vhd");
     parser.AddArgument(Utf8String(shell), L"--shell");
-    parser.AddArgument(reinterpret_cast<bool&>(settings.EnableDnsTunneling), L"--dns-tunneling");
-    parser.AddArgument(Integer(settings.MemoryMb), L"--memory");
-    parser.AddArgument(Integer(settings.CpuCount), L"--cpu");
-    parser.AddArgument(Integer(reinterpret_cast<int&>(settings.NetworkingMode)), L"--networking-mode");
-    parser.AddArgument(Utf8String(fsType), L"--fstype");
-    parser.AddArgument(containerRootVhd, L"--container-vhd");
+    parser.AddArgument(
+        SetFlag<int, WslaFeatureFlagsDnsTunneling>(reinterpret_cast<int&>(sessionSettings.FeatureFlags)), L"--dns-tunneling");
+    parser.AddArgument(Integer(sessionSettings.MemoryMb), L"--memory");
+    parser.AddArgument(Integer(sessionSettings.CpuCount), L"--cpu");
+    parser.AddArgument(Utf8String(rootVhdTypeOverride), L"--fstype");
+    parser.AddArgument(storagePath, L"--storage");
+    parser.AddArgument(Integer(reinterpret_cast<int&>(sessionSettings.NetworkingMode)), L"--networking-mode");
     parser.AddArgument(Utf8String(containerImage), L"--image");
     parser.AddArgument(debugShell, L"--debug-shell");
     parser.AddArgument(help, L"--help");
@@ -1577,7 +1572,7 @@ int WslaShell(_In_ std::wstring_view commandLine)
         return 1;
     }
 
-    switch (settings.NetworkingMode)
+    switch (sessionSettings.NetworkingMode)
     {
     case WSLANetworkingMode::WSLANetworkingModeNone:
     case WSLANetworkingMode::WSLANetworkingModeNAT:
@@ -1587,46 +1582,41 @@ int WslaShell(_In_ std::wstring_view commandLine)
         THROW_HR(E_INVALIDARG);
     }
 
-    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"WSLA Test Session"};
     wil::com_ptr<IWSLASession> session;
-    settings.RootVhd = vhd.c_str();
-    settings.RootVhdType = fsType.c_str();
+
+    if (!rootVhdOverride.empty())
+    {
+        if (rootVhdTypeOverride.empty())
+        {
+            wprintf(L"--fstype required when --vhd is passed\n");
+            return 1;
+        }
+
+        sessionSettings.RootVhdOverride = rootVhdOverride.c_str();
+        sessionSettings.RootVhdTypeOverride = rootVhdTypeOverride.c_str();
+    }
+
+    if (!storagePath.empty())
+    {
+        storagePath = std::filesystem::weakly_canonical(storagePath).wstring();
+        sessionSettings.StoragePath = storagePath.c_str();
+    }
 
     if (!debugShell.empty())
     {
         THROW_IF_FAILED(userSession->OpenSessionByName(debugShell.c_str(), &session));
     }
     else
     {
-        THROW_IF_FAILED(userSession->CreateSession(&sessionSettings, &settings, &session));
+        THROW_IF_FAILED(userSession->CreateSession(&sessionSettings, &session));
         THROW_IF_FAILED(session->GetVirtualMachine(&virtualMachine));
 
         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);
-        }
     }
 
     std::optional<wil::com_ptr<IWSLAContainer>> container;

src/windows/wslaservice/exe/WSLASession.cpp

@@ -17,23 +17,85 @@ Module Name:
 #include "WSLAUserSession.h"
 #include "WSLAContainer.h"
 #include "ServiceProcessLauncher.h"
+#include "WslCoreFilesystem.h"
 
 using wsl::windows::service::wsla::WSLASession;
+using wsl::windows::service::wsla::WSLAVirtualMachine;
 
-WSLASession::WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl, const VIRTUAL_MACHINE_SETTINGS& VmSettings) :
-    m_sessionSettings(Settings),
-    m_userSession(&userSessionImpl),
-    m_virtualMachine(wil::MakeOrThrow<WSLAVirtualMachine>(VmSettings, userSessionImpl.GetUserSid(), &userSessionImpl)),
-    m_displayName(Settings.DisplayName)
+WSLASession::WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl) :
+    m_sessionSettings(Settings), m_userSession(&userSessionImpl), m_displayName(Settings.DisplayName)
 {
     WSL_LOG("SessionCreated", TraceLoggingValue(m_displayName.c_str(), "DisplayName"));
 
+    m_virtualMachine = wil::MakeOrThrow<WSLAVirtualMachine>(CreateVmSettings(Settings), userSessionImpl.GetUserSid());
+
     if (Settings.TerminationCallback != nullptr)
     {
         m_virtualMachine->RegisterCallback(Settings.TerminationCallback);
     }
 
     m_virtualMachine->Start();
+
+    ConfigureStorage(Settings);
+
+    // Launch the init script.
+    // TODO: Replace with something more robust once the final VHD is ready.
+    try
+    {
+        ServiceProcessLauncher launcher{"/bin/sh", {"/bin/sh", "-c", "/etc/lsw-init.sh"}};
+        auto result = launcher.Launch(*m_virtualMachine.Get()).WaitAndCaptureOutput();
+
+        THROW_HR_IF_MSG(E_FAIL, result.Code != 0, "Init script failed: %hs", launcher.FormatResult(result).c_str());
+    }
+    catch (...)
+    {
+        // Ignore issues launching the init script with custom root VHD's, for convenience.
+        // TODO: Remove once the final VHD is ready.
+        if (Settings.RootVhdOverride == nullptr)
+        {
+            throw;
+        }
+    }
+}
+
+WSLAVirtualMachine::Settings WSLASession::CreateVmSettings(const WSLA_SESSION_SETTINGS& Settings)
+{
+    WSLAVirtualMachine::Settings vmSettings{};
+    vmSettings.CpuCount = Settings.CpuCount;
+    vmSettings.MemoryMb = Settings.MemoryMb;
+    vmSettings.NetworkingMode = Settings.NetworkingMode;
+    vmSettings.BootTimeoutMs = Settings.BootTimeoutMs;
+    vmSettings.FeatureFlags = static_cast<WSLAFeatureFlags>(Settings.FeatureFlags);
+    vmSettings.DisplayName = Settings.DisplayName;
+
+    if (Settings.RootVhdOverride != nullptr)
+    {
+        THROW_HR_IF(E_INVALIDARG, Settings.RootVhdTypeOverride == nullptr);
+
+        vmSettings.RootVhd = Settings.RootVhdOverride;
+        vmSettings.RootVhdType = Settings.RootVhdTypeOverride;
+    }
+    else
+    {
+
+#ifdef WSLA_TEST_DISTRO_PATH
+
+        vmSettings.RootVhd = TEXT(WSLA_TEST_DISTRO_PATH);
+
+#else
+        vmSettings.RootVhd = std::filesystem::path(common::wslutil::GetMsiPackagePath().value()) / L"wslarootfs.vhd";
+
+#endif
+
+        vmSettings.RootVhdType = "squashfs";
+    }
+
+    if (Settings.DmesgOutput != 0)
+    {
+        vmSettings.DmesgHandle.reset(wsl::windows::common::wslutil::DuplicateHandleFromCallingProcess(ULongToHandle(Settings.DmesgOutput)));
+    }
+
+    return vmSettings;
 }
 
 WSLASession::~WSLASession()
@@ -54,6 +116,71 @@ WSLASession::~WSLASession()
     }
 }
 
+void WSLASession::ConfigureStorage(const WSLA_SESSION_SETTINGS& Settings)
+{
+    if (Settings.StoragePath == nullptr)
+    {
+        // If no storage path is specified, use a tmpfs for convenience.
+        m_virtualMachine->Mount("", "/root", "tmpfs", "", 0);
+        return;
+    }
+
+    std::filesystem::path storagePath{Settings.StoragePath};
+    THROW_HR_IF_MSG(E_INVALIDARG, !storagePath.is_absolute(), "Storage path is not absolute: %ls", storagePath.c_str());
+
+    m_storageVhdPath = storagePath / "storage.vhdx";
+
+    std::string diskDevice;
+    std::optional<ULONG> diskLun{};
+    bool vhdCreated = false;
+
+    auto deleteVhdOnFailure = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [&]() {
+        if (vhdCreated)
+        {
+            if (diskLun.has_value())
+            {
+                m_virtualMachine->DetachDisk(diskLun.value());
+            }
+
+            auto runAsUser = wil::CoImpersonateClient();
+            LOG_IF_WIN32_BOOL_FALSE(DeleteFileW(m_storageVhdPath.c_str()));
+        }
+    });
+
+    auto result =
+        wil::ResultFromException([&]() { diskDevice = m_virtualMachine->AttachDisk(m_storageVhdPath.c_str(), false).second; });
+
+    if (FAILED(result))
+    {
+        THROW_HR_IF_MSG(
+            result,
+            result != HRESULT_FROM_WIN32(ERROR_PATH_NOT_FOUND) && result != HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND),
+            "Failed to attach vhd: %ls",
+            m_storageVhdPath.c_str());
+
+        // If the VHD wasn't found, create it.
+        WSL_LOG("CreateStorageVhd", TraceLoggingValue(m_storageVhdPath.c_str(), "StorageVhdPath"));
+
+        auto runAsUser = wil::CoImpersonateClient();
+
+        std::filesystem::create_directories(storagePath);
+        wsl::core::filesystem::CreateVhd(
+            m_storageVhdPath.c_str(), Settings.MaximumStorageSizeMb * _1MB, m_userSession->GetUserSid(), false, false);
+        vhdCreated = true;
+
+        // Then attach the new disk.
+        std::tie(diskLun, diskDevice) = m_virtualMachine->AttachDisk(m_storageVhdPath.c_str(), false);
+
+        // Then format it.
+        Ext4Format(diskDevice);
+    }
+
+    // Mount the device to /root.
+    m_virtualMachine->Mount(diskDevice.c_str(), "/root", "ext4", "", 0);
+
+    deleteVhdOnFailure.release();
+}
+
 HRESULT WSLASession::GetDisplayName(LPWSTR* DisplayName)
 {
     *DisplayName = wil::make_unique_string<wil::unique_cotaskmem_string>(m_displayName.c_str()).release();
@@ -135,6 +262,15 @@ try
 }
 CATCH_RETURN();
 
+void WSLASession::Ext4Format(const std::string& Device)
+{
+    constexpr auto mkfsPath = "/usr/sbin/mkfs.ext4";
+    ServiceProcessLauncher launcher(mkfsPath, {mkfsPath, Device});
+    auto result = launcher.Launch(*m_virtualMachine.Get()).WaitAndCaptureOutput();
+
+    THROW_HR_IF_MSG(E_FAIL, result.Code != 0, "%hs", launcher.FormatResult(result).c_str());
+}
+
 HRESULT WSLASession::FormatVirtualDisk(LPCWSTR Path)
 try
 {
@@ -150,11 +286,7 @@ try
     auto detachDisk = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [this, lun]() { m_virtualMachine->DetachDisk(lun); });
 
     // Format it to ext4.
-    constexpr auto mkfsPath = "/usr/sbin/mkfs.ext4";
-    ServiceProcessLauncher launcher(mkfsPath, {mkfsPath, device});
-    auto result = launcher.Launch(*m_virtualMachine.Get()).WaitAndCaptureOutput();
-
-    THROW_HR_IF_MSG(E_FAIL, result.Code != 0, "%hs", launcher.FormatResult(result).c_str());
+    Ext4Format(device);
 
     return S_OK;
 }

src/windows/wslaservice/exe/WSLASession.h

@@ -23,7 +23,7 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     : public Microsoft::WRL::RuntimeClass<Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>, IWSLASession, IFastRundown>
 {
 public:
-    WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl, const VIRTUAL_MACHINE_SETTINGS& VmSettings);
+    WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl);
     ~WSLASession();
 
     IFACEMETHOD(GetDisplayName)(LPWSTR* DisplayName) override;
@@ -52,10 +52,16 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     void OnUserSessionTerminating();
 
 private:
+    static WSLAVirtualMachine::Settings CreateVmSettings(const WSLA_SESSION_SETTINGS& Settings);
+
+    void ConfigureStorage(const WSLA_SESSION_SETTINGS& Settings);
+    void Ext4Format(const std::string& Device);
+
     WSLA_SESSION_SETTINGS m_sessionSettings; // TODO: Revisit to see if we should have session settings as a member or not
     WSLAUserSessionImpl* m_userSession = nullptr;
     Microsoft::WRL::ComPtr<WSLAVirtualMachine> m_virtualMachine;
     std::wstring m_displayName;
+    std::filesystem::path m_storageVhdPath;
     std::mutex m_lock;
 
     // TODO: Add container tracking here. Could reuse m_lock for that.

src/windows/wslaservice/exe/WSLAUserSession.cpp

@@ -47,9 +47,9 @@ PSID WSLAUserSessionImpl::GetUserSid() const
     return m_tokenInfo->User.Sid;
 }
 
-HRESULT WSLAUserSessionImpl::CreateSession(const WSLA_SESSION_SETTINGS* Settings, const VIRTUAL_MACHINE_SETTINGS* VmSettings, IWSLASession** WslaSession)
+HRESULT WSLAUserSessionImpl::CreateSession(const WSLA_SESSION_SETTINGS* Settings, IWSLASession** WslaSession)
 {
-    auto session = wil::MakeOrThrow<WSLASession>(*Settings, *this, *VmSettings);
+    auto session = wil::MakeOrThrow<WSLASession>(*Settings, *this);
 
     std::lock_guard lock(m_wslaSessionsLock);
     auto it = m_sessions.emplace(session.Get());
@@ -94,14 +94,13 @@ HRESULT wsl::windows::service::wsla::WSLAUserSession::GetVersion(_Out_ WSLA_VERS
     return S_OK;
 }
 
-HRESULT wsl::windows::service::wsla::WSLAUserSession::CreateSession(
-    const WSLA_SESSION_SETTINGS* Settings, const VIRTUAL_MACHINE_SETTINGS* VmSettings, IWSLASession** WslaSession)
+HRESULT wsl::windows::service::wsla::WSLAUserSession::CreateSession(const WSLA_SESSION_SETTINGS* Settings, IWSLASession** WslaSession)
 try
 {
     auto session = m_session.lock();
     RETURN_HR_IF(RPC_E_DISCONNECTED, !session);
 
-    return session->CreateSession(Settings, VmSettings, WslaSession);
+    return session->CreateSession(Settings, WslaSession);
 }
 CATCH_RETURN();