Make the estimated size in ARP a little more accurate.

Fixes 4039

7 files changed, 86 insertions, 58 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 9ee7b58c..69106f65 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -79,7 +79,8 @@ static HRESULT WINAPI AuthenticationRequired(
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BOOL fLog,
-    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType,
+    __inout DWORD64* pqwEstimatedSize
     );
 static HRESULT ExecuteDependentRegistrationActions(
     __in HANDLE hPipe,
@@ -422,8 +423,9 @@ extern "C" HRESULT ApplyRegister(
     HRESULT hr = S_OK;
     LPWSTR sczEngineWorkingPath = NULL;
     BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+    DWORD64 qwEstimatedSize = 0;
 
-    CalculateKeepRegistration(pEngineState, FALSE, &registrationType);
+    CalculateKeepRegistration(pEngineState, FALSE, &registrationType, &qwEstimatedSize);
 
     hr = UserExperienceOnRegisterBegin(&pEngineState->userExperience, &registrationType);
     ExitOnRootFailure(hr, "BA aborted register begin.");
@@ -451,12 +453,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->registration.fDetectedForeignProviderKeyBundleId, pEngineState->plan.qwEstimatedSize, registrationType);
+            hr = ElevationSessionBegin(pEngineState->companionConnection.hPipe, sczEngineWorkingPath, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->registration.fDetectedForeignProviderKeyBundleId, qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session in per-machine process.");
         }
         else
         {
-            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->cache, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.qwEstimatedSize, registrationType);
+            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->cache, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session.");
         }
     }
@@ -491,6 +493,7 @@ extern "C" HRESULT ApplyUnregister(
     BURN_RESUME_MODE resumeMode = BURN_RESUME_MODE_NONE;
     BOOTSTRAPPER_REGISTRATION_TYPE defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
     BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
+    DWORD64 qwEstimatedSize = 0;
 
     // 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.
@@ -513,7 +516,7 @@ extern "C" HRESULT ApplyUnregister(
         defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
     }
 
-    CalculateKeepRegistration(pEngineState, TRUE, &defaultRegistrationType);
+    CalculateKeepRegistration(pEngineState, TRUE, &defaultRegistrationType, &qwEstimatedSize);
 
     registrationType = defaultRegistrationType;
 
@@ -547,12 +550,12 @@ extern "C" HRESULT ApplyUnregister(
 
     if (pEngineState->registration.fPerMachine)
     {
-        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->registration.fDetectedForeignProviderKeyBundleId, registrationType);
+        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->registration.fDetectedForeignProviderKeyBundleId, qwEstimatedSize, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-machine process.");
     }
     else
     {
-        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->cache, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, registrationType);
+        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->cache, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, qwEstimatedSize, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-user process.");
     }
 
@@ -880,9 +883,12 @@ extern "C" void ApplyClean(
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BOOL fLog,
-    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType,
+    __inout DWORD64* pqwEstimatedSize
     )
 {
+    DWORD64 qwEstimatedSize = pEngineState->section.qwBundleSize;
+
     if (fLog)
     {
         LogId(REPORT_STANDARD, MSG_POST_APPLY_CALCULATE_REGISTRATION);
@@ -911,16 +917,29 @@ static void CalculateKeepRegistration(
         {
             *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_FULL;
 
-            if (!fLog)
+            if (BURN_PACKAGE_TYPE_MSP == pPackage->type)
             {
-                break;
+                qwEstimatedSize += pPackage->qwSize;
             }
+
+            qwEstimatedSize += pPackage->qwInstallSize;
         }
-        else if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState && BOOTSTRAPPER_REGISTRATION_TYPE_NONE == *pRegistrationType)
+
+        if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState)
         {
-            *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+            if (BOOTSTRAPPER_REGISTRATION_TYPE_NONE == *pRegistrationType)
+            {
+                *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+            }
+
+            qwEstimatedSize += pPackage->qwSize;
         }
     }
+
+    if (pqwEstimatedSize)
+    {
+        *pqwEstimatedSize = qwEstimatedSize;
+    }
 }
 
 static HRESULT ExecuteDependentRegistrationActions(
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 56a62955..3305cf62 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -636,6 +636,7 @@ extern "C" HRESULT ElevationSessionEnd(
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BOOL fDetectedForeignProviderKeyBundleId,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
@@ -654,6 +655,9 @@ extern "C" HRESULT ElevationSessionEnd(
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)fDetectedForeignProviderKeyBundleId);
     ExitOnFailure(hr, "Failed to write dependency registration action to message buffer.");
 
+    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.");
 
@@ -2580,6 +2584,7 @@ static HRESULT OnSessionEnd(
     SIZE_T iData = 0;
     DWORD dwResumeMode = 0;
     DWORD dwRestart = 0;
+    DWORD64 qwEstimatedSize = 0;
     DWORD dwRegistrationType = 0;
 
     // Deserialize message data.
@@ -2592,11 +2597,14 @@ static HRESULT OnSessionEnd(
     hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&pRegistration->fDetectedForeignProviderKeyBundleId);
     ExitOnFailure(hr, "Failed to read dependency registration action.");
 
+    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.");
 
     // suspend session in per-machine process
-    hr = RegistrationSessionEnd(pRegistration, pCache, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
+    hr = RegistrationSessionEnd(pRegistration, pCache, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, qwEstimatedSize, (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 3484057e..a74b6027 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -44,6 +44,7 @@ HRESULT ElevationSessionEnd(
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BOOL fDetectedForeignProviderKeyBundleId,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT ElevationSaveState(
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index 47da22c0..68c1d2ba 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -1178,31 +1178,6 @@ extern "C" HRESULT PlanExecutePackage(
         pPackage->rollback = BOOTSTRAPPER_ACTION_STATE_NONE;
     }
 
-    // Add the cache and install size to estimated size if it will be on the machine at the end of the install
-    if (BOOTSTRAPPER_REQUEST_STATE_PRESENT == pPackage->requested ||
-        fRequestedCache ||
-        (BOOTSTRAPPER_PACKAGE_STATE_PRESENT == pPackage->currentState && BOOTSTRAPPER_REQUEST_STATE_ABSENT < pPackage->requested)
-       )
-    {
-        // If the package will remain in the cache, add the package size to the estimated size
-        if (BOOTSTRAPPER_CACHE_TYPE_REMOVE < pPackage->cacheType)
-        {
-            pPlan->qwEstimatedSize += pPackage->qwSize;
-        }
-
-        // If the package will end up installed on the machine, add the install size to the estimated size.
-        if (BOOTSTRAPPER_REQUEST_STATE_CACHE < pPackage->requested)
-        {
-            // MSP packages get cached automatically by windows installer with any embedded cabs, so include that in the size as well
-            if (BURN_PACKAGE_TYPE_MSP == pPackage->type)
-            {
-                pPlan->qwEstimatedSize += pPackage->qwSize;
-            }
-
-            pPlan->qwEstimatedSize += pPackage->qwInstallSize;
-        }
-    }
-
     // Add execute actions.
     switch (pPackage->type)
     {
@@ -3079,7 +3054,6 @@ extern "C" void PlanDump(
     LogStringLine(PlanDumpLevel, "     disallow-removal: %hs", LoggingTrueFalseToString(pPlan->fDisallowRemoval));
     LogStringLine(PlanDumpLevel, "     downgrade: %hs", LoggingTrueFalseToString(pPlan->fDowngrade));
     LogStringLine(PlanDumpLevel, "     registration options: %hs", LoggingRegistrationOptionsToString(pPlan->dwRegistrationOperations));
-    LogStringLine(PlanDumpLevel, "     estimated size: %llu", pPlan->qwEstimatedSize);
     if (pPlan->sczLayoutDirectory)
     {
         LogStringLine(PlanDumpLevel, "     layout directory: %ls", pPlan->sczLayoutDirectory);
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index fbb7a5a1..2b04e097 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -267,8 +267,6 @@ typedef struct _BURN_PLAN
 
     DWORD64 qwCacheSizeTotal;
 
-    DWORD64 qwEstimatedSize;
-
     DWORD cExecutePackagesTotal;
     DWORD cOverallProgressTicksTotal;
 
diff --git a/src/burn/engine/registration.cpp b/src/burn/engine/registration.cpp
index 0ffbc89e..9953be83 100644
--- a/src/burn/engine/registration.cpp
+++ b/src/burn/engine/registration.cpp
@@ -102,6 +102,10 @@ static HRESULT UpdateBundleNameRegistration(
     __in HKEY hkRegistration,
     __in BOOL fInProgressRegistration
     );
+static HRESULT UpdateEstimatedSize(
+    __in HKEY hkRegistration,
+    __in DWORD64 qwEstimatedSize
+    );
 static BOOL IsWuRebootPending();
 static BOOL IsRegistryRebootPending();
 
@@ -600,8 +604,8 @@ extern "C" HRESULT RegistrationSessionBegin(
     )
 {
     HRESULT hr = S_OK;
-    DWORD dwSize = 0;
     HKEY hkRegistration = NULL;
+    BOOL fCreated = FALSE;
     LPWSTR sczPublisher = NULL;
 
     AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE != registrationType, "Registration type can't be NONE");
@@ -620,7 +624,7 @@ extern "C" HRESULT RegistrationSessionBegin(
     }
 
     // create registration key
-    hr = RegCreate(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, &hkRegistration);
+    hr = RegCreateEx(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, REG_KEY_DEFAULT, FALSE, NULL, &hkRegistration, &fCreated);
     ExitOnFailure(hr, "Failed to create registration key.");
 
     // Write any ARP values and software tags.
@@ -797,22 +801,12 @@ extern "C" HRESULT RegistrationSessionBegin(
         ExitOnFailure(hr, "Failed to write update registration.");
     }
 
-    // Update estimated size.
-    qwEstimatedSize /= 1024; // Convert bytes to KB
-    if (0 < qwEstimatedSize)
+    // Only set estimated size here for the first time.
+    // It will always get updated at the end of the session.
+    if (fCreated)
     {
-        if (DWORD_MAX < qwEstimatedSize)
-        {
-            // ARP doesn't support QWORDs here
-            dwSize = DWORD_MAX;
-        }
-        else
-        {
-            dwSize = static_cast<DWORD>(qwEstimatedSize);
-        }
-
-        hr = RegWriteNumber(hkRegistration, REGISTRY_BUNDLE_ESTIMATED_SIZE, dwSize);
-        ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_ESTIMATED_SIZE);
+        hr = UpdateEstimatedSize(hkRegistration, qwEstimatedSize);
+        ExitOnFailure(hr, "Failed to update estimated size.");
     }
 
     // Register the bundle dependency key.
@@ -879,6 +873,7 @@ extern "C" HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
@@ -921,6 +916,9 @@ extern "C" HRESULT RegistrationSessionEnd(
         // update display name
         hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
         ExitOnFailure(hr, "Failed to update name and publisher.");
+
+        hr = UpdateEstimatedSize(hkRegistration, qwEstimatedSize);
+        ExitOnFailure(hr, "Failed to update estimated size.");
     }
 
     // Update resume mode.
@@ -1604,6 +1602,35 @@ LExit:
     return hr;
 }
 
+static HRESULT UpdateEstimatedSize(
+    __in HKEY hkRegistration,
+    __in DWORD64 qwEstimatedSize
+    )
+{
+    HRESULT hr = S_OK;
+    DWORD dwSize = 0;
+
+    qwEstimatedSize /= 1024; // Convert bytes to KB
+    if (0 < qwEstimatedSize)
+    {
+        if (DWORD_MAX < qwEstimatedSize)
+        {
+            // ARP doesn't support QWORDs here
+            dwSize = DWORD_MAX;
+        }
+        else
+        {
+            dwSize = static_cast<DWORD>(qwEstimatedSize);
+        }
+
+        hr = RegWriteNumber(hkRegistration, REGISTRY_BUNDLE_ESTIMATED_SIZE, dwSize);
+        ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_ESTIMATED_SIZE);
+    }
+
+LExit:
+    return hr;
+}
+
 static BOOL IsWuRebootPending()
 {
     HRESULT hr = S_OK;
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index c07be962..cb5349a3 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -213,6 +213,7 @@ HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT RegistrationSaveState(