Only block shutdown during Apply.

5 files changed, 76 insertions, 11 deletions

diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index b4d74341..5f492efd 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -531,15 +531,6 @@ namespace WixToolset.Test.BA
             this.ShutdownUiThread();
         }
 
-        protected override void OnSystemShutdown(SystemShutdownEventArgs args)
-        {
-            // Always prevent shutdown.
-            this.Log("Disallowed system request to shut down the bootstrapper application.");
-            args.Cancel = true;
-
-            this.ShutdownUiThread();
-        }
-
         protected override void OnUnregisterBegin(UnregisterBeginEventArgs args)
         {
             if (this.forceKeepRegistration && args.RegistrationType == RegistrationType.None)
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index 7c533fea..3cf994fc 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -88,6 +88,56 @@ public: // IBAFunctions
         return hr;
     }
 
+    virtual STDMETHODIMP WndProc(
+        __in HWND hWnd,
+        __in UINT uMsg,
+        __in WPARAM /*wParam*/,
+        __in LPARAM lParam,
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
+        )
+    {
+        switch (uMsg)
+        {
+        case WM_QUERYENDSESSION:
+            if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+            {
+                DWORD dwEndSession = static_cast<DWORD>(lParam);
+                if (ENDSESSION_CRITICAL & dwEndSession)
+                {
+                    // Return false to get the WM_ENDSESSION message so that critical shutdowns can be delayed.
+                    *plResult = FALSE;
+                    *pfProcessed = TRUE;
+                }
+            }
+            break;
+        case WM_ENDSESSION:
+            if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+            {
+                ::MessageBoxW(hWnd, L"WM_ENDSESSION", L"BAFunctions WndProc", MB_OK);
+            }
+            break;
+        }
+        return S_OK;
+    }
+
+public: //IBootstrapperApplication
+    virtual STDMETHODIMP OnExecuteBegin(
+        __in DWORD /*cExecutingPackages*/,
+        __inout BOOL* pfCancel
+        )
+    {
+        if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+        {
+            if (IDCANCEL == ::MessageBoxW(m_hwndParent, L"Shutdown requests should be denied right now.", L"OnExecuteBegin", MB_OKCANCEL))
+            {
+                *pfCancel = TRUE;
+            }
+        }
+
+        return S_OK;
+    }
+
 private:
     HRESULT OnShowTheme()
     {
@@ -138,6 +188,8 @@ private:
 
         ThemeInitializeWindowClass(m_pBafTheme, &wc, CBafThmUtilTesting::TestingWndProc, m_hModule, BAFTHMUTILTESTING_WINDOW_CLASS);
 
+        Assert(wc.lpszClassName);
+
         // If the theme did not provide an icon, try using the icon from the bundle engine.
         if (!wc.hIcon)
         {
diff --git a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
index 20706b6a..44abc645 100644
--- a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
+++ b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
@@ -1,4 +1,4 @@
-<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
 
 
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
@@ -8,7 +8,10 @@
       <Payload SourceFile="$(var.BafThmUtilTesting.TargetPath)" bal:BAFunctions="yes" />
     </BootstrapperApplication>
     <PackageGroup Id="BundlePackages">
-      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)" />
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)">
+        <MsiProperty Name="FORCERESTARTCA" Value="[FORCERESTARTCA]" />
+      </MsiPackage>
     </PackageGroup>
+    <Variable Name="FORCERESTARTCA" bal:Overridable="yes" />
   </Fragment>
 </Wix>
diff --git a/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj b/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
index c9c81c3a..61dc27ab 100644
--- a/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
+++ b/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
@@ -2,8 +2,12 @@
 <Project Sdk="WixToolset.Sdk">
   <PropertyGroup>
     <UpgradeCode>{0D803A6E-8090-4174-8DAC-810ECC2B1BBF}</UpgradeCode>
+    <ProductComponentsRef>true</ProductComponentsRef>
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="..\..\Templates\Package.wxs" Link="Package.wxs" />
   </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Util.wixext" />
+  </ItemGroup>
 </Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/Manual/PackageA/ProductA.wxs b/src/test/burn/TestData/Manual/PackageA/ProductA.wxs
new file mode 100644
index 00000000..bf1ef152
--- /dev/null
+++ b/src/test/burn/TestData/Manual/PackageA/ProductA.wxs
@@ -0,0 +1,15 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:util="http://wixtoolset.org/schemas/v4/wxs/util">
+    <Fragment>
+        <ComponentGroup Id="ProductComponents" />
+
+        <Property Id="FORCERESTARTCA" Secure="yes" />
+        <Property Id="ForceRestartCA" Value="&quot;shutdown.exe&quot; -r -f -t 0" />
+        <CustomAction Id="ForceRestartCA" DllEntry="WixQuietExec" BinaryRef="Wix4UtilCA_X86" Execute="deferred" Return="ignore" />
+
+        <InstallExecuteSequence>
+            <Custom Action="ForceRestartCA" After="InstallFiles" Condition="FORCERESTARTCA = 1" />
+        </InstallExecuteSequence>
+    </Fragment>
+</Wix>