3 files changed, 20 insertions, 13 deletions

diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index ef9c6e97..e30a05bb 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -189,7 +189,8 @@ static HRESULT ElevatedProcessDetect(
     );
 static HRESULT OnApplyUninitialize(
     __in HANDLE* phLock,
-    __in BOOL* pfApplying
+    __in BOOL* pfApplying,
+    __in BOOL* pfDisabledAutomaticUpdates
     );
 static HRESULT OnSessionBegin(
     __in BURN_CACHE* pCache,
@@ -1543,7 +1544,6 @@ extern "C" HRESULT ElevationChildPumpMessages(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __out HANDLE* phLock,
-    __out BOOL* pfDisabledAutomaticUpdates,
     __out DWORD* pdwChildExitCode,
     __out BOOL* pfRestart,
     __out BOOL* pfApplying
@@ -1554,6 +1554,7 @@ extern "C" HRESULT ElevationChildPumpMessages(
     BURN_ELEVATION_CHILD_MESSAGE_CONTEXT context = { };
     HANDLE hCacheThread = NULL;
     BURN_PIPE_RESULT result = { };
+    BOOL fDisabledAutomaticUpdates = FALSE;
 
     cacheContext.dwLoggingTlsId = dwLoggingTlsId;
     cacheContext.hPipe = hCachePipe;
@@ -1568,7 +1569,7 @@ extern "C" HRESULT ElevationChildPumpMessages(
     context.dwLoggingTlsId = dwLoggingTlsId;
     context.hPipe = hPipe;
     context.phLock = phLock;
-    context.pfDisabledAutomaticUpdates = pfDisabledAutomaticUpdates;
+    context.pfDisabledAutomaticUpdates = &fDisabledAutomaticUpdates;
     context.pfApplying = pfApplying;
     context.pApprovedExes = pApprovedExes;
     context.pCache = pCache;
@@ -1595,6 +1596,11 @@ extern "C" HRESULT ElevationChildPumpMessages(
 LExit:
     ReleaseHandle(hCacheThread);
 
+    if (fDisabledAutomaticUpdates)
+    {
+        ElevationChildResumeAutomaticUpdates();
+    }
+
     return hr;
 }
 
@@ -2128,7 +2134,7 @@ static HRESULT ProcessElevatedChildMessage(
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_UNINITIALIZE:
-        hrResult = OnApplyUninitialize(pContext->phLock, pContext->pfApplying);
+        hrResult = OnApplyUninitialize(pContext->phLock, pContext->pfApplying, pContext->pfDisabledAutomaticUpdates);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_SESSION_BEGIN:
@@ -2433,7 +2439,8 @@ LExit:
 
 static HRESULT OnApplyUninitialize(
     __in HANDLE* phLock,
-    __in BOOL* pfApplying
+    __in BOOL* pfApplying,
+    __in BOOL* pfDisabledAutomaticUpdates
     )
 {
     Assert(phLock);
@@ -2442,6 +2449,13 @@ static HRESULT OnApplyUninitialize(
 
     *pfApplying = FALSE;
 
+    if (*pfDisabledAutomaticUpdates)
+    {
+        *pfDisabledAutomaticUpdates = FALSE;
+
+        ElevationChildResumeAutomaticUpdates();
+    }
+
     if (*phLock)
     {
         ::ReleaseMutex(*phLock);
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index d08d2edc..bda0fdef 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -178,7 +178,6 @@ HRESULT ElevationChildPumpMessages(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __out HANDLE* phLock,
-    __out BOOL* pfDisabledAutomaticUpdates,
     __out DWORD* pdwChildExitCode,
     __out BOOL* pfRestart,
     __out BOOL* pfApplying
diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index 334ea23b..2d2da1d5 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -638,7 +638,6 @@ static HRESULT RunElevated(
 {
     HRESULT hr = S_OK;
     HANDLE hLock = NULL;
-    BOOL fDisabledAutomaticUpdates = FALSE;
 
     // connect to per-user process
     hr = PipeChildConnect(&pEngineState->companionConnection, TRUE);
@@ -666,7 +665,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, &pEngineState->plan.fApplying);
+    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, &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.");
 
@@ -676,11 +675,6 @@ LExit:
     // If the message window is still around, close it.
     UiCloseMessageWindow(pEngineState);
 
-    if (fDisabledAutomaticUpdates)
-    {
-        ElevationChildResumeAutomaticUpdates();
-    }
-
     if (hLock)
     {
         ::ReleaseMutex(hLock);