1 files changed, 139 insertions, 27 deletions

diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 6c4a775f..355b4a34 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -19,6 +19,7 @@ typedef enum _BURN_ELEVATION_MESSAGE_TYPE
     BURN_ELEVATION_MESSAGE_TYPE_CACHE_VERIFY_PAYLOAD,
     BURN_ELEVATION_MESSAGE_TYPE_CACHE_CLEANUP,
     BURN_ELEVATION_MESSAGE_TYPE_PROCESS_DEPENDENT_REGISTRATION,
+    BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_RELATED_BUNDLE,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_EXE_PACKAGE,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSI_PACKAGE,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSP_PACKAGE,
@@ -230,11 +231,18 @@ static HRESULT OnProcessDependentRegistration(
     __in BYTE* pbData,
     __in SIZE_T cbData
     );
+static HRESULT OnExecuteRelatedBundle(
+    __in HANDLE hPipe,
+    __in BURN_CACHE* pCache,
+    __in BURN_RELATED_BUNDLES* pRelatedBundles,
+    __in BURN_VARIABLES* pVariables,
+    __in BYTE* pbData,
+    __in SIZE_T cbData
+    );
 static HRESULT OnExecuteExePackage(
     __in HANDLE hPipe,
     __in BURN_CACHE* pCache,
     __in BURN_PACKAGES* pPackages,
-    __in BURN_RELATED_BUNDLES* pRelatedBundles,
     __in BURN_VARIABLES* pVariables,
     __in BYTE* pbData,
     __in SIZE_T cbData
@@ -818,10 +826,10 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteExePackage - 
+ ElevationExecuteRelatedBundle - 
 
 *******************************************************************/
-extern "C" HRESULT ElevationExecuteExePackage(
+extern "C" HRESULT ElevationExecuteRelatedBundle(
     __in HANDLE hPipe,
     __in BURN_EXECUTE_ACTION* pExecuteAction,
     __in BURN_VARIABLES* pVariables,
@@ -838,22 +846,22 @@ extern "C" HRESULT ElevationExecuteExePackage(
     DWORD dwResult = 0;
 
     // serialize message data
-    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->exePackage.pPackage->sczId);
+    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->relatedBundle.pRelatedBundle->package.sczId);
     ExitOnFailure(hr, "Failed to write package id to message buffer.");
 
-    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)pExecuteAction->exePackage.action);
+    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)pExecuteAction->relatedBundle.action);
     ExitOnFailure(hr, "Failed to write action to message buffer.");
 
     hr = BuffWriteNumber(&pbData, &cbData, fRollback);
     ExitOnFailure(hr, "Failed to write rollback.");
 
-    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->exePackage.sczIgnoreDependencies);
+    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->relatedBundle.sczIgnoreDependencies);
     ExitOnFailure(hr, "Failed to write the list of dependencies to ignore to the message buffer.");
 
-    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->exePackage.sczAncestors);
+    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->relatedBundle.sczAncestors);
     ExitOnFailure(hr, "Failed to write the list of ancestors to the message buffer.");
 
-    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->exePackage.sczEngineWorkingDirectory);
+    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->relatedBundle.sczEngineWorkingDirectory);
     ExitOnFailure(hr, "Failed to write the custom working directory to the message buffer.");
 
     hr = VariableSerialize(pVariables, FALSE, &pbData, &cbData);
@@ -863,6 +871,54 @@ extern "C" HRESULT ElevationExecuteExePackage(
     context.pfnGenericMessageHandler = pfnGenericMessageHandler;
     context.pvContext = pvContext;
 
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_RELATED_BUNDLE, pbData, cbData, ProcessGenericExecuteMessages, &context, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_RELATED_BUNDLE message to per-machine process.");
+
+    hr = ProcessResult(dwResult, pRestart);
+
+LExit:
+    ReleaseBuffer(pbData);
+
+    return hr;
+}
+
+/*******************************************************************
+ ElevationExecuteExePackage - 
+
+*******************************************************************/
+extern "C" HRESULT ElevationExecuteExePackage(
+    __in HANDLE hPipe,
+    __in BURN_EXECUTE_ACTION* pExecuteAction,
+    __in BURN_VARIABLES* pVariables,
+    __in BOOL fRollback,
+    __in PFN_GENERICMESSAGEHANDLER pfnGenericMessageHandler,
+    __in LPVOID pvContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
+    )
+{
+    HRESULT hr = S_OK;
+    BYTE* pbData = NULL;
+    SIZE_T cbData = 0;
+    BURN_ELEVATION_GENERIC_MESSAGE_CONTEXT context = { };
+    DWORD dwResult = 0;
+
+    // serialize message data
+    hr = BuffWriteString(&pbData, &cbData, pExecuteAction->exePackage.pPackage->sczId);
+    ExitOnFailure(hr, "Failed to write package id to message buffer.");
+
+    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)pExecuteAction->exePackage.action);
+    ExitOnFailure(hr, "Failed to write action to message buffer.");
+
+    hr = BuffWriteNumber(&pbData, &cbData, fRollback);
+    ExitOnFailure(hr, "Failed to write rollback.");
+
+    hr = VariableSerialize(pVariables, FALSE, &pbData, &cbData);
+    ExitOnFailure(hr, "Failed to write variables.");
+
+    // send message
+    context.pfnGenericMessageHandler = pfnGenericMessageHandler;
+    context.pvContext = pvContext;
+
     hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_EXE_PACKAGE, pbData, cbData, ProcessGenericExecuteMessages, &context, &dwResult);
     ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_EXE_PACKAGE message to per-machine process.");
 
@@ -1911,8 +1967,12 @@ static HRESULT ProcessElevatedChildMessage(
         hrResult = OnProcessDependentRegistration(pContext->pRegistration, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
 
+    case BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_RELATED_BUNDLE:
+        hrResult = OnExecuteRelatedBundle(pContext->hPipe, pContext->pCache, &pContext->pRegistration->relatedBundles, pContext->pVariables, (BYTE*)pMsg->pvData, pMsg->cbData);
+        break;
+
     case BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_EXE_PACKAGE:
-        hrResult = OnExecuteExePackage(pContext->hPipe, pContext->pCache, pContext->pPackages, &pContext->pRegistration->relatedBundles, pContext->pVariables, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnExecuteExePackage(pContext->hPipe, pContext->pCache, pContext->pPackages, pContext->pVariables, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSI_PACKAGE:
@@ -2473,10 +2533,9 @@ LExit:
     return hr;
 }
 
-static HRESULT OnExecuteExePackage(
+static HRESULT OnExecuteRelatedBundle(
     __in HANDLE hPipe,
     __in BURN_CACHE* pCache,
-    __in BURN_PACKAGES* pPackages,
     __in BURN_RELATED_BUNDLES* pRelatedBundles,
     __in BURN_VARIABLES* pVariables,
     __in BYTE* pbData,
@@ -2491,15 +2550,15 @@ static HRESULT OnExecuteExePackage(
     LPWSTR sczIgnoreDependencies = NULL;
     LPWSTR sczAncestors = NULL;
     LPWSTR sczEngineWorkingDirectory = NULL;
-    BOOTSTRAPPER_APPLY_RESTART exeRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+    BOOTSTRAPPER_APPLY_RESTART bundleRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;
 
-    executeAction.type = BURN_EXECUTE_ACTION_TYPE_EXE_PACKAGE;
+    executeAction.type = BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE;
 
     // Deserialize message data.
     hr = BuffReadString(pbData, cbData, &iData, &sczPackage);
-    ExitOnFailure(hr, "Failed to read EXE package id.");
+    ExitOnFailure(hr, "Failed to read related bundle id.");
 
-    hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&executeAction.exePackage.action);
+    hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&executeAction.relatedBundle.action);
     ExitOnFailure(hr, "Failed to read action.");
 
     hr = BuffReadNumber(pbData, cbData, &iData, &dwRollback);
@@ -2517,36 +2576,32 @@ static HRESULT OnExecuteExePackage(
     hr = VariableDeserialize(pVariables, FALSE, pbData, cbData, &iData);
     ExitOnFailure(hr, "Failed to read variables.");
 
-    hr = PackageFindById(pPackages, sczPackage, &executeAction.exePackage.pPackage);
-    if (E_NOTFOUND == hr)
-    {
-        hr = PackageFindRelatedById(pRelatedBundles, sczPackage, &executeAction.exePackage.pPackage);
-    }
-    ExitOnFailure(hr, "Failed to find package: %ls", sczPackage);
+    hr = RelatedBundleFindById(pRelatedBundles, sczPackage, &executeAction.relatedBundle.pRelatedBundle);
+    ExitOnFailure(hr, "Failed to find related bundle: %ls", sczPackage);
 
     // Pass the list of dependencies to ignore, if any, to the related bundle.
     if (sczIgnoreDependencies && *sczIgnoreDependencies)
     {
-        hr = StrAllocString(&executeAction.exePackage.sczIgnoreDependencies, sczIgnoreDependencies, 0);
+        hr = StrAllocString(&executeAction.relatedBundle.sczIgnoreDependencies, sczIgnoreDependencies, 0);
         ExitOnFailure(hr, "Failed to allocate the list of dependencies to ignore.");
     }
 
     // Pass the list of ancestors, if any, to the related bundle.
     if (sczAncestors && *sczAncestors)
     {
-        hr = StrAllocString(&executeAction.exePackage.sczAncestors, sczAncestors, 0);
+        hr = StrAllocString(&executeAction.relatedBundle.sczAncestors, sczAncestors, 0);
         ExitOnFailure(hr, "Failed to allocate the list of ancestors.");
     }
 
     if (sczEngineWorkingDirectory && *sczEngineWorkingDirectory)
     {
-        hr = StrAllocString(&executeAction.exePackage.sczEngineWorkingDirectory, sczEngineWorkingDirectory, 0);
+        hr = StrAllocString(&executeAction.relatedBundle.sczEngineWorkingDirectory, sczEngineWorkingDirectory, 0);
         ExitOnFailure(hr, "Failed to allocate the custom working directory.");
     }
 
-    // Execute EXE package.
-    hr = ExeEngineExecutePackage(&executeAction, pCache, pVariables, static_cast<BOOL>(dwRollback), GenericExecuteMessageHandler, hPipe, &exeRestart);
-    ExitOnFailure(hr, "Failed to execute EXE package.");
+    // Execute related bundle.
+    hr = BundlePackageEngineExecuteRelatedBundle(&executeAction, pCache, pVariables, static_cast<BOOL>(dwRollback), GenericExecuteMessageHandler, hPipe, &bundleRestart);
+    ExitOnFailure(hr, "Failed to execute related bundle.");
 
 LExit:
     ReleaseStr(sczEngineWorkingDirectory);
@@ -2557,6 +2612,63 @@ LExit:
 
     if (SUCCEEDED(hr))
     {
+        if (BOOTSTRAPPER_APPLY_RESTART_REQUIRED == bundleRestart)
+        {
+            hr = HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED);
+        }
+        else if (BOOTSTRAPPER_APPLY_RESTART_INITIATED == bundleRestart)
+        {
+            hr = HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_INITIATED);
+        }
+    }
+
+    return hr;
+}
+
+static HRESULT OnExecuteExePackage(
+    __in HANDLE hPipe,
+    __in BURN_CACHE* pCache,
+    __in BURN_PACKAGES* pPackages,
+    __in BURN_VARIABLES* pVariables,
+    __in BYTE* pbData,
+    __in SIZE_T cbData
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T iData = 0;
+    LPWSTR sczPackage = NULL;
+    DWORD dwRollback = 0;
+    BURN_EXECUTE_ACTION executeAction = { };
+    BOOTSTRAPPER_APPLY_RESTART exeRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+
+    executeAction.type = BURN_EXECUTE_ACTION_TYPE_EXE_PACKAGE;
+
+    // Deserialize message data.
+    hr = BuffReadString(pbData, cbData, &iData, &sczPackage);
+    ExitOnFailure(hr, "Failed to read EXE package id.");
+
+    hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&executeAction.exePackage.action);
+    ExitOnFailure(hr, "Failed to read action.");
+
+    hr = BuffReadNumber(pbData, cbData, &iData, &dwRollback);
+    ExitOnFailure(hr, "Failed to read rollback.");
+
+    hr = VariableDeserialize(pVariables, FALSE, pbData, cbData, &iData);
+    ExitOnFailure(hr, "Failed to read variables.");
+
+    hr = PackageFindById(pPackages, sczPackage, &executeAction.exePackage.pPackage);
+    ExitOnFailure(hr, "Failed to find package: %ls", sczPackage);
+
+    // Execute EXE package.
+    hr = ExeEngineExecutePackage(&executeAction, pCache, pVariables, static_cast<BOOL>(dwRollback), GenericExecuteMessageHandler, hPipe, &exeRestart);
+    ExitOnFailure(hr, "Failed to execute EXE package.");
+
+LExit:
+    ReleaseStr(sczPackage);
+    PlanUninitializeExecuteAction(&executeAction);
+
+    if (SUCCEEDED(hr))
+    {
         if (BOOTSTRAPPER_APPLY_RESTART_REQUIRED == exeRestart)
         {
             hr = HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED);