From 237bdc6a97c0dd455da99c16e6b3b7cac4c79e86 Mon Sep 17 00:00:00 2001 From: Sean Hall Date: Fri, 2 Sep 2022 16:12:42 -0500 Subject: Add WixCanRestart and make wixstdba elevate for restart when required. Fixes 5499 --- src/burn/engine/EngineForApplication.h | 16 ---------------- src/burn/engine/core.cpp | 7 ++++--- src/burn/engine/core.h | 1 + src/burn/engine/elevation.cpp | 7 ++++++- src/burn/engine/elevation.h | 1 + src/burn/engine/engine.cpp | 23 +++-------------------- src/burn/engine/platform.h | 16 ++++++++++++++++ src/burn/engine/variable.cpp | 25 +++++++++++++++++++++++++ src/burn/test/BurnUnitTest/ElevationTest.cpp | 2 +- src/burn/test/BurnUnitTest/ExitCodeTest.cpp | 2 +- src/burn/test/BurnUnitTest/VariableTest.cpp | 1 + 11 files changed, 59 insertions(+), 42 deletions(-) (limited to 'src/burn') 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(pmsg->lParam)); + hr = CoreElevate(pEngineState, WM_BURN_ELEVATE, reinterpret_cast(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(Environment::GetFolderPath(Environment::SpecialFolder::ApplicationData) + "\\", VariableGetStringHelper(&variables, L"AppDataFolder")); -- cgit v1.2.3-55-g6feb