Always remove all dependencies the bundle could have registered.

#6297

7 files changed, 37 insertions, 13 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
diff --git a/src/test/BurnUnitTest/RegistrationTest.cpp b/src/test/BurnUnitTest/RegistrationTest.cpp
index 9c7bf4ce..1687385b 100644
--- a/src/test/BurnUnitTest/RegistrationTest.cpp
+++ b/src/test/BurnUnitTest/RegistrationTest.cpp
@@ -71,6 +71,7 @@ namespace Bootstrapper
             BOOTSTRAPPER_COMMAND command = { };
             BURN_REGISTRATION registration = { };
             BURN_LOGGING logging = { };
+            BURN_PACKAGES packages = { };
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"));
             try
             {
@@ -121,7 +122,7 @@ namespace Bootstrapper
                 Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)(Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr)));
 
                 // end session
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -159,6 +160,7 @@ namespace Bootstrapper
             BOOTSTRAPPER_COMMAND command = { };
             BURN_REGISTRATION registration = { };
             BURN_LOGGING logging = { };
+            BURN_PACKAGES packages = { };
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"));
             try
             {
@@ -210,7 +212,7 @@ namespace Bootstrapper
                 Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
 
                 // complete registration
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was updated
@@ -232,7 +234,7 @@ namespace Bootstrapper
                 Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
 
                 // delete registration
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -269,6 +271,7 @@ namespace Bootstrapper
             BOOTSTRAPPER_COMMAND command = { };
             BURN_REGISTRATION registration = { };
             BURN_LOGGING logging = { };
+            BURN_PACKAGES packages = { };
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"));
             try
             {
@@ -322,7 +325,7 @@ namespace Bootstrapper
                 Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
 
                 // finish registration
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was updated
@@ -355,7 +358,7 @@ namespace Bootstrapper
                 Assert::Equal<String^>(String::Concat(L"\"", Path::Combine(cacheDirectory, gcnew String(L"setup.exe")), L"\" /burn.runonce"), (String^)Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
 
                 // delete registration
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -392,6 +395,7 @@ namespace Bootstrapper
             BOOTSTRAPPER_COMMAND command = { };
             BURN_REGISTRATION registration = { };
             BURN_LOGGING logging = { };
+            BURN_PACKAGES packages = { };
             BYTE rgbData[256] = { };
             BOOTSTRAPPER_RESUME_TYPE resumeType = BOOTSTRAPPER_RESUME_TYPE_NONE;
             BYTE* pbBuffer = NULL;
@@ -459,7 +463,7 @@ namespace Bootstrapper
                 Assert::Equal((int)BOOTSTRAPPER_RESUME_TYPE_INTERRUPTED, (int)resumeType);
 
                 // suspend session
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_SUSPEND, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_SUSPEND, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_REGISTER);
                 TestThrowOnFailure(hr, L"Failed to suspend session.");
 
                 // verify that run key was removed
@@ -486,7 +490,7 @@ namespace Bootstrapper
                 Assert::NotEqual((Object^)nullptr, Registry::GetValue(gcnew String(TEST_RUN_KEY), gcnew String(L"{D54F896D-1952-43e6-9C67-B5652240618C}"), nullptr));
 
                 // end session
-                hr = RegistrationSessionEnd(&registration, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
+                hr = RegistrationSessionEnd(&registration, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BURN_DEPENDENCY_REGISTRATION_ACTION_UNREGISTER);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // read resume type after session