Include bundle reboot-pending in RebootPending variable.

Fixes https://github.com/wixtoolset/issues/issues/5332

3 files changed, 53 insertions, 54 deletions

diff --git a/src/engine/core.h b/src/engine/core.h
index 75a61614..b4e0e0d3 100644
--- a/src/engine/core.h
+++ b/src/engine/core.h
@@ -46,6 +46,7 @@ const LPCWSTR BURN_BUNDLE_SOURCE_PROCESS_FOLDER = L"WixBundleSourceProcessFolder
 const LPCWSTR BURN_BUNDLE_TAG = L"WixBundleTag";
 const LPCWSTR BURN_BUNDLE_UILEVEL = L"WixBundleUILevel";
 const LPCWSTR BURN_BUNDLE_VERSION = L"WixBundleVersion";
+const LPCWSTR BURN_REBOOT_PENDING = L"RebootPending";
 
 // The following constants must stay in sync with src\wix\Binder.cs
 const LPCWSTR BURN_BUNDLE_NAME = L"WixBundleName";
diff --git a/src/engine/registration.cpp b/src/engine/registration.cpp
index dc4b88bf..fc5ae627 100644
--- a/src/engine/registration.cpp
+++ b/src/engine/registration.cpp
@@ -92,6 +92,10 @@ static HRESULT UpdateBundleNameRegistration(
     __in BURN_VARIABLES* pVariables,
     __in HKEY hkRegistration
     );
+static BOOL IsWuRebootPending();
+static BOOL IsBundleRebootPending(
+    __in BURN_REGISTRATION* pRegistration
+);
 
 // function definitions
 
@@ -443,7 +447,10 @@ extern "C" HRESULT RegistrationSetVariables(
     ExitOnFailure(hr, "Failed to overwrite the bundle tag built-in variable.");
 
     hr = VariableSetVersion(pVariables, BURN_BUNDLE_VERSION, pRegistration->pVersion, TRUE);
-    ExitOnFailure(hr, "Failed to overwrite the bundle tag built-in variable.");
+    ExitOnFailure(hr, "Failed to overwrite the bundle version built-in variable.");
+
+    hr = VariableSetNumeric(pVariables, BURN_REBOOT_PENDING, IsBundleRebootPending(pRegistration) || IsWuRebootPending(), TRUE);
+    ExitOnFailure(hr, "Failed to overwrite the bundle reboot-pending built-in variable.");
 
 LExit:
     ReleaseStr(sczBundleManufacturer);
@@ -491,17 +498,10 @@ extern "C" HRESULT RegistrationDetectResumeType(
     )
 {
     HRESULT hr = S_OK;
-    LPWSTR sczRebootRequiredKey = NULL;
-    HKEY hkRebootRequired = NULL;
     HKEY hkRegistration = NULL;
     DWORD dwResume = 0;
 
-    // Check to see if a restart is pending for this bundle.
-    hr = StrAllocFormatted(&sczRebootRequiredKey, REGISTRY_REBOOT_PENDING_FORMAT, pRegistration->sczRegistrationKey);
-    ExitOnFailure(hr, "Failed to format pending restart registry key to read.");
-
-    hr = RegOpen(pRegistration->hkRoot, sczRebootRequiredKey, KEY_QUERY_VALUE, &hkRebootRequired);
-    if (SUCCEEDED(hr))
+    if (IsBundleRebootPending(pRegistration))
     {
         *pResumeType = BOOTSTRAPPER_RESUME_TYPE_REBOOT_PENDING;
         ExitFunction1(hr = S_OK);
@@ -554,8 +554,6 @@ extern "C" HRESULT RegistrationDetectResumeType(
 
 LExit:
     ReleaseRegKey(hkRegistration);
-    ReleaseRegKey(hkRebootRequired);
-    ReleaseStr(sczRebootRequiredKey);
 
     return hr;
 }
@@ -1591,3 +1589,46 @@ LExit:
 
     return hr;
 }
+
+static BOOL IsWuRebootPending()
+{
+    HRESULT hr = S_OK;
+    BOOL fRebootPending = FALSE;
+
+    // Do a best effort to ask WU if a reboot is required. If anything goes
+    // wrong then let's pretend a reboot is not required.
+    hr = ::CoInitialize(NULL);
+    if (SUCCEEDED(hr) || RPC_E_CHANGED_MODE == hr)
+    {
+        hr = WuaRestartRequired(&fRebootPending);
+        if (FAILED(hr))
+        {
+            fRebootPending = FALSE;
+        }
+
+        ::CoUninitialize();
+    }
+
+    return fRebootPending;
+}
+
+static BOOL IsBundleRebootPending(BURN_REGISTRATION* pRegistration)
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczRebootRequiredKey = NULL;
+    HKEY hkRebootRequired = NULL;
+    BOOL fBundleRebootPending = FALSE;
+
+    // Check to see if a restart is pending for this bundle.
+    hr = StrAllocFormatted(&sczRebootRequiredKey, REGISTRY_REBOOT_PENDING_FORMAT, pRegistration->sczRegistrationKey);
+    ExitOnFailure(hr, "Failed to format pending restart registry key to read.");
+
+    hr = RegOpen(pRegistration->hkRoot, sczRebootRequiredKey, KEY_QUERY_VALUE, &hkRebootRequired);
+    fBundleRebootPending = SUCCEEDED(hr);
+
+LExit:
+    ReleaseStr(sczRebootRequiredKey);
+    ReleaseRegKey(hkRebootRequired);
+
+    return fBundleRebootPending;
+}
diff --git a/src/engine/variable.cpp b/src/engine/variable.cpp
index fed23151..d0c67504 100644
--- a/src/engine/variable.cpp
+++ b/src/engine/variable.cpp
@@ -133,10 +133,6 @@ static HRESULT InitializeVariablePrivileged(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue
     );
-static HRESULT InitializeVariableRebootPending(
-    __in DWORD_PTR dwpData,
-    __inout BURN_VARIANT* pValue
-    );
 static HRESULT InitializeSystemLanguageID(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue
@@ -247,7 +243,6 @@ extern "C" HRESULT VariableInitialize(
 #endif
         {L"ProgramFiles6432Folder", InitializeVariable6432Folder, CSIDL_PROGRAM_FILES},
         {L"ProgramMenuFolder", InitializeVariableCsidlFolder, CSIDL_PROGRAMS},
-        {L"RebootPending", InitializeVariableRebootPending, 0},
         {L"SendToFolder", InitializeVariableCsidlFolder, CSIDL_SENDTO},
         {L"ServicePackLevel", InitializeVariableVersionNT, OS_INFO_VARIABLE_ServicePackLevel},
         {L"StartMenuFolder", InitializeVariableCsidlFolder, CSIDL_STARTMENU},
@@ -2031,44 +2026,6 @@ LExit:
     return hr;
 }
 
-static HRESULT InitializeVariableRebootPending(
-    __in DWORD_PTR dwpData,
-    __inout BURN_VARIANT* pValue
-    )
-{
-    UNREFERENCED_PARAMETER(dwpData);
-
-    HRESULT hr = S_OK;
-    BOOL fRebootPending = FALSE;
-    BOOL fComInitialized = FALSE;
-
-    // Do a best effort to ask WU if a reboot is required. If anything goes
-    // wrong then let's pretend a reboot is not required.
-    hr = ::CoInitialize(NULL);
-    if (SUCCEEDED(hr) || RPC_E_CHANGED_MODE == hr)
-    {
-        fComInitialized = TRUE;
-
-        hr = WuaRestartRequired(&fRebootPending);
-        if (FAILED(hr))
-        {
-            fRebootPending = FALSE;
-            hr = S_OK;
-        }
-    }
-
-    hr = BVariantSetNumeric(pValue, fRebootPending);
-    ExitOnFailure(hr, "Failed to set reboot pending variant value.");
-
-LExit:
-    if (fComInitialized)
-    {
-        ::CoUninitialize();
-    }
-
-    return hr;
-}
-
 static HRESULT InitializeSystemLanguageID(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue