1 files changed, 138 insertions, 101 deletions

diff --git a/src/libs/dutil/WixToolset.DUtil/butil.cpp b/src/libs/dutil/WixToolset.DUtil/butil.cpp
index 4c96dfc1..4262d573 100644
--- a/src/libs/dutil/WixToolset.DUtil/butil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/butil.cpp
@@ -55,6 +55,7 @@ static HRESULT OpenBundleKey(
     __in_z LPCWSTR wzBundleId,
     __in BUNDLE_INSTALL_CONTEXT context,
     __in_opt LPCWSTR wzSubKey,
+    __in REG_KEY_BITNESS kbKeyBitness,
     __inout HKEY* phKey
     );
 static HRESULT CopyStringToBuffer(
@@ -62,6 +63,14 @@ static HRESULT CopyStringToBuffer(
     __in_z_opt LPWSTR wzBuffer,
     __inout SIZE_T* pcchBuffer
     );
+static HRESULT DoBundleEnumRelatedBundle(
+    __in HKEY hkRoot,
+    __in REG_KEY_BITNESS kbKeyBitness,
+    __in_z LPCWSTR wzUpgradeCode,
+    __inout PDWORD pdwStartIndex,
+    __deref_out_z LPWSTR* psczBundleId
+    );
+
 
 DAPI_(HRESULT) BundleGetBundleInfo(
     __in_z LPCWSTR wzBundleId,
@@ -153,113 +162,26 @@ DAPI_(HRESULT) BundleEnumRelatedBundle(
 {
     HRESULT hr = S_OK;
     HKEY hkRoot = BUNDLE_INSTALL_CONTEXT_USER == context ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE;
-    HKEY hkUninstall = NULL;
-    HKEY hkBundle = NULL;
-    LPWSTR sczUninstallSubKey = NULL;
-    LPWSTR sczUninstallSubKeyPath = NULL;
-    LPWSTR sczValue = NULL;
-    DWORD dwType = 0;
-
-    LPWSTR* rgsczBundleUpgradeCodes = NULL;
-    DWORD cBundleUpgradeCodes = 0;
-    BOOL fUpgradeCodeFound = FALSE;
 
     if (!wzUpgradeCode || !pdwStartIndex)
     {
         ButilExitOnFailure(hr = E_INVALIDARG, "An invalid parameter was passed to the function.");
     }
 
-    hr = RegOpen(hkRoot, BUNDLE_REGISTRATION_REGISTRY_UNINSTALL_KEY, KEY_READ, &hkUninstall);
-    ButilExitOnFailure(hr, "Failed to open bundle uninstall key path.");
-
-    for (DWORD dwIndex = *pdwStartIndex; !fUpgradeCodeFound; dwIndex++)
+    hr = DoBundleEnumRelatedBundle(hkRoot, REG_KEY_DEFAULT, wzUpgradeCode, pdwStartIndex, psczBundleId);
+    ButilExitOnFailure(hr, "Failed to enumerate default-bitness bundles.");
+    if (S_FALSE == hr)
     {
-        hr = RegKeyEnum(hkUninstall, dwIndex, &sczUninstallSubKey);
-        ButilExitOnFailure(hr, "Failed to enumerate bundle uninstall key path.");
-
-        hr = StrAllocFormatted(&sczUninstallSubKeyPath, L"%ls\\%ls", BUNDLE_REGISTRATION_REGISTRY_UNINSTALL_KEY, sczUninstallSubKey);
-        ButilExitOnFailure(hr, "Failed to allocate bundle uninstall key path.");
-        
-        hr = RegOpen(hkRoot, sczUninstallSubKeyPath, KEY_READ, &hkBundle);
-        ButilExitOnFailure(hr, "Failed to open uninstall key path.");
-
-        // If it's a bundle, it should have a BundleUpgradeCode value of type REG_SZ (old) or REG_MULTI_SZ
-        hr = RegGetType(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &dwType);
-        if (FAILED(hr))
-        {
-            ReleaseRegKey(hkBundle);
-            ReleaseNullStr(sczUninstallSubKey);
-            ReleaseNullStr(sczUninstallSubKeyPath);
-            // Not a bundle
-            continue;
-        }
-
-        switch (dwType)
-        {
-            case REG_SZ:
-                hr = RegReadString(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &sczValue);
-                ButilExitOnFailure(hr, "Failed to read BundleUpgradeCode string property.");
-
-                if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, sczValue, -1, wzUpgradeCode, -1))
-                {
-                    *pdwStartIndex = dwIndex;
-                    fUpgradeCodeFound = TRUE;
-                    break;
-                }
-
-                ReleaseNullStr(sczValue);
-
-                break;
-            case REG_MULTI_SZ:
-                hr = RegReadStringArray(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &rgsczBundleUpgradeCodes, &cBundleUpgradeCodes);
-                ButilExitOnFailure(hr, "Failed to read BundleUpgradeCode  multi-string property.");
-
-                for (DWORD i = 0; i < cBundleUpgradeCodes; i++)
-                {
-                    LPWSTR wzBundleUpgradeCode = rgsczBundleUpgradeCodes[i];
-                    if (wzBundleUpgradeCode && *wzBundleUpgradeCode)
-                    {
-                        if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, wzBundleUpgradeCode, -1, wzUpgradeCode, -1))
-                        {
-                            *pdwStartIndex = dwIndex;
-                            fUpgradeCodeFound = TRUE;
-                            break;
-                        }
-                    }
-                }
-                ReleaseNullStrArray(rgsczBundleUpgradeCodes, cBundleUpgradeCodes);
-
-                break;
-
-            default:
-                ButilExitWithRootFailure(hr, E_NOTIMPL, "BundleUpgradeCode of type 0x%x not implemented.", dwType);
-        }
-
-        if (fUpgradeCodeFound)
-        {
-            if (psczBundleId)
-            {
-                *psczBundleId = sczUninstallSubKey;
-                sczUninstallSubKey = NULL;
-            }
-
-            break;
-        }
-
-        // Cleanup before next iteration
-        ReleaseRegKey(hkBundle);
-        ReleaseNullStr(sczUninstallSubKey);
-        ReleaseNullStr(sczUninstallSubKeyPath);
+#if defined(_WIN64)
+        hr = DoBundleEnumRelatedBundle(hkRoot, REG_KEY_32BIT, wzUpgradeCode, pdwStartIndex, psczBundleId);
+        ButilExitOnFailure(hr, "Failed to enumerate 32-bit bundles.");
+#else
+        hr = DoBundleEnumRelatedBundle(hkRoot, REG_KEY_64BIT, wzUpgradeCode, pdwStartIndex, psczBundleId);
+        ButilExitOnFailure(hr, "Failed to enumerate 64-bit bundles.");
+#endif
     }
 
 LExit:
-    ReleaseStr(sczValue);
-    ReleaseStr(sczUninstallSubKey);
-    ReleaseStr(sczUninstallSubKeyPath);
-    ReleaseRegKey(hkBundle);
-    ReleaseRegKey(hkUninstall);
-    ReleaseStrArray(rgsczBundleUpgradeCodes, cBundleUpgradeCodes);
-
     return hr;
 }
 
@@ -276,7 +198,7 @@ DAPI_(HRESULT) BundleEnumRelatedBundleFixed(
     size_t cchValue = 0;
 
     hr = BundleEnumRelatedBundle(wzUpgradeCode, context, pdwStartIndex, &sczValue);
-    if (SUCCEEDED(hr) && wzBundleId)
+    if (S_OK == hr && wzBundleId)
     {
         hr = ::StringCchLengthW(sczValue, STRSAFE_MAX_CCH, &cchValue);
         ButilExitOnRootFailure(hr, "Failed to calculate length of string.");
@@ -380,8 +302,9 @@ static HRESULT LocateAndQueryBundleValue(
 
     *pStatus = INTERNAL_BUNDLE_STATUS_SUCCESS;
 
-    if (FAILED(hr = OpenBundleKey(wzBundleId, BUNDLE_INSTALL_CONTEXT_MACHINE, wzSubKey, phKey)) &&
-        FAILED(hr = OpenBundleKey(wzBundleId, BUNDLE_INSTALL_CONTEXT_USER, wzSubKey, phKey)))
+    if (FAILED(hr = OpenBundleKey(wzBundleId, BUNDLE_INSTALL_CONTEXT_MACHINE, wzSubKey, REG_KEY_32BIT, phKey)) &&
+        FAILED(hr = OpenBundleKey(wzBundleId, BUNDLE_INSTALL_CONTEXT_MACHINE, wzSubKey, REG_KEY_64BIT, phKey)) &&
+        FAILED(hr = OpenBundleKey(wzBundleId, BUNDLE_INSTALL_CONTEXT_USER, wzSubKey, REG_KEY_DEFAULT, phKey)))
     {
         if (E_FILENOTFOUND == hr)
         {
@@ -413,6 +336,7 @@ static HRESULT OpenBundleKey(
     __in_z LPCWSTR wzBundleId,
     __in BUNDLE_INSTALL_CONTEXT context,
     __in_opt LPCWSTR wzSubKey,
+    __in REG_KEY_BITNESS kbKeyBitness,
     __inout HKEY* phKey
     )
 {
@@ -433,7 +357,7 @@ static HRESULT OpenBundleKey(
     }
     ButilExitOnFailure(hr, "Failed to allocate bundle uninstall key path.");
 
-    hr = RegOpen(hkRoot, sczKeypath, KEY_READ, phKey);
+    hr = RegOpenEx(hkRoot, sczKeypath, KEY_READ, kbKeyBitness, phKey);
     ButilExitOnFailure(hr, "Failed to open bundle uninstall key path.");
 
 LExit:
@@ -472,3 +396,116 @@ static HRESULT CopyStringToBuffer(
 
     return hr;
 }
+
+static HRESULT DoBundleEnumRelatedBundle(
+    __in HKEY hkRoot,
+    __in REG_KEY_BITNESS kbKeyBitness,
+    __in_z LPCWSTR wzUpgradeCode,
+    __inout PDWORD pdwStartIndex,
+    __deref_out_z LPWSTR* psczBundleId
+)
+{
+    HRESULT hr = S_OK;
+    BOOL fUpgradeCodeFound = FALSE;
+    HKEY hkUninstall = NULL;
+    HKEY hkBundle = NULL;
+    LPWSTR sczUninstallSubKey = NULL;
+    LPWSTR sczUninstallSubKeyPath = NULL;
+    LPWSTR sczValue = NULL;
+    DWORD dwType = 0;
+    LPWSTR* rgsczBundleUpgradeCodes = NULL;
+    DWORD cBundleUpgradeCodes = 0;
+
+    hr = RegOpenEx(hkRoot, BUNDLE_REGISTRATION_REGISTRY_UNINSTALL_KEY, KEY_READ, kbKeyBitness, &hkUninstall);
+    ButilExitOnFailure(hr, "Failed to open bundle uninstall key path.");
+
+    for (DWORD dwIndex = *pdwStartIndex; !fUpgradeCodeFound; dwIndex++)
+    {
+        hr = RegKeyEnum(hkUninstall, dwIndex, &sczUninstallSubKey);
+        ButilExitOnFailure(hr, "Failed to enumerate bundle uninstall key path.");
+
+        hr = StrAllocFormatted(&sczUninstallSubKeyPath, L"%ls\\%ls", BUNDLE_REGISTRATION_REGISTRY_UNINSTALL_KEY, sczUninstallSubKey);
+        ButilExitOnFailure(hr, "Failed to allocate bundle uninstall key path.");
+
+        hr = RegOpenEx(hkRoot, sczUninstallSubKeyPath, KEY_READ, kbKeyBitness, &hkBundle);
+        ButilExitOnFailure(hr, "Failed to open uninstall key path.");
+
+        // If it's a bundle, it should have a BundleUpgradeCode value of type REG_SZ (old) or REG_MULTI_SZ
+        hr = RegGetType(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &dwType);
+        if (FAILED(hr))
+        {
+            ReleaseRegKey(hkBundle);
+            ReleaseNullStr(sczUninstallSubKey);
+            ReleaseNullStr(sczUninstallSubKeyPath);
+            // Not a bundle
+            continue;
+        }
+
+        switch (dwType)
+        {
+        case REG_SZ:
+            hr = RegReadString(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &sczValue);
+            ButilExitOnFailure(hr, "Failed to read BundleUpgradeCode string property.");
+
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, sczValue, -1, wzUpgradeCode, -1))
+            {
+                *pdwStartIndex = dwIndex;
+                fUpgradeCodeFound = TRUE;
+                break;
+            }
+
+            ReleaseNullStr(sczValue);
+
+            break;
+        case REG_MULTI_SZ:
+            hr = RegReadStringArray(hkBundle, BUNDLE_REGISTRATION_REGISTRY_BUNDLE_UPGRADE_CODE, &rgsczBundleUpgradeCodes, &cBundleUpgradeCodes);
+            ButilExitOnFailure(hr, "Failed to read BundleUpgradeCode multi-string property.");
+
+            for (DWORD i = 0; i < cBundleUpgradeCodes; i++)
+            {
+                LPWSTR wzBundleUpgradeCode = rgsczBundleUpgradeCodes[i];
+                if (wzBundleUpgradeCode && *wzBundleUpgradeCode)
+                {
+                    if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, wzBundleUpgradeCode, -1, wzUpgradeCode, -1))
+                    {
+                        *pdwStartIndex = dwIndex;
+                        fUpgradeCodeFound = TRUE;
+                        break;
+                    }
+                }
+            }
+            ReleaseNullStrArray(rgsczBundleUpgradeCodes, cBundleUpgradeCodes);
+
+            break;
+
+        default:
+            ButilExitWithRootFailure(hr, E_NOTIMPL, "BundleUpgradeCode of type 0x%x not implemented.", dwType);
+        }
+
+        if (fUpgradeCodeFound)
+        {
+            if (psczBundleId)
+            {
+                *psczBundleId = sczUninstallSubKey;
+                sczUninstallSubKey = NULL;
+            }
+
+            break;
+        }
+
+        // Cleanup before next iteration
+        ReleaseRegKey(hkBundle);
+        ReleaseNullStr(sczUninstallSubKey);
+        ReleaseNullStr(sczUninstallSubKeyPath);
+    }
+
+LExit:
+    ReleaseStr(sczValue);
+    ReleaseStr(sczUninstallSubKey);
+    ReleaseStr(sczUninstallSubKeyPath);
+    ReleaseRegKey(hkBundle);
+    ReleaseRegKey(hkUninstall);
+    ReleaseStrArray(rgsczBundleUpgradeCodes, cBundleUpgradeCodes);
+
+    return FAILED(hr) ? hr : fUpgradeCodeFound ? S_OK : S_FALSE;
+}