6 files changed, 26 insertions, 6 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 44ff9429..f57a56fe 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -435,7 +435,7 @@ extern "C" HRESULT ApplyUnregister(
     }
     else
     {
-        hr = RegistrationSessionEnd(&pEngineState->registration, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction);
+        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->packages, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction);
         ExitOnFailure(hr, "Failed to end session in per-user process.");
     }
 
diff --git a/src/engine/dependency.cpp b/src/engine/dependency.cpp
index 3d978740..4833de94 100644
--- a/src/engine/dependency.cpp
+++ b/src/engine/dependency.cpp
@@ -680,10 +680,12 @@ LExit:
 }
 
 extern "C" void DependencyUnregisterBundle(
-    __in const BURN_REGISTRATION* pRegistration
+    __in const BURN_REGISTRATION* pRegistration,
+    __in const BURN_PACKAGES* pPackages
     )
 {
     HRESULT hr = S_OK;
+    LPCWSTR wzDependentProviderKey = pRegistration->sczId;
 
     // Remove the bundle provider key.
     hr = DepUnregisterDependency(pRegistration->hkRoot, pRegistration->sczProviderKey);
@@ -695,6 +697,19 @@ extern "C" void DependencyUnregisterBundle(
     {
         LogId(REPORT_VERBOSE, MSG_DEPENDENCY_BUNDLE_UNREGISTERED_FAILED, pRegistration->sczProviderKey, hr);
     }
+
+    // Best effort to make sure this bundle is not registered as a dependent for anything.
+    for (DWORD i = 0; i < pPackages->cPackages; ++i)
+    {
+        const BURN_PACKAGE* pPackage = pPackages->rgPackages + i;
+        UnregisterPackageDependency(pPackage->fPerMachine, pPackage, wzDependentProviderKey);
+    }
+
+    for (DWORD i = 0; i < pRegistration->relatedBundles.cRelatedBundles; ++i)
+    {
+        const BURN_PACKAGE* pPackage = &pRegistration->relatedBundles.rgRelatedBundles[i].package;
+        UnregisterPackageDependency(pPackage->fPerMachine, pPackage, wzDependentProviderKey);
+    }
 }
 
 // internal functions
diff --git a/src/engine/dependency.h b/src/engine/dependency.h
index 5390bede..efb9f2f2 100644
--- a/src/engine/dependency.h
+++ b/src/engine/dependency.h
@@ -169,7 +169,8 @@ HRESULT DependencyProcessDependentRegistration(
  Note: Does not check for existing dependents before removing the key.
 *********************************************************************/
 void DependencyUnregisterBundle(
-    __in const BURN_REGISTRATION* pRegistration
+    __in const BURN_REGISTRATION* pRegistration,
+    __in const BURN_PACKAGES* pPackages
     );
 
 #if defined(__cplusplus)
diff --git a/src/engine/elevation.cpp b/src/engine/elevation.cpp
index c9ed9810..7c5dae4b 100644
--- a/src/engine/elevation.cpp
+++ b/src/engine/elevation.cpp
@@ -154,6 +154,7 @@ static HRESULT OnSessionResume(
     __in DWORD cbData
     );
 static HRESULT OnSessionEnd(
+    __in BURN_PACKAGES* pPackages,
     __in BURN_REGISTRATION* pRegistration,
     __in BYTE* pbData,
     __in DWORD cbData
@@ -1637,7 +1638,7 @@ static HRESULT ProcessElevatedChildMessage(
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_SESSION_END:
-        hrResult = OnSessionEnd(pContext->pRegistration, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnSessionEnd(pContext->pPackages, pContext->pRegistration, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_SAVE_STATE:
@@ -1954,6 +1955,7 @@ LExit:
 }
 
 static HRESULT OnSessionEnd(
+    __in BURN_PACKAGES* pPackages,
     __in BURN_REGISTRATION* pRegistration,
     __in BYTE* pbData,
     __in DWORD cbData
@@ -1976,7 +1978,7 @@ static HRESULT OnSessionEnd(
     ExitOnFailure(hr, "Failed to read dependency registration action.");
 
     // suspend session in per-machine process
-    hr = RegistrationSessionEnd(pRegistration, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction);
+    hr = RegistrationSessionEnd(pRegistration, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction);
     ExitOnFailure(hr, "Failed to suspend registration session.");
 
 LExit:
diff --git a/src/engine/registration.cpp b/src/engine/registration.cpp
index 9c821422..dc4b88bf 100644
--- a/src/engine/registration.cpp
+++ b/src/engine/registration.cpp
@@ -869,6 +869,7 @@ LExit:
  *******************************************************************/
 extern "C" HRESULT RegistrationSessionEnd(
     __in BURN_REGISTRATION* pRegistration,
+    __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction
@@ -912,7 +913,7 @@ extern "C" HRESULT RegistrationSessionEnd(
             BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER == dependencyRegistrationAction)
         {
             // Remove the bundle dependency key.
-            DependencyUnregisterBundle(pRegistration);
+            DependencyUnregisterBundle(pRegistration, pPackages);
         }
 
         // Delete update registration key.
diff --git a/src/engine/registration.h b/src/engine/registration.h
index 56bcb1f0..4aca5a05 100644
--- a/src/engine/registration.h
+++ b/src/engine/registration.h
@@ -196,6 +196,7 @@ HRESULT RegistrationSessionResume(
     );
 HRESULT RegistrationSessionEnd(
     __in BURN_REGISTRATION* pRegistration,
+    __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction