20 files changed, 336 insertions, 76 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
diff --git a/src/burn/test/BurnUnitTest/ElevationTest.cpp b/src/burn/test/BurnUnitTest/ElevationTest.cpp
index 8d4cc7ff..713d9e07 100644
--- a/src/burn/test/BurnUnitTest/ElevationTest.cpp
+++ b/src/burn/test/BurnUnitTest/ElevationTest.cpp
@@ -68,7 +68,7 @@ namespace Bootstrapper
                 //
                 // per-user side setup
                 //
-                hr = ElevationElevate(&engineState, NULL);
+                hr = ElevationElevate(&engineState, WM_BURN_ELEVATE, NULL);
                 TestThrowOnFailure(hr, L"Failed to elevate.");
 
                 // post execute message
diff --git a/src/burn/test/BurnUnitTest/ExitCodeTest.cpp b/src/burn/test/BurnUnitTest/ExitCodeTest.cpp
index d7d91d06..465ee04b 100644
--- a/src/burn/test/BurnUnitTest/ExitCodeTest.cpp
+++ b/src/burn/test/BurnUnitTest/ExitCodeTest.cpp
@@ -118,7 +118,7 @@ static void LoadEngineState(
                 //
                 LoadEngineState(&engineState);
 
-                hr = ElevationElevate(&engineState, NULL);
+                hr = ElevationElevate(&engineState, WM_BURN_ELEVATE, NULL);
                 TestThrowOnFailure(hr, L"Failed to elevate.");
 
                 for (DWORD i = 0; i < countof(rgExitCodeItems); ++i)
diff --git a/src/burn/test/BurnUnitTest/VariableTest.cpp b/src/burn/test/BurnUnitTest/VariableTest.cpp
index 8ee6e179..b373ae8e 100644
--- a/src/burn/test/BurnUnitTest/VariableTest.cpp
+++ b/src/burn/test/BurnUnitTest/VariableTest.cpp
@@ -564,6 +564,7 @@ namespace Bootstrapper
                 VariableGetNumericHelper(&variables, L"TerminalServer");
                 VariableGetNumericHelper(&variables, L"UserUILanguageID");
                 VariableGetNumericHelper(&variables, L"UserLanguageID");
+                VariableGetNumericHelper(&variables, L"WixCanRestart");
 
                 // known folders
                 Assert::Equal<String^>(Environment::GetFolderPath(Environment::SpecialFolder::ApplicationData) + "\\", VariableGetStringHelper(&variables, L"AppDataFolder"));
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 8c4b0b35..50de2c7f 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -4,6 +4,7 @@
 #include "BalBaseBootstrapperApplicationProc.h"
 #include "BalBaseBootstrapperApplication.h"
 
+static const LPCWSTR WIXBUNDLE_VARIABLE_CANRESTART = L"WixCanRestart";
 static const LPCWSTR WIXBUNDLE_VARIABLE_ELEVATED = L"WixBundleElevated";
 
 static const LPCWSTR WIXSTDBA_WINDOW_CLASS = L"WixStdBA";
@@ -1289,6 +1290,28 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+
+    virtual STDMETHODIMP OnElevateComplete(
+        __in HRESULT hrStatus
+        )
+    {
+        if (m_fElevatingForRestart)
+        {
+            m_fElevatingForRestart = FALSE;
+
+            if (SUCCEEDED(hrStatus))
+            {
+                m_fAllowRestart = TRUE;
+
+                ::SendMessageW(m_hWnd, WM_CLOSE, 0, 0);
+            }
+            // else if failed then OnError showed the user an error message box
+        }
+
+        return __super::OnElevateComplete(hrStatus);
+    }
+
+
     virtual STDMETHODIMP_(void) BAProcFallback(
         __in BOOTSTRAPPER_APPLICATION_MESSAGE message,
         __in const LPVOID pvArgs,
@@ -3734,14 +3757,16 @@ private:
 
             if (dwOldPageId != dwNewPageId)
             {
+                LONGLONG llCanRestart = 0;
+                LONGLONG llElevated = 0;
+
+                BalGetNumericVariable(WIXBUNDLE_VARIABLE_CANRESTART, &llCanRestart);
+                BalGetNumericVariable(WIXBUNDLE_VARIABLE_ELEVATED, &llElevated);
+                m_fRestartRequiresElevation = !llCanRestart && !llElevated;
+
                 // Enable disable controls per-page.
                 if (m_rgdwPageIds[WIXSTDBA_PAGE_INSTALL] == dwNewPageId) // on the "Install" page, ensure the install button is enabled/disabled correctly.
                 {
-                    LONGLONG llElevated = 0;
-                    if (m_Bundle.fPerMachine)
-                    {
-                        BalGetNumericVariable(WIXBUNDLE_VARIABLE_ELEVATED, &llElevated);
-                    }
                     ThemeControlElevates(m_pControlInstallButton, (m_Bundle.fPerMachine && !llElevated));
 
                     // If the EULA control exists, show it only if a license URL is provided as well.
@@ -3757,12 +3782,18 @@ private:
                 }
                 else if (m_rgdwPageIds[WIXSTDBA_PAGE_MODIFY] == dwNewPageId)
                 {
+                    ThemeControlElevates(m_pControlRepairButton, (m_Bundle.fPerMachine && !llElevated));
+                    ThemeControlElevates(m_pControlUninstallButton, (m_Bundle.fPerMachine && !llElevated));
+
                     ThemeControlEnable(m_pControlRepairButton, !m_fSuppressRepair);
                 }
                 else if (m_rgdwPageIds[WIXSTDBA_PAGE_SUCCESS] == dwNewPageId) // on the "Success" page, check if the restart or launch button should be enabled.
                 {
                     BOOL fEnableRestartButton = FALSE;
                     BOOL fLaunchTargetExists = FALSE;
+
+                    ThemeControlElevates(m_pControlSuccessRestartButton, m_fRestartRequiresElevation);
+
                     if (m_fShouldRestart)
                     {
                         if (BAL_INFO_RESTART_PROMPT == m_BalInfoCommand.restart)
@@ -3784,6 +3815,8 @@ private:
                     BOOL fShowErrorMessage = FALSE;
                     BOOL fEnableRestartButton = FALSE;
 
+                    ThemeControlElevates(m_pControlFailureRestartButton, m_fRestartRequiresElevation);
+
                     if (FAILED(m_hrFinal))
                     {
                         // If we know the failure message, use that.
@@ -4156,8 +4189,20 @@ private:
     {
         AssertSz(m_fRestartRequired, "Restart must be requested to be able to click on the restart button.");
 
-        m_fAllowRestart = TRUE;
-        ::SendMessageW(m_hWnd, WM_CLOSE, 0, 0);
+        if (m_fRestartRequiresElevation)
+        {
+            m_fElevatingForRestart = TRUE;
+            ThemeControlEnable(m_pControlFailureRestartButton, FALSE);
+            ThemeControlEnable(m_pControlSuccessRestartButton, FALSE);
+
+            m_pEngine->Elevate(m_hWnd);
+        }
+        else
+        {
+            m_fAllowRestart = TRUE;
+
+            ::SendMessageW(m_hWnd, WM_CLOSE, 0, 0);
+        }
     }
 
 
@@ -4649,6 +4694,8 @@ public:
         m_fRestartRequired = FALSE;
         m_fShouldRestart = FALSE;
         m_fAllowRestart = FALSE;
+        m_fRestartRequiresElevation = FALSE;
+        m_fElevatingForRestart = FALSE;
 
         m_sczLicenseFile = NULL;
         m_sczLicenseUrl = NULL;
@@ -4956,6 +5003,8 @@ private:
     BOOL m_fRestartRequired;
     BOOL m_fShouldRestart;
     BOOL m_fAllowRestart;
+    BOOL m_fRestartRequiresElevation;
+    BOOL m_fElevatingForRestart;
 
     LPWSTR m_sczLicenseFile;
     LPWSTR m_sczLicenseUrl;
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/procutil.h b/src/libs/dutil/WixToolset.DUtil/inc/procutil.h
index d61d91b5..6a641a5b 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/procutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/procutil.h
@@ -23,9 +23,21 @@ HRESULT DAPI ProcSystem(
     __out BOOL* pfSystem
     );
 
-HRESULT DAPI ProcTokenUser(
+HRESULT DAPI ProcGetTokenInformation(
     __in HANDLE hProcess,
-    __out TOKEN_USER** ppTokenUser
+    __in TOKEN_INFORMATION_CLASS tokenInformationClass,
+    __out LPVOID* ppvTokenInformation
+    );
+
+HRESULT DAPI ProcHasPrivilege(
+    __in HANDLE hProcess,
+    __in LPCWSTR wzPrivilegeName,
+    __out BOOL* pfHasPrivilege
+    );
+
+HRESULT DAPI ProcEnablePrivilege(
+    __in HANDLE hProcess,
+    __in LPCWSTR wzPrivilegeName
     );
 
 HRESULT DAPI ProcWow64(
diff --git a/src/libs/dutil/WixToolset.DUtil/procutil.cpp b/src/libs/dutil/WixToolset.DUtil/procutil.cpp
index 29f575ae..376aec6d 100644
--- a/src/libs/dutil/WixToolset.DUtil/procutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/procutil.cpp
@@ -85,7 +85,7 @@ extern "C" HRESULT DAPI ProcSystem(
     HRESULT hr = S_OK;
     TOKEN_USER* pTokenUser = NULL;
 
-    hr = ProcTokenUser(hProcess, &pTokenUser);
+    hr = ProcGetTokenInformation(hProcess, TokenUser, reinterpret_cast<LPVOID*>(&pTokenUser));
     ProcExitOnFailure(hr, "Failed to get TokenUser from process token.");
 
     *pfSystem = ::IsWellKnownSid(pTokenUser->User.Sid, WinLocalSystemSid);
@@ -96,15 +96,16 @@ LExit:
     return hr;
 }
 
-extern "C" HRESULT DAPI ProcTokenUser(
+extern "C" HRESULT DAPI ProcGetTokenInformation(
     __in HANDLE hProcess,
-    __out TOKEN_USER** ppTokenUser
+    __in TOKEN_INFORMATION_CLASS tokenInformationClass,
+    __out LPVOID* ppvTokenInformation
     )
 {
     HRESULT hr = S_OK;
     DWORD er = ERROR_SUCCESS;
     HANDLE hToken = NULL;
-    TOKEN_USER* pTokenUser = NULL;
+    LPVOID pvTokenInformation = NULL;
     DWORD cbToken = 0;
 
     if (!::OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
@@ -112,33 +113,104 @@ extern "C" HRESULT DAPI ProcTokenUser(
         ProcExitWithLastError(hr, "Failed to open process token.");
     }
 
-    if (::GetTokenInformation(hToken, TokenUser, pTokenUser, 0, &cbToken))
-    {
-        er = ERROR_SUCCESS;
-    }
-    else
+    if (!::GetTokenInformation(hToken, tokenInformationClass, pvTokenInformation, 0, &cbToken))
     {
         er = ::GetLastError();
     }
 
     if (er != ERROR_INSUFFICIENT_BUFFER)
     {
-        ProcExitOnWin32Error(er, hr, "Failed to get user from process token size.");
+        ProcExitOnWin32Error(er, hr, "Failed to get information from process token size.");
     }
 
-    pTokenUser = reinterpret_cast<TOKEN_USER*>(MemAlloc(cbToken, TRUE));
-    ProcExitOnNull(pTokenUser, hr, E_OUTOFMEMORY, "Failed to allocate token information.");
+    pvTokenInformation = MemAlloc(cbToken, TRUE);
+    ProcExitOnNull(pvTokenInformation, hr, E_OUTOFMEMORY, "Failed to allocate token information.");
 
-    if (!::GetTokenInformation(hToken, TokenUser, pTokenUser, cbToken, &cbToken))
+    if (!::GetTokenInformation(hToken, tokenInformationClass, pvTokenInformation, cbToken, &cbToken))
     {
-        ProcExitWithLastError(hr, "Failed to get user from process token.");
+        ProcExitWithLastError(hr, "Failed to get information from process token.");
     }
 
-    *ppTokenUser = pTokenUser;
-    pTokenUser = NULL;
+    *ppvTokenInformation = pvTokenInformation;
+    pvTokenInformation = NULL;
+
+LExit:
+    ReleaseMem(pvTokenInformation);
+    ReleaseHandle(hToken);
+
+    return hr;
+}
+
+extern "C" HRESULT DAPI ProcHasPrivilege(
+    __in HANDLE hProcess,
+    __in LPCWSTR wzPrivilegeName,
+    __out BOOL* pfHasPrivilege
+    )
+{
+    HRESULT hr = S_OK;
+    TOKEN_PRIVILEGES* pTokenPrivileges = NULL;
+    LUID luidPrivilege = { };
+
+    *pfHasPrivilege = FALSE;
+
+    if (!::LookupPrivilegeValueW(NULL, wzPrivilegeName, &luidPrivilege))
+    {
+        ProcExitWithLastError(hr, "Failed to get privilege LUID: %ls", wzPrivilegeName);
+    }
+
+    hr = ProcGetTokenInformation(hProcess, TokenPrivileges, reinterpret_cast<LPVOID*>(&pTokenPrivileges));
+    ProcExitOnFailure(hr, "Failed to get token privilege information.");
+
+    for (DWORD i = 0; i < pTokenPrivileges->PrivilegeCount; ++i)
+    {
+        LUID* pTokenLuid = &pTokenPrivileges->Privileges[i].Luid;
+
+        if (luidPrivilege.LowPart == pTokenLuid->LowPart && luidPrivilege.HighPart == pTokenLuid->HighPart)
+        {
+            *pfHasPrivilege = TRUE;
+            break;
+        }
+    }
+
+LExit:
+    ReleaseMem(pTokenPrivileges);
+
+    return hr;
+}
+
+extern "C" HRESULT DAPI ProcEnablePrivilege(
+    __in HANDLE hProcess,
+    __in LPCWSTR wzPrivilegeName
+    )
+{
+    HRESULT hr = S_OK;
+    HANDLE hToken = NULL;
+    TOKEN_PRIVILEGES priv = { };
+
+    priv.PrivilegeCount = 1;
+    priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+
+    if (!::LookupPrivilegeValueW(NULL, wzPrivilegeName, &priv.Privileges[0].Luid))
+    {
+        ProcExitWithLastError(hr, "Failed to get privilege LUID: %ls", wzPrivilegeName);
+    }
+
+    if (!::OpenProcessToken(hProcess, TOKEN_ADJUST_PRIVILEGES, &hToken))
+    {
+        ProcExitWithLastError(hr, "Failed to get process token to adjust privileges.");
+    }
+
+    if (!::AdjustTokenPrivileges(hToken, FALSE, &priv, sizeof(TOKEN_PRIVILEGES), NULL, 0))
+    {
+        ProcExitWithLastError(hr, "Failed to adjust token to add privilege: %ls", wzPrivilegeName);
+    }
+
+    if (ERROR_NOT_ALL_ASSIGNED == ::GetLastError())
+    {
+        hr = S_FALSE;
+    }
 
 LExit:
-    ReleaseMem(pTokenUser);
     ReleaseHandle(hToken);
 
     return hr;
diff --git a/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp b/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp
index 297d90f4..8672ed38 100644
--- a/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp
+++ b/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp
@@ -6,6 +6,7 @@ using namespace System;
 using namespace System::Security::Principal;
 using namespace Xunit;
 using namespace WixBuildTools::TestSupport;
+using namespace WixBuildTools::TestSupport::XunitExtensions;
 
 namespace DutilTests
 {
@@ -13,7 +14,7 @@ namespace DutilTests
     {
     public:
         [Fact]
-        void ProcTokenUserTest()
+        void ProcGetTokenInformationTest()
         {
             HRESULT hr = S_OK;
             TOKEN_USER* pTokenUser = NULL;
@@ -21,7 +22,7 @@ namespace DutilTests
 
             try
             {
-                hr = ProcTokenUser(::GetCurrentProcess(), &pTokenUser);
+                hr = ProcGetTokenInformation(::GetCurrentProcess(), TokenUser, reinterpret_cast<LPVOID*>(&pTokenUser));
                 NativeAssert::Succeeded(hr, "Failed to get TokenUser for current process.");
 
                 if (!::ConvertSidToStringSidW(pTokenUser->User.Sid, &sczSid))
@@ -38,5 +39,28 @@ namespace DutilTests
                 ReleaseStr(sczSid);
             }
         }
+
+        [SkippableFact]
+        void ProcHasPrivilegeTest()
+        {
+            HRESULT hr = S_OK;
+            BOOL fHasPrivilege = FALSE;
+
+            hr = ProcHasPrivilege(::GetCurrentProcess(), SE_CREATE_TOKEN_NAME, &fHasPrivilege);
+            NativeAssert::Succeeded(hr, "Failed to check privilege for current process.");
+
+            if (fHasPrivilege)
+            {
+                WixAssert::Skip("Didn't expect process to have SE_CREATE_TOKEN_NAME privilege");
+            }
+
+            hr = ProcHasPrivilege(::GetCurrentProcess(), SE_INC_WORKING_SET_NAME, &fHasPrivilege);
+            NativeAssert::Succeeded(hr, "Failed to check privilege for current process.");
+
+            if (!fHasPrivilege)
+            {
+                WixAssert::Skip("Expected process to have SE_INC_WORKING_SET_NAME privilege");
+            }
+        }
     };
 }
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index c619dbd6..a2b8e041 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -477,7 +477,7 @@ static HRESULT LogUserSid()
     TOKEN_USER* pTokenUser = NULL;
     LPWSTR sczSid = NULL;
 
-    hr = ProcTokenUser(::GetCurrentProcess(), &pTokenUser);
+    hr = ProcGetTokenInformation(::GetCurrentProcess(), TokenUser, reinterpret_cast<LPVOID*>(&pTokenUser));
     BalExitOnFailure(hr, "Failed to get user from process token.");
 
     if (!::ConvertSidToStringSidW(pTokenUser->User.Sid, &sczSid))
diff --git a/src/test/burn/TestData/Manual/BundleA/ManualTests.txt b/src/test/burn/TestData/Manual/BundleA/ManualTests.txt
index d432f94a..b8c6e74d 100644
--- a/src/test/burn/TestData/Manual/BundleA/ManualTests.txt
+++ b/src/test/burn/TestData/Manual/BundleA/ManualTests.txt
@@ -94,11 +94,12 @@ CanRestartFromUnelevatedPerMachineBundleWithoutShutdownPrivilege
 2. Click Install.
 3. Verify that the UAC prompt came up and accept elevation.
 4. Click OK on the OnExecuteBegin message box.
-5. Click Restart.
-6. The machine should restart.
-7. Login to the machine.
-8. Verify that the bundle did not automatically start running.
-9. Look in the bundle log and verify that the restart request didn't cause any errors, and that it logged messages that look like (the process id for w005 must match the elevated i400 and i401):
+5. Verify the Restart button doesn't have the UAC shield and then click it.
+6. Verify that there was no UAC prompt.
+7. The machine should restart.
+8. Login to the machine.
+9. Verify that the bundle did not automatically start running.
+10. Look in the bundle log and verify that the restart request didn't cause any errors, and that it logged messages that look like (the process id for w005 must match the elevated i400 and i401):
 
 [0DDC:0448]w005: Restarting computer...
 [1228:18CC]i400: Received system request to shut down the process: allowed: Yes, elevated: No, critical: No, logoff: No, close app: No
@@ -107,4 +108,47 @@ CanRestartFromUnelevatedPerMachineBundleWithoutShutdownPrivilege
 [0DDC:0954]i401: Received result of system request to shut down the process: closing: Yes, elevated: Yes, critical: No, logoff: No, close app: No
 =======================================
 
+(11. Uninstall the bundle)
+
+CanRestartFromUnelevatedPerUserBundleWithoutShutdownPrivilege
+
+    Note: Requires different User Rights Assignment settings from CanLogRestartFailure - Only Administrators should have the shutdown privilege. Users should not have it.
+
+1. Run BundleC.exe unelevated with the command line EXEEXITCODE=3010.
+2. Click Install.
+3. Verify the Restart button has the UAC shield and then click it.
+4. Verify that the UAC prompt came up and accept elevation.
+5. The machine should restart.
+6. Login to the machine.
+7. Verify that the bundle did not automatically start running.
+8. Look in the bundle log and verify that the restart request didn't cause any errors, and that it logged messages that look like (the process id for w005 must match the elevated i400 and i401):
+
+[0DDC:0448]w005: Restarting computer...
+[1228:18CC]i400: Received system request to shut down the process: allowed: Yes, elevated: No, critical: No, logoff: No, close app: No
+[1228:18CC]i401: Received result of system request to shut down the process: closing: Yes, elevated: No, critical: No, logoff: No, close app: No
+[0DDC:0954]i400: Received system request to shut down the process: allowed: Yes, elevated: Yes, critical: No, logoff: No, close app: No
+[0DDC:0954]i401: Received result of system request to shut down the process: closing: Yes, elevated: Yes, critical: No, logoff: No, close app: No
+=======================================
+
+(9. Uninstall the bundle)
+
+CanRestartFromUnelevatedPerUserBundleWithShutdownPrivilege
+
+    Note: Requires different User Rights Assignment settings from CanLogRestartFailure - Administrators and Users should have the shutdown privilege.
+
+1. Run BundleC.exe unelevated with the command line EXEEXITCODE=3010.
+2. Click Install.
+3. Verify that there was no UAC prompt.
+4. Verify the Restart button doesn't have the UAC shield and then click it.
+5. Verify that there was no UAC prompt.
+6. The machine should restart.
+7. Login to the machine.
+8. Verify that the bundle did not automatically start running.
+9. Look in the bundle log and verify that the restart request didn't cause any errors, and that it logged messages that look like (there should be no elevated i400 and i401 and the process id for w005 must match):
+
+[1228:0448]w005: Restarting computer...
+[1228:18CC]i400: Received system request to shut down the process: allowed: Yes, elevated: No, critical: No, logoff: No, close app: No
+[1228:18CC]i401: Received result of system request to shut down the process: closing: Yes, elevated: No, critical: No, logoff: No, close app: No
+=======================================
+
 (10. Uninstall the bundle)
diff --git a/src/test/burn/TestData/Manual/BundleC/BundleC.wixproj b/src/test/burn/TestData/Manual/BundleC/BundleC.wixproj
new file mode 100644
index 00000000..eabd529c
--- /dev/null
+++ b/src/test/burn/TestData/Manual/BundleC/BundleC.wixproj
@@ -0,0 +1,18 @@
+<!-- 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. -->
+<Project Sdk="WixToolset.Sdk">
+  <PropertyGroup>
+    <OutputType>Bundle</OutputType>
+    <BA>hyperlinkLicense</BA>
+    <UpgradeCode>{5CAE82BB-A2F3-4994-9BA8-4ACEAFAE7738}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\Bundle.wxs" Link="Bundle.wxs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="..\PackagePerUser\PackagePerUser.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Bal.wixext" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/Manual/BundleC/BundleC.wxs b/src/test/burn/TestData/Manual/BundleC/BundleC.wxs
new file mode 100644
index 00000000..aadcd8b2
--- /dev/null
+++ b/src/test/burn/TestData/Manual/BundleC/BundleC.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:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackagePerUser" SourceFile="$(var.PackagePerUser.TargetPath)" />
+      <ExePackage Id="ExeA" Cache="remove" PerMachine="no"
+                  DetectCondition="" Permanent="yes" InstallArguments="/ec [EXEEXITCODE]">
+        <PayloadGroupRef Id="TestExePayloads" />
+      </ExePackage>
+    </PackageGroup>
+    <Variable Name="EXEEXITCODE" bal:Overridable="yes" Value="0" />
+  </Fragment>
+</Wix>
diff --git a/src/test/burn/TestData/Manual/PackagePerUser/PackagePerUser.wixproj b/src/test/burn/TestData/Manual/PackagePerUser/PackagePerUser.wixproj
new file mode 100644
index 00000000..0efa4d63
--- /dev/null
+++ b/src/test/burn/TestData/Manual/PackagePerUser/PackagePerUser.wixproj
@@ -0,0 +1,9 @@
+<!-- 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. -->
+<Project Sdk="WixToolset.Sdk">
+  <PropertyGroup>
+    <UpgradeCode>{F7FBAC90-07A6-4DEB-A9C5-E267EDDA28EF}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\PackagePerUser.wxs" Link="PackagePerUser.wxs" />
+  </ItemGroup>
+</Project>
\ No newline at end of file