8 files changed, 56 insertions, 40 deletions

diff --git a/src/burn/engine/EngineForApplication.h b/src/burn/engine/EngineForApplication.h
index e5e8f6d7..d25a7e51 100644
--- a/src/burn/engine/EngineForApplication.h
+++ b/src/burn/engine/EngineForApplication.h
@@ -6,22 +6,6 @@
 extern "C" {
 #endif
 
-// constants
-
-enum WM_BURN
-{
-    WM_BURN_FIRST = WM_APP + 0xFFF, // this enum value must always be first.
-
-    WM_BURN_DETECT,
-    WM_BURN_PLAN,
-    WM_BURN_ELEVATE,
-    WM_BURN_APPLY,
-    WM_BURN_LAUNCH_APPROVED_EXE,
-    WM_BURN_QUIT,
-
-    WM_BURN_LAST, // this enum value must always be last.
-};
-
 // structs
 
 typedef struct _BOOTSTRAPPER_ENGINE_CONTEXT
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index c8dce17b..93b9c002 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -593,6 +593,7 @@ LExit:
 
 extern "C" HRESULT CoreElevate(
     __in BURN_ENGINE_STATE* pEngineState,
+    __in WM_BURN reason,
     __in_opt HWND hwndParent
     )
 {
@@ -608,7 +609,7 @@ extern "C" HRESULT CoreElevate(
             ExitOnFailure(hr, "Failed to cache engine to working directory.");
         }
 
-        hr = ElevationElevate(pEngineState, hwndParent);
+        hr = ElevationElevate(pEngineState, reason, hwndParent);
         if (E_SUSPECTED_AV_INTERFERENCE == hr && 1 > cAVRetryAttempts)
         {
             ++cAVRetryAttempts;
@@ -720,7 +721,7 @@ extern "C" HRESULT CoreApply(
     // Elevate.
     if (pEngineState->plan.fPerMachine)
     {
-        hr = CoreElevate(pEngineState, pEngineState->userExperience.hwndApply);
+        hr = CoreElevate(pEngineState, WM_BURN_APPLY, pEngineState->userExperience.hwndApply);
         ExitOnFailure(hr, "Failed to elevate.");
 
         hr = ElevationApplyInitialize(pEngineState->companionConnection.hPipe, &pEngineState->userExperience, &pEngineState->variables, &pEngineState->plan);
@@ -872,7 +873,7 @@ extern "C" HRESULT CoreLaunchApprovedExe(
     ExitOnRootFailure(hr, "BA aborted LaunchApprovedExe begin.");
 
     // Elevate.
-    hr = CoreElevate(pEngineState, pLaunchApprovedExe->hwndParent);
+    hr = CoreElevate(pEngineState, WM_BURN_LAUNCH_APPROVED_EXE, pLaunchApprovedExe->hwndParent);
     ExitOnFailure(hr, "Failed to elevate.");
 
     // Launch.
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 812b40b1..3e636640 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -251,6 +251,7 @@ HRESULT CorePlan(
     );
 HRESULT CoreElevate(
     __in BURN_ENGINE_STATE* pEngineState,
+    __in WM_BURN reason,
     __in_opt HWND hwndParent
     );
 HRESULT CoreApply(
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 4d2e8544..7fd372b0 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -402,6 +402,7 @@ static HRESULT ElevatedOnExecuteActionComplete(
 
 extern "C" HRESULT ElevationElevate(
     __in BURN_ENGINE_STATE* pEngineState,
+    __in WM_BURN reason,
     __in_opt HWND hwndParent
     )
 {
@@ -450,7 +451,11 @@ extern "C" HRESULT ElevationElevate(
         else if (HRESULT_FROM_WIN32(ERROR_CANCELLED) == hr)
         {
             // The user clicked "Cancel" on the elevation prompt or the elevation prompt timed out, provide the notification with the option to retry.
-            hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT);
+            if (WM_BURN_APPLY == reason)
+            {
+                hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT);
+            }
+
             nResult = UserExperienceSendError(&pEngineState->userExperience, BOOTSTRAPPER_ERROR_TYPE_ELEVATE, NULL, hr, NULL, MB_ICONERROR | MB_RETRYCANCEL, IDNOACTION);
         }
     } while (IDRETRY == nResult);
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index 0d15b470..810287a3 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -10,6 +10,7 @@ extern "C" {
 // Parent (per-user process) side functions.
 HRESULT ElevationElevate(
     __in BURN_ENGINE_STATE* pEngineState,
+    __in WM_BURN reason,
     __in_opt HWND hwndParent
     );
 HRESULT ElevationApplyInitialize(
diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index 89082a88..48196655 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -893,7 +893,7 @@ static HRESULT ProcessMessage(
         break;
 
     case WM_BURN_ELEVATE:
-        hr = CoreElevate(pEngineState, reinterpret_cast<HWND>(pmsg->lParam));
+        hr = CoreElevate(pEngineState, WM_BURN_ELEVATE, reinterpret_cast<HWND>(pmsg->lParam));
         break;
 
     case WM_BURN_APPLY:
@@ -1069,26 +1069,10 @@ static HRESULT Restart(
     )
 {
     HRESULT hr = S_OK;
-    HANDLE hProcessToken = NULL;
-    TOKEN_PRIVILEGES priv = { };
     DWORD dwRetries = 0;
 
-    if (!::OpenProcessToken(::GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hProcessToken))
-    {
-        ExitWithLastError(hr, "Failed to get process token.");
-    }
-
-    priv.PrivilegeCount = 1;
-    priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
-    if (!::LookupPrivilegeValueW(NULL, L"SeShutdownPrivilege", &priv.Privileges[0].Luid))
-    {
-        ExitWithLastError(hr, "Failed to get shutdown privilege LUID.");
-    }
-
-    if (!::AdjustTokenPrivileges(hProcessToken, FALSE, &priv, sizeof(TOKEN_PRIVILEGES), NULL, 0))
-    {
-        ExitWithLastError(hr, "Failed to adjust token to add shutdown privileges.");
-    }
+    hr = ProcEnablePrivilege(::GetCurrentProcess(), SE_SHUTDOWN_NAME);
+    ExitOnFailure(hr, "Failed to enable shutdown privilege in process token.");
 
     pEngineState->fRestarting = TRUE;
     CoreUpdateRestartState(pEngineState, BURN_RESTART_STATE_REQUESTING);
@@ -1147,7 +1131,6 @@ static HRESULT Restart(
     }
 
 LExit:
-    ReleaseHandle(hProcessToken);
     return hr;
 }
 
diff --git a/src/burn/engine/platform.h b/src/burn/engine/platform.h
index 49d8f3e9..60184c23 100644
--- a/src/burn/engine/platform.h
+++ b/src/burn/engine/platform.h
@@ -8,6 +8,22 @@
 extern "C" {
 #endif
 
+// constants
+
+enum WM_BURN
+{
+    WM_BURN_FIRST = WM_APP + 0xFFF, // this enum value must always be first.
+
+    WM_BURN_DETECT,
+    WM_BURN_PLAN,
+    WM_BURN_ELEVATE,
+    WM_BURN_APPLY,
+    WM_BURN_LAUNCH_APPROVED_EXE,
+    WM_BURN_QUIT,
+
+    WM_BURN_LAST, // this enum value must always be last.
+};
+
 // forward declare
 
 enum BURN_MODE;
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