Remove extension data after extensions get PreDecompileTables callback

Also update Util.wixext decompiler extension to handle the new
elements that compile into CustomActions with no additional table data.
These exposed the weakness in the decompiler where extension data was
removed before extensions got a chance to pre-decompile.

Fixes 7151

1 files changed, 50 insertions, 8 deletions

diff --git a/src/ext/Util/wixext/UtilDecompiler.cs b/src/ext/Util/wixext/UtilDecompiler.cs
index 1201fdd5..0a78201e 100644
--- a/src/ext/Util/wixext/UtilDecompiler.cs
+++ b/src/ext/Util/wixext/UtilDecompiler.cs
@@ -3,16 +3,15 @@
 namespace WixToolset.Util
 {
     using System;
+    using System.Collections;
+    using System.Collections.Generic;
     using System.IO;
+    using System.Linq;
     using System.Text;
-    using System.Collections;
-    using System.Diagnostics;
-
+    using System.Xml.Linq;
     using WixToolset.Data;
-    using WixToolset.Extensibility;
     using WixToolset.Data.WindowsInstaller;
-    using System.Collections.Generic;
-    using System.Xml.Linq;
+    using WixToolset.Extensibility;
     using WixToolset.Util.Symbols;
 
     /// <summary>
@@ -22,16 +21,47 @@ namespace WixToolset.Util
     {
         public override IReadOnlyCollection<TableDefinition> TableDefinitions => UtilTableDefinitions.All;
 
+        private static readonly Dictionary<string, XName> CustomActionMapping = new Dictionary<string, XName>()
+        {
+            {  "Wix4BroadcastEnvironmentChange_X86", UtilConstants.BroadcastEnvironmentChange },
+            {  "Wix4BroadcastEnvironmentChange_X64", UtilConstants.BroadcastEnvironmentChange },
+            {  "Wix4BroadcastEnvironmentChange_ARM64", UtilConstants.BroadcastEnvironmentChange },
+            {  "Wix4BroadcastSettingChange_X86", UtilConstants.BroadcastSettingChange },
+            {  "Wix4BroadcastSettingChange_X64", UtilConstants.BroadcastSettingChange },
+            {  "Wix4BroadcastSettingChange_ARM64", UtilConstants.BroadcastSettingChange },
+            {  "Wix4CheckRebootRequired_X86", UtilConstants.CheckRebootRequired },
+            {  "Wix4CheckRebootRequired_X64", UtilConstants.CheckRebootRequired },
+            {  "Wix4CheckRebootRequired_ARM64", UtilConstants.CheckRebootRequired },
+            {  "Wix4QueryNativeMachine_X86", UtilConstants.QueryNativeMachine },
+            {  "Wix4QueryNativeMachine_X64", UtilConstants.QueryNativeMachine },
+            {  "Wix4QueryNativeMachine_ARM64", UtilConstants.QueryNativeMachine },
+            {  "Wix4QueryOsDriverInfo_X86", UtilConstants.QueryWindowsDriverInfo },
+            {  "Wix4QueryOsDriverInfo_X64", UtilConstants.QueryWindowsDriverInfo },
+            {  "Wix4QueryOsDriverInfo_ARM64", UtilConstants.QueryWindowsDriverInfo },
+            {  "Wix4QueryOsInfo_X86", UtilConstants.QueryWindowsSuiteInfo },
+            {  "Wix4QueryOsInfo_X64", UtilConstants.QueryWindowsSuiteInfo },
+            {  "Wix4QueryOsInfo_ARM64", UtilConstants.QueryWindowsSuiteInfo },
+        };
+
+        private IReadOnlyCollection<string> customActionNames;
+
         /// <summary>
         /// Called at the beginning of the decompilation of a database.
         /// </summary>
         /// <param name="tables">The collection of all tables.</param>
         public override void PreDecompileTables(TableIndexedCollection tables)
         {
+            this.RememberCustomActionNames(tables);
             this.CleanupSecureCustomProperties(tables);
             this.CleanupInternetShortcutRemoveFileTables(tables);
         }
 
+        private void RememberCustomActionNames(TableIndexedCollection tables)
+        {
+            var customActionTable = tables["CustomAction"];
+            this.customActionNames = customActionTable?.Rows.Select(r => r.GetPrimaryKey()).Distinct().ToList() ?? (IReadOnlyCollection<string>)Array.Empty<string>();
+        }
+
         /// <summary>
         /// Decompile the SecureCustomProperties field to PropertyRefs for known extension properties.
         /// </summary>
@@ -195,6 +225,7 @@ namespace WixToolset.Util
         /// <param name="tables">The collection of all tables.</param>
         public override void PostDecompileTables(TableIndexedCollection tables)
         {
+            this.FinalizeCustomActions();
             this.FinalizePerfmonTable(tables);
             this.FinalizePerfmonManifestTable(tables);
             this.FinalizeSecureObjectsTable(tables);
@@ -223,7 +254,7 @@ namespace WixToolset.Util
                     AttributeIfNotNull("RebootPrompt", 0x2 == (attribute & 0x2)),
                     AttributeIfNotNull("ElevatedCloseMessage", 0x4 == (attribute & 0x4)),
                     NumericAttributeIfNotNull("Sequence", row, 5),
-                    NumericAttributeIfNotNull("Property", row, 6)
+                    AttributeIfNotNull("Property", row, 6)
                     );
             }
         }
@@ -379,7 +410,7 @@ namespace WixToolset.Util
                 if (this.DecompilerHelper.TryGetIndexedElement("Wix4FileShare", row.FieldAsString(0), out var fileShare) ||
                     this.DecompilerHelper.TryGetIndexedElement("FileShare", row.FieldAsString(0), out fileShare))
                 {
-                        fileShare.Add(fileSharePermission);
+                    fileShare.Add(fileSharePermission);
                 }
                 else
                 {
@@ -695,6 +726,17 @@ namespace WixToolset.Util
             }
         }
 
+        private void FinalizeCustomActions()
+        {
+            foreach (var customActionName in this.customActionNames)
+            {
+                if (CustomActionMapping.TryGetValue(customActionName, out var elementName))
+                {
+                    this.DecompilerHelper.AddElementToRoot(elementName);
+                }
+            }
+        }
+
         /// <summary>
         /// Finalize the Perfmon table.
         /// </summary>