9 files changed, 98 insertions, 16 deletions

diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index dc0af8ce..5483e46f 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -469,7 +469,7 @@ extern "C" HRESULT CorePlan(
     pEngineState->plan.fDisableRollback = pEngineState->fDisableRollback || BOOTSTRAPPER_ACTION_UNSAFE_UNINSTALL == pEngineState->plan.action;
     pEngineState->plan.fPlanPackageCacheRollback = BOOTSTRAPPER_REGISTRATION_TYPE_NONE == pEngineState->registration.detectedRegistrationType;
 
-    hr = PlanPackagesAndBundleScope(pEngineState->packages.rgPackages, pEngineState->packages.cPackages, pEngineState->plan.plannedScope, pEngineState->registration.scope, pEngineState->command.commandLineScope, pEngineState->registration.detectedScope, &pEngineState->plan.plannedScope, &pEngineState->registration.fPerMachine);
+    hr = PlanPackagesAndBundleScope(pEngineState->packages.rgPackages, pEngineState->packages.cPackages, pEngineState->registration.sczPrimaryUpgradeCode, pEngineState->registration.relatedBundles.rgRelatedBundles, pEngineState->registration.relatedBundles.cRelatedBundles, pEngineState->plan.plannedScope, pEngineState->registration.scope, pEngineState->command.commandLineScope, pEngineState->registration.detectedScope, &pEngineState->plan.plannedScope, &pEngineState->registration.fPerMachine);
     ExitOnFailure(hr, "Failed to determine packages and bundle scope.");
 
     if (BOOTSTRAPPER_PACKAGE_SCOPE_PER_MACHINE_OR_PER_USER == pEngineState->registration.scope || BOOTSTRAPPER_PACKAGE_SCOPE_PER_USER_OR_PER_MACHINE == pEngineState->registration.scope)
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 520fdcd5..2ad9111f 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -235,11 +235,6 @@ HRESULT CoreSerializeEngineState(
 HRESULT CoreQueryRegistration(
     __in BURN_ENGINE_STATE* pEngineState
     );
-//HRESULT CoreDeserializeEngineState(
-//    __in BURN_ENGINE_STATE* pEngineState,
-//    __in_bcount(cbBuffer) BYTE* pbBuffer,
-//    __in SIZE_T cbBuffer
-//    );
 HRESULT CoreDetect(
     __in BURN_ENGINE_STATE* pEngineState,
     __in_opt HWND hwndParent
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index f7ec0359..edafef97 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -418,7 +418,14 @@ MessageId=227
 Severity=Success
 SymbolicName=MSG_PLAN_INSTALLED_SCOPE
 Language=English
-Bundle was already installed with scope: %1!hs!
+Bundle was already installed with scope: %1!hs!. Scope cannot change during maintenance.
+.
+
+MessageId=228
+Severity=Success
+SymbolicName=MSG_PLAN_UPGRADE_SCOPE
+Language=English
+Upgraded bundle %1!ls! was already installed with scope: %2!hs!. Scope cannot change during upgrade.
 .
 
 MessageId=201
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index 889ad68c..2257b28c 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -144,6 +144,12 @@ static BOOL ForceCache(
     __in BURN_PLAN* pPlan,
     __in BURN_PACKAGE* pPackage
     );
+static HRESULT GetUpgradedBundleScope(
+    __in DWORD cRelatedBundles,
+    __in BURN_RELATED_BUNDLE* rgRelatedBundles,
+    __in_z LPCWSTR wzUpgradeCode,
+    __inout BOOTSTRAPPER_SCOPE* pScope
+);
 
 // function definitions
 
@@ -823,6 +829,9 @@ LExit:
 extern "C" HRESULT PlanPackagesAndBundleScope(
     __in BURN_PACKAGE* rgPackages,
     __in DWORD cPackages,
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BURN_RELATED_BUNDLE* rgRelatedBundles,
+    __in DWORD cRelatedBundles,
     __in BOOTSTRAPPER_SCOPE scope,
     __in BOOTSTRAPPER_PACKAGE_SCOPE authoredScope,
     __in BOOTSTRAPPER_SCOPE commandLineScope,
@@ -855,7 +864,11 @@ extern "C" HRESULT PlanPackagesAndBundleScope(
         }
     }
 
-    if (BOOTSTRAPPER_SCOPE_DEFAULT != detectedScope)
+    // If we're upgrading, lock the scope to that of the upgraded bundle.
+    hr = GetUpgradedBundleScope(cRelatedBundles, rgRelatedBundles, wzUpgradeCode, &scope);
+
+    // If we're in maintenance mode instead, lock the scope to the original scope.
+    if (S_FALSE == hr && BOOTSTRAPPER_SCOPE_DEFAULT != detectedScope)
     {
         scope = detectedScope;
 
@@ -888,6 +901,48 @@ extern "C" HRESULT PlanPackagesAndBundleScope(
 }
 
 
+static HRESULT GetUpgradedBundleScope(
+    __in DWORD cRelatedBundles,
+    __in BURN_RELATED_BUNDLE* rgRelatedBundles,
+    __in_z LPCWSTR wzUpgradeCode,
+    __inout BOOTSTRAPPER_SCOPE* pScope
+    )
+{
+    HRESULT hr = S_OK;
+
+    for (DWORD i = 0; i < cRelatedBundles; ++i)
+    {
+        BURN_RELATED_BUNDLE* pRelatedBundle = rgRelatedBundles + i;
+
+        if (BOOTSTRAPPER_RELATION_UPGRADE == pRelatedBundle->detectRelationType)
+        {
+            for (DWORD j = 0; j < pRelatedBundle->package.Bundle.cUpgradeCodes; ++j)
+            {
+                LPCWSTR wzRelatedUpgradeCode = *(pRelatedBundle->package.Bundle.rgsczUpgradeCodes + j);
+
+                // Is the related bundle's upgrade code the same as ours?
+                // If so, lock our scope to the "original" bundle's scope.
+                if (CSTR_EQUAL == ::CompareStringOrdinal(wzRelatedUpgradeCode, -1, wzUpgradeCode, -1, FALSE))
+                if (CSTR_EQUAL == ::CompareStringOrdinal(wzRelatedUpgradeCode, -1, wzUpgradeCode, -1, TRUE))
+                {
+                    *pScope = pRelatedBundle->detectedScope;
+
+                    LogId(REPORT_STANDARD, MSG_PLAN_UPGRADE_SCOPE, pRelatedBundle->package.Bundle.sczBundleCode, LoggingBundleScopeToString(*pScope));
+
+                    ExitFunction();
+                }
+            }
+        }
+    }
+
+    // No upgrade codes or none match, which is fine. But note it via S_FALSE
+    // so we can distinguish the upgrade case from the maintenance case.
+    hr = S_FALSE;
+
+LExit:
+    return hr;
+}
+
 static HRESULT PlanPackagesHelper(
     __in BURN_PACKAGE* rgPackages,
     __in DWORD cPackages,
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index f9996ac0..09918440 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -485,6 +485,9 @@ void PlanDump(
 HRESULT PlanPackagesAndBundleScope(
     __in BURN_PACKAGE* rgPackages,
     __in DWORD cPackages,
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BURN_RELATED_BUNDLE* rgRelatedBundles,
+    __in DWORD cRelatedBundles,
     __in BOOTSTRAPPER_SCOPE scope,
     __in BOOTSTRAPPER_PACKAGE_SCOPE authoredScope,
     __in BOOTSTRAPPER_SCOPE commandLineScope,
diff --git a/src/burn/engine/registration.cpp b/src/burn/engine/registration.cpp
index e588bef0..7c9ca182 100644
--- a/src/burn/engine/registration.cpp
+++ b/src/burn/engine/registration.cpp
@@ -29,9 +29,10 @@ const LPCWSTR REGISTRY_BUNDLE_UNINSTALL_STRING = L"UninstallString";
 const LPCWSTR REGISTRY_BUNDLE_RESUME_COMMAND_LINE = L"BundleResumeCommandLine";
 const LPCWSTR REGISTRY_BUNDLE_VERSION_MAJOR = L"VersionMajor";
 const LPCWSTR REGISTRY_BUNDLE_VERSION_MINOR = L"VersionMinor";
-const LPCWSTR REGISTRY_BUNDLE_SCOPE = L"BundleScope";
 const LPCWSTR SWIDTAG_FOLDER = L"swidtag";
 const LPCWSTR REGISTRY_BUNDLE_VARIABLE_KEY = L"variables";
+const LPCWSTR REGISTRY_BUNDLE_INSTALLED = L"Installed";
+
 
 // internal function declarations
 
@@ -146,6 +147,10 @@ extern "C" HRESULT RegistrationParseFromXml(
     hr = XmlGetAttributeEx(pixnRegistrationNode, L"Tag", &pRegistration->sczTag);
     ExitOnRequiredXmlQueryFailure(hr, "Failed to get @Tag.");
 
+    // @PrimaryUpgradeCode
+    hr = XmlGetAttributeEx(pixnRegistrationNode, L"PrimaryUpgradeCode", &pRegistration->sczPrimaryUpgradeCode);
+    ExitOnOptionalXmlQueryFailure(hr, fFoundXml, "Failed to get @PrimaryUpgradeCode.");
+
     hr = BundlePackageEngineParseRelatedCodes(pixnBundle, &pRegistration->rgsczDetectCodes, &pRegistration->cDetectCodes, &pRegistration->rgsczUpgradeCodes, &pRegistration->cUpgradeCodes, &pRegistration->rgsczAddonCodes, &pRegistration->cAddonCodes, &pRegistration->rgsczPatchCodes, &pRegistration->cPatchCodes);
     ExitOnFailure(hr, "Failed to parse related bundles");
 
@@ -327,7 +332,8 @@ extern "C" void RegistrationUninitialize(
 {
     ReleaseStr(pRegistration->sczCode);
     ReleaseStr(pRegistration->sczTag);
-
+    ReleaseStr(pRegistration->sczPrimaryUpgradeCode);
+    
     for (DWORD i = 0; i < pRegistration->cDetectCodes; ++i)
     {
         ReleaseStr(pRegistration->rgsczDetectCodes[i]);
@@ -666,8 +672,8 @@ extern "C" HRESULT RegistrationSessionBegin(
     hr = RegWriteStringFormatted(hkRegistration, REGISTRY_BUNDLE_DISPLAY_ICON, L"%s,0", pRegistration->sczCacheExecutablePath);
     ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_DISPLAY_ICON);
 
-    hr = RegWriteNumber(hkRegistration, REGISTRY_BUNDLE_SCOPE, pRegistration->fPerMachine ? BOOTSTRAPPER_SCOPE_PER_MACHINE : BOOTSTRAPPER_SCOPE_PER_USER);
-    ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_SCOPE);
+    hr = RegWriteNumber(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE, pRegistration->fPerMachine ? BOOTSTRAPPER_SCOPE_PER_MACHINE : BOOTSTRAPPER_SCOPE_PER_USER);
+    ExitOnFailure(hr, "Failed to write %ls value.", BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE);
 
     // update display name
     hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
@@ -1780,8 +1786,8 @@ static HRESULT DetectInstalled(
 
         pRegistration->detectedRegistrationType = (1 == dwInstalled) ? BOOTSTRAPPER_REGISTRATION_TYPE_FULL : BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
 
-        hr = RegReadNumber(hkRegistration, REGISTRY_BUNDLE_SCOPE, &dwScope);
-        ExitOnFailure(hr, "Failed to read registration %ls@%ls.", pRegistration->sczRegistrationKey, REGISTRY_BUNDLE_SCOPE);
+        hr = RegReadNumber(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE, &dwScope);
+        ExitOnFailure(hr, "Failed to read registration %ls@%ls.", pRegistration->sczRegistrationKey, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE);
 
         pRegistration->detectedScope = static_cast<BOOTSTRAPPER_SCOPE>(dwScope);
     }
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index f340999a..97c6951b 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -21,8 +21,7 @@ const LPCWSTR BURN_REGISTRATION_REGISTRY_ENGINE_VERSION = L"EngineVersion";
 const LPCWSTR BURN_REGISTRATION_REGISTRY_ENGINE_PROTOCOL_VERSION = L"EngineProtocolVersion";
 const LPCWSTR BURN_REGISTRATION_REGISTRY_BUNDLE_PROVIDER_KEY = L"BundleProviderKey";
 const LPCWSTR BURN_REGISTRATION_REGISTRY_BUNDLE_TAG = L"BundleTag";
-
-const LPCWSTR REGISTRY_BUNDLE_INSTALLED = L"Installed";
+const LPCWSTR BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE = L"BundleScope";
 
 enum BURN_RESUME_MODE
 {
@@ -69,6 +68,7 @@ typedef struct _BURN_RELATED_BUNDLE
     BOOTSTRAPPER_REQUEST_STATE defaultRequestedRestore;
     BOOTSTRAPPER_REQUEST_STATE requestedRestore;
     BOOTSTRAPPER_ACTION_STATE restore;
+    BOOTSTRAPPER_SCOPE detectedScope;
 } BURN_RELATED_BUNDLE;
 
 typedef struct _BURN_RELATED_BUNDLES
@@ -106,6 +106,7 @@ typedef struct _BURN_REGISTRATION
     BOOTSTRAPPER_PACKAGE_SCOPE scope;
     LPWSTR sczCode;
     LPWSTR sczTag;
+    LPWSTR sczPrimaryUpgradeCode;
 
     LPWSTR *rgsczDetectCodes;
     DWORD cDetectCodes;
diff --git a/src/burn/engine/relatedbundle.cpp b/src/burn/engine/relatedbundle.cpp
index d0b97af0..166c4ae8 100644
--- a/src/burn/engine/relatedbundle.cpp
+++ b/src/burn/engine/relatedbundle.cpp
@@ -325,6 +325,7 @@ static HRESULT LoadRelatedBundleFromKey(
     BOOL fExists = FALSE;
     BURN_DEPENDENCY_PROVIDER dependencyProvider = { };
     BURN_DEPENDENCY_PROVIDER* pBundleDependencyProvider = NULL;
+    DWORD dwScope = 0;
 
     // Only support progress from engines that are compatible.
     hr = RegReadNumber(hkBundleCode, BURN_REGISTRATION_REGISTRY_ENGINE_PROTOCOL_VERSION, &dwEngineProtocolVersion);
@@ -355,6 +356,11 @@ static HRESULT LoadRelatedBundleFromKey(
         LogId(REPORT_WARNING, MSG_RELATED_PACKAGE_INVALID_VERSION, wzRelatedBundleCode, sczBundleVersion);
     }
 
+    hr = RegReadNumber(hkBundleCode, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE, &dwScope);
+    ExitOnFailure(hr, "Failed to read registration %ls for bundle %ls.", wzRelatedBundleCode, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE);
+
+    pRelatedBundle->detectedScope = static_cast<BOOTSTRAPPER_SCOPE>(dwScope);
+
     hr = RegReadString(hkBundleCode, BURN_REGISTRATION_REGISTRY_BUNDLE_CACHE_PATH, &sczCachePath);
     ExitOnFailure(hr, "Failed to read cache path from registry for bundle: %ls", wzRelatedBundleCode);
 
@@ -390,6 +396,12 @@ static HRESULT LoadRelatedBundleFromKey(
 
     pRelatedBundle->detectRelationType = relationType;
 
+    hr = StrAllocString(&pRelatedBundle->package.Bundle.sczBundleCode, wzRelatedBundleCode, 0);
+    ExitOnFailure(hr, "Failed to bundle code to related bundle.");
+
+    hr = RegReadStringArray(hkBundleCode, BURN_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &pRelatedBundle->package.Bundle.rgsczUpgradeCodes, &pRelatedBundle->package.Bundle.cUpgradeCodes);
+    ExitOnFailure(hr, "Failed to read upgrade codes.");
+
     hr = PseudoBundleInitializeRelated(&pRelatedBundle->package, fSupportsBurnProtocol, fPerMachine, wzRelatedBundleCode,
 #ifdef DEBUG
                                        pRelatedBundle->detectRelationType,
diff --git a/src/burn/test/BurnUnitTest/RelatedBundleTest.cpp b/src/burn/test/BurnUnitTest/RelatedBundleTest.cpp
index dcccc589..38c99121 100644
--- a/src/burn/test/BurnUnitTest/RelatedBundleTest.cpp
+++ b/src/burn/test/BurnUnitTest/RelatedBundleTest.cpp
@@ -178,6 +178,9 @@ namespace Bootstrapper
 
                 hr = RegWriteString(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_VERSION, wzVersion);
                 NativeAssert::Succeeded(hr, "Failed to write %ls value.", BURN_REGISTRATION_REGISTRY_BUNDLE_VERSION);
+
+                hr = RegWriteNumber(hkRegistration, BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE, fPerMachine ? BOOTSTRAPPER_SCOPE_PER_MACHINE : BOOTSTRAPPER_SCOPE_PER_USER);
+                NativeAssert::Succeeded(hr, "Failed to write %ls value.", BURN_REGISTRATION_REGISTRY_BUNDLE_SCOPE);
             }
             finally
             {