Add WixCanRestart and make wixstdba elevate for restart when required.

Fixes 5499

1 files changed, 25 insertions, 0 deletions

diff --git a/src/burn/engine/variable.cpp b/src/burn/engine/variable.cpp
index 81fa31b7..3947e29e 100644
--- a/src/burn/engine/variable.cpp
+++ b/src/burn/engine/variable.cpp
@@ -153,6 +153,10 @@ static HRESULT InitializeVariablePrivileged(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue
     );
+static HRESULT InitializeVariableProcessTokenPrivilege(
+    __in DWORD_PTR dwpData,
+    __inout BURN_VARIANT* pValue
+    );
 static HRESULT InitializeSystemLanguageID(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue
@@ -284,6 +288,7 @@ extern "C" HRESULT VariableInitialize(
         {L"WindowsFolder", InitializeVariableCsidlFolder, CSIDL_WINDOWS},
         {L"WindowsVolume", InitializeVariableWindowsVolumeFolder, 0},
         {BURN_BUNDLE_ACTION, InitializeVariableNumeric, 0, FALSE, TRUE},
+        {L"WixCanRestart", InitializeVariableProcessTokenPrivilege, (DWORD_PTR)SE_SHUTDOWN_NAME},
         {BURN_BUNDLE_COMMAND_LINE_ACTION, InitializeVariableNumeric, 0, FALSE, TRUE},
         {BURN_BUNDLE_EXECUTE_PACKAGE_CACHE_FOLDER, InitializeVariableString, NULL, FALSE, TRUE},
         {BURN_BUNDLE_EXECUTE_PACKAGE_ACTION, InitializeVariableString, NULL, FALSE, TRUE},
@@ -2165,6 +2170,26 @@ LExit:
     return hr;
 }
 
+static HRESULT InitializeVariableProcessTokenPrivilege(
+    __in DWORD_PTR dwpData,
+    __inout BURN_VARIANT* pValue
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fHasPrivilege = FALSE;
+    LPCWSTR wzPrivilegeName = (LPCWSTR)dwpData;
+
+    hr = ProcHasPrivilege(::GetCurrentProcess(), wzPrivilegeName, &fHasPrivilege);
+    ExitOnFailure(hr, "Failed to check if process token has privilege: %ls.", wzPrivilegeName);
+
+    // set value
+    hr = BVariantSetNumeric(pValue, fHasPrivilege);
+    ExitOnFailure(hr, "Failed to set variant value.");
+
+LExit:
+    return hr;
+}
+
 static HRESULT InitializeSystemLanguageID(
     __in DWORD_PTR dwpData,
     __inout BURN_VARIANT* pValue