Clean up synchronization between the engine and the BA.

10 files changed, 80 insertions, 109 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index dd4932aa..76831461 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -366,7 +366,7 @@ extern "C" HRESULT ApplyRegister(
     hr = CoreSaveEngineState(pEngineState);
     if (FAILED(hr))
     {
-        LogErrorId(hr, MSG_STATE_NOT_SAVED, NULL, NULL, NULL);
+        LogErrorId(hr, MSG_STATE_NOT_SAVED);
         hr = S_OK;
     }
 
@@ -1348,8 +1348,6 @@ static HRESULT PromptForSource(
     HRESULT hr = S_OK;
     BOOTSTRAPPER_RESOLVESOURCE_ACTION action = BOOTSTRAPPER_RESOLVESOURCE_ACTION_NONE;
 
-    UserExperienceDeactivateEngine(pUX);
-
     hr = UserExperienceOnResolveSource(pUX, wzPackageOrContainerId, wzPayloadId, wzLocalSource, wzDownloadSource, &action);
     if (FAILED(hr))
     {
@@ -1376,7 +1374,6 @@ static HRESULT PromptForSource(
     }
 
 LExit:
-    UserExperienceActivateEngine(pUX, NULL);
     return hr;
 }
 
diff --git a/src/engine/core.cpp b/src/engine/core.cpp
index ae09ea65..a644d377 100644
--- a/src/engine/core.cpp
+++ b/src/engine/core.cpp
@@ -236,16 +236,12 @@ extern "C" HRESULT CoreDetect(
     )
 {
     HRESULT hr = S_OK;
-    BOOL fActivated = FALSE;
     BOOL fDetectBegan = FALSE;
     BURN_PACKAGE* pPackage = NULL;
     HRESULT hrFirstPackageFailure = S_OK;
 
     LogId(REPORT_STANDARD, MSG_DETECT_BEGIN, pEngineState->packages.cPackages);
 
-    hr = UserExperienceActivateEngine(&pEngineState->userExperience, &fActivated);
-    ExitOnFailure(hr, "Engine cannot start detect because it is busy with another action.");
-
     // Detect if bundle installed state has changed since start up. This
     // only happens if Apply() changed the state of bundle (installed or
     // uninstalled). In that case, Detect() can be used here to reset
@@ -369,11 +365,6 @@ LExit:
         hr = hrFirstPackageFailure;
     }
 
-    if (fActivated)
-    {
-        UserExperienceDeactivateEngine(&pEngineState->userExperience);
-    }
-
     if (fDetectBegan)
     {
         UserExperienceOnDetectComplete(&pEngineState->userExperience, hr);
@@ -392,7 +383,6 @@ extern "C" HRESULT CorePlan(
     )
 {
     HRESULT hr = S_OK;
-    BOOL fActivated = FALSE;
     BOOL fPlanBegan = FALSE;
     LPWSTR sczLayoutDirectory = NULL;
     HANDLE hSyncpointEvent = NULL;
@@ -401,9 +391,6 @@ extern "C" HRESULT CorePlan(
 
     LogId(REPORT_STANDARD, MSG_PLAN_BEGIN, pEngineState->packages.cPackages, LoggingBurnActionToString(action));
 
-    hr = UserExperienceActivateEngine(&pEngineState->userExperience, &fActivated);
-    ExitOnFailure(hr, "Engine cannot start plan because it is busy with another action.");
-
     fPlanBegan = TRUE;
     hr = UserExperienceOnPlanBegin(&pEngineState->userExperience, pEngineState->packages.cPackages);
     ExitOnRootFailure(hr, "BA aborted plan begin.");
@@ -496,11 +483,6 @@ extern "C" HRESULT CorePlan(
     PlanDump(&pEngineState->plan);
 
 LExit:
-    if (fActivated)
-    {
-        UserExperienceDeactivateEngine(&pEngineState->userExperience);
-    }
-
     if (fPlanBegan)
     {
         UserExperienceOnPlanComplete(&pEngineState->userExperience, hr);
@@ -551,10 +533,10 @@ extern "C" HRESULT CoreApply(
     )
 {
     HRESULT hr = S_OK;
-    BOOL fActivated = FALSE;
     HANDLE hLock = NULL;
     DWORD cOverallProgressTicks = 0;
     HANDLE hCacheThread = NULL;
+    BOOL fApplyInitialize = FALSE;
     BOOL fElevated = FALSE;
     BOOL fRegistered = FALSE;
     BOOL fKeepRegistration = pEngineState->plan.fKeepRegistrationDefault;
@@ -567,9 +549,6 @@ extern "C" HRESULT CoreApply(
 
     LogId(REPORT_STANDARD, MSG_APPLY_BEGIN);
 
-    hr = UserExperienceActivateEngine(&pEngineState->userExperience, &fActivated);
-    ExitOnFailure(hr, "Engine cannot start apply because it is busy with another action.");
-
     // Ensure any previous attempts to execute are reset.
     ApplyReset(&pEngineState->userExperience, &pEngineState->packages);
 
@@ -599,6 +578,7 @@ extern "C" HRESULT CoreApply(
     ExitOnFailure(hr, "Another per-user setup is already executing.");
 
     // Initialize only after getting a lock.
+    fApplyInitialize = TRUE;
     ApplyInitialize();
 
     pEngineState->userExperience.hwndApply = hwndParent;
@@ -627,7 +607,7 @@ extern "C" HRESULT CoreApply(
         ExitOnFailure(hr, "Failed to elevate.");
 
         hr = ElevationApplyInitialize(pEngineState->companionConnection.hPipe, &pEngineState->userExperience, &pEngineState->variables, pEngineState->plan.action, pEngineState->automaticUpdates, !pEngineState->fDisableSystemRestore);
-        ExitOnFailure(hr, "Another per-machine setup is already executing.");
+        ExitOnFailure(hr, "Failed to initialize apply in elevated process.");
 
         fElevated = TRUE;
     }
@@ -704,7 +684,10 @@ LExit:
 
     pEngineState->userExperience.hwndApply = NULL;
 
-    ApplyUninitialize();
+    if (fApplyInitialize)
+    {
+        ApplyUninitialize();
+    }
 
     if (hLock)
     {
@@ -712,11 +695,6 @@ LExit:
         ::CloseHandle(hLock);
     }
 
-    if (fActivated)
-    {
-        UserExperienceDeactivateEngine(&pEngineState->userExperience);
-    }
-
     ReleaseHandle(hCacheThread);
 
     UserExperienceOnApplyComplete(&pEngineState->userExperience, hr, restart, &applyCompleteAction);
@@ -736,14 +714,10 @@ extern "C" HRESULT CoreLaunchApprovedExe(
     )
 {
     HRESULT hr = S_OK;
-    BOOL fActivated = FALSE;
     DWORD dwProcessId = 0;
 
     LogId(REPORT_STANDARD, MSG_LAUNCH_APPROVED_EXE_BEGIN, pLaunchApprovedExe->sczId);
 
-    hr = UserExperienceActivateEngine(&pEngineState->userExperience, &fActivated);
-    ExitOnFailure(hr, "Engine cannot start LaunchApprovedExe because it is busy with another action.");
-
     hr = UserExperienceOnLaunchApprovedExeBegin(&pEngineState->userExperience);
     ExitOnRootFailure(hr, "BA aborted LaunchApprovedExe begin.");
 
@@ -755,11 +729,6 @@ extern "C" HRESULT CoreLaunchApprovedExe(
     hr = ElevationLaunchApprovedExe(pEngineState->companionConnection.hPipe, pLaunchApprovedExe, &dwProcessId);
 
 LExit:
-    if (fActivated)
-    {
-        UserExperienceDeactivateEngine(&pEngineState->userExperience);
-    }
-
     UserExperienceOnLaunchApprovedExeComplete(&pEngineState->userExperience, hr, dwProcessId);
 
     LogId(REPORT_STANDARD, MSG_LAUNCH_APPROVED_EXE_COMPLETE, hr, dwProcessId);
@@ -782,7 +751,7 @@ extern "C" HRESULT CoreQuit(
         hr = CoreSaveEngineState(pEngineState);
         if (FAILED(hr))
         {
-            LogErrorId(hr, MSG_STATE_NOT_SAVED, NULL, NULL, NULL);
+            LogErrorId(hr, MSG_STATE_NOT_SAVED);
             hr = S_OK;
         }
     }
diff --git a/src/engine/core.h b/src/engine/core.h
index 0161eaa0..fae4bfe5 100644
--- a/src/engine/core.h
+++ b/src/engine/core.h
@@ -77,11 +77,6 @@ enum BURN_AU_PAUSE_ACTION
 
 typedef struct _BURN_ENGINE_STATE
 {
-    // synchronization
-    CRITICAL_SECTION csActive; // Any call from the UX that reads or alters the engine state
-                               // needs to be syncronized through this critical section.
-                               // Note: The engine must never do a UX callback while in this critical section.
-
     // UX flow control
     //BOOL fSuspend;             // Is TRUE when UX made Suspend() call on core.
     //BOOL fForcedReboot;        // Is TRUE when UX made Reboot() call on core.
diff --git a/src/engine/detect.cpp b/src/engine/detect.cpp
index 3b8d63e2..ff6ba4d7 100644
--- a/src/engine/detect.cpp
+++ b/src/engine/detect.cpp
@@ -9,7 +9,7 @@ typedef struct _DETECT_AUTHENTICATION_REQUIRED_DATA
 } DETECT_AUTHENTICATION_REQUIRED_DATA;
 
 // internal function definitions
-static HRESULT AuthenticationRequired(
+static HRESULT WINAPI AuthenticationRequired(
     __in LPVOID pData,
     __in HINTERNET hUrl,
     __in long lHttpCode,
@@ -288,7 +288,7 @@ LExit:
     return hr;
 }
 
-static HRESULT AuthenticationRequired(
+static HRESULT WINAPI AuthenticationRequired(
     __in LPVOID pData,
     __in HINTERNET hUrl,
     __in long lHttpCode,
diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp
index 9eee7fde..2c6bad03 100644
--- a/src/engine/engine.cpp
+++ b/src/engine/engine.cpp
@@ -330,7 +330,6 @@ static HRESULT InitializeEngineState(
 
     pEngineState->automaticUpdates = BURN_AU_PAUSE_ACTION_IFELEVATED;
     pEngineState->dwElevatedLoggingTlsId = TLS_OUT_OF_INDEXES;
-    ::InitializeCriticalSection(&pEngineState->csActive);
     ::InitializeCriticalSection(&pEngineState->userExperience.csEngineActive);
     PipeConnectionInitialize(&pEngineState->companionConnection);
     PipeConnectionInitialize(&pEngineState->embeddedConnection);
@@ -418,8 +417,6 @@ static void UninitializeEngineState(
         ::TlsFree(pEngineState->dwElevatedLoggingTlsId);
     }
 
-    ::DeleteCriticalSection(&pEngineState->csActive);
-
     // clear struct
     memset(pEngineState, 0, sizeof(BURN_ENGINE_STATE));
 }
@@ -805,6 +802,8 @@ static HRESULT ProcessMessage(
 {
     HRESULT hr = S_OK;
 
+    UserExperienceActivateEngine(&pEngineState->userExperience);
+
     switch (pmsg->message)
     {
     case WM_BURN_DETECT:
@@ -832,6 +831,8 @@ static HRESULT ProcessMessage(
         break;
     }
 
+    UserExperienceDeactivateEngine(&pEngineState->userExperience);
+
     return hr;
 }
 
diff --git a/src/engine/engine.mc b/src/engine/engine.mc
index d6e9751a..b120c5bb 100644
--- a/src/engine/engine.mc
+++ b/src/engine/engine.mc
@@ -938,7 +938,7 @@ MessageId=501
 Severity=Warning
 SymbolicName=MSG_STATE_NOT_SAVED
 Language=English
-The state file could not be saved, error: 0x%1!x!. Continuing...
+The state file could not be saved. Continuing...
 .
 
 MessageId=600
diff --git a/src/engine/externalengine.cpp b/src/engine/externalengine.cpp
index 39878c8b..f9a06437 100644
--- a/src/engine/externalengine.cpp
+++ b/src/engine/externalengine.cpp
@@ -276,7 +276,9 @@ HRESULT ExternalEngineSetUpdate(
     WCHAR wzGuid[39];
     RPC_STATUS rs = RPC_S_OK;
 
-    ::EnterCriticalSection(&pEngineState->csActive);
+    ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
+    hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
+    ExitOnFailure(hr, "Engine is active, cannot change engine state.");
 
     if ((!wzLocalSource || !*wzLocalSource) && (!wzDownloadSource || !*wzDownloadSource))
     {
@@ -332,7 +334,7 @@ HRESULT ExternalEngineSetUpdate(
     }
 
 LExit:
-    ::LeaveCriticalSection(&pEngineState->csActive);
+    ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
 
     ReleaseStr(sczCommandline);
     ReleaseStr(sczLocalSource);
@@ -351,7 +353,7 @@ HRESULT ExternalEngineSetLocalSource(
     BURN_CONTAINER* pContainer = NULL;
     BURN_PAYLOAD* pPayload = NULL;
 
-    ::EnterCriticalSection(&pEngineState->csActive);
+    ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
     hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
     ExitOnFailure(hr, "Engine is active, cannot change engine state.");
 
@@ -387,7 +389,7 @@ HRESULT ExternalEngineSetLocalSource(
     }
 
 LExit:
-    ::LeaveCriticalSection(&pEngineState->csActive);
+    ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
 
     return hr;
 }
@@ -406,7 +408,7 @@ HRESULT ExternalEngineSetDownloadSource(
     BURN_PAYLOAD* pPayload = NULL;
     DOWNLOAD_SOURCE* pDownloadSource = NULL;
 
-    ::EnterCriticalSection(&pEngineState->csActive);
+    ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
     hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
     ExitOnFailure(hr, "Engine is active, cannot change engine state.");
 
@@ -470,7 +472,7 @@ HRESULT ExternalEngineSetDownloadSource(
     }
 
 LExit:
-    ::LeaveCriticalSection(&pEngineState->csActive);
+    ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
 
     return hr;
 }
@@ -686,7 +688,7 @@ HRESULT ExternalEngineLaunchApprovedExe(
     pLaunchApprovedExe = (BURN_LAUNCH_APPROVED_EXE*)MemAlloc(sizeof(BURN_LAUNCH_APPROVED_EXE), TRUE);
     ExitOnNull(pLaunchApprovedExe, hr, E_OUTOFMEMORY, "Failed to alloc BURN_LAUNCH_APPROVED_EXE");
 
-    ::EnterCriticalSection(&pEngineState->csActive);
+    ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
     fLeaveCriticalSection = TRUE;
     hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
     ExitOnFailure(hr, "Engine is active, cannot change engine state.");
@@ -699,9 +701,6 @@ HRESULT ExternalEngineLaunchApprovedExe(
     hr = ApprovedExesFindById(&pEngineState->approvedExes, wzApprovedExeForElevationId, &pApprovedExe);
     ExitOnFailure(hr, "BA requested unknown approved exe with id: %ls", wzApprovedExeForElevationId);
 
-    ::LeaveCriticalSection(&pEngineState->csActive);
-    fLeaveCriticalSection = FALSE;
-
     hr = StrAllocString(&pLaunchApprovedExe->sczId, wzApprovedExeForElevationId, NULL);
     ExitOnFailure(hr, "Failed to copy the id.");
 
@@ -723,7 +722,7 @@ HRESULT ExternalEngineLaunchApprovedExe(
 LExit:
     if (fLeaveCriticalSection)
     {
-        ::LeaveCriticalSection(&pEngineState->csActive);
+        ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
     }
 
     if (FAILED(hr))
diff --git a/src/engine/userexperience.cpp b/src/engine/userexperience.cpp
index ea8b8a19..3a36cab6 100644
--- a/src/engine/userexperience.cpp
+++ b/src/engine/userexperience.cpp
@@ -24,6 +24,13 @@ static HRESULT SendBAMessage(
     __inout LPVOID pvResults
     );
 
+static HRESULT SendBAMessageFromInactiveEngine(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in BOOTSTRAPPER_APPLICATION_MESSAGE message,
+    __in const LPVOID pvArgs,
+    __inout LPVOID pvResults
+    );
+
 
 // function definitions
 
@@ -230,51 +237,31 @@ extern "C" int UserExperienceSendError(
     return nResult;
 }
 
-extern "C" HRESULT UserExperienceActivateEngine(
-    __in BURN_USER_EXPERIENCE* pUserExperience,
-    __out_opt BOOL* pfActivated
+extern "C" void UserExperienceActivateEngine(
+    __in BURN_USER_EXPERIENCE* pUserExperience
     )
 {
-    HRESULT hr = S_OK;
-    BOOL fActivated;
-
     ::EnterCriticalSection(&pUserExperience->csEngineActive);
-    if (InterlockedCompareExchange(reinterpret_cast<LONG*>(&pUserExperience->fEngineActive), TRUE, FALSE))
-    {
-        AssertSz(FALSE, "Engine should have been deactivated before activating it.");
-
-        fActivated = FALSE;
-        hr = HRESULT_FROM_WIN32(ERROR_INVALID_STATE);
-    }
-    else
-    {
-        fActivated = TRUE;
-    }
+    AssertSz(!pUserExperience->fEngineActive, "Engine should have been deactivated before activating it.");
+    pUserExperience->fEngineActive = TRUE;
     ::LeaveCriticalSection(&pUserExperience->csEngineActive);
-
-    if (pfActivated)
-    {
-        *pfActivated = fActivated;
-    }
-    ExitOnRootFailure(hr, "Engine active cannot be changed because it was already in that state.");
-
-LExit:
-    return hr;
 }
 
 extern "C" void UserExperienceDeactivateEngine(
     __in BURN_USER_EXPERIENCE* pUserExperience
     )
 {
-    BOOL fActive = InterlockedExchange(reinterpret_cast<LONG*>(&pUserExperience->fEngineActive), FALSE);
-    fActive = fActive; // prevents warning in "ship" build.
-    AssertSz(fActive, "Engine should have be active before deactivating it.");
+    ::EnterCriticalSection(&pUserExperience->csEngineActive);
+    AssertSz(pUserExperience->fEngineActive, "Engine should have been active before deactivating it.");
+    pUserExperience->fEngineActive = FALSE;
+    ::LeaveCriticalSection(&pUserExperience->csEngineActive);
 }
 
 extern "C" HRESULT UserExperienceEnsureEngineInactive(
     __in BURN_USER_EXPERIENCE* pUserExperience
     )
 {
+    // Make a slight optimization here by ignoring the critical section, because all callers should have needed to enter it for their operation anyway.
     HRESULT hr = pUserExperience->fEngineActive ? HRESULT_FROM_WIN32(ERROR_BUSY) : S_OK;
     ExitOnRootFailure(hr, "Engine is active, cannot proceed.");
 
@@ -346,7 +333,7 @@ EXTERN_C BAAPI UserExperienceOnApplyComplete(
     results.cbSize = sizeof(results);
     results.action = *pAction;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONAPPLYCOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONAPPLYCOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnApplyComplete failed.");
 
     *pAction = results.action;
@@ -761,7 +748,7 @@ EXTERN_C BAAPI UserExperienceOnDetectComplete(
 
     results.cbSize = sizeof(results);
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTCOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTCOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnDetectComplete failed.");
 
 LExit:
@@ -990,7 +977,7 @@ LExit:
 
 EXTERN_C BAAPI UserExperienceOnDetectUpdate(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in_z LPCWSTR wzUpdateLocation,
+    __in_z_opt LPCWSTR wzUpdateLocation,
     __in DWORD64 dw64Size,
     __in VERUTIL_VERSION* pVersion,
     __in_z_opt LPCWSTR wzTitle,
@@ -1016,7 +1003,7 @@ EXTERN_C BAAPI UserExperienceOnDetectUpdate(
     results.cbSize = sizeof(results);
     results.fStopProcessingUpdates = *pfStopProcessingUpdates;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATE, &args, &results);
     ExitOnFailure(hr, "BA OnDetectUpdate failed.");
 
     if (results.fCancel)
@@ -1045,7 +1032,7 @@ EXTERN_C BAAPI UserExperienceOnDetectUpdateBegin(
     results.cbSize = sizeof(results);
     results.fSkip = *pfSkip;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATEBEGIN, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATEBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnDetectUpdateBegin failed.");
 
     if (results.fCancel)
@@ -1074,7 +1061,7 @@ EXTERN_C BAAPI UserExperienceOnDetectUpdateComplete(
     results.cbSize = sizeof(results);
     results.fIgnoreError = *pfIgnoreError;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATECOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATECOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnDetectUpdateComplete failed.");
 
     if (FAILED(hrStatus))
@@ -1124,7 +1111,7 @@ EXTERN_C BAAPI UserExperienceOnElevateComplete(
 
     results.cbSize = sizeof(results);
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONELEVATECOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONELEVATECOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnElevateComplete failed.");
 
 LExit:
@@ -1452,7 +1439,7 @@ EXTERN_C BAAPI UserExperienceOnLaunchApprovedExeComplete(
 
     results.cbSize = sizeof(results);
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONLAUNCHAPPROVEDEXECOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONLAUNCHAPPROVEDEXECOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnLaunchApprovedExeComplete failed.");
 
 LExit:
@@ -1571,7 +1558,7 @@ EXTERN_C BAAPI UserExperienceOnPlanComplete(
 
     results.cbSize = sizeof(results);
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPLETE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnPlanComplete failed.");
 
 LExit:
@@ -1831,7 +1818,7 @@ EXTERN_C BAAPI UserExperienceOnResolveSource(
     results.cbSize = sizeof(results);
     results.action = *pAction;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONRESOLVESOURCE, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONRESOLVESOURCE, &args, &results);
     ExitOnFailure(hr, "BA OnResolveSource failed.");
 
     if (results.fCancel)
@@ -2317,3 +2304,21 @@ static HRESULT SendBAMessage(
 
     return hr;
 }
+
+static HRESULT SendBAMessageFromInactiveEngine(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in BOOTSTRAPPER_APPLICATION_MESSAGE message,
+    __in const LPVOID pvArgs,
+    __inout LPVOID pvResults
+    )
+{
+    HRESULT hr = S_OK;
+
+    UserExperienceDeactivateEngine(pUserExperience);
+
+    hr = SendBAMessage(pUserExperience, message, pvArgs, pvResults);
+
+    UserExperienceActivateEngine(pUserExperience);
+
+    return hr;
+}
diff --git a/src/engine/userexperience.h b/src/engine/userexperience.h
index 3176e762..f51c09ff 100644
--- a/src/engine/userexperience.h
+++ b/src/engine/userexperience.h
@@ -82,13 +82,17 @@ int UserExperienceSendError(
     __in DWORD uiFlags,
     __in int nRecommendation
     );
-HRESULT UserExperienceActivateEngine(
-    __in BURN_USER_EXPERIENCE* pUserExperience,
-    __out_opt BOOL* pfActivated
+void UserExperienceActivateEngine(
+    __in BURN_USER_EXPERIENCE* pUserExperience
     );
 void UserExperienceDeactivateEngine(
     __in BURN_USER_EXPERIENCE* pUserExperience
     );
+/********************************************************************
+ UserExperienceEnsureEngineInactive - Verifies the engine is inactive.
+   The caller MUST enter the csActive critical section before calling.
+
+*********************************************************************/
 HRESULT UserExperienceEnsureEngineInactive(
     __in BURN_USER_EXPERIENCE* pUserExperience
     );
@@ -240,7 +244,7 @@ BAAPI UserExperienceOnDetectTargetMsiPackage(
     );
 BAAPI UserExperienceOnDetectUpdate(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in_z LPCWSTR wzUpdateLocation,
+    __in_z_opt LPCWSTR wzUpdateLocation,
     __in DWORD64 dw64Size,
     __in VERUTIL_VERSION* pVersion,
     __in_z_opt LPCWSTR wzTitle,
diff --git a/src/test/BurnUnitTest/PlanTest.cpp b/src/test/BurnUnitTest/PlanTest.cpp
index fb4ca246..10b12e7b 100644
--- a/src/test/BurnUnitTest/PlanTest.cpp
+++ b/src/test/BurnUnitTest/PlanTest.cpp
@@ -529,7 +529,6 @@ namespace Bootstrapper
             HRESULT hr = S_OK;
             LPWSTR sczFilePath = NULL;
 
-            ::InitializeCriticalSection(&pEngineState->csActive);
             ::InitializeCriticalSection(&pEngineState->userExperience.csEngineActive);
 
             hr = VariableInitialize(&pEngineState->variables);
@@ -567,6 +566,8 @@ namespace Bootstrapper
 
             hr = DepDependencyArrayAlloc(&pRegistration->rgIgnoredDependencies, &pRegistration->cIgnoredDependencies, pRegistration->sczProviderKey, NULL);
             NativeAssert::Succeeded(hr, "Failed to add the bundle provider key to the list of dependencies to ignore.");
+
+            pEngineState->userExperience.fEngineActive = TRUE;
         }
 
         void DetectAttachedContainerAsAttached(BURN_ENGINE_STATE* pEngineState)