13 files changed, 281 insertions, 7 deletions

diff --git a/src/burn/engine/registration.cpp b/src/burn/engine/registration.cpp
index 85c006f7..966fc849 100644
--- a/src/burn/engine/registration.cpp
+++ b/src/burn/engine/registration.cpp
@@ -808,13 +808,6 @@ extern "C" HRESULT RegistrationSessionBegin(
         ExitOnFailure(hr, "Failed to write software tags.");
     }
 
-    // Update registration.
-    if (pRegistration->update.fRegisterUpdate)
-    {
-        hr = WriteUpdateRegistration(pRegistration, pVariables);
-        ExitOnFailure(hr, "Failed to write update registration.");
-    }
-
     // Only set install date and initial estimated size here for the first time.
     // Estimated size will always get updated at the end of the session.
     if (fCreated)
@@ -904,6 +897,13 @@ extern "C" HRESULT RegistrationSessionEnd(
 
         hr = UpdateEstimatedSize(hkRegistration, qwEstimatedSize);
         ExitOnFailure(hr, "Failed to update estimated size.");
+
+        // Update registration.
+        if (pRegistration->update.fRegisterUpdate)
+        {
+            hr = WriteUpdateRegistration(pRegistration, pVariables);
+            ExitOnFailure(hr, "Failed to write update registration.");
+        }
     }
 
     // Update resume mode.
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/Bundle.wxs b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/Bundle.wxs
new file mode 100644
index 00000000..2e920297
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/Bundle.wxs
@@ -0,0 +1,58 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+<?ifndef TestVersion?>
+<?define TestVersion = 1.0.0.0?>
+<?endif?>
+<?ifndef BundleLogDirectory?>
+<?define BundleLogDirectory = .?>
+<?endif?>
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
+  <Bundle Name="~$(TestGroupName)" Version="$(TestVersion)" Manufacturer="Acme" UpgradeCode="$(UpgradeCode)" Compressed="yes">
+    <Log Prefix="$(BundleLogDirectory)\~$(TestGroupName)_$(BundleName)" />
+
+    <Variable Name="TestGroupName" Value="$(TestGroupName)" />
+
+    <OptionalUpdateRegistration ProductFamily="GreetingsAndFelicitations" Department="Setup Geeks" Classification="Bundle" />
+
+    <?ifdef SoftwareTag?>
+    <SoftwareTag Regid="regid.1995-08.com.example" InstallPath="[CommonAppDataFolder]TestingSwidTags" />
+    <?endif?>
+
+    <?ifndef BA?>
+      <!-- pulled in through the PackageGroupRef below -->
+    <?elseif $(BA) = "TestBA_x64"?>
+      <!-- pulled in through the PackageGroupRef below -->
+    <?elseif $(BA) = "WixBA"?>
+      <!-- pulled in through the PackageGroupRef below -->
+    <?elseif $(BA) = "hyperlinkLicense"?>
+      <BootstrapperApplication>
+        <bal:WixStandardBootstrapperApplication LicenseUrl="" Theme="hyperlinkLicense" />
+        <PayloadGroupRef Id="ExtraPayloads" />
+      </BootstrapperApplication>
+    <?elseif $(BA) = "iui"?>
+      <BootstrapperApplication>
+        <bal:WixInternalUIBootstrapperApplication />
+      </BootstrapperApplication>
+    <?else?>
+      <BootstrapperApplicationRef Id="$(BA)" />
+    <?endif?>
+
+    <Chain>
+      <?ifndef BA?>
+        <PackageGroupRef Id="TestBA" />
+      <?elseif $(BA) = "TestBA_x64"?>
+        <PackageGroupRef Id="TestBA_x64" />
+      <?elseif $(BA) = "WixBA"?>
+        <PackageGroupRef Id="WixBA" />
+      <?endif?>
+
+      <PackageGroupRef Id="BundlePackages" />
+    </Chain>
+  </Bundle>
+
+  <Fragment>
+    <PayloadGroup Id="virtual ExtraPayloads" />
+  </Fragment>
+
+</Wix>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleA.props b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleA.props
new file mode 100644
index 00000000..6d91ba06
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleA.props
@@ -0,0 +1,8 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project>
+  <PropertyGroup>
+    <OutputType>Bundle</OutputType>
+    <BA>hyperlinkLicense</BA>
+    <UpgradeCode>{9510DDE7-CD4F-45B3-9D57-3F7EA04DB33D}</UpgradeCode>
+  </PropertyGroup>
+</Project>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wixproj b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wixproj
new file mode 100644
index 00000000..0dea1e40
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wixproj
@@ -0,0 +1,12 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="BundleA.props" />
+  <ItemGroup>
+    <ProjectReference Include="..\PackageAv1\PackageAv1.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.BootstrapperApplications.wixext" />
+    <PackageReference Include="WixToolset.NetFx.wixext" />
+  </ItemGroup>
+</Project>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wxs b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wxs
new file mode 100644
index 00000000..7bf16212
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv1/BundleAv1.wxs
@@ -0,0 +1,10 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs">
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageAv1.TargetPath)" />
+    </PackageGroup>
+  </Fragment>
+</Wix>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wixproj b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wixproj
new file mode 100644
index 00000000..271bf54b
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wixproj
@@ -0,0 +1,18 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="..\BundleAv1\BundleA.props" />
+  <PropertyGroup>
+    <TestVersion>2.0.0.0</TestVersion>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\BundleAv1\Bundle.wxs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="..\PackageAv2\PackageAv2.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.BootstrapperApplications.wixext" />
+    <PackageReference Include="WixToolset.NetFx.wixext" />
+  </ItemGroup>
+</Project>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wxs b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wxs
new file mode 100644
index 00000000..5cbee5a8
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/BundleAv2/BundleAv2.wxs
@@ -0,0 +1,10 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs">
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageAv2.TargetPath)" />
+    </PackageGroup>
+  </Fragment>
+</Wix>
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageA.props b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageA.props
new file mode 100644
index 00000000..25ae3f42
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageA.props
@@ -0,0 +1,9 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project>
+  <PropertyGroup>
+    <UpgradeCode>{30746E93-9A4A-48EC-BAEA-3093A2530E73}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\Package.wxs" Link="Package.wxs" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageAv1.wixproj b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageAv1.wixproj
new file mode 100644
index 00000000..45d3b2c8
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv1/PackageAv1.wixproj
@@ -0,0 +1,4 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="PackageA.props" />
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv2/PackageAv2.wixproj b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv2/PackageAv2.wixproj
new file mode 100644
index 00000000..d4455f2b
--- /dev/null
+++ b/src/test/burn/TestData/OptionalUpdateRegistrationTests/PackageAv2/PackageAv2.wixproj
@@ -0,0 +1,7 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="..\PackageAv1\PackageA.props" />
+  <PropertyGroup>
+    <TestVersion>2.0.0.0</TestVersion>
+  </PropertyGroup>
+</Project>
diff --git a/src/test/burn/WixTestTools/BundleUpdateRegistration.cs b/src/test/burn/WixTestTools/BundleUpdateRegistration.cs
new file mode 100644
index 00000000..2ce68122
--- /dev/null
+++ b/src/test/burn/WixTestTools/BundleUpdateRegistration.cs
@@ -0,0 +1,57 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixTestTools
+{
+    using System;
+    using Microsoft.Win32;
+
+    public class BundleUpdateRegistration
+    {
+        public const string BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PACKAGE_NAME = "PackageName";
+        public const string BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PACKAGE_VERSION = "PackageVersion";
+        public const string BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PUBLISHER = "Publisher";
+        public const string BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PUBLISHING_GROUP = "PublishingGroup";
+
+        public string PackageName { get; set; }
+
+        public string PackageVersion { get; set; }
+
+        public string Publisher { get; set; }
+
+        public string PublishingGroup { get; set; }
+
+        public static bool TryGetPerMachineBundleUpdateRegistration(string manufacturer, string productFamily, string name, bool x64, out BundleUpdateRegistration registration)
+        {
+            return TryGetUpdateRegistration(manufacturer, productFamily, name, x64, perUser: false, out registration);
+        }
+
+        public static bool TryGetPerUserBundleUpdateRegistration(string manufacturer, string productFamily, string name, out BundleUpdateRegistration registration)
+        {
+            return TryGetUpdateRegistration(manufacturer, productFamily, name, x64: true, perUser: true, out registration);
+        }
+
+        private static bool TryGetUpdateRegistration(string manufacturer, string productFamily, string name, bool x64, bool perUser, out BundleUpdateRegistration registration)
+        {
+            var baseKey = perUser ? Registry.CurrentUser : Registry.LocalMachine;
+            var baseKeyPath = x64 ? @$"SOFTWARE\{manufacturer}\Updates\{productFamily}\{name}"
+                : @$"SOFTWARE\WOW6432Node\{manufacturer}\Updates\{productFamily}\{name}";
+            using var idKey = baseKey.OpenSubKey(baseKeyPath);
+
+            if (idKey == null)
+            {
+                registration = null;
+                return false;
+            }
+
+            registration = new BundleUpdateRegistration()
+            {
+                PackageName = idKey.GetValue(BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PACKAGE_NAME) as string,
+                PackageVersion = idKey.GetValue(BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PACKAGE_VERSION) as string,
+                Publisher = idKey.GetValue(BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PUBLISHER) as string,
+                PublishingGroup = idKey.GetValue(BURN_UPDATE_REGISTRATION_REGISTRY_BUNDLE_PUBLISHING_GROUP) as string,
+            };
+
+            return true;
+        }
+    }
+}
diff --git a/src/test/burn/WixTestTools/BundleVerifier.cs b/src/test/burn/WixTestTools/BundleVerifier.cs
index da61d96e..b6181047 100644
--- a/src/test/burn/WixTestTools/BundleVerifier.cs
+++ b/src/test/burn/WixTestTools/BundleVerifier.cs
@@ -6,6 +6,7 @@ namespace WixTestTools
     using System.IO;
     using System.Linq;
     using System.Text;
+    using System.Xml.Linq;
     using Microsoft.Win32;
     using WixInternal.TestSupport;
     using WixToolset.Data;
@@ -23,6 +24,8 @@ namespace WixTestTools
 
         private WixBundleSymbol BundleSymbol { get; set; }
 
+        private WixUpdateRegistrationSymbol UpdateRegistrationSymbol { get; set; }
+
         private WixBundleSymbol GetBundleSymbol()
         {
             if (this.BundleSymbol == null)
@@ -36,6 +39,19 @@ namespace WixTestTools
             return this.BundleSymbol;
         }
 
+        private WixUpdateRegistrationSymbol GetUpdateRegistrationSymbol()
+        {
+            if (this.UpdateRegistrationSymbol == null)
+            {
+                using var wixOutput = WixOutput.Read(this.BundlePdb);
+                var intermediate = Intermediate.Load(wixOutput);
+                var section = intermediate.Sections.Single();
+                this.UpdateRegistrationSymbol = section.Symbols.OfType<WixUpdateRegistrationSymbol>().Single();
+            }
+
+            return this.UpdateRegistrationSymbol;
+        }
+
         public string GetFullBurnPolicyRegistryPath()
         {
             var bundleSymbol = this.GetBundleSymbol();
@@ -118,6 +134,27 @@ namespace WixTestTools
             }
         }
 
+        public bool TryGetUpdateRegistration(out BundleUpdateRegistration registration)
+        {
+            var bundleSymbol = this.GetBundleSymbol();
+            var x64 = bundleSymbol.Platform != Platform.X86;
+
+            var updateRegistrationSymbol = this.GetUpdateRegistrationSymbol();
+            var manufacturer = updateRegistrationSymbol.Manufacturer;
+            var productFamily = updateRegistrationSymbol.ProductFamily;
+            var name = updateRegistrationSymbol.Name;
+
+
+            if (bundleSymbol.PerMachine)
+            {
+                return BundleUpdateRegistration.TryGetPerMachineBundleUpdateRegistration(manufacturer, productFamily, name, x64, out registration);
+            }
+            else
+            {
+                return BundleUpdateRegistration.TryGetPerUserBundleUpdateRegistration(manufacturer, productFamily, name, out registration);
+            }
+        }
+
         public BundleRegistration VerifyRegisteredAndInPackageCache(int? expectedSystemComponent = null)
         {
             Assert.True(this.TryGetRegistration(out var registration));
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs
new file mode 100644
index 00000000..3731b0cc
--- /dev/null
+++ b/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs
@@ -0,0 +1,44 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolsetTest.BurnE2E
+{
+    using WixTestTools;
+    using Xunit;
+    using Xunit.Abstractions;
+
+    public class OptionalUpdateRegistrationTests : BurnE2ETests
+    {
+        public OptionalUpdateRegistrationTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
+
+        [RuntimeFact]
+        public void BundleUpdateRegistrationIsStickyAndAccurateAcrossUpgrades()
+        {
+            var packageAv1 = this.CreatePackageInstaller("PackageAv1");
+            var packageAv2 = this.CreatePackageInstaller("PackageAv2");
+            var bundleAv1 = this.CreateBundleInstaller("BundleAv1");
+            var bundleAv2 = this.CreateBundleInstaller("BundleAv2");
+
+            bundleAv1.Install();
+            bundleAv1.VerifyRegisteredAndInPackageCache();
+            var gotV1Registration = bundleAv1.TryGetUpdateRegistration(out var v1Registration);
+
+            bundleAv2.Install();
+            bundleAv2.VerifyRegisteredAndInPackageCache();
+            var gotV2Registration = bundleAv2.TryGetUpdateRegistration(out var v2Registration);
+
+            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache();
+
+            Assert.True(gotV1Registration, "Missing update registration after v1 install.");
+            Assert.True(gotV2Registration, "Missing update registration after v2 upgrade.");
+
+            Assert.Equal("Acme", v1Registration.Publisher);
+            Assert.Equal("Acme", v2Registration.Publisher);
+            Assert.Equal("Setup Geeks", v1Registration.PublishingGroup);
+            Assert.Equal("Setup Geeks", v2Registration.PublishingGroup);
+            Assert.Equal("~OptionalUpdateRegistrationTests", v1Registration.PackageName);
+            Assert.Equal("~OptionalUpdateRegistrationTests", v2Registration.PackageName);
+            Assert.Equal("1.0.0.0", v1Registration.PackageVersion);
+            Assert.Equal("2.0.0.0", v2Registration.PackageVersion);
+        }
+    }
+}