4 files changed, 163 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
diff --git a/src/test/BurnUnitTest/RegistrationTest.cpp b/src/test/BurnUnitTest/RegistrationTest.cpp
index 1687385b..883a1258 100644
--- a/src/test/BurnUnitTest/RegistrationTest.cpp
+++ b/src/test/BurnUnitTest/RegistrationTest.cpp
@@ -73,6 +73,7 @@ namespace Bootstrapper
             BURN_LOGGING logging = { };
             BURN_PACKAGES packages = { };
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"));
+
             try
             {
                 // set mock API's
@@ -261,6 +262,115 @@ namespace Bootstrapper
         }
 
         [Fact]
+        void RegisterVariablesTest()
+        {
+            HRESULT hr = S_OK;
+            IXMLDOMElement* pixeBundle = NULL;
+            LPWSTR sczCurrentProcess = NULL;
+            BURN_VARIABLES variables = { };
+            BURN_USER_EXPERIENCE userExperience = { };
+            BOOTSTRAPPER_COMMAND command = { };
+            BURN_REGISTRATION registration = { };
+            BURN_LOGGING logging = { };
+            BURN_PACKAGES packages = { };
+            String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"));
+            try
+            {
+                // set mock API's
+                RegFunctionOverride(RegistrationTest_RegCreateKeyExW, RegistrationTest_RegOpenKeyExW, RegistrationTest_RegDeleteKeyExW, NULL, NULL, NULL, NULL, NULL, NULL);
+
+                Registry::CurrentUser->CreateSubKey(gcnew String(HKCU_PATH));
+
+                logging.sczPath = L"BurnUnitTest.txt";
+
+                LPCWSTR wzDocument =
+                    L"<Bundle>"
+                    L"    <UX>"
+                    L"        <Payload Id='ux.dll' FilePath='ux.dll' Packaging='embedded' SourcePath='ux.dll' Hash='000000000000' />"
+                    L"    </UX>"
+                    L"    <Registration Id='{D54F896D-1952-43e6-9C67-B5652240618C}' UpgradeCode='{D54F896D-1952-43e6-9C67-B5652240618C}' Tag='foo' ProviderKey='bar' Version='1.0.0.0' ExecutableName='setup.exe' PerMachine='no'>"
+                    L"        <Arp Register='yes' Publisher='WiX Toolset' DisplayName='Product1' DisplayVersion='1.0.0.0' />"
+                    L"    </Registration>"
+                    L"</Bundle>";
+
+                // load XML document
+                LoadBundleXmlHelper(wzDocument, &pixeBundle);
+
+                hr = VariableInitialize(&variables);
+                TestThrowOnFailure(hr, L"Failed to initialize variables.");
+
+                hr = UserExperienceParseFromXml(&userExperience, pixeBundle);
+                TestThrowOnFailure(hr, L"Failed to parse UX from XML.");
+
+                hr = RegistrationParseFromXml(&registration, pixeBundle);
+                TestThrowOnFailure(hr, L"Failed to parse registration from XML.");
+
+                hr = PlanSetResumeCommand(&registration, BOOTSTRAPPER_ACTION_INSTALL, &command, &logging);
+                TestThrowOnFailure(hr, L"Failed to set registration resume command.");
+
+                hr = PathForCurrentProcess(&sczCurrentProcess, NULL);
+                TestThrowOnFailure(hr, L"Failed to get current process path.");
+
+                //
+                // install
+                //
+
+                // write registration
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &variables, &userExperience, BURN_REGISTRATION_ACTION_OPERATIONS_WRITE_REGISTRATION, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER, 0);
+                TestThrowOnFailure(hr, L"Failed to register bundle.");
+
+                // verify that registration was created
+                Assert::Equal(Int32(BURN_RESUME_MODE_ACTIVE), (Int32)Registry::GetValue(gcnew String(TEST_UNINSTALL_KEY), gcnew String(L"Resume"), nullptr));
+                Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
+
+                // complete registration
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_REQUIRED, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
+                TestThrowOnFailure(hr, L"Failed to unregister bundle.");
+
+                // verify that registration variables were updated
+                registration.fInstalled = TRUE;
+
+                hr = RegistrationSetVariables(&registration, &variables);
+                TestThrowOnFailure(hr, L"Failed to set registration variables.");
+
+                Assert::Equal(1ll, VariableGetNumericHelper(&variables, BURN_BUNDLE_INSTALLED));
+                Assert::Equal(1ll, VariableGetNumericHelper(&variables, BURN_REBOOT_PENDING));
+                Assert::Equal<String^>(gcnew String(L"foo"), VariableGetStringHelper(&variables, BURN_BUNDLE_TAG));
+                Assert::Equal<String^>(gcnew String(L"bar"), VariableGetStringHelper(&variables, BURN_BUNDLE_PROVIDER_KEY));
+                Assert::Equal<String^>(gcnew String(L"1.0.0.0"), VariableGetVersionHelper(&variables, BURN_BUNDLE_VERSION));
+
+                //
+                // uninstall
+                //
+
+                // delete registration
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
+                TestThrowOnFailure(hr, L"Failed to unregister bundle.");
+
+                // verify that registration was removed
+                Assert::Equal((Object^)nullptr, Registry::GetValue(gcnew String(TEST_UNINSTALL_KEY), gcnew String(L"Resume"), nullptr));
+                Assert::Equal((Object^)nullptr, Registry::GetValue(gcnew String(TEST_UNINSTALL_KEY), gcnew String(L"Installed"), nullptr));
+                Assert::Equal((Object^)nullptr, Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
+            }
+            finally
+            {
+                ReleaseStr(sczCurrentProcess);
+                ReleaseObject(pixeBundle);
+                UserExperienceUninitialize(&userExperience);
+                RegistrationUninitialize(&registration);
+                VariablesUninitialize(&variables);
+
+                Registry::CurrentUser->DeleteSubKeyTree(gcnew String(ROOT_PATH));
+                if (Directory::Exists(cacheDirectory))
+                {
+                    Directory::Delete(cacheDirectory, true);
+                }
+
+                RegFunctionOverride(NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+            }
+        }
+
+        [Fact]
         void RegisterArpFullTest()
         {
             HRESULT hr = S_OK;