6 files changed, 90 insertions, 3 deletions

diff --git a/src/api/burn/WixToolset.Mba.Core/IPackageInfo.cs b/src/api/burn/WixToolset.Mba.Core/IPackageInfo.cs
index ee3be820..c6285f03 100644
--- a/src/api/burn/WixToolset.Mba.Core/IPackageInfo.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IPackageInfo.cs
@@ -68,6 +68,11 @@ namespace WixToolset.Mba.Core
         string PrereqLicenseUrl { get; }
 
         /// <summary>
+        /// See <see cref="PrimaryPackageType"/>
+        /// </summary>
+        PrimaryPackageType PrimaryPackageType { get; }
+
+        /// <summary>
         /// 
         /// </summary>
         string ProductCode { get; }
diff --git a/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs b/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
index e400cbe4..3fa1b49e 100644
--- a/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
+++ b/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
@@ -64,6 +64,37 @@ namespace WixToolset.Mba.Core
     }
 
     /// <summary>
+    /// Metadata for BAs like WixInternalUIBootstrapperApplication that only support one main package.
+    /// </summary>
+    public enum PrimaryPackageType
+    {
+        /// <summary>
+        /// Not a primary package.
+        /// </summary>
+        None,
+
+        /// <summary>
+        /// The default package if no architecture specific package is available for the current architecture.
+        /// </summary>
+        Default,
+
+        /// <summary>
+        /// The package to use on x86 machines.
+        /// </summary>
+        X86,
+
+        /// <summary>
+        /// The package to use on x64 machines.
+        /// </summary>
+        X64,
+
+        /// <summary>
+        /// The package to use on ARM64 machines.
+        /// </summary>
+        ARM64,
+    }
+
+    /// <summary>
     /// Default implementation of <see cref="IPackageInfo"/>.
     /// </summary>
     public class PackageInfo : IPackageInfo
@@ -117,6 +148,9 @@ namespace WixToolset.Mba.Core
         public string PrereqLicenseUrl { get; internal set; }
 
         /// <inheritdoc/>
+        public PrimaryPackageType PrimaryPackageType { get; internal set; }
+
+        /// <inheritdoc/>
         public object CustomData { get; set; }
 
         internal PackageInfo() { }
diff --git a/src/api/burn/balutil/balinfo.cpp b/src/api/burn/balutil/balinfo.cpp
index 5832281f..52a7f911 100644
--- a/src/api/burn/balutil/balinfo.cpp
+++ b/src/api/burn/balutil/balinfo.cpp
@@ -529,6 +529,33 @@ static HRESULT ParseBalPackageInfoFromXml(
         hr = XmlGetAttributeEx(pNode, L"DisplayInternalUICondition", &pPackage->sczDisplayInternalUICondition);
         ExitOnOptionalXmlQueryFailure(hr, fXmlFound, "Failed to get DisplayInternalUICondition setting for package.");
 
+        hr = XmlGetAttributeEx(pNode, L"PrimaryPackageType", &scz);
+        ExitOnOptionalXmlQueryFailure(hr, fXmlFound, "Failed to get PrimaryPackageType setting for package.");
+
+        if (fXmlFound)
+        {
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"default", -1))
+            {
+                pPackage->primaryPackageType = BAL_INFO_PRIMARY_PACKAGE_TYPE_DEFAULT;
+            }
+            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"x86", -1))
+            {
+                pPackage->primaryPackageType = BAL_INFO_PRIMARY_PACKAGE_TYPE_X86;
+            }
+            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"x64", -1))
+            {
+                pPackage->primaryPackageType = BAL_INFO_PRIMARY_PACKAGE_TYPE_X64;
+            }
+            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"arm64", -1))
+            {
+                pPackage->primaryPackageType = BAL_INFO_PRIMARY_PACKAGE_TYPE_ARM64;
+            }
+            else
+            {
+                ExitWithRootFailure(hr, E_INVALIDARG, "Invalid value for WixBalPackageInfo/@PrimaryPackageType: %ls", scz);
+            }
+        }
+
         ReleaseNullObject(pNode);
     }
     ExitOnFailure(hr, "Failed to parse all WixBalPackageInfo elements.");
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index aa1ca56f..c3223bee 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -451,8 +451,8 @@ public: // IBootstrapperApplication
         __in BOOTSTRAPPER_ERROR_TYPE errorType,
         __in_z LPCWSTR wzPackageId,
         __in DWORD dwCode,
-        __in_z LPCWSTR /*wzError*/,
-        __in DWORD /*dwUIHint*/,
+        __in_z LPCWSTR wzError,
+        __in DWORD dwUIHint,
         __in DWORD /*cData*/,
         __in_ecount_z_opt(cData) LPCWSTR* /*rgwzData*/,
         __in int /*nRecommendation*/,
@@ -461,7 +461,15 @@ public: // IBootstrapperApplication
     {
         BalRetryErrorOccurred(wzPackageId, dwCode);
 
-        if (CheckCanceled())
+        if (BOOTSTRAPPER_DISPLAY_EMBEDDED == m_display)
+        {
+            HRESULT hr = m_pEngine->SendEmbeddedError(dwCode, wzError, dwUIHint, pResult);
+            if (FAILED(hr))
+            {
+                *pResult = IDERROR;
+            }
+        }
+        else if (CheckCanceled())
         {
             *pResult = IDCANCEL;
         }
diff --git a/src/api/burn/balutil/inc/balinfo.h b/src/api/burn/balutil/inc/balinfo.h
index cd61f34d..818ff5ef 100644
--- a/src/api/burn/balutil/inc/balinfo.h
+++ b/src/api/burn/balutil/inc/balinfo.h
@@ -20,6 +20,15 @@ typedef enum BAL_INFO_PACKAGE_TYPE
     BAL_INFO_PACKAGE_TYPE_BUNDLE_CHAIN,
 } BAL_INFO_PACKAGE_TYPE;
 
+typedef enum _BAL_INFO_PRIMARY_PACKAGE_TYPE
+{
+    BAL_INFO_PRIMARY_PACKAGE_TYPE_NONE,
+    BAL_INFO_PRIMARY_PACKAGE_TYPE_DEFAULT,
+    BAL_INFO_PRIMARY_PACKAGE_TYPE_X86,
+    BAL_INFO_PRIMARY_PACKAGE_TYPE_X64,
+    BAL_INFO_PRIMARY_PACKAGE_TYPE_ARM64,
+} BAL_INFO_PRIMARY_PACKAGE_TYPE;
+
 typedef enum _BAL_INFO_RESTART
 {
     BAL_INFO_RESTART_UNKNOWN,
@@ -54,6 +63,7 @@ typedef struct _BAL_INFO_PACKAGE
     BOOL fPrereqPackage;
     LPWSTR sczPrereqLicenseFile;
     LPWSTR sczPrereqLicenseUrl;
+    BAL_INFO_PRIMARY_PACKAGE_TYPE primaryPackageType;
     LPVOID pvCustomData;
 } BAL_INFO_PACKAGE;
 
diff --git a/src/api/wix/WixToolset.Data/Burn/BurnConstants.cs b/src/api/wix/WixToolset.Data/Burn/BurnConstants.cs
index 1ecccbd2..8eeb030b 100644
--- a/src/api/wix/WixToolset.Data/Burn/BurnConstants.cs
+++ b/src/api/wix/WixToolset.Data/Burn/BurnConstants.cs
@@ -14,6 +14,9 @@ namespace WixToolset.Data.Burn
         public const string BundleExtensionDataWixOutputStreamName = "wix-bextdata";
         public const string BootstrapperApplicationDataWixOutputStreamName = "wix-badata.xml";
 
+        public const string BootstrapperApplicationDataNamespace = "http://wixtoolset.org/schemas/v4/BootstrapperApplicationData";
+        public const string BundleExtensionDataNamespace = "http://wixtoolset.org/schemas/v4/BundleExtensionData";
+
         public const string BootstrapperApplicationDataSymbolDefinitionTag = "WixBootstrapperApplicationData";
         public const string BundleExtensionSearchSymbolDefinitionTag = "WixBundleExtensionSearch";