10 files changed, 181 insertions, 93 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 3ad22e9b..73b5b396 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -3326,6 +3326,14 @@ static int GenericExecuteMessageHandler(
         }
         break;
 
+    case GENERIC_EXECUTE_MESSAGE_PROCESS_CANCEL:
+        {
+            BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION action = BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION_ABANDON;
+            UserExperienceOnExecuteProcessCancel(pContext->pUX, pContext->wzExecutingPackageId, pMessage->processCancel.dwProcessId, &action); // ignore return value.
+            nResult = BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION_WAIT == action ? IDRETRY : IDIGNORE;
+        }
+        break;
+
     case GENERIC_EXECUTE_MESSAGE_ERROR:
         UserExperienceOnError(pContext->pUX, BOOTSTRAPPER_ERROR_TYPE_EXE_PACKAGE, pContext->wzExecutingPackageId, pMessage->error.dwErrorCode, pMessage->error.wzMessage, pMessage->dwUIHint, 0, NULL, &nResult); // ignore return value.
         break;
diff --git a/src/burn/engine/apply.h b/src/burn/engine/apply.h
index 1717a71a..47f0ece6 100644
--- a/src/burn/engine/apply.h
+++ b/src/burn/engine/apply.h
@@ -13,6 +13,7 @@ enum GENERIC_EXECUTE_MESSAGE_TYPE
     GENERIC_EXECUTE_MESSAGE_ERROR,
     GENERIC_EXECUTE_MESSAGE_PROGRESS,
     GENERIC_EXECUTE_MESSAGE_NETFX_FILES_IN_USE,
+    GENERIC_EXECUTE_MESSAGE_PROCESS_CANCEL,
 };
 
 typedef struct _APPLY_AUTHENTICATION_REQUIRED_DATA
@@ -43,6 +44,10 @@ typedef struct _GENERIC_EXECUTE_MESSAGE
             DWORD cFiles;
             LPCWSTR* rgwzFiles;
         } filesInUse;
+        struct
+        {
+            DWORD dwProcessId;
+        } processCancel;
     };
 } GENERIC_EXECUTE_MESSAGE;
 
diff --git a/src/burn/engine/bundlepackageengine.cpp b/src/burn/engine/bundlepackageengine.cpp
index 88a00f5e..0bee054f 100644
--- a/src/burn/engine/bundlepackageengine.cpp
+++ b/src/burn/engine/bundlepackageengine.cpp
@@ -251,7 +251,6 @@ extern "C" HRESULT BundlePackageEngineExecuteRelatedBundle(
     )
 {
     HRESULT hr = S_OK;
-    int nResult = IDNOACTION;
     LPCWSTR wzArguments = NULL;
     LPWSTR sczArguments = NULL;
     LPWSTR sczArgumentsFormatted = NULL;
@@ -420,31 +419,10 @@ extern "C" HRESULT BundlePackageEngineExecuteRelatedBundle(
         hr = EmbeddedRunBundle(sczExecutablePath, sczCommand, pfnGenericMessageHandler, pvContext, &dwExitCode);
         ExitOnFailure(hr, "Failed to run bundle as embedded from path: %ls", sczExecutablePath);
     }
-    else // create and wait for the executable process while sending fake progress to allow cancel.
+    else
     {
-        // Make the cache location of the executable the current directory to help those executables
-        // that expect stuff to be relative to them.
-        si.cb = sizeof(si);
-        if (!::CreateProcessW(sczExecutablePath, sczCommand, NULL, NULL, TRUE, CREATE_NO_WINDOW, NULL, sczCachedDirectory, &si, &pi))
-        {
-            ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", sczExecutablePath);
-        }
-
-        do
-        {
-            message.type = GENERIC_EXECUTE_MESSAGE_PROGRESS;
-            message.dwUIHint = MB_OKCANCEL;
-            message.progress.dwPercentage = 50;
-            nResult = pfnGenericMessageHandler(&message, pvContext);
-            hr = (IDOK == nResult || IDNOACTION == nResult) ? S_OK : IDCANCEL == nResult ? HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) : HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE);
-            ExitOnRootFailure(hr, "Bootstrapper application aborted during BUNDLE progress.");
-
-            hr = ProcWaitForCompletion(pi.hProcess, 500, &dwExitCode);
-            if (HRESULT_FROM_WIN32(WAIT_TIMEOUT) != hr)
-            {
-                ExitOnFailure(hr, "Failed to wait for executable to complete: %ls", sczExecutablePath);
-            }
-        } while (HRESULT_FROM_WIN32(WAIT_TIMEOUT) == hr);
+        hr = ExeEngineRunProcess(pfnGenericMessageHandler, pvContext, pPackage, sczExecutablePath, sczCommand, sczCachedDirectory, &dwExitCode);
+        ExitOnFailure(hr, "Failed to run BUNDLE process");
     }
 
     hr = ExeEngineHandleExitCode(pPackage->Bundle.rgExitCodes, pPackage->Bundle.cExitCodes, dwExitCode, pRestart);
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 636d67ce..3c2872f1 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -43,6 +43,7 @@ typedef enum _BURN_ELEVATION_MESSAGE_TYPE
     BURN_ELEVATION_MESSAGE_TYPE_BURN_CACHE_COMPLETE,
     BURN_ELEVATION_MESSAGE_TYPE_BURN_CACHE_SUCCESS,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROGRESS,
+    BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROCESS_CANCEL,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_ERROR,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSI_MESSAGE,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSI_FILES_IN_USE,
@@ -1812,7 +1813,14 @@ static HRESULT ProcessGenericExecuteMessages(
 
         // read message parameters
         hr = BuffReadNumber((BYTE*)pMsg->pvData, pMsg->cbData, &iData, &message.progress.dwPercentage);
-        ExitOnFailure(hr, "Failed to progress.");
+        ExitOnFailure(hr, "Failed to read progress.");
+        break;
+
+    case BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROCESS_CANCEL:
+        message.type = GENERIC_EXECUTE_MESSAGE_PROCESS_CANCEL;
+
+        hr = BuffReadNumber((BYTE*)pMsg->pvData, pMsg->cbData, &iData, &message.processCancel.dwProcessId);
+        ExitOnFailure(hr, "Failed to read processId.");
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_ERROR:
@@ -3450,6 +3458,13 @@ static int GenericExecuteMessageHandler(
         dwMessage = BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROGRESS;
         break;
 
+    case GENERIC_EXECUTE_MESSAGE_PROCESS_CANCEL:
+        hr = BuffWriteNumber(&pbData, &cbData, pMessage->processCancel.dwProcessId);
+        ExitOnFailure(hr, "Failed to write progress percentage to message buffer.");
+
+        dwMessage = BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROCESS_CANCEL;
+        break;
+
     case GENERIC_EXECUTE_MESSAGE_ERROR:
         // serialize message data
         hr = BuffWriteNumber(&pbData, &cbData, pMessage->error.dwErrorCode);
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index 53e6b256..9e139661 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -933,6 +933,13 @@ Language=English
 Could not create system restore point, error: 0x%1!x!. Continuing...
 .
 
+MessageId=364
+Severity=Success
+SymbolicName=MSG_EXECUTE_PROCESS_DELAYED_CANCEL_REQUESTED
+Language=English
+Bootstrapper application requested delayed cancel during package process progress, id: %1!ls!. Waiting...
+.
+
 MessageId=370
 Severity=Success
 SymbolicName=MSG_SESSION_BEGIN
diff --git a/src/burn/engine/exeengine.cpp b/src/burn/engine/exeengine.cpp
index c984f5a7..4c3c6fb0 100644
--- a/src/burn/engine/exeengine.cpp
+++ b/src/burn/engine/exeengine.cpp
@@ -317,7 +317,6 @@ extern "C" HRESULT ExeEngineExecutePackage(
     )
 {
     HRESULT hr = S_OK;
-    int nResult = IDNOACTION;
     LPCWSTR wzArguments = NULL;
     LPWSTR sczArguments = NULL;
     LPWSTR sczArgumentsFormatted = NULL;
@@ -327,10 +326,7 @@ extern "C" HRESULT ExeEngineExecutePackage(
     LPWSTR sczCommand = NULL;
     LPWSTR sczCommandObfuscated = NULL;
     HANDLE hExecutableFile = INVALID_HANDLE_VALUE;
-    STARTUPINFOW si = { };
-    PROCESS_INFORMATION pi = { };
     DWORD dwExitCode = 0;
-    GENERIC_EXECUTE_MESSAGE message = { };
     BURN_PACKAGE* pPackage = pExecuteAction->exePackage.pPackage;
     BURN_PAYLOAD* pPackagePayload = pPackage->payloads.rgItems[0].pPayload;
 
@@ -442,37 +438,10 @@ extern "C" HRESULT ExeEngineExecutePackage(
         hr = NetFxRunChainer(sczExecutablePath, sczCommand, pfnGenericMessageHandler, pvContext, &dwExitCode);
         ExitOnFailure(hr, "Failed to run netfx chainer: %ls", sczExecutablePath);
     }
-    else // create and wait for the executable process while sending fake progress to allow cancel.
+    else
     {
-        // Make the cache location of the executable the current directory to help those executables
-        // that expect stuff to be relative to them.
-        si.cb = sizeof(si);
-        if (!::CreateProcessW(sczExecutablePath, sczCommand, NULL, NULL, TRUE, CREATE_NO_WINDOW, NULL, sczCachedDirectory, &si, &pi))
-        {
-            ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", sczExecutablePath);
-        }
-
-        if (pPackage->Exe.fFireAndForget)
-        {
-            ::WaitForInputIdle(pi.hProcess, 5000);
-            ExitFunction();
-        }
-
-        do
-        {
-            message.type = GENERIC_EXECUTE_MESSAGE_PROGRESS;
-            message.dwUIHint = MB_OKCANCEL;
-            message.progress.dwPercentage = 50;
-            nResult = pfnGenericMessageHandler(&message, pvContext);
-            hr = (IDOK == nResult || IDNOACTION == nResult) ? S_OK : IDCANCEL == nResult ? HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) : HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE);
-            ExitOnRootFailure(hr, "Bootstrapper application aborted during EXE progress.");
-
-            hr = ProcWaitForCompletion(pi.hProcess, 500, &dwExitCode);
-            if (HRESULT_FROM_WIN32(WAIT_TIMEOUT) != hr)
-            {
-                ExitOnFailure(hr, "Failed to wait for executable to complete: %ls", sczExecutablePath);
-            }
-        } while (HRESULT_FROM_WIN32(WAIT_TIMEOUT) == hr);
+        hr = ExeEngineRunProcess(pfnGenericMessageHandler, pvContext, pPackage, sczExecutablePath, sczCommand, sczCachedDirectory, &dwExitCode);
+        ExitOnFailure(hr, "Failed to run EXE process");
     }
 
     hr = ExeEngineHandleExitCode(pPackage->Exe.rgExitCodes, pPackage->Exe.cExitCodes, dwExitCode, pRestart);
@@ -487,8 +456,6 @@ LExit:
     StrSecureZeroFreeString(sczCommand);
     ReleaseStr(sczCommandObfuscated);
 
-    ReleaseHandle(pi.hThread);
-    ReleaseHandle(pi.hProcess);
     ReleaseFileHandle(hExecutableFile);
 
     // Best effort to clear the execute package cache folder and action variables.
@@ -498,6 +465,99 @@ LExit:
     return hr;
 }
 
+extern "C" HRESULT ExeEngineRunProcess(
+    __in PFN_GENERICMESSAGEHANDLER pfnGenericMessageHandler,
+    __in LPVOID pvContext,
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzExecutablePath,
+    __in_z LPWSTR wzCommand,
+    __in_z_opt LPCWSTR wzCachedDirectory,
+    __inout DWORD* pdwExitCode
+    )
+{
+    HRESULT hr = S_OK;
+    STARTUPINFOW si = { };
+    PROCESS_INFORMATION pi = { };
+    GENERIC_EXECUTE_MESSAGE message = { };
+    int nResult = IDNOACTION;
+    DWORD dwProcessId = 0;
+    BOOL fDelayedCancel = FALSE;
+    BOOL fFireAndForget = BURN_PACKAGE_TYPE_EXE == pPackage->type && pPackage->Exe.fFireAndForget;
+    BOOL fInheritHandles = BURN_PACKAGE_TYPE_BUNDLE == pPackage->type;
+
+    // Make the cache location of the executable the current directory to help those executables
+    // that expect stuff to be relative to them.
+    si.cb = sizeof(si);
+    if (!::CreateProcessW(wzExecutablePath, wzCommand, NULL, NULL, fInheritHandles, CREATE_NO_WINDOW, NULL, wzCachedDirectory, &si, &pi))
+    {
+        ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", wzExecutablePath);
+    }
+
+    if (fFireAndForget)
+    {
+        ::WaitForInputIdle(pi.hProcess, 5000);
+        ExitFunction();
+    }
+
+    dwProcessId = ::GetProcessId(pi.hProcess);
+
+    // Wait for the executable process while sending fake progress to allow cancel.
+    do
+    {
+        message.type = GENERIC_EXECUTE_MESSAGE_PROGRESS;
+        message.dwUIHint = MB_OKCANCEL;
+        message.progress.dwPercentage = 50;
+        nResult = pfnGenericMessageHandler(&message, pvContext);
+
+        if (IDCANCEL == nResult)
+        {
+            memset(&message, 0, sizeof(message));
+            message.type = GENERIC_EXECUTE_MESSAGE_PROCESS_CANCEL;
+            message.dwUIHint = MB_ABORTRETRYIGNORE;
+            message.processCancel.dwProcessId = dwProcessId;
+            nResult = pfnGenericMessageHandler(&message, pvContext);
+
+            if (IDIGNORE == nResult) // abandon
+            {
+                nResult = IDCANCEL;
+                fDelayedCancel = FALSE;
+            }
+            //else if (IDABORT == nResult) // kill
+            else // wait
+            {
+                if (!fDelayedCancel)
+                {
+                    fDelayedCancel = TRUE;
+
+                    LogId(REPORT_STANDARD, MSG_EXECUTE_PROCESS_DELAYED_CANCEL_REQUESTED, pPackage->sczId);
+                }
+
+                nResult = IDNOACTION;
+            }
+        }
+
+        hr = (IDOK == nResult || IDNOACTION == nResult) ? S_OK : IDCANCEL == nResult ? HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) : HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE);
+        ExitOnRootFailure(hr, "Bootstrapper application aborted during package process progress.");
+
+        hr = ProcWaitForCompletion(pi.hProcess, 500, pdwExitCode);
+        if (HRESULT_FROM_WIN32(WAIT_TIMEOUT) != hr)
+        {
+            ExitOnFailure(hr, "Failed to wait for executable to complete: %ls", wzExecutablePath);
+        }
+    } while (HRESULT_FROM_WIN32(WAIT_TIMEOUT) == hr);
+
+    if (fDelayedCancel)
+    {
+        ExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT), "Bootstrapper application cancelled during package process progress, exit code: 0x%x", *pdwExitCode);
+    }
+
+LExit:
+    ReleaseHandle(pi.hThread);
+    ReleaseHandle(pi.hProcess);
+
+    return hr;
+}
+
 extern "C" void ExeEngineUpdateInstallRegistrationState(
     __in BURN_EXECUTE_ACTION* pAction,
     __in HRESULT hrExecute
diff --git a/src/burn/engine/exeengine.h b/src/burn/engine/exeengine.h
index 743621b7..636988f1 100644
--- a/src/burn/engine/exeengine.h
+++ b/src/burn/engine/exeengine.h
@@ -42,6 +42,15 @@ HRESULT ExeEngineExecutePackage(
     __in LPVOID pvContext,
     __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
+HRESULT ExeEngineRunProcess(
+    __in PFN_GENERICMESSAGEHANDLER pfnGenericMessageHandler,
+    __in LPVOID pvContext,
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzExecutablePath,
+    __in_z LPWSTR wzCommand,
+    __in_z_opt LPCWSTR wzCachedDirectory,
+    __inout DWORD* pdwExitCode
+    );
 void ExeEngineUpdateInstallRegistrationState(
     __in BURN_EXECUTE_ACTION* pAction,
     __in HRESULT hrExecute
diff --git a/src/burn/engine/msuengine.cpp b/src/burn/engine/msuengine.cpp
index 091bbe62..2f1fb61c 100644
--- a/src/burn/engine/msuengine.cpp
+++ b/src/burn/engine/msuengine.cpp
@@ -264,7 +264,6 @@ extern "C" HRESULT MsuEngineExecutePackage(
     )
 {
     HRESULT hr = S_OK;
-    int nResult = IDNOACTION;
     LPWSTR sczCachedDirectory = NULL;
     LPWSTR sczMsuPath = NULL;
     LPWSTR sczWindowsPath = NULL;
@@ -350,35 +349,8 @@ extern "C" HRESULT MsuEngineExecutePackage(
     hr = EnsureWUServiceEnabled(fStopWusaService, &schWu, &fWuWasDisabled);
     ExitOnFailure(hr, "Failed to ensure WU service was enabled to install MSU package.");
 
-    // create process
-    si.cb = sizeof(si);
-    if (!::CreateProcessW(sczWusaPath, sczCommand, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi))
-    {
-        ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", sczWusaPath);
-    }
-
-    do
-    {
-        message.type = GENERIC_EXECUTE_MESSAGE_PROGRESS;
-        message.dwUIHint = MB_OKCANCEL;
-        message.progress.dwPercentage = 50;
-        nResult = pfnGenericMessageHandler(&message, pvContext);
-        hr = (IDOK == nResult || IDNOACTION == nResult) ? S_OK : IDCANCEL == nResult ? HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) : HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE);
-        ExitOnRootFailure(hr, "Bootstrapper application aborted during MSU progress.");
-
-        // wait for process to terminate
-        hr = ProcWaitForCompletion(pi.hProcess, 500, &dwExitCode);
-        if (HRESULT_FROM_WIN32(WAIT_TIMEOUT) != hr)
-        {
-            ExitOnFailure(hr, "Failed to wait for executable to complete: %ls", sczWusaPath);
-        }
-    } while (HRESULT_FROM_WIN32(WAIT_TIMEOUT) == hr);
-
-    // get process exit code
-    if (!::GetExitCodeProcess(pi.hProcess, &dwExitCode))
-    {
-        ExitWithLastError(hr, "Failed to get process exit code.");
-    }
+    hr = ExeEngineRunProcess(pfnGenericMessageHandler, pvContext, pPackage, sczWusaPath, sczCommand, NULL, &dwExitCode);
+    ExitOnFailure(hr, "Failed to run MSU process");
 
     // We'll normalize the restart required error code from wusa.exe just in case. Most likely
     // that on reboot we'll actually get WU_S_REBOOT_REQUIRED.
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 81ce8bb9..06f87363 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -104,7 +104,7 @@ extern "C" HRESULT UserExperienceLoad(
     args.pCommand = pCommand;
     args.pfnBootstrapperEngineProc = EngineForApplicationProc;
     args.pvBootstrapperEngineProcContext = pEngineContext;
-    args.qwEngineAPIVersion = MAKEQWORDVERSION(2022, 3, 14, 0);
+    args.qwEngineAPIVersion = MAKEQWORDVERSION(2022, 3, 17, 0);
 
     results.cbSize = sizeof(BOOTSTRAPPER_CREATE_RESULTS);
 
@@ -1701,6 +1701,34 @@ LExit:
     return hr;
 }
 
+BAAPI UserExperienceOnExecuteProcessCancel(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzPackageId,
+    __in DWORD dwProcessId,
+    __inout BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION* pAction
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONEXECUTEPROCESSCANCEL_ARGS args = { };
+    BA_ONEXECUTEPROCESSCANCEL_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.wzPackageId = wzPackageId;
+    args.dwProcessId = dwProcessId;
+    args.recommendation = *pAction;
+
+    results.cbSize = sizeof(results);
+    results.action = *pAction;
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONEXECUTEPROCESSCANCEL, &args, &results);
+    ExitOnFailure(hr, "BA OnExecuteProcessCancel failed.");
+
+    *pAction = results.action;
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnExecuteProgress(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzPackageId,
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index 2f18acdd..de558ad5 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -398,6 +398,12 @@ BAAPI UserExperienceOnExecutePatchTarget(
     __in_z LPCWSTR wzPackageId,
     __in_z LPCWSTR wzTargetProductCode
     );
+BAAPI UserExperienceOnExecuteProcessCancel(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzPackageId,
+    __in DWORD dwProcessId,
+    __inout BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION* pAction
+    );
 BAAPI UserExperienceOnExecuteProgress(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzPackageId,