From 29f7e00586412163a20e298fbf84505f8a917425 Mon Sep 17 00:00:00 2001 From: Sean Hall Date: Tue, 3 May 2022 15:30:50 -0500 Subject: Only block shutdown during Apply. --- src/burn/engine/core.cpp | 4 ++++ src/burn/engine/elevation.cpp | 21 ++++++++++++++++----- src/burn/engine/elevation.h | 3 ++- src/burn/engine/engine.cpp | 2 +- src/burn/engine/plan.h | 1 + src/burn/engine/uithread.cpp | 20 +++++++------------- src/burn/engine/userexperience.cpp | 25 ------------------------- src/burn/engine/userexperience.h | 5 ----- 8 files changed, 31 insertions(+), 50 deletions(-) (limited to 'src/burn') diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp index ec996c4c..19b739f9 100644 --- a/src/burn/engine/core.cpp +++ b/src/burn/engine/core.cpp @@ -681,6 +681,8 @@ extern "C" HRESULT CoreApply( hr = ApplyLock(FALSE, &hLock); ExitOnFailure(hr, "Another per-user setup is already executing."); + pEngineState->plan.fApplying = TRUE; + // Initialize only after getting a lock. fApplyInitialize = TRUE; ApplyInitialize(); @@ -814,6 +816,8 @@ LExit: ApplyUninitialize(); } + pEngineState->plan.fApplying = FALSE; + if (hLock) { ::ReleaseMutex(hLock); diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp index d2c5c8c6..75b24ec3 100644 --- a/src/burn/engine/elevation.cpp +++ b/src/burn/engine/elevation.cpp @@ -94,6 +94,7 @@ typedef struct _BURN_ELEVATION_CHILD_MESSAGE_CONTEXT HANDLE hPipe; HANDLE* phLock; BOOL* pfDisabledAutomaticUpdates; + BOOL* pfApplying; BURN_APPROVED_EXES* pApprovedExes; BURN_CACHE* pCache; BURN_CONTAINERS* pContainers; @@ -176,6 +177,7 @@ static HRESULT OnApplyInitialize( __in BURN_PACKAGES* pPackages, __in HANDLE* phLock, __in BOOL* pfDisabledWindowsUpdate, + __in BOOL* pfApplying, __in BYTE* pbData, __in SIZE_T cbData ); @@ -185,7 +187,8 @@ static HRESULT ElevatedProcessDetect( __in BURN_PACKAGES* pPackages ); static HRESULT OnApplyUninitialize( - __in HANDLE* phLock + __in HANDLE* phLock, + __in BOOL* pfApplying ); static HRESULT OnSessionBegin( __in BURN_CACHE* pCache, @@ -1538,7 +1541,8 @@ extern "C" HRESULT ElevationChildPumpMessages( __out HANDLE* phLock, __out BOOL* pfDisabledAutomaticUpdates, __out DWORD* pdwChildExitCode, - __out BOOL* pfRestart + __out BOOL* pfRestart, + __out BOOL* pfApplying ) { HRESULT hr = S_OK; @@ -1561,6 +1565,7 @@ extern "C" HRESULT ElevationChildPumpMessages( context.hPipe = hPipe; context.phLock = phLock; context.pfDisabledAutomaticUpdates = pfDisabledAutomaticUpdates; + context.pfApplying = pfApplying; context.pApprovedExes = pApprovedExes; context.pCache = pCache; context.pContainers = pContainers; @@ -2114,11 +2119,11 @@ static HRESULT ProcessElevatedChildMessage( break; case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE: - hrResult = OnApplyInitialize(pContext->hPipe, pContext->pVariables, pContext->pRegistration, pContext->pPackages, pContext->phLock, pContext->pfDisabledAutomaticUpdates, (BYTE*)pMsg->pvData, pMsg->cbData); + hrResult = OnApplyInitialize(pContext->hPipe, pContext->pVariables, pContext->pRegistration, pContext->pPackages, pContext->phLock, pContext->pfDisabledAutomaticUpdates, pContext->pfApplying, (BYTE*)pMsg->pvData, pMsg->cbData); break; case BURN_ELEVATION_MESSAGE_TYPE_APPLY_UNINITIALIZE: - hrResult = OnApplyUninitialize(pContext->phLock); + hrResult = OnApplyUninitialize(pContext->phLock, pContext->pfApplying); break; case BURN_ELEVATION_MESSAGE_TYPE_SESSION_BEGIN: @@ -2267,6 +2272,7 @@ static HRESULT OnApplyInitialize( __in BURN_PACKAGES* pPackages, __in HANDLE* phLock, __in BOOL* pfDisabledWindowsUpdate, + __in BOOL* pfApplying, __in BYTE* pbData, __in SIZE_T cbData ) @@ -2296,6 +2302,8 @@ static HRESULT OnApplyInitialize( hr = ApplyLock(TRUE, phLock); ExitOnFailure(hr, "Failed to acquire lock due to setup in other session."); + *pfApplying = TRUE; + // Detect. hr = ElevatedProcessDetect(pRegistration, pVariables, pPackages); ExitOnFailure(hr, "Failed to run detection in elevated process."); @@ -2419,13 +2427,16 @@ LExit: } static HRESULT OnApplyUninitialize( - __in HANDLE* phLock + __in HANDLE* phLock, + __in BOOL* pfApplying ) { Assert(phLock); // TODO: end system restore point. + *pfApplying = FALSE; + if (*phLock) { ::ReleaseMutex(*phLock); diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h index 0dc189c5..d08d2edc 100644 --- a/src/burn/engine/elevation.h +++ b/src/burn/engine/elevation.h @@ -180,7 +180,8 @@ HRESULT ElevationChildPumpMessages( __out HANDLE* phLock, __out BOOL* pfDisabledAutomaticUpdates, __out DWORD* pdwChildExitCode, - __out BOOL* pfRestart + __out BOOL* pfRestart, + __out BOOL* pfApplying ); HRESULT ElevationChildResumeAutomaticUpdates(); diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp index 821626f8..ee848acf 100644 --- a/src/burn/engine/engine.cpp +++ b/src/burn/engine/engine.cpp @@ -653,7 +653,7 @@ static HRESULT RunElevated( SrpInitialize(TRUE); // Pump messages from parent process. - hr = ElevationChildPumpMessages(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe, pEngineState->companionConnection.hCachePipe, &pEngineState->approvedExes, &pEngineState->cache, &pEngineState->containers, &pEngineState->packages, &pEngineState->payloads, &pEngineState->variables, &pEngineState->registration, &pEngineState->userExperience, &hLock, &fDisabledAutomaticUpdates, &pEngineState->userExperience.dwExitCode, &pEngineState->fRestart); + hr = ElevationChildPumpMessages(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe, pEngineState->companionConnection.hCachePipe, &pEngineState->approvedExes, &pEngineState->cache, &pEngineState->containers, &pEngineState->packages, &pEngineState->payloads, &pEngineState->variables, &pEngineState->registration, &pEngineState->userExperience, &hLock, &fDisabledAutomaticUpdates, &pEngineState->userExperience.dwExitCode, &pEngineState->fRestart, &pEngineState->plan.fApplying); LogRedirect(NULL, NULL); // reset logging so the next failure gets written to "log buffer" for the failure log. ExitOnFailure(hr, "Failed to pump messages from parent process."); diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h index 2b04e097..46a9a363 100644 --- a/src/burn/engine/plan.h +++ b/src/burn/engine/plan.h @@ -264,6 +264,7 @@ typedef struct _BURN_PLAN LPWSTR sczLayoutDirectory; BOOL fPlanPackageCacheRollback; BOOL fDowngrade; + BOOL fApplying; DWORD64 qwCacheSizeTotal; diff --git a/src/burn/engine/uithread.cpp b/src/burn/engine/uithread.cpp index bb1e5972..cdd8613c 100644 --- a/src/burn/engine/uithread.cpp +++ b/src/burn/engine/uithread.cpp @@ -17,8 +17,7 @@ struct UITHREAD_CONTEXT struct UITHREAD_INFO { BOOL fElevatedEngine; - BURN_USER_EXPERIENCE* pUserExperience; - BOOL* pfCriticalShutdownInitiated; + BURN_ENGINE_STATE* pEngineState; }; @@ -132,8 +131,7 @@ static DWORD WINAPI ThreadProc( fRegistered = TRUE; info.fElevatedEngine = fElevatedEngine; - info.pUserExperience = &pEngineState->userExperience; - info.pfCriticalShutdownInitiated = &pEngineState->fCriticalShutdownInitiated; + info.pEngineState = pEngineState; // Create the window to handle reboots without activating it. hWnd = ::CreateWindowExW(WS_EX_NOACTIVATE, wc.lpszClassName, NULL, WS_POPUP, 0, 0, 0, 0, HWND_DESKTOP, NULL, pContext->hInstance, &info); @@ -196,21 +194,17 @@ static LRESULT CALLBACK WndProc( { DWORD dwEndSession = static_cast(lParam); BOOL fCritical = ENDSESSION_CRITICAL & dwEndSession; - BOOL fCancel = TRUE; + BOOL fCancel = FALSE; BOOL fRet = FALSE; - // Always block shutdown in the elevated process, but ask the BA in the non-elevated. + // Always block shutdown during apply. UITHREAD_INFO* pInfo = reinterpret_cast(::GetWindowLongPtrW(hWnd, GWLP_USERDATA)); - if (!pInfo->fElevatedEngine) + if (pInfo->pEngineState->plan.fApplying) { - // TODO: instead of recommending canceling all non-critical shutdowns, maybe we should only recommend cancel - // when the engine is doing work? - fCancel = !fCritical; - // TODO: There's a race condition here where the BA may not have been loaded, or already was unloaded. - UserExperienceOnSystemShutdown(pInfo->pUserExperience, dwEndSession, &fCancel); + fCancel = TRUE; } - *pInfo->pfCriticalShutdownInitiated |= fCritical; + pInfo->pEngineState->fCriticalShutdownInitiated |= fCritical; fRet = !fCancel; LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet)); diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp index a3cbcd4a..4325a6ee 100644 --- a/src/burn/engine/userexperience.cpp +++ b/src/burn/engine/userexperience.cpp @@ -2631,31 +2631,6 @@ LExit: return hr; } -EXTERN_C BAAPI UserExperienceOnSystemShutdown( - __in BURN_USER_EXPERIENCE* pUserExperience, - __in DWORD dwEndSession, - __inout BOOL* pfCancel - ) -{ - HRESULT hr = S_OK; - BA_ONSYSTEMSHUTDOWN_ARGS args = { }; - BA_ONSYSTEMSHUTDOWN_RESULTS results = { }; - - args.cbSize = sizeof(args); - args.dwEndSession = dwEndSession; - - results.cbSize = sizeof(results); - results.fCancel = *pfCancel; - - hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN, &args, &results); - ExitOnFailure(hr, "BA OnSystemShutdown failed."); - - *pfCancel = results.fCancel; - -LExit: - return hr; -} - EXTERN_C BAAPI UserExperienceOnUnregisterBegin( __in BURN_USER_EXPERIENCE* pUserExperience, __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h index 2059b521..e7489710 100644 --- a/src/burn/engine/userexperience.h +++ b/src/burn/engine/userexperience.h @@ -583,11 +583,6 @@ BAAPI UserExperienceOnSystemRestorePointComplete( __in BURN_USER_EXPERIENCE* pUserExperience, __in HRESULT hrStatus ); -BAAPI UserExperienceOnSystemShutdown( - __in BURN_USER_EXPERIENCE* pUserExperience, - __in DWORD dwEndSession, - __inout BOOL* pfCancel - ); BAAPI UserExperienceOnUnregisterBegin( __in BURN_USER_EXPERIENCE* pUserExperience, __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType -- cgit v1.2.3-55-g6feb