WSLC SDK install API update

src/windows/WslcSDK/winrt/CMakeLists.txt

@@ -9,6 +9,7 @@ set(SOURCES
     ContainerVolume.cpp
     ImageInfo.cpp
     ImageProgress.cpp
+    InstallOptions.cpp
     InstallProgress.cpp
     Process.cpp
     ProcessCrashInformation.cpp
@@ -33,6 +34,7 @@ set(HEADERS
     Helpers.h
     ImageInfo.h
     ImageProgress.h
+    InstallOptions.h
     InstallProgress.h
     Process.h
     ProcessCrashInformation.h

src/windows/WslcSDK/winrt/InstallOptions.cpp

@@ -0,0 +1,41 @@
+/*++
+
+Copyright (c) Microsoft. All rights reserved.
+
+Module Name:
+
+    InstallOptions.cpp
+
+Abstract:
+
+    This file contains the implementation of the WinRT wrapper for the WSLC SDK InstallOptions class.
+
+--*/
+
+#include "precomp.h"
+#include "InstallOptions.h"
+#include "Microsoft.WSL.Containers.InstallOptions.g.cpp"
+
+namespace winrt::Microsoft::WSL::Containers::implementation {
+
+winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> InstallOptions::Components()
+{
+    return m_components;
+}
+
+void InstallOptions::Components(winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> value)
+{
+    m_components = std::move(value);
+}
+
+bool InstallOptions::Repair()
+{
+    return m_repair;
+}
+
+void InstallOptions::Repair(bool value)
+{
+    m_repair = value;
+}
+
+} // namespace winrt::Microsoft::WSL::Containers::implementation

src/windows/WslcSDK/winrt/InstallOptions.h

@@ -0,0 +1,41 @@
+/*++
+
+Copyright (c) Microsoft. All rights reserved.
+
+Module Name:
+
+    InstallOptions.h
+
+Abstract:
+
+    This file contains the definition of the WinRT wrapper for the WSLC SDK InstallOptions class.
+
+--*/
+
+#pragma once
+#include "Microsoft.WSL.Containers.InstallOptions.g.h"
+#include "Helpers.h"
+
+namespace winrt::Microsoft::WSL::Containers::implementation {
+struct InstallOptions : InstallOptionsT<InstallOptions>
+{
+    InstallOptions() = default;
+
+    winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> Components();
+    void Components(winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> value);
+    bool Repair();
+    void Repair(bool value);
+
+private:
+    winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> m_components = nullptr;
+    bool m_repair = false;
+};
+} // namespace winrt::Microsoft::WSL::Containers::implementation
+
+namespace winrt::Microsoft::WSL::Containers::factory_implementation {
+struct InstallOptions : InstallOptionsT<InstallOptions, implementation::InstallOptions>
+{
+};
+} // namespace winrt::Microsoft::WSL::Containers::factory_implementation
+
+DEFINE_TYPE_HELPERS(InstallOptions);

src/windows/WslcSDK/winrt/WslcService.cpp

@@ -24,6 +24,60 @@ using namespace winrt::Windows::Foundation::Collections;
 namespace winrt::Microsoft::WSL::Containers::implementation {
 
 namespace {
+    WslcComponentFlags GetComponentsForInstall(const InstallOptions& options)
+    {
+        WslcComponentFlags result = WslcComponentFlags::WSLC_COMPONENT_FLAG_NONE;
+        bool shouldCheckMissingComponents = true;
+
+        if (options)
+        {
+            auto components = options.Components();
+            if (components)
+            {
+                shouldCheckMissingComponents = false;
+
+                for (const auto& component : components)
+                {
+                    switch (component)
+                    {
+                    case Component::VirtualMachinePlatform:
+                        result |= WslcComponentFlags::WSLC_COMPONENT_FLAG_VIRTUAL_MACHINE_PLATFORM;
+                        break;
+                    case Component::WslPackage:
+                        result |= WslcComponentFlags::WSLC_COMPONENT_FLAG_WSL_PACKAGE;
+                        break;
+                    case Component::SdkNeedsUpdate:
+                        THROW_HR(WSLC_E_SDK_UPDATE_NEEDED);
+                    default:
+                        THROW_HR(E_INVALIDARG);
+                    }
+                }
+            }
+        }
+
+        if (shouldCheckMissingComponents)
+        {
+            winrt::check_hresult(WslcGetMissingComponents(&result));
+        }
+
+        return result;
+    }
+
+    WslcInstallOptions GetOptionsForInstall(const InstallOptions& options)
+    {
+        WslcInstallOptions result = WslcInstallOptions::WSLC_INSTALL_OPTION_NONE;
+
+        if (options)
+        {
+            if (options.Repair())
+            {
+                result |= WslcInstallOptions::WSLC_INSTALL_OPTION_REPAIR;
+            }
+        }
+
+        return result;
+    }
+
     void CALLBACK InstallProgressCallback(WslcComponentFlags component, uint32_t progressSteps, uint32_t totalSteps, PVOID context) noexcept
     {
         try
@@ -64,16 +118,22 @@ winrt::Microsoft::WSL::Containers::ServiceVersion WslcService::GetVersion()
     return winrt::make<implementation::ServiceVersion>(version.major, version.minor, version.revision);
 }
 
-IAsyncActionWithProgress<winrt::Microsoft::WSL::Containers::InstallProgress> WslcService::InstallWithDependenciesAsync()
+IAsyncActionWithProgress<winrt::Microsoft::WSL::Containers::InstallProgress> WslcService::InstallWithDependenciesAsync(
+    winrt::Microsoft::WSL::Containers::InstallOptions options)
 {
+    auto components = GetComponentsForInstall(options);
+    auto wslcOptions = GetOptionsForInstall(options);
+
     co_await winrt::resume_background();
 
     auto context = ProgressCallbackHelper<winrt::Microsoft::WSL::Containers::InstallProgress>{co_await winrt::get_progress_token()};
-    winrt::check_hresult(WslcInstallWithDependencies(InstallProgressCallback, &context));
+    winrt::check_hresult(WslcInstallWithDependencies(components, wslcOptions, InstallProgressCallback, &context));
 }
 
-void WslcService::InstallWithDependencies()
+void WslcService::InstallWithDependencies(winrt::Microsoft::WSL::Containers::InstallOptions options)
 {
-    winrt::check_hresult(WslcInstallWithDependencies(nullptr, nullptr));
+    auto components = GetComponentsForInstall(options);
+    auto wslcOptions = GetOptionsForInstall(options);
+    winrt::check_hresult(WslcInstallWithDependencies(components, wslcOptions, nullptr, nullptr));
 }
 } // namespace winrt::Microsoft::WSL::Containers::implementation

src/windows/WslcSDK/winrt/WslcService.h

@@ -23,8 +23,9 @@ struct WslcService
 
     static winrt::Windows::Foundation::Collections::IVectorView<winrt::Microsoft::WSL::Containers::Component> GetMissingComponents();
     static winrt::Microsoft::WSL::Containers::ServiceVersion GetVersion();
-    static void InstallWithDependencies();
-    static winrt::Windows::Foundation::IAsyncActionWithProgress<winrt::Microsoft::WSL::Containers::InstallProgress> InstallWithDependenciesAsync();
+    static void InstallWithDependencies(winrt::Microsoft::WSL::Containers::InstallOptions options);
+    static winrt::Windows::Foundation::IAsyncActionWithProgress<winrt::Microsoft::WSL::Containers::InstallProgress> InstallWithDependenciesAsync(
+        winrt::Microsoft::WSL::Containers::InstallOptions options);
 };
 } // namespace winrt::Microsoft::WSL::Containers::implementation
 namespace winrt::Microsoft::WSL::Containers::factory_implementation {

src/windows/WslcSDK/winrt/wslcsdk.idl

@@ -249,6 +249,14 @@ namespace Microsoft.WSL.Containers
         UInt32 Revision { get; };
     };
 
+    runtimeclass InstallOptions
+    {
+        InstallOptions();
+
+        IVectorView<Component> Components;
+        Boolean Repair;
+    };
+
     runtimeclass InstallProgress
     {
         Component Component { get; };
@@ -260,8 +268,8 @@ namespace Microsoft.WSL.Containers
     {
         static IVectorView<Component> GetMissingComponents();
         static ServiceVersion GetVersion();
-        static void InstallWithDependencies();
-        static Windows.Foundation.IAsyncActionWithProgress<InstallProgress> InstallWithDependenciesAsync();
+        static void InstallWithDependencies(InstallOptions options);
+        static Windows.Foundation.IAsyncActionWithProgress<InstallProgress> InstallWithDependenciesAsync(InstallOptions options);
     };
 
     enum VhdType

src/windows/WslcSDK/wslcsdk.cpp

@@ -1664,33 +1664,36 @@ try
 }
 CATCH_RETURN();
 
-STDAPI WslcInstallWithDependencies(_In_opt_ WslcInstallCallback progressCallback, _In_opt_ PVOID context)
+STDAPI WslcInstallWithDependencies(
+    _In_ WslcComponentFlags components, _In_ WslcInstallOptions options, _In_opt_ WslcInstallCallback progressCallback, _In_opt_ PVOID context)
 try
 {
+    // This API cannot update the SDK that the client is using.
+    RETURN_HR_IF(WSLC_E_SDK_UPDATE_NEEDED, WI_IsFlagSet(components, WSLC_COMPONENT_FLAG_SDK_NEEDS_UPDATE));
+
     HRESULT result = S_OK;
-    bool needsVirtualMachine = NeedsVirtualMachineServicesInstalled();
-    auto runtimeResult = CreateSessionManagerRaw().second;
 
-    if (!needsVirtualMachine && SUCCEEDED(runtimeResult))
+    if (components == WSLC_COMPONENT_FLAG_NONE)
     {
         return result;
     }
 
-    THROW_HR_IF(runtimeResult, runtimeResult != REGDB_E_CLASSNOTREG && runtimeResult != WSLC_E_SDK_UPDATE_NEEDED);
-
-    // Installing these components requires elevation.
+    // Installing components requires elevation.
     RETURN_HR_IF(
         HRESULT_FROM_WIN32(ERROR_ELEVATION_REQUIRED),
         !wsl::windows::common::security::IsTokenElevated(GetCurrentThreadEffectiveToken()) &&
             !wsl::windows::common::security::IsTokenLocalSystem(nullptr));
 
-    if (needsVirtualMachine)
+    bool isRepair = WI_IsFlagSet(options, WSLC_INSTALL_OPTION_REPAIR);
+
+    if (WI_IsFlagSet(components, WSLC_COMPONENT_FLAG_VIRTUAL_MACHINE_PLATFORM))
     {
         if (progressCallback)
         {
             progressCallback(WSLC_COMPONENT_FLAG_VIRTUAL_MACHINE_PLATFORM, 0, 1, context);
         }
 
+        // No difference between install and repair, just let DISM attempt to enable the feature.
         auto exitCode = WslInstall::InstallOptionalComponent(WslInstall::c_optionalFeatureNameVmp, false);
         if (exitCode == ERROR_SUCCESS_REBOOT_REQUIRED)
         {
@@ -1707,9 +1710,15 @@ try
         {
             progressCallback(WSLC_COMPONENT_FLAG_VIRTUAL_MACHINE_PLATFORM, 1, 1, context);
         }
+
+        // If a reboot is required, we need the reboot to happen before attempting to install the WSL package.
+        if (result == HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED))
+        {
+            return result;
+        }
     }
 
-    if (!SUCCEEDED(runtimeResult))
+    if (WI_IsFlagSet(components, WSLC_COMPONENT_FLAG_WSL_PACKAGE))
     {
         std::function<void(uint32_t)> callback;
         if (progressCallback)
@@ -1719,14 +1728,16 @@ try
             };
         }
 
-        wsl::windows::common::WindowsUpdateContext wuContext;
-        wuContext.RunUpdateFlow(true, callback);
+        using WindowsUpdateContext = wsl::windows::common::WindowsUpdateContext;
+        WindowsUpdateContext wuContext;
+        wuContext.RunUpdateFlow(
+            isRepair ? WindowsUpdateContext::UpdateOptions::ResetProductRegistration : WindowsUpdateContext::UpdateOptions::EnsureProductRegistration,
+            callback);
 
-        // Because we do a forced install here, we expect an update.
         if (wuContext.GetUpdateCount() == 0)
         {
             // During the preview period, the package may not be published yet, so fall back to getting it from GH.
-            // When moving to GA, change this to a hard error to indicate a service configuration issue.
+            // When moving to GA, change this to an error like WSL_E_NO_UPDATE_AVAILABLE or similar.
             if (callback)
             {
                 callback(0);

src/windows/WslcSDK/wslcsdk.h

@@ -624,8 +624,18 @@ STDAPI WslcGetVersion(_Out_writes_(1) WslcVersion* version);
 typedef __callback void(CALLBACK* WslcInstallCallback)(
     _In_ WslcComponentFlags component, _In_ uint32_t progressSteps, _In_ uint32_t totalSteps, _In_opt_ PVOID context);
 
+typedef enum WslcInstallOptions
+{
+    WSLC_INSTALL_OPTION_NONE = 0,
+    // Allows components to be reinstalled.
+    WSLC_INSTALL_OPTION_REPAIR = 1,
+} WslcInstallOptions;
+
+DEFINE_ENUM_FLAG_OPERATORS(WslcInstallOptions);
+
 // Callbacks will only be made for components that are actively installed by this call.
-// That list can be acquired prior to this call with `WslcCanRun`.
-STDAPI WslcInstallWithDependencies(_In_opt_ WslcInstallCallback progressCallback, _In_opt_ PVOID context);
+// The list of required components can be acquired prior to this call with `WslcGetMissingComponents`.
+STDAPI WslcInstallWithDependencies(
+    _In_ WslcComponentFlags components, _In_ WslcInstallOptions options, _In_opt_ WslcInstallCallback progressCallback, _In_opt_ PVOID context);
 
 EXTERN_C_END

src/windows/common/WindowsUpdateIntegration.cpp

@@ -126,18 +126,12 @@ namespace anon {
     };
 } // namespace anon
 
-WindowsUpdateContext::WindowsUpdateContext() :
-    WindowsUpdateContext(std::make_unique<anon::DefaultWindowsUpdateClassFactory>(), WslProductIdentifier())
+WindowsUpdateContext::WindowsUpdateContext() : WindowsUpdateContext(std::make_unique<anon::DefaultWindowsUpdateClassFactory>())
 {
 }
 
-WindowsUpdateContext::WindowsUpdateContext(std::wstring product) :
-    WindowsUpdateContext(std::make_unique<anon::DefaultWindowsUpdateClassFactory>(), std::move(product))
-{
-}
-
-WindowsUpdateContext::WindowsUpdateContext(std::unique_ptr<WindowsUpdateClassFactory> factory, std::wstring product) :
-    m_factory(std::move(factory)), m_product(std::move(product))
+WindowsUpdateContext::WindowsUpdateContext(std::unique_ptr<WindowsUpdateClassFactory> factory) :
+    m_factory(std::move(factory)), m_product(WslProductIdentifier())
 {
     m_session = m_factory->CreateUpdateSession();
 
@@ -158,9 +152,12 @@ std::wstring WindowsUpdateContext::WslProductIdentifier()
     return STRING_TO_WIDE_STRING(DCAT_PRODUCT_NAME);
 }
 
-void WindowsUpdateContext::EnsureProductRegistryEntry() const
+void WindowsUpdateContext::EnsureProductRegistryEntry(bool reset) const
 {
-    wsl::windows::common::helpers::RegisterWithDcat(false);
+    if (reset || !wsl::windows::common::helpers::VersionRegisteredWithDcat())
+    {
+        wsl::windows::common::helpers::RegisterWithDcat(false);
+    }
 }
 
 size_t WindowsUpdateContext::SearchForUpdates()
@@ -336,14 +333,15 @@ void WindowsUpdateContext::InstallUpdates(const std::function<void(uint32_t)>& p
     THROW_IF_FAILED(installationHResult);
 }
 
-void WindowsUpdateContext::RunUpdateFlow(bool forceInstall, const std::function<void(uint32_t)>& progress)
+void WindowsUpdateContext::RunUpdateFlow(UpdateOptions options, const std::function<void(uint32_t)>& progress)
 {
     TraceLoggingWriteTagged(
         *m_activity,
         "RunUpdateFlow",
         TraceLoggingKeyword(MICROSOFT_KEYWORD_MEASURES),
         TelemetryPrivacyDataTag(PDT_ProductAndServiceUsage),
-        TraceLoggingBool(forceInstall, "forceInstall"));
+        TraceLoggingBool(options == UpdateOptions::ResetProductRegistration, "forceInstall"),
+        TraceLoggingBool(options == UpdateOptions::EnsureProductRegistration, "ensureInstall"));
 
     static_assert(
         DownloadProgressPercent + InstallProgressPercent == 100, "Download and Install progress values must add up to 100.");
@@ -353,9 +351,9 @@ void WindowsUpdateContext::RunUpdateFlow(bool forceInstall, const std::function<
         progress(0);
     }
 
-    if (forceInstall)
+    if (options != UpdateOptions::None)
     {
-        EnsureProductRegistryEntry();
+        EnsureProductRegistryEntry(options == UpdateOptions::ResetProductRegistration);
     }
 
     size_t updateCount = SearchForUpdates();