Fix signaling the app shutdown event running as admin (#3874)

There's bug in the OS where the WM_QUERYENDSESSION message is not send
when a package is being updated if the package is running elevated. This
makes `winget configure --enable` to stop at 95% and eventually get
terminated by the system for an update. The update is successful, but
the experience is not the best. It also makes it harder to use
`IConfigurationStatics::IConfigurationStatics` in an elevated context.

The fix is to use the `PackageCatalog` APIs and register for
`PackageUpdating` events when winget is running elevated and in a
package context and signal the competition event if the progress is more
than 0. Otherwise create the window and listen to messages.

Enable test where a new update is being registered since now it will
work and move test to force the WM_QUERYENDSESSION to UTs

src/AppInstallerCLICore/Commands/TestCommand.cpp

@@ -5,6 +5,7 @@
 #ifndef AICLI_DISABLE_TEST_HOOKS
 
 #include "TestCommand.h"
+#include "AppInstallerRuntime.h"
 
 namespace AppInstaller::CLI
 {
@@ -15,6 +16,60 @@ namespace AppInstaller::CLI
             AICLI_LOG(CLI, Info, << message);
             context.Reporter.Info() << message << std::endl;
         }
+
+        HRESULT WaitForShutdown(Execution::Context& context)
+        {
+            LogAndReport(context, "Waiting for app shutdown event");
+            if (!Execution::WaitForAppShutdownEvent())
+            {
+                LogAndReport(context, "Failed getting app shutdown event");
+                return APPINSTALLER_CLI_ERROR_INTERNAL_ERROR;
+            }
+
+            LogAndReport(context, "Succeeded waiting for app shutdown event");
+            return S_OK;
+        }
+
+        HRESULT AppShutdownWindowMessage(Execution::Context& context)
+        {
+            auto windowHandle = Execution::GetWindowHandle();
+
+            if (windowHandle == NULL)
+            {
+                LogAndReport(context, "Window was not created");
+                return APPINSTALLER_CLI_ERROR_INTERNAL_ERROR;
+            }
+
+            if (context.Args.Contains(Execution::Args::Type::Force))
+            {
+                LogAndReport(context, "Sending WM_QUERYENDSESSION message");
+                THROW_LAST_ERROR_IF(!SendMessageTimeout(
+                    windowHandle,
+                    WM_QUERYENDSESSION,
+                    NULL,
+                    ENDSESSION_CLOSEAPP,
+                    (SMTO_ABORTIFHUNG | SMTO_ERRORONEXIT),
+                    5000,
+                    NULL));
+            }
+
+            HRESULT hr = WaitForShutdown(context);
+
+            if (context.Args.Contains(Execution::Args::Type::Force))
+            {
+                LogAndReport(context, "Sending WM_ENDSESSION message");
+                THROW_LAST_ERROR_IF(!SendMessageTimeout(
+                    windowHandle,
+                    WM_ENDSESSION,
+                    NULL,
+                    ENDSESSION_CLOSEAPP,
+                    (SMTO_ABORTIFHUNG | SMTO_ERRORONEXIT),
+                    5000,
+                    NULL));
+            }
+
+            return hr;
+        }
     }
 
     std::vector<std::unique_ptr<Command>> TestCommand::GetCommands() const
@@ -49,36 +104,19 @@ namespace AppInstaller::CLI
 
     void TestAppShutdownCommand::ExecuteInternal(Execution::Context& context) const
     {
-        auto windowHandle = Execution::GetWindowHandle();
+        HRESULT hr = E_FAIL;
 
-        if (windowHandle == NULL)
+        // Only package context and admin won't create the window message.
+        if (!Runtime::IsRunningInPackagedContext() || !Runtime::IsRunningAsAdmin())
         {
-            LogAndReport(context, "Window was not created");
-            AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_INTERNAL_ERROR);
+            hr = AppShutdownWindowMessage(context);
         }
-
-        if (context.Args.Contains(Execution::Args::Type::Force))
-        {
-            LogAndReport(context, "Sending WM_QUERYENDSESSION message");
-            THROW_LAST_ERROR_IF(!SendMessageTimeout(
-                windowHandle,
-                WM_QUERYENDSESSION,
-                NULL,
-                ENDSESSION_CLOSEAPP,
-                (SMTO_ABORTIFHUNG | SMTO_ERRORONEXIT),
-                5000,
-                NULL));
-        }
-
-        LogAndReport(context, "Waiting for app shutdown event");
-        bool result = Execution::WaitForAppShutdownEvent();
-        if (!result)
+        else
         {
-            LogAndReport(context, "Failed getting app shutdown event");
-            AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_INTERNAL_ERROR);
+            hr = WaitForShutdown(context);
         }
 
-        LogAndReport(context, "Succeeded waiting for app shutdown event");
+        AICLI_TERMINATE_CONTEXT(hr);
     }
 
     Resource::LocString TestAppShutdownCommand::ShortDescription() const

src/AppInstallerCLICore/ExecutionContext.cpp

@@ -67,12 +67,31 @@ namespace AppInstaller::CLI::Execution
         private:
             SignalTerminationHandler()
             {
-                // Create message only window.
-                m_messageQueueReady.create();
-                m_windowThread = std::thread(&SignalTerminationHandler::CreateWindowAndStartMessageLoop, this);
-                if (!m_messageQueueReady.wait(100))
+                if (Runtime::IsRunningAsAdmin() && Runtime::IsRunningInPackagedContext())
                 {
-                    AICLI_LOG(CLI, Warning, << "Timeout creating winget window");
+                    m_catalog = winrt::Windows::ApplicationModel::PackageCatalog::OpenForCurrentPackage();
+                    m_updatingEvent = m_catalog.PackageUpdating(
+                        winrt::auto_revoke, [this](winrt::Windows::ApplicationModel::PackageCatalog, winrt::Windows::ApplicationModel::PackageUpdatingEventArgs args)
+                        {
+                            // There are 3 events being hit with 0%, 1% and 38%
+                            // Typically the window message is received between the first two.
+                            constexpr double minProgress = 0;
+                            auto progress = args.Progress();
+                            if (progress > minProgress)
+                            {
+                                SignalTerminationHandler::Instance().StartAppShutdown();
+                            }
+                        });
+                }
+                else
+                {
+                    // Create message only window.
+                    m_messageQueueReady.create();
+                    m_windowThread = std::thread(&SignalTerminationHandler::CreateWindowAndStartMessageLoop, this);
+                    if (!m_messageQueueReady.wait(100))
+                    {
+                        AICLI_LOG(CLI, Warning, << "Timeout creating winget window");
+                    }
                 }
 
                 // Set up ctrl-c handler.
@@ -236,6 +255,8 @@ namespace AppInstaller::CLI::Execution
             wil::unique_event m_messageQueueReady;
             wil::unique_hwnd m_windowHandle;
             std::thread m_windowThread;
+            winrt::Windows::ApplicationModel::PackageCatalog m_catalog = nullptr;
+            decltype(winrt::Windows::ApplicationModel::PackageCatalog{ nullptr }.PackageUpdating(winrt::auto_revoke, nullptr)) m_updatingEvent;
         };
 
         void SetSignalTerminationHandlerContext(bool add, Context* context)

src/AppInstallerCLIE2ETests/AppShutdownTests.cs

@@ -23,14 +23,18 @@ public class AppShutdownTests
         /// Runs winget test appshutdown and register the application to force a WM_QUERYENDSESSION message.
         /// </summary>
         [Test]
-        [Ignore("This test won't work on Window Server")]
         public void RegisterApplicationTest()
         {
             if (!TestSetup.Parameters.PackagedContext)
             {
                 Assert.Ignore("Not packaged context.");
             }
 
+            if (!TestCommon.ExecutingAsAdministrator && TestCommon.IsCIEnvironment)
+            {
+                Assert.Ignore("This test won't work on Window Server as non-admin");
+            }
+
             if (string.IsNullOrEmpty(TestSetup.Parameters.AICLIPackagePath))
             {
                 throw new NullReferenceException("AICLIPackagePath");
@@ -91,26 +95,5 @@ public void RegisterApplicationTest()
             // Look for the output.
             Assert.True(testCmdTask.Result.StdOut.Contains("Succeeded waiting for app shutdown event"));
         }
-
-        /// <summary>
-        /// Runs winget test appshutdown --force.
-        /// </summary>
-        [Test]
-        public void RegisterApplicationTest_Force()
-        {
-            if (!TestSetup.Parameters.PackagedContext)
-            {
-                Assert.Ignore("Not packaged context.");
-            }
-
-            if (string.IsNullOrEmpty(TestSetup.Parameters.AICLIPackagePath))
-            {
-                throw new NullReferenceException("AICLIPackagePath");
-            }
-
-            var result = TestCommon.RunAICLICommand("test", "appshutdown --force", timeOut: 300000, throwOnTimeout: false);
-            TestContext.Out.Write(result.StdOut);
-            Assert.True(result.StdOut.Contains("Succeeded waiting for app shutdown event"));
-        }
     }
 }

src/AppInstallerCLIE2ETests/Helpers/TestCommon.cs

@@ -13,6 +13,7 @@ namespace AppInstallerCLIE2ETests.Helpers
     using System.Linq;
     using System.Management.Automation;
     using System.Reflection;
+    using System.Security.Principal;
     using System.Text;
     using System.Threading;
     using AppInstallerCLIE2ETests;
@@ -78,6 +79,31 @@ public enum TestModuleLocation
             Default,
         }
 
+        /// <summary>
+        /// Gets a value indicating whether the current assembly is executing in an administrative context.
+        /// </summary>
+        [System.Diagnostics.CodeAnalysis.SuppressMessage("Interoperability", "CA1416:Validate platform compatibility", Justification = "Windows only API")]
+        public static bool ExecutingAsAdministrator
+        {
+            get
+            {
+                WindowsIdentity identity = WindowsIdentity.GetCurrent();
+                WindowsPrincipal principal = new (identity);
+                return principal.IsInRole(WindowsBuiltInRole.Administrator);
+            }
+        }
+
+        /// <summary>
+        /// Gets a value indicating whether the test is running in the CI build.
+        /// </summary>
+        public static bool IsCIEnvironment
+        {
+            get
+            {
+                return Environment.GetEnvironmentVariable("BUILD_BUILDNUMBER") != null;
+            }
+        }
+
         /// <summary>
         /// Run winget command.
         /// </summary>

src/AppInstallerCLITests/AppInstallerCLITests.vcxproj

@@ -194,6 +194,7 @@
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="AdminSettings.cpp" />
+    <ClCompile Include="AppShutdown.cpp" />
     <ClCompile Include="Archive.cpp" />
     <ClCompile Include="Argument.cpp" />
     <ClCompile Include="ARPChanges.cpp" />

src/AppInstallerCLITests/AppInstallerCLITests.vcxproj.filters

@@ -329,6 +329,9 @@
     <ClCompile Include="ArpHelper.cpp">
       <Filter>Source Files\Repository</Filter>
     </ClCompile>
+    <ClCompile Include="AppShutdown.cpp">
+      <Filter>Source Files\Common</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <None Include="PropertySheet.props" />

src/AppInstallerCLITests/AppShutdown.cpp

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#include "pch.h"
+#include <winget/Runtime.h>
+#include "Commands/TestCommand.h"
+
+using namespace AppInstaller::CLI;
+
+TEST_CASE("AppShutdown_WindowMessage", "[appShutdown]")
+{
+    if (AppInstaller::Runtime::IsRunningAsAdmin() && AppInstaller::Runtime::IsRunningInPackagedContext())
+    {
+        WARN("Test can't run as admin in package context");
+        return;
+    }
+
+    std::ostringstream output;
+    Execution::Context context{ output, std::cin };
+    context.Args.AddArg(Execution::Args::Type::Force);
+
+    TestAppShutdownCommand appShutdownCmd({});
+    appShutdownCmd.Execute(context);
+
+    REQUIRE(context.IsTerminated());
+    REQUIRE(S_OK == context.GetTerminationHR());
+}