Automatically uninstall the bundle after Quit if eligible.

For now, the requirements are:
* The bundle is installed and
* The bundle is per-user or has already elevated and
* No non-permanent packages are installed and
* No non-permanent packages are cached and
* No related bundle would run by default during uninstall and
* The bundle didn't Uninstall/Cache/Install/Modify/Repair and
* The BA didn't opt out of this behavior

10 files changed, 131 insertions, 12 deletions

diff --git a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index c0baa958..48bd813d 100644
--- a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -221,6 +221,9 @@ enum BOOTSTRAPPER_SHUTDOWN_ACTION
     // restart the engine which will load the bootstrapper application again.
     // Typically used to switch from a native bootstrapper application to a managed one.
     BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER,
+    // Opts out of the engine behavior of trying to uninstall itself
+    // when no non-permanent packages are installed.
+    BOOTSTRAPPER_SHUTDOWN_ACTION_SKIP_CLEANUP,
 };
 
 enum BURN_MSI_PROPERTY
@@ -470,6 +473,7 @@ struct BA_ONDETECTCOMPLETE_ARGS
 {
     DWORD cbSize;
     HRESULT hrStatus;
+    BOOL fEligibleForCleanup;
 };
 
 struct BA_ONDETECTCOMPLETE_RESULTS
diff --git a/src/engine/core.cpp b/src/engine/core.cpp
index 36471e93..1503f8d8 100644
--- a/src/engine/core.cpp
+++ b/src/engine/core.cpp
@@ -300,7 +300,7 @@ extern "C" HRESULT CoreDetect(
     ExitOnFailure(hr, "Failed to detect provider key bundle id.");
 
     // Report the related bundles.
-    hr = DetectReportRelatedBundles(&pEngineState->userExperience, &pEngineState->registration, pEngineState->command.relationType, pEngineState->command.action);
+    hr = DetectReportRelatedBundles(&pEngineState->userExperience, &pEngineState->registration, pEngineState->command.relationType, pEngineState->command.action, &pEngineState->registration.fEligibleForCleanup);
     ExitOnFailure(hr, "Failed to report detected related bundles.");
 
     // Do update detection.
@@ -344,6 +344,14 @@ extern "C" HRESULT CoreDetect(
     {
         pPackage = pEngineState->packages.rgPackages + iPackage;
 
+        // If any packages that can affect registration are present, then the bundle should not automatically be uninstalled.
+        if (pEngineState->registration.fEligibleForCleanup && pPackage->fCanAffectRegistration &&
+            (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState ||
+             BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->installRegistrationState))
+        {
+            pEngineState->registration.fEligibleForCleanup = FALSE;
+        }
+
         LogId(REPORT_STANDARD, MSG_DETECTED_PACKAGE, pPackage->sczId, LoggingPackageStateToString(pPackage->currentState), LoggingCacheStateToString(pPackage->cache), LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->installRegistrationState), LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->cacheRegistrationState));
 
         if (BURN_PACKAGE_TYPE_MSI == pPackage->type)
@@ -377,12 +385,12 @@ LExit:
 
     if (fDetectBegan)
     {
-        UserExperienceOnDetectComplete(&pEngineState->userExperience, hr);
+        UserExperienceOnDetectComplete(&pEngineState->userExperience, hr, pEngineState->registration.fEligibleForCleanup);
     }
 
     pEngineState->userExperience.hwndDetect = NULL;
 
-    LogId(REPORT_STANDARD, MSG_DETECT_COMPLETE, hr);
+    LogId(REPORT_STANDARD, MSG_DETECT_COMPLETE, hr, !fDetectBegan ? "(failed)" : LoggingBoolToString(pEngineState->registration.fInstalled), FAILED(hr) ? "(failed)" : LoggingBoolToString(pEngineState->registration.fEligibleForCleanup));
 
     return hr;
 }
@@ -1053,6 +1061,54 @@ LExit:
     return hr;
 }
 
+extern "C" HRESULT CoreCleanup(
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    HRESULT hr = S_OK;
+    LONGLONG llValue = 0;
+    BOOL fNeedsElevation = pEngineState->registration.fPerMachine && INVALID_HANDLE_VALUE == pEngineState->companionConnection.hPipe;
+
+    if (fNeedsElevation)
+    {
+        hr = VariableGetNumeric(&pEngineState->variables, BURN_BUNDLE_ELEVATED, &llValue);
+        ExitOnFailure(hr, "Failed to get value of WixBundleElevated variable during cleanup");
+
+        if (llValue)
+        {
+            fNeedsElevation = FALSE;
+        }
+    }
+
+    if (pEngineState->fApplied && BOOTSTRAPPER_ACTION_LAYOUT < pEngineState->plan.action && BOOTSTRAPPER_ACTION_UPDATE_REPLACE > pEngineState->plan.action ||
+        fNeedsElevation)
+    {
+        ExitFunction();
+    }
+
+    if (!pEngineState->fDetected)
+    {
+        hr = CoreDetect(pEngineState, pEngineState->hMessageWindow);
+        ExitOnFailure(hr, "Detect during cleanup failed");
+    }
+
+    if (!pEngineState->registration.fEligibleForCleanup)
+    {
+        ExitFunction();
+    }
+
+    hr = CorePlan(pEngineState, BOOTSTRAPPER_ACTION_UNINSTALL);
+    ExitOnFailure(hr, "Plan during cleanup failed");
+
+    hr = CoreApply(pEngineState, pEngineState->hMessageWindow);
+    ExitOnFailure(hr, "Apply during cleanup failed");
+
+    // Need to think about cache=always
+
+LExit:
+    return hr;
+}
+
 // internal helper functions
 
 static HRESULT ParseCommandLine(
diff --git a/src/engine/core.h b/src/engine/core.h
index fd7311e3..47cfd559 100644
--- a/src/engine/core.h
+++ b/src/engine/core.h
@@ -204,6 +204,9 @@ HRESULT CoreAppendFileHandleSelfToCommandLine(
     __deref_inout_z LPWSTR* psczCommandLine,
     __deref_inout_z_opt LPWSTR* psczObfuscatedCommandLine
     );
+HRESULT CoreCleanup(
+    __in BURN_ENGINE_STATE* pEngineState
+    );
 
 #if defined(__cplusplus)
 }
diff --git a/src/engine/detect.cpp b/src/engine/detect.cpp
index 159df3d0..b702306e 100644
--- a/src/engine/detect.cpp
+++ b/src/engine/detect.cpp
@@ -42,6 +42,7 @@ extern "C" void DetectReset(
     pRegistration->fEnabledForwardCompatibleBundle = FALSE;
     PackageUninitialize(&pRegistration->forwardCompatibleBundle);
     pRegistration->fSelfRegisteredAsDependent = FALSE;
+    pRegistration->fEligibleForCleanup = FALSE;
 
     if (pRegistration->rgIgnoredDependencies)
     {
@@ -184,11 +185,14 @@ extern "C" HRESULT DetectReportRelatedBundles(
     __in BURN_USER_EXPERIENCE* pUX,
     __in BURN_REGISTRATION* pRegistration,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
-    __in BOOTSTRAPPER_ACTION action
+    __in BOOTSTRAPPER_ACTION action,
+    __out BOOL* pfEligibleForCleanup
     )
 {
     HRESULT hr = S_OK;
     int nCompareResult = 0;
+    BOOTSTRAPPER_REQUEST_STATE uninstallRequestState = BOOTSTRAPPER_REQUEST_STATE_NONE;
+    *pfEligibleForCleanup = pRegistration->fInstalled;
 
     for (DWORD iRelatedBundle = 0; iRelatedBundle < pRegistration->relatedBundles.cRelatedBundles; ++iRelatedBundle)
     {
@@ -201,7 +205,7 @@ extern "C" HRESULT DetectReportRelatedBundles(
             if (BOOTSTRAPPER_RELATION_UPGRADE != relationType && BOOTSTRAPPER_ACTION_UNINSTALL < action)
             {
                 hr = VerCompareParsedVersions(pRegistration->pVersion, pRelatedBundle->pVersion, &nCompareResult);
-                ExitOnFailure(hr, "Failed to compare bundle version '%ls' to related bundle version '%ls'", pRegistration->pVersion, pRelatedBundle->pVersion);
+                ExitOnFailure(hr, "Failed to compare bundle version '%ls' to related bundle version '%ls'", pRegistration->pVersion->sczVersion, pRelatedBundle->pVersion->sczVersion);
 
                 if (nCompareResult < 0)
                 {
@@ -244,6 +248,19 @@ extern "C" HRESULT DetectReportRelatedBundles(
 
         hr = UserExperienceOnDetectRelatedBundle(pUX, pRelatedBundle->package.sczId, pRelatedBundle->relationType, pRelatedBundle->sczTag, pRelatedBundle->package.fPerMachine, pRelatedBundle->pVersion, operation);
         ExitOnRootFailure(hr, "BA aborted detect related bundle.");
+
+        // For now, if any related bundles will be executed during uninstall by default then never automatically clean up the bundle.
+        if (*pfEligibleForCleanup)
+        {
+            uninstallRequestState = BOOTSTRAPPER_REQUEST_STATE_NONE;
+            hr = PlanDefaultRelatedBundleRequestState(relationType, pRelatedBundle->relationType, BOOTSTRAPPER_ACTION_UNINSTALL, pRegistration->pVersion, pRelatedBundle->pVersion, &uninstallRequestState);
+            ExitOnFailure(hr, "Failed to get the default request state for related bundle for calculating fEligibleForCleanup");
+
+            if (BOOTSTRAPPER_REQUEST_STATE_NONE != uninstallRequestState)
+            {
+                *pfEligibleForCleanup = FALSE;
+            }
+        }
     }
 
 LExit:
diff --git a/src/engine/detect.h b/src/engine/detect.h
index 01488f1a..7989c9dd 100644
--- a/src/engine/detect.h
+++ b/src/engine/detect.h
@@ -30,7 +30,8 @@ HRESULT DetectReportRelatedBundles(
     __in BURN_USER_EXPERIENCE* pUX,
     __in BURN_REGISTRATION* pRegistration,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
-    __in BOOTSTRAPPER_ACTION action
+    __in BOOTSTRAPPER_ACTION action,
+    __out BOOL* pfEligibleForCleanup
     );
 
 HRESULT DetectUpdate(
diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp
index e3ace592..bc27cb14 100644
--- a/src/engine/engine.cpp
+++ b/src/engine/engine.cpp
@@ -39,7 +39,8 @@ static HRESULT RunRunOnce(
     );
 static HRESULT RunApplication(
     __in BURN_ENGINE_STATE* pEngineState,
-    __out BOOL* pfReloadApp
+    __out BOOL* pfReloadApp,
+    __out BOOL* pfSkipCleanup
     );
 static HRESULT ProcessMessage(
     __in BURN_ENGINE_STATE* pEngineState,
@@ -529,6 +530,7 @@ static HRESULT RunNormal(
     HANDLE hPipesCreatedEvent = NULL;
     BOOL fContinueExecution = TRUE;
     BOOL fReloadApp = FALSE;
+    BOOL fSkipCleanup = FALSE;
     BURN_EXTENSION_ENGINE_CONTEXT extensionEngineContext = { };
 
     // Initialize logging.
@@ -584,11 +586,18 @@ static HRESULT RunNormal(
     do
     {
         fReloadApp = FALSE;
+        pEngineState->fQuit = FALSE;
 
-        hr = RunApplication(pEngineState, &fReloadApp);
+        hr = RunApplication(pEngineState, &fReloadApp, &fSkipCleanup);
         ExitOnFailure(hr, "Failed while running ");
     } while (fReloadApp);
 
+    if (!fSkipCleanup)
+    {
+        hr = CoreCleanup(pEngineState);
+        ExitOnFailure(hr, "Failed to cleanup before shutting down");
+    }
+
 LExit:
     BurnExtensionUnload(&pEngineState->extensions);
 
@@ -732,7 +741,8 @@ LExit:
 
 static HRESULT RunApplication(
     __in BURN_ENGINE_STATE* pEngineState,
-    __out BOOL* pfReloadApp
+    __out BOOL* pfReloadApp,
+    __out BOOL* pfSkipCleanup
     )
 {
     HRESULT hr = S_OK;
@@ -787,6 +797,11 @@ LExit:
             LogId(REPORT_STANDARD, MSG_BA_REQUESTED_RELOAD);
             *pfReloadApp = TRUE;
         }
+        else if (BOOTSTRAPPER_SHUTDOWN_ACTION_SKIP_CLEANUP == shutdownAction)
+        {
+            LogId(REPORT_STANDARD, MSG_BA_REQUESTED_SKIP_CLEANUP);
+            *pfSkipCleanup = TRUE;
+        }
     }
 
     // Unload BA.
diff --git a/src/engine/engine.mc b/src/engine/engine.mc
index 59a05676..c90f08e3 100644
--- a/src/engine/engine.mc
+++ b/src/engine/engine.mc
@@ -121,6 +121,13 @@ Language=English
 The manifest contains an invalid version string: '%1!ls!'
 .
 
+MessageId=14
+Severity=Success
+SymbolicName=MSG_BA_REQUESTED_SKIP_CLEANUP
+Language=English
+Bootstrapper application opted out of any engine behavior to automatically uninstall the bundle during shutdown.
+.
+
 MessageId=51
 Severity=Error
 SymbolicName=MSG_FAILED_PARSE_CONDITION
@@ -286,7 +293,7 @@ MessageId=199
 Severity=Success
 SymbolicName=MSG_DETECT_COMPLETE
 Language=English
-Detect complete, result: 0x%1!x!
+Detect complete, result: 0x%1!x!, installed: %2!hs!, eligible for cleanup: %3!hs!
 .
 
 MessageId=200
diff --git a/src/engine/registration.h b/src/engine/registration.h
index c1e52ac9..55d5a4c8 100644
--- a/src/engine/registration.h
+++ b/src/engine/registration.h
@@ -146,6 +146,7 @@ typedef struct _BURN_REGISTRATION
     UINT cDependents;                    // Only valid after detect.
     LPCWSTR wzSelfDependent;             // Only valid after detect.
     BOOL fSelfRegisteredAsDependent;     // Only valid after detect.
+    BOOL fEligibleForCleanup;            // Only valid after detect.
 
     LPWSTR sczDetectedProviderKeyBundleId;
     LPWSTR sczAncestors;
diff --git a/src/engine/userexperience.cpp b/src/engine/userexperience.cpp
index 3a36cab6..a0fb341d 100644
--- a/src/engine/userexperience.cpp
+++ b/src/engine/userexperience.cpp
@@ -736,7 +736,8 @@ LExit:
 
 EXTERN_C BAAPI UserExperienceOnDetectComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in HRESULT hrStatus
+    __in HRESULT hrStatus,
+    __in BOOL fEligibleForCleanup
     )
 {
     HRESULT hr = S_OK;
@@ -745,6 +746,7 @@ EXTERN_C BAAPI UserExperienceOnDetectComplete(
 
     args.cbSize = sizeof(args);
     args.hrStatus = hrStatus;
+    args.fEligibleForCleanup = fEligibleForCleanup;
 
     results.cbSize = sizeof(results);
 
@@ -2296,12 +2298,18 @@ static HRESULT SendBAMessage(
 {
     HRESULT hr = S_OK;
 
+    if (!pUserExperience->hUXModule)
+    {
+        ExitFunction();
+    }
+
     hr = pUserExperience->pfnBAProc(message, pvArgs, pvResults, pUserExperience->pvBAProcContext);
     if (hr == E_NOTIMPL)
     {
         hr = S_OK;
     }
 
+LExit:
     return hr;
 }
 
@@ -2314,11 +2322,17 @@ static HRESULT SendBAMessageFromInactiveEngine(
 {
     HRESULT hr = S_OK;
 
+    if (!pUserExperience->hUXModule)
+    {
+        ExitFunction();
+    }
+
     UserExperienceDeactivateEngine(pUserExperience);
 
     hr = SendBAMessage(pUserExperience, message, pvArgs, pvResults);
 
     UserExperienceActivateEngine(pUserExperience);
 
+LExit:
     return hr;
 }
diff --git a/src/engine/userexperience.h b/src/engine/userexperience.h
index f51c09ff..363c0f06 100644
--- a/src/engine/userexperience.h
+++ b/src/engine/userexperience.h
@@ -191,7 +191,8 @@ BAAPI UserExperienceOnDetectBegin(
     );
 BAAPI UserExperienceOnDetectComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in HRESULT hrStatus
+    __in HRESULT hrStatus,
+    __in BOOL fEligibleForCleanup
     );
 BAAPI UserExperienceOnDetectForwardCompatibleBundle(
     __in BURN_USER_EXPERIENCE* pUserExperience,