11 files changed, 249 insertions, 82 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index aad9e6eb..c32f4c84 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -75,7 +75,8 @@ static HRESULT WINAPI AuthenticationRequired(
 
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
-    __inout BOOL* pfKeepRegistration
+    __in BOOL fLog,
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     );
 static HRESULT ExecuteDependentRegistrationActions(
     __in HANDLE hPipe,
@@ -376,8 +377,11 @@ extern "C" HRESULT ApplyRegister(
 {
     HRESULT hr = S_OK;
     LPWSTR sczEngineWorkingPath = NULL;
+    BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
 
-    hr = UserExperienceOnRegisterBegin(&pEngineState->userExperience);
+    CalculateKeepRegistration(pEngineState, FALSE, &registrationType);
+
+    hr = UserExperienceOnRegisterBegin(&pEngineState->userExperience, &registrationType);
     ExitOnRootFailure(hr, "BA aborted register begin.");
 
     // If we have a resume mode that suggests the bundle is on the machine.
@@ -386,12 +390,12 @@ extern "C" HRESULT ApplyRegister(
         // resume previous session
         if (pEngineState->registration.fPerMachine)
         {
-            hr = ElevationSessionResume(pEngineState->companionConnection.hPipe, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables);
+            hr = ElevationSessionResume(pEngineState->companionConnection.hPipe, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, registrationType);
             ExitOnFailure(hr, "Failed to resume registration session in per-machine process.");
         }
         else
         {
-            hr = RegistrationSessionResume(&pEngineState->registration, &pEngineState->variables);
+            hr = RegistrationSessionResume(&pEngineState->registration, &pEngineState->variables, registrationType);
             ExitOnFailure(hr, "Failed to resume registration session.");
         }
     }
@@ -403,12 +407,12 @@ extern "C" HRESULT ApplyRegister(
         // begin new session
         if (pEngineState->registration.fPerMachine)
         {
-            hr = ElevationSessionBegin(pEngineState->companionConnection.hPipe, sczEngineWorkingPath, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.dependencyRegistrationAction, pEngineState->plan.qwEstimatedSize);
+            hr = ElevationSessionBegin(pEngineState->companionConnection.hPipe, sczEngineWorkingPath, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.dependencyRegistrationAction, pEngineState->plan.qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session in per-machine process.");
         }
         else
         {
-            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.dependencyRegistrationAction, pEngineState->plan.qwEstimatedSize);
+            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.dependencyRegistrationAction, pEngineState->plan.qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session.");
         }
     }
@@ -441,17 +445,11 @@ extern "C" HRESULT ApplyUnregister(
 {
     HRESULT hr = S_OK;
     BURN_RESUME_MODE resumeMode = BURN_RESUME_MODE_NONE;
-    BOOL fKeepRegistration = pEngineState->plan.fDisallowRemoval;
-
-    CalculateKeepRegistration(pEngineState, &fKeepRegistration);
-
-    hr = UserExperienceOnUnregisterBegin(&pEngineState->userExperience, &fKeepRegistration);
-    ExitOnRootFailure(hr, "BA aborted unregister begin.");
+    BOOTSTRAPPER_REGISTRATION_TYPE defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
+    BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
 
-    // Calculate the correct resume mode. If a restart has been initiated, that trumps all other
+    // Calculate special cases for the resume mode. If a restart has been initiated, that trumps all other
     // modes. If the user chose to suspend the install then we'll use that as the resume mode.
-    // Barring those special cases, if it was determined that we should keep the registration then
-    // do that, otherwise the resume mode was initialized to none and registration will be removed.
     if (BOOTSTRAPPER_APPLY_RESTART_INITIATED == restart)
     {
         resumeMode = BURN_RESUME_MODE_REBOOT_PENDING;
@@ -460,28 +458,50 @@ extern "C" HRESULT ApplyUnregister(
     {
         resumeMode = BURN_RESUME_MODE_SUSPEND;
     }
-    else if (fKeepRegistration)
+    else if (pEngineState->plan.fDisallowRemoval)
+    {
+        resumeMode = BURN_RESUME_MODE_ARP;
+    }
+
+    // If there was a special case, make sure the registration is kept.
+    if (BURN_RESUME_MODE_NONE < resumeMode)
+    {
+        defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+    }
+
+    CalculateKeepRegistration(pEngineState, TRUE, &defaultRegistrationType);
+
+    registrationType = defaultRegistrationType;
+
+    hr = UserExperienceOnUnregisterBegin(&pEngineState->userExperience, &registrationType);
+    ExitOnRootFailure(hr, "BA aborted unregister begin.");
+
+    // Barring the special cases, if it was determined that we should keep the registration then
+    // do that, otherwise the resume mode is NONE and registration will be removed.
+    if (BURN_RESUME_MODE_NONE == resumeMode && BOOTSTRAPPER_REGISTRATION_TYPE_NONE < registrationType)
     {
         resumeMode = BURN_RESUME_MODE_ARP;
     }
 
     // If apply failed in any way and we're going to be keeping the bundle registered then
     // execute any rollback dependency registration actions.
-    if (fFailed && fKeepRegistration)
+    if (fFailed && BURN_RESUME_MODE_NONE < resumeMode)
     {
         // Execute any rollback registration actions.
         HRESULT hrRegistrationRollback = ExecuteDependentRegistrationActions(pEngineState->companionConnection.hPipe, &pEngineState->registration, pEngineState->plan.rgRollbackRegistrationActions, pEngineState->plan.cRollbackRegistrationActions);
-        UNREFERENCED_PARAMETER(hrRegistrationRollback);
+        IgnoreRollbackError(hrRegistrationRollback, "Dependent registration actions failed");
     }
 
+    LogId(REPORT_STANDARD, MSG_SESSION_END, pEngineState->registration.sczRegistrationKey, LoggingResumeModeToString(resumeMode), LoggingRestartToString(restart), LoggingBoolToString(pEngineState->registration.fDisableResume), LoggingRegistrationTypeToString(defaultRegistrationType), LoggingRegistrationTypeToString(registrationType));
+
     if (pEngineState->registration.fPerMachine)
     {
-        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction);
+        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-machine process.");
     }
     else
     {
-        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction);
+        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, pEngineState->plan.dependencyRegistrationAction, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-user process.");
     }
 
@@ -751,10 +771,14 @@ extern "C" void ApplyClean(
 
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
-    __inout BOOL* pfKeepRegistration
+    __in BOOL fLog,
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     )
 {
-    LogId(REPORT_STANDARD, MSG_POST_APPLY_CALCULATE_REGISTRATION);
+    if (fLog)
+    {
+        LogId(REPORT_STANDARD, MSG_POST_APPLY_CALCULATE_REGISTRATION);
+    }
 
     for (DWORD i = 0; i < pEngineState->packages.cPackages; ++i)
     {
@@ -765,17 +789,28 @@ static void CalculateKeepRegistration(
             MspEngineFinalizeInstallRegistrationState(pPackage);
         }
 
-        LogId(REPORT_STANDARD, MSG_POST_APPLY_PACKAGE, pPackage->sczId, LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->installRegistrationState), LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->cacheRegistrationState));
+        if (fLog)
+        {
+            LogId(REPORT_STANDARD, MSG_POST_APPLY_PACKAGE, pPackage->sczId, LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->installRegistrationState), LoggingPackageRegistrationStateToString(pPackage->fCanAffectRegistration, pPackage->cacheRegistrationState));
+        }
 
         if (!pPackage->fCanAffectRegistration)
         {
             continue;
         }
 
-        if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->installRegistrationState ||
-            BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState)
+        if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->installRegistrationState)
+        {
+            *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_FULL;
+
+            if (!fLog)
+            {
+                break;
+            }
+        }
+        else if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState && BOOTSTRAPPER_REGISTRATION_TYPE_NONE == *pRegistrationType)
         {
-            *pfKeepRegistration = TRUE;
+            *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
         }
     }
 }
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 27cb2e25..3a28188b 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -48,8 +48,9 @@ const LPCWSTR BURN_BUNDLE_UILEVEL = L"WixBundleUILevel";
 const LPCWSTR BURN_BUNDLE_VERSION = L"WixBundleVersion";
 const LPCWSTR BURN_REBOOT_PENDING = L"RebootPending";
 
-// The following constants must stay in sync with src\wix\Binder.cs
+// The following constants must stay in sync with src\api\wix\WixToolset.Data\Burn\BurnConstants.cs
 const LPCWSTR BURN_BUNDLE_NAME = L"WixBundleName";
+const LPCWSTR BURN_BUNDLE_INPROGRESS_NAME = L"WixBundleInProgressName";
 const LPCWSTR BURN_BUNDLE_ORIGINAL_SOURCE = L"WixBundleOriginalSource";
 const LPCWSTR BURN_BUNDLE_ORIGINAL_SOURCE_FOLDER = L"WixBundleOriginalSourceFolder";
 const LPCWSTR BURN_BUNDLE_LAST_USED_SOURCE = L"WixBundleLastUsedSource";
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 85e4b692..db2a82a6 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -477,7 +477,8 @@ extern "C" HRESULT ElevationSessionBegin(
     __in BURN_VARIABLES* pVariables,
     __in DWORD dwRegistrationOperations,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
-    __in DWORD64 qwEstimatedSize
+    __in DWORD64 qwEstimatedSize,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
@@ -504,6 +505,9 @@ extern "C" HRESULT ElevationSessionBegin(
     hr = BuffWriteNumber64(&pbData, &cbData, qwEstimatedSize);
     ExitOnFailure(hr, "Failed to write estimated size to message buffer.");
 
+    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)registrationType);
+    ExitOnFailure(hr, "Failed to write registration type to message buffer.");
+
     hr = VariableSerialize(pVariables, FALSE, &pbData, &cbData);
     ExitOnFailure(hr, "Failed to write variables.");
 
@@ -527,7 +531,8 @@ extern "C" HRESULT ElevationSessionResume(
     __in HANDLE hPipe,
     __in_z LPCWSTR wzResumeCommandLine,
     __in BOOL fDisableResume,
-    __in BURN_VARIABLES* pVariables
+    __in BURN_VARIABLES* pVariables,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
@@ -542,6 +547,9 @@ extern "C" HRESULT ElevationSessionResume(
     hr = BuffWriteNumber(&pbData, &cbData, fDisableResume);
     ExitOnFailure(hr, "Failed to write resume flag.");
 
+    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)registrationType);
+    ExitOnFailure(hr, "Failed to write registration type to message buffer.");
+
     hr = VariableSerialize(pVariables, FALSE, &pbData, &cbData);
     ExitOnFailure(hr, "Failed to write variables.");
 
@@ -565,7 +573,8 @@ extern "C" HRESULT ElevationSessionEnd(
     __in HANDLE hPipe,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
-    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction
+    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
@@ -583,6 +592,9 @@ extern "C" HRESULT ElevationSessionEnd(
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)dependencyRegistrationAction);
     ExitOnFailure(hr, "Failed to write dependency registration action to message buffer.");
 
+    hr = BuffWriteNumber(&pbData, &cbData, (DWORD)registrationType);
+    ExitOnFailure(hr, "Failed to write registration type to message buffer.");
+
     // send message
     hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_SESSION_END, pbData, cbData, NULL, NULL, &dwResult);
     ExitOnFailure(hr, "Failed to send message to per-machine process.");
@@ -2080,6 +2092,7 @@ static HRESULT OnSessionBegin(
     DWORD dwRegistrationOperations = 0;
     DWORD dwDependencyRegistrationAction = 0;
     DWORD64 qwEstimatedSize = 0;
+    DWORD dwRegistrationType = 0;
 
     // Deserialize message data.
     hr = BuffReadString(pbData, cbData, &iData, &sczEngineWorkingPath);
@@ -2100,11 +2113,14 @@ static HRESULT OnSessionBegin(
     hr = BuffReadNumber64(pbData, cbData, &iData, &qwEstimatedSize);
     ExitOnFailure(hr, "Failed to read estimated size.");
 
+    hr = BuffReadNumber(pbData, cbData, &iData, &dwRegistrationType);
+    ExitOnFailure(hr, "Failed to read dependency registration action.");
+
     hr = VariableDeserialize(pVariables, FALSE, pbData, cbData, &iData);
     ExitOnFailure(hr, "Failed to read variables.");
 
     // Begin session in per-machine process.
-    hr = RegistrationSessionBegin(sczEngineWorkingPath, pRegistration, pVariables, dwRegistrationOperations, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction, qwEstimatedSize);
+    hr = RegistrationSessionBegin(sczEngineWorkingPath, pRegistration, pVariables, dwRegistrationOperations, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction, qwEstimatedSize, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
     ExitOnFailure(hr, "Failed to begin registration session.");
 
 LExit:
@@ -2122,6 +2138,7 @@ static HRESULT OnSessionResume(
 {
     HRESULT hr = S_OK;
     SIZE_T iData = 0;
+    DWORD dwRegistrationType = 0;
 
     // Deserialize message data.
     hr = BuffReadString(pbData, cbData, &iData, &pRegistration->sczResumeCommandLine);
@@ -2130,11 +2147,14 @@ static HRESULT OnSessionResume(
     hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&pRegistration->fDisableResume);
     ExitOnFailure(hr, "Failed to read resume flag.");
 
+    hr = BuffReadNumber(pbData, cbData, &iData, &dwRegistrationType);
+    ExitOnFailure(hr, "Failed to read dependency registration action.");
+
     hr = VariableDeserialize(pVariables, FALSE, pbData, cbData, &iData);
     ExitOnFailure(hr, "Failed to read variables.");
 
     // resume session in per-machine process
-    hr = RegistrationSessionResume(pRegistration, pVariables);
+    hr = RegistrationSessionResume(pRegistration, pVariables, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
     ExitOnFailure(hr, "Failed to resume registration session.");
 
 LExit:
@@ -2154,6 +2174,7 @@ static HRESULT OnSessionEnd(
     DWORD dwResumeMode = 0;
     DWORD dwRestart = 0;
     DWORD dwDependencyRegistrationAction = 0;
+    DWORD dwRegistrationType = 0;
 
     // Deserialize message data.
     hr = BuffReadNumber(pbData, cbData, &iData, &dwResumeMode);
@@ -2165,8 +2186,11 @@ static HRESULT OnSessionEnd(
     hr = BuffReadNumber(pbData, cbData, &iData, &dwDependencyRegistrationAction);
     ExitOnFailure(hr, "Failed to read dependency registration action.");
 
+    hr = BuffReadNumber(pbData, cbData, &iData, &dwRegistrationType);
+    ExitOnFailure(hr, "Failed to read dependency registration action.");
+
     // suspend session in per-machine process
-    hr = RegistrationSessionEnd(pRegistration, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction);
+    hr = RegistrationSessionEnd(pRegistration, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BURN_DEPENDENCY_REGISTRATION_ACTION)dwDependencyRegistrationAction, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
     ExitOnFailure(hr, "Failed to suspend registration session.");
 
 LExit:
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index af2dec8b..00dca8dc 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -31,19 +31,22 @@ HRESULT ElevationSessionBegin(
     __in BURN_VARIABLES* pVariables,
     __in DWORD dwRegistrationOperations,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
-    __in DWORD64 qwEstimatedSize
+    __in DWORD64 qwEstimatedSize,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT ElevationSessionResume(
     __in HANDLE hPipe,
     __in_z LPCWSTR wzResumeCommandLine,
     __in BOOL fDisableResume,
-    __in BURN_VARIABLES* pVariables
+    __in BURN_VARIABLES* pVariables,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT ElevationSessionEnd(
     __in HANDLE hPipe,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
-    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction
+    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT ElevationSaveState(
     __in HANDLE hPipe,
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index 929bf67a..ad9d676f 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -889,7 +889,7 @@ MessageId=372
 Severity=Success
 SymbolicName=MSG_SESSION_END
 Language=English
-Session end, registration key: %1!ls!, resume: %2!hs!, restart: %3!hs!, disable resume: %4!hs!
+Session end, registration key: %1!ls!, resume: %2!hs!, restart: %3!hs!, disable resume: %4!hs!, default registration: %5!hs!, ba requested registration: %6!hs!
 .
 
 MessageId=373
diff --git a/src/burn/engine/logging.cpp b/src/burn/engine/logging.cpp
index 7ee1ec85..2db7defd 100644
--- a/src/burn/engine/logging.cpp
+++ b/src/burn/engine/logging.cpp
@@ -561,6 +561,23 @@ extern "C" LPCSTR LoggingPerMachineToString(
     return "PerUser";
 }
 
+extern "C" LPCSTR LoggingRegistrationTypeToString(
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
+    )
+{
+    switch (registrationType)
+    {
+    case BOOTSTRAPPER_REGISTRATION_TYPE_NONE:
+        return "None";
+    case BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS:
+        return "InProgress";
+    case BOOTSTRAPPER_REGISTRATION_TYPE_FULL:
+        return "Full";
+    default:
+        return "Invalid";
+    }
+}
+
 extern "C" LPCSTR LoggingRestartToString(
     __in BOOTSTRAPPER_APPLY_RESTART restart
     )
diff --git a/src/burn/engine/logging.h b/src/burn/engine/logging.h
index 909ce591..21ea6297 100644
--- a/src/burn/engine/logging.h
+++ b/src/burn/engine/logging.h
@@ -123,6 +123,10 @@ LPCSTR LoggingPerMachineToString(
     __in BOOL fPerMachine
     );
 
+LPCSTR LoggingRegistrationTypeToString(
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
+    );
+
 LPCSTR LoggingRestartToString(
     __in BOOTSTRAPPER_APPLY_RESTART restart
     );
diff --git a/src/burn/engine/registration.cpp b/src/burn/engine/registration.cpp
index 5b246112..eed1fee2 100644
--- a/src/burn/engine/registration.cpp
+++ b/src/burn/engine/registration.cpp
@@ -46,12 +46,23 @@ static HRESULT SetPaths(
 static HRESULT GetBundleManufacturer(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __out LPWSTR* psczBundleManufacturer
+    __out_z LPWSTR* psczBundleManufacturer
+    );
+static HRESULT GetBundleInProgressName(
+    __in BURN_REGISTRATION* pRegistration,
+    __in BURN_VARIABLES* pVariables,
+    __out_z LPWSTR* psczBundleName
     );
 static HRESULT GetBundleName(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __out LPWSTR* psczBundleName
+    __out_z LPWSTR* psczBundleName
+    );
+static HRESULT EnsureRegistrationVariable(
+    __in BURN_VARIABLES* pVariables,
+    __in_z LPCWSTR wzVariable,
+    __in_z LPCWSTR wzDefaultValue,
+    __out_z LPWSTR* psczValue
     );
 static HRESULT UpdateResumeMode(
     __in BURN_REGISTRATION* pRegistration,
@@ -91,7 +102,8 @@ static HRESULT RegWriteStringVariable(
 static HRESULT UpdateBundleNameRegistration(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __in HKEY hkRegistration
+    __in HKEY hkRegistration,
+    __in BOOL fInProgressRegistration
     );
 static BOOL IsWuRebootPending();
 static BOOL IsBundleRebootPending(
@@ -176,6 +188,13 @@ extern "C" HRESULT RegistrationParseFromXml(
             ExitOnFailure(hr, "Failed to get @DisplayName.");
         }
 
+        // @InProgressDisplayName
+        hr = XmlGetAttributeEx(pixnArpNode, L"InProgressDisplayName", &pRegistration->sczInProgressDisplayName);
+        if (E_NOTFOUND != hr)
+        {
+            ExitOnFailure(hr, "Failed to get @InProgressDisplayName.");
+        }
+
         // @DisplayVersion
         hr = XmlGetAttributeEx(pixnArpNode, L"DisplayVersion", &pRegistration->sczDisplayVersion);
         if (E_NOTFOUND != hr)
@@ -372,6 +391,7 @@ extern "C" void RegistrationUninitialize(
     ReleaseStr(pRegistration->sczStateFile);
 
     ReleaseStr(pRegistration->sczDisplayName);
+    ReleaseStr(pRegistration->sczInProgressDisplayName);
     ReleaseStr(pRegistration->sczDisplayVersion);
     ReleaseStr(pRegistration->sczPublisher);
     ReleaseStr(pRegistration->sczHelpLink);
@@ -421,8 +441,7 @@ extern "C" HRESULT RegistrationSetVariables(
     )
 {
     HRESULT hr = S_OK;
-    LPWSTR sczBundleManufacturer = NULL;
-    LPWSTR sczBundleName = NULL;
+    LPWSTR scz = NULL;
 
     if (pRegistration->fInstalled)
     {
@@ -431,10 +450,13 @@ extern "C" HRESULT RegistrationSetVariables(
     }
 
     // Ensure the registration bundle name is updated.
-    hr = GetBundleName(pRegistration, pVariables, &sczBundleName);
+    hr = GetBundleInProgressName(pRegistration, pVariables, &scz);
     ExitOnFailure(hr, "Failed to initialize bundle name.");
 
-    hr = GetBundleManufacturer(pRegistration, pVariables, &sczBundleName);
+    hr = GetBundleName(pRegistration, pVariables, &scz);
+    ExitOnFailure(hr, "Failed to initialize bundle name.");
+
+    hr = GetBundleManufacturer(pRegistration, pVariables, &scz);
     ExitOnFailure(hr, "Failed to initialize bundle manufacturer.");
 
     if (pRegistration->sczActiveParent && *pRegistration->sczActiveParent)
@@ -456,8 +478,7 @@ extern "C" HRESULT RegistrationSetVariables(
     ExitOnFailure(hr, "Failed to overwrite the bundle reboot-pending built-in variable.");
 
 LExit:
-    ReleaseStr(sczBundleManufacturer);
-    ReleaseStr(sczBundleName);
+    ReleaseStr(scz);
 
     return hr;
 }
@@ -595,7 +616,8 @@ extern "C" HRESULT RegistrationSessionBegin(
     __in BURN_VARIABLES* pVariables,
     __in DWORD dwRegistrationOptions,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
-    __in DWORD64 qwEstimatedSize
+    __in DWORD64 qwEstimatedSize,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
@@ -603,6 +625,8 @@ extern "C" HRESULT RegistrationSessionBegin(
     HKEY hkRegistration = NULL;
     LPWSTR sczPublisher = NULL;
 
+    AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE != registrationType, "Registration type can't be NONE");
+
     LogId(REPORT_VERBOSE, MSG_SESSION_BEGIN, pRegistration->sczRegistrationKey, dwRegistrationOptions, LoggingBoolToString(pRegistration->fDisableResume));
 
     // Cache bundle executable.
@@ -668,7 +692,7 @@ extern "C" HRESULT RegistrationSessionBegin(
         ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_DISPLAY_ICON);
 
         // update display name
-        hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration);
+        hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
         ExitOnFailure(hr, "Failed to update name and publisher.");
 
         // DisplayVersion: provided by UI
@@ -841,12 +865,15 @@ LExit:
 *******************************************************************/
 extern "C" HRESULT RegistrationSessionResume(
     __in BURN_REGISTRATION* pRegistration,
-    __in BURN_VARIABLES* pVariables
+    __in BURN_VARIABLES* pVariables,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
     HKEY hkRegistration = NULL;
 
+    AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE != registrationType, "Registration type can't be NONE");
+
     // open registration key
     hr = RegOpen(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, &hkRegistration);
     ExitOnFailure(hr, "Failed to open registration key.");
@@ -856,7 +883,7 @@ extern "C" HRESULT RegistrationSessionResume(
     ExitOnFailure(hr, "Failed to update resume mode.");
 
     // update display name
-    hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration);
+    hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
     ExitOnFailure(hr, "Failed to update name and publisher.");
 
 LExit:
@@ -876,7 +903,8 @@ extern "C" HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
-    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction
+    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
     HRESULT hr = S_OK;
@@ -884,8 +912,6 @@ extern "C" HRESULT RegistrationSessionEnd(
     HKEY hkRebootRequired = NULL;
     HKEY hkRegistration = NULL;
 
-    LogId(REPORT_STANDARD, MSG_SESSION_END, pRegistration->sczRegistrationKey, LoggingResumeModeToString(resumeMode), LoggingRestartToString(restart), LoggingBoolToString(pRegistration->fDisableResume));
-
     // If a restart is required for any reason, write a volatile registry key to track of
     // of that fact until the reboot has taken place.
     if (BOOTSTRAPPER_APPLY_RESTART_NONE != restart)
@@ -910,6 +936,8 @@ extern "C" HRESULT RegistrationSessionEnd(
     // If no resume mode, then remove the bundle registration.
     if (BURN_RESUME_MODE_NONE == resumeMode)
     {
+        AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE == registrationType, "Registration type must be NONE if resume mode is NONE");
+
         // If we just registered the bundle dependency but something went wrong and caused us to not
         // keep the bundle registration (like rollback) or we are supposed to unregister the bundle
         // dependency when unregistering the bundle, do so.
@@ -939,9 +967,15 @@ extern "C" HRESULT RegistrationSessionEnd(
     }
     else // the mode needs to be updated so open the registration key.
     {
+        AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE != registrationType, "Registration type must not be NONE if resume mode is not NONE");
+
         // Open registration key.
         hr = RegOpen(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, &hkRegistration);
         ExitOnFailure(hr, "Failed to open registration key.");
+
+        // update display name
+        hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
+        ExitOnFailure(hr, "Failed to update name and publisher.");
     }
 
     // Update resume mode.
@@ -1133,44 +1167,69 @@ LExit:
 static HRESULT GetBundleManufacturer(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __out LPWSTR* psczBundleManufacturer
+    __out_z LPWSTR* psczBundleManufacturer
     )
 {
     HRESULT hr = S_OK;
     LPCWSTR wzPublisher = pRegistration->sczPublisher ? pRegistration->sczPublisher : L"";
 
-    hr = VariableGetString(pVariables, BURN_BUNDLE_MANUFACTURER, psczBundleManufacturer);
-    if (E_NOTFOUND == hr)
-    {
-        hr = VariableSetString(pVariables, BURN_BUNDLE_MANUFACTURER, wzPublisher, FALSE, FALSE);
-        ExitOnFailure(hr, "Failed to set bundle manufacturer.");
-
-        hr = StrAllocString(psczBundleManufacturer, wzPublisher, 0);
-    }
+    hr = EnsureRegistrationVariable(pVariables, BURN_BUNDLE_MANUFACTURER, wzPublisher, psczBundleManufacturer);
     ExitOnFailure(hr, "Failed to get bundle manufacturer.");
 
 LExit:
     return hr;
 }
 
+static HRESULT GetBundleInProgressName(
+    __in BURN_REGISTRATION* pRegistration,
+    __in BURN_VARIABLES* pVariables,
+    __out_z LPWSTR* psczInProgressBundleName
+    )
+{
+    HRESULT hr = S_OK;
+    LPCWSTR wzInProgressDisplayName = pRegistration->sczInProgressDisplayName ? pRegistration->sczInProgressDisplayName : L"";
+
+    hr = EnsureRegistrationVariable(pVariables, BURN_BUNDLE_INPROGRESS_NAME, wzInProgressDisplayName, psczInProgressBundleName);
+    ExitOnFailure(hr, "Failed to ensure in-progress bundle name.");
+
+LExit:
+    return hr;
+}
+
 static HRESULT GetBundleName(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __out LPWSTR* psczBundleName
+    __out_z LPWSTR* psczBundleName
     )
 {
     HRESULT hr = S_OK;
     LPCWSTR wzDisplayName = pRegistration->sczDisplayName ? pRegistration->sczDisplayName : L"";
 
-    hr = VariableGetString(pVariables, BURN_BUNDLE_NAME, psczBundleName);
+    hr = EnsureRegistrationVariable(pVariables, BURN_BUNDLE_NAME, wzDisplayName, psczBundleName);
+    ExitOnFailure(hr, "Failed to ensure bundle name.");
+
+LExit:
+    return hr;
+}
+
+static HRESULT EnsureRegistrationVariable(
+    __in BURN_VARIABLES* pVariables,
+    __in_z LPCWSTR wzVariable,
+    __in_z LPCWSTR wzDefaultValue,
+    __out_z LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    hr = VariableGetString(pVariables, wzVariable, psczValue);
     if (E_NOTFOUND == hr)
     {
-        hr = VariableSetString(pVariables, BURN_BUNDLE_NAME, wzDisplayName, FALSE, FALSE);
-        ExitOnFailure(hr, "Failed to set bundle name.");
+        hr = VariableSetString(pVariables, wzVariable, wzDefaultValue, FALSE, FALSE);
+        ExitOnFailure(hr, "Failed to set registration variable.");
 
-        hr = StrAllocString(psczBundleName, wzDisplayName, 0);
+        hr = StrAllocString(psczValue, wzDefaultValue, 0);
     }
-    ExitOnFailure(hr, "Failed to get bundle name.");
+    ExitOnFailure(hr, "Failed to get registration variable.");
 
 LExit:
     return hr;
@@ -1584,15 +1643,26 @@ LExit:
 static HRESULT UpdateBundleNameRegistration(
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_VARIABLES* pVariables,
-    __in HKEY hkRegistration
+    __in HKEY hkRegistration,
+    __in BOOL fInProgressRegistration
     )
 {
     HRESULT hr = S_OK;
     LPWSTR sczDisplayName = NULL;
 
-    // DisplayName: provided by UI
-    hr = GetBundleName(pRegistration, pVariables, &sczDisplayName);
-    hr = RegWriteString(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_DISPLAY_NAME, SUCCEEDED(hr) ? sczDisplayName : pRegistration->sczDisplayName);
+    if (fInProgressRegistration)
+    {
+        hr = GetBundleInProgressName(pRegistration, pVariables, &sczDisplayName);
+        ExitOnFailure(hr, "Failed to get bundle in-progress name.");
+    }
+
+    if (!sczDisplayName || !*sczDisplayName)
+    {
+        hr = GetBundleName(pRegistration, pVariables, &sczDisplayName);
+        ExitOnFailure(hr, "Failed to get bundle name.");
+    }
+
+    hr = RegWriteString(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_DISPLAY_NAME, sczDisplayName);
     ExitOnFailure(hr, "Failed to write %ls value.", BURN_REGISTRATION_REGISTRY_BUNDLE_DISPLAY_NAME);
 
 LExit:
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index 6d8a6d2a..936e5966 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -123,6 +123,7 @@ typedef struct _BURN_REGISTRATION
 
     // ARP registration
     LPWSTR sczDisplayName;
+    LPWSTR sczInProgressDisplayName;
     LPWSTR sczDisplayVersion;
     LPWSTR sczPublisher;
     LPWSTR sczHelpLink;
@@ -190,11 +191,13 @@ HRESULT RegistrationSessionBegin(
     __in BURN_VARIABLES* pVariables,
     __in DWORD dwRegistrationOptions,
     __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
-    __in DWORD64 qwEstimatedSize
+    __in DWORD64 qwEstimatedSize,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT RegistrationSessionResume(
     __in BURN_REGISTRATION* pRegistration,
-    __in BURN_VARIABLES* pVariables
+    __in BURN_VARIABLES* pVariables,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT RegistrationSessionEnd(
     __in BURN_REGISTRATION* pRegistration,
@@ -202,7 +205,8 @@ HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
-    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction
+    __in BURN_DEPENDENCY_REGISTRATION_ACTION dependencyRegistrationAction,
+    __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT RegistrationSaveState(
     __in BURN_REGISTRATION* pRegistration,
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 5a54f26e..07f4b831 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2119,7 +2119,8 @@ EXTERN_C BAAPI UserExperienceOnProgress(
 }
 
 EXTERN_C BAAPI UserExperienceOnRegisterBegin(
-    __in BURN_USER_EXPERIENCE* pUserExperience
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     )
 {
     HRESULT hr = S_OK;
@@ -2127,8 +2128,10 @@ EXTERN_C BAAPI UserExperienceOnRegisterBegin(
     BA_ONREGISTERBEGIN_RESULTS results = { };
 
     args.cbSize = sizeof(args);
+    args.recommendedRegistrationType = *pRegistrationType;
 
     results.cbSize = sizeof(results);
+    results.registrationType = *pRegistrationType;
 
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONREGISTERBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnRegisterBegin failed.");
@@ -2137,6 +2140,10 @@ EXTERN_C BAAPI UserExperienceOnRegisterBegin(
     {
         hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT);
     }
+    else if (BOOTSTRAPPER_REGISTRATION_TYPE_NONE < results.registrationType && BOOTSTRAPPER_REGISTRATION_TYPE_FULL >= results.registrationType)
+    {
+        *pRegistrationType = results.registrationType;
+    }
 
 LExit:
     return hr;
@@ -2316,7 +2323,7 @@ LExit:
 
 EXTERN_C BAAPI UserExperienceOnUnregisterBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __inout BOOL* pfKeepRegistration
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     )
 {
     HRESULT hr = S_OK;
@@ -2324,16 +2331,17 @@ EXTERN_C BAAPI UserExperienceOnUnregisterBegin(
     BA_ONUNREGISTERBEGIN_RESULTS results = { };
 
     args.cbSize = sizeof(args);
-    args.fKeepRegistration = *pfKeepRegistration;
+    args.recommendedRegistrationType = *pRegistrationType;
 
     results.cbSize = sizeof(results);
+    results.registrationType = *pRegistrationType;
 
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONUNREGISTERBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnUnregisterBegin failed.");
 
-    if (!args.fKeepRegistration && results.fForceKeepRegistration)
+    if (BOOTSTRAPPER_REGISTRATION_TYPE_NONE < results.registrationType && BOOTSTRAPPER_REGISTRATION_TYPE_FULL >= results.registrationType)
     {
-        *pfKeepRegistration = TRUE;
+        *pRegistrationType = results.registrationType;
     }
 
 LExit:
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index f2453dca..56bc3239 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -486,7 +486,8 @@ BAAPI UserExperienceOnProgress(
     __in DWORD dwOverallPercentage
     );
 BAAPI UserExperienceOnRegisterBegin(
-    __in BURN_USER_EXPERIENCE* pUserExperience
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     );
 BAAPI UserExperienceOnRegisterComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
@@ -522,7 +523,7 @@ BAAPI UserExperienceOnSystemShutdown(
     );
 BAAPI UserExperienceOnUnregisterBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
-    __inout BOOL* pfKeepRegistration
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
     );
 BAAPI UserExperienceOnUnregisterComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,