Provide Bundle Icon in BA container and update wixstdba to use itHEADmain

Fixes 8104

11 files changed, 126 insertions, 21 deletions

diff --git a/src/ext/Bal/stdbas/WixInternalUIBootstrapperApplication.cpp b/src/ext/Bal/stdbas/WixInternalUIBootstrapperApplication.cpp
index 270ce40c..bae098d4 100644
--- a/src/ext/Bal/stdbas/WixInternalUIBootstrapperApplication.cpp
+++ b/src/ext/Bal/stdbas/WixInternalUIBootstrapperApplication.cpp
@@ -557,14 +557,19 @@ private:
     HRESULT CreateMainWindow()
     {
         HRESULT hr = S_OK;
-        WNDCLASSW wc = { };
+        WNDCLASSEXW wc = { };
         DWORD dwWindowStyle = WS_POPUP;
 
+        LoadBundleIcon(m_hModule, &m_hIcon, &m_hSmallIcon);
+
+        wc.cbSize = sizeof(WNDCLASSEXW);
         wc.lpfnWndProc = CWixInternalUIBootstrapperApplication::WndProc;
         wc.hInstance = m_hModule;
         wc.lpszClassName = WIXIUIBA_WINDOW_CLASS;
+        wc.hIcon = m_hIcon;
+        wc.hIconSm = m_hSmallIcon;
 
-        if (!::RegisterClassW(&wc))
+        if (!::RegisterClassExW(&wc))
         {
             ExitWithLastError(hr, "Failed to register window.");
         }
@@ -601,6 +606,18 @@ private:
             ::UnregisterClassW(WIXIUIBA_WINDOW_CLASS, m_hModule);
             m_fRegistered = FALSE;
         }
+
+        if (m_hIcon)
+        {
+            ::DestroyIcon(m_hIcon);
+            m_hIcon = NULL;
+        }
+
+        if (m_hSmallIcon)
+        {
+            ::DestroyIcon(m_hSmallIcon);
+            m_hSmallIcon = NULL;
+        }
     }
 
     //
@@ -807,6 +824,8 @@ public:
         m_sczFailedMessage = NULL;
 
         m_hUiThread = NULL;
+        m_hIcon = NULL;
+        m_hSmallIcon = NULL;
         m_fRegistered = FALSE;
         m_hWnd = NULL;
 
@@ -847,6 +866,8 @@ private:
     LPWSTR m_sczConfirmCloseMessage;
 
     HANDLE m_hUiThread;
+    HICON m_hIcon;
+    HICON m_hSmallIcon;
     BOOL m_fRegistered;
     HWND m_hWnd;
 
diff --git a/src/ext/Bal/stdbas/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/stdbas/WixStandardBootstrapperApplication.cpp
index ddb2d1c5..c1e8b2b3 100644
--- a/src/ext/Bal/stdbas/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/stdbas/WixStandardBootstrapperApplication.cpp
@@ -3301,26 +3301,24 @@ private:
     HRESULT CreateMainWindow()
     {
         HRESULT hr = S_OK;
-        WNDCLASSW wc = { };
+        WNDCLASSEXW wc = { sizeof(WNDCLASSEXW)};
         DWORD dwWindowStyle = 0;
         int x = CW_USEDEFAULT;
         int y = CW_USEDEFAULT;
         POINT ptCursor = { };
 
-        ThemeInitializeWindowClass(m_pTheme, &wc, CWixStandardBootstrapperApplication::WndProc, m_hModule, WIXSTDBA_WINDOW_CLASS);
+        ThemeInitializeWindowClassEx(m_pTheme, &wc, CWixStandardBootstrapperApplication::WndProc, m_hModule, WIXSTDBA_WINDOW_CLASS);
 
-        // If the theme did not provide an icon, try using the icon from the bundle engine.
+        // If the theme did not provide an icon, try using the icon from the bundle then fallback to the bundle engine.
         if (!wc.hIcon)
         {
-            HMODULE hBootstrapperEngine = ::GetModuleHandleW(NULL);
-            if (hBootstrapperEngine)
-            {
-                wc.hIcon = ::LoadIconW(hBootstrapperEngine, MAKEINTRESOURCEW(1));
-            }
+            LoadBundleIcon(m_hModule, &m_hIcon, &m_hSmallIcon);
+            wc.hIcon = m_hIcon;
+            wc.hIconSm = m_hSmallIcon;
         }
 
         // Register the window class and create the window.
-        if (!::RegisterClassW(&wc))
+        if (!::RegisterClassExW(&wc))
         {
             ExitWithLastError(hr, "Failed to register window.");
         }
@@ -3358,7 +3356,6 @@ private:
         return hr;
     }
 
-
     //
     // InitializeTaskbarButton - initializes taskbar button for progress.
     //
@@ -3397,6 +3394,18 @@ private:
             ::UnregisterClassW(WIXSTDBA_WINDOW_CLASS, m_hModule);
             m_fRegistered = FALSE;
         }
+
+        if (m_hIcon)
+        {
+            ::DestroyIcon(m_hIcon);
+            m_hIcon = NULL;
+        }
+
+        if (m_hSmallIcon)
+        {
+            ::DestroyIcon(m_hSmallIcon);
+            m_hSmallIcon = NULL;
+        }
     }
 
 
@@ -4822,6 +4831,8 @@ public:
         m_pTheme = NULL;
         memset(m_rgdwPageIds, 0, sizeof(m_rgdwPageIds));
         m_hUiThread = NULL;
+        m_hIcon = NULL;
+        m_hSmallIcon = NULL;
         m_fRegistered = FALSE;
         m_hWnd = NULL;
 
@@ -5110,6 +5121,8 @@ private:
     THEME_ASSIGN_CONTROL_ID m_rgInitControls[LAST_WIXSTDBA_CONTROL - WIXSTDBA_FIRST_ASSIGN_CONTROL_ID];
     DWORD m_rgdwPageIds[countof(vrgwzPageNames)];
     HANDLE m_hUiThread;
+    HICON m_hIcon;
+    HICON m_hSmallIcon;
     BOOL m_fRegistered;
     HWND m_hWnd;
 
diff --git a/src/ext/Bal/stdbas/precomp.h b/src/ext/Bal/stdbas/precomp.h
index a6c4cfe4..0faf1548 100644
--- a/src/ext/Bal/stdbas/precomp.h
+++ b/src/ext/Bal/stdbas/precomp.h
@@ -42,3 +42,9 @@
 
 #include "stdbas.messages.h"
 #include "WixStandardBootstrapperApplication.h"
+
+HRESULT LoadBundleIcon(
+    __in HMODULE hModule,
+    __out HICON* phIcon,
+    __out HICON* phSmallIcon
+);
diff --git a/src/ext/Bal/stdbas/stdbas.cpp b/src/ext/Bal/stdbas/stdbas.cpp
new file mode 100644
index 00000000..909bdc23
--- /dev/null
+++ b/src/ext/Bal/stdbas/stdbas.cpp
@@ -0,0 +1,50 @@
+// 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.
+
+#include "precomp.h"
+
+static const LPCWSTR WIX_BUNDLE_ICON_FILENAME = L"WixBundle.ico";
+
+
+//
+// LoadBundleIcon - loads the icon that was (optionally) authored in the bundle otherwise use the one embedded in the bootstrapper application.
+//
+HRESULT LoadBundleIcon(
+    __in HMODULE hModule,
+    __out HICON* phIcon,
+    __out HICON* phSmallIcon
+)
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczIconPath = NULL;
+    int nIconWidth = ::GetSystemMetrics(SM_CXICON);
+    int nIconHeight = ::GetSystemMetrics(SM_CYICON);
+    int nSmallIconWidth = ::GetSystemMetrics(SM_CXSMICON);
+    int nSmallIconHeight = ::GetSystemMetrics(SM_CYSMICON);
+    HICON hIcon = NULL;
+    HICON hSmallIcon = NULL;
+
+    // First look for the optional authored bundle icon.
+    hr = PathRelativeToModule(&sczIconPath, WIX_BUNDLE_ICON_FILENAME, hModule);
+    ExitOnFailure(hr, "Failed to get path to bundle icon: %ls", WIX_BUNDLE_ICON_FILENAME);
+
+    if (FileExistsEx(sczIconPath, NULL))
+    {
+        hIcon = reinterpret_cast<HICON>(::LoadImageW(NULL, sczIconPath, IMAGE_ICON, nIconWidth, nIconHeight, LR_LOADFROMFILE));
+
+        hSmallIcon = reinterpret_cast<HICON>(::LoadImageW(NULL, sczIconPath, IMAGE_ICON, nSmallIconWidth, nSmallIconHeight, LR_LOADFROMFILE));
+    }
+    else // fallback to the first icon resource in the bootstrapper application.
+    {
+        hIcon = reinterpret_cast<HICON>(::LoadImageW(hModule, MAKEINTRESOURCEW(1), IMAGE_ICON, nIconWidth, nIconHeight, LR_DEFAULTCOLOR));
+
+        hSmallIcon = reinterpret_cast<HICON>(::LoadImageW(hModule, MAKEINTRESOURCEW(1), IMAGE_ICON, nSmallIconWidth, nSmallIconHeight, LR_DEFAULTCOLOR));
+    }
+
+    *phIcon = hIcon;
+    *phSmallIcon = hSmallIcon;
+
+LExit:
+    ReleaseStr(sczIconPath);
+
+    return hr;
+}
diff --git a/src/ext/Bal/stdbas/stdbas.vcxproj b/src/ext/Bal/stdbas/stdbas.vcxproj
index fb2b0682..ba990af7 100644
--- a/src/ext/Bal/stdbas/stdbas.vcxproj
+++ b/src/ext/Bal/stdbas/stdbas.vcxproj
@@ -60,6 +60,7 @@
     <ClCompile Include="precomp.cpp">
       <PrecompiledHeader>Create</PrecompiledHeader>
     </ClCompile>
+    <ClCompile Include="stdbas.cpp" />
     <ClCompile Include="WixInternalUIBootstrapperApplication.cpp" />
     <ClCompile Include="WixStandardBootstrapperApplication.cpp" />
   </ItemGroup>
@@ -80,4 +81,4 @@ rc.exe -fo "$(OutDir)stdbas.res" "$(IntDir)stdbas.messages.rc"</Command>
   </ItemGroup>
 
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
-</Project>
+</Project>
\ No newline at end of file
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/BalExtensionFixture.cs b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/BalExtensionFixture.cs
index 576e7c6c..e599b8e2 100644
--- a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/BalExtensionFixture.cs
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/BalExtensionFixture.cs
@@ -24,7 +24,8 @@ namespace WixToolsetTest.BootstrapperApplications
             {
                 var baseFolder = fs.GetFolder();
                 var bundleFile = Path.Combine(baseFolder, "bin", "test.exe");
-                var bundleSourceFolder = TestData.Get(@"TestData\WixStdBa");
+                var bundleSourceFolder = TestData.Get(@"TestData", "WixStdBa");
+                var dataFolder = TestData.Get(@"TestData", ".Data");
                 var intermediateFolder = Path.Combine(baseFolder, "obj");
                 var baFolderPath = Path.Combine(baseFolder, "ba");
                 var extractFolderPath = Path.Combine(baseFolder, "extract");
@@ -36,6 +37,7 @@ namespace WixToolsetTest.BootstrapperApplications
                     "-ext", TestData.Get(@"WixToolset.BootstrapperApplications.wixext.dll"),
                     "-intermediateFolder", intermediateFolder,
                     "-bindpath", Path.Combine(bundleSourceFolder, "data"),
+                    "-bindpath", dataFolder,
                     "-o", bundleFile,
                 });
                 compileResult.AssertSuccess();
@@ -62,7 +64,8 @@ namespace WixToolsetTest.BootstrapperApplications
             {
                 var baseFolder = fs.GetFolder();
                 var bundleFile = Path.Combine(baseFolder, "bin", "test.exe");
-                var bundleSourceFolder = TestData.Get(@"TestData\WixStdBa");
+                var bundleSourceFolder = TestData.Get(@"TestData", "WixStdBa");
+                var dataFolder = TestData.Get(@"TestData", ".Data");
                 var intermediateFolder = Path.Combine(baseFolder, "obj");
                 var baFolderPath = Path.Combine(baseFolder, "ba");
                 var extractFolderPath = Path.Combine(baseFolder, "extract");
@@ -74,6 +77,7 @@ namespace WixToolsetTest.BootstrapperApplications
                     "-ext", TestData.Get(@"WixToolset.BootstrapperApplications.wixext.dll"),
                     "-intermediateFolder", intermediateFolder,
                     "-bindpath", Path.Combine(bundleSourceFolder, "data"),
+                    "-bindpath", dataFolder,
                     "-o", bundleFile,
                 });
                 compileResult.AssertSuccess();
@@ -101,6 +105,7 @@ namespace WixToolsetTest.BootstrapperApplications
                 var baseFolder = fs.GetFolder();
                 var bundleFile = Path.Combine(baseFolder, "bin", "test.exe");
                 var bundleSourceFolder = TestData.Get("TestData", "WixStdBa");
+                var dataFolder = TestData.Get(@"TestData", ".Data");
                 var intermediateFolder = Path.Combine(baseFolder, "obj");
                 var baFolderPath = Path.Combine(baseFolder, "ba");
                 var extractFolderPath = Path.Combine(baseFolder, "extract");
@@ -112,6 +117,7 @@ namespace WixToolsetTest.BootstrapperApplications
                     "-ext", TestData.Get(@"WixToolset.BootstrapperApplications.wixext.dll"),
                     "-intermediateFolder", intermediateFolder,
                     "-bindpath", Path.Combine(bundleSourceFolder, "data"),
+                    "-bindpath", dataFolder,
                     "-o", bundleFile,
                 });
                 compileResult.AssertSuccess();
@@ -141,7 +147,8 @@ namespace WixToolsetTest.BootstrapperApplications
             {
                 var baseFolder = fs.GetFolder();
                 var bundleFile = Path.Combine(baseFolder, "bin", "test.exe");
-                var bundleSourceFolder = TestData.Get(@"TestData\Overridable");
+                var bundleSourceFolder = TestData.Get(@"TestData", "Overridable");
+                var dataFolder = TestData.Get(@"TestData", ".Data");
                 var intermediateFolder = Path.Combine(baseFolder, "obj");
                 var baFolderPath = Path.Combine(baseFolder, "ba");
                 var extractFolderPath = Path.Combine(baseFolder, "extract");
@@ -152,6 +159,7 @@ namespace WixToolsetTest.BootstrapperApplications
                     Path.Combine(bundleSourceFolder, "Bundle.wxs"),
                     "-ext", TestData.Get(@"WixToolset.BootstrapperApplications.wixext.dll"),
                     "-intermediateFolder", intermediateFolder,
+                    "-bindpath", dataFolder,
                     "-o", bundleFile,
                 });
                 compileResult.AssertSuccess();
@@ -182,7 +190,8 @@ namespace WixToolsetTest.BootstrapperApplications
             {
                 var baseFolder = fs.GetFolder();
                 var bundleFile = Path.Combine(baseFolder, "bin", "test.exe");
-                var bundleSourceFolder = TestData.Get(@"TestData\WixStdBa");
+                var bundleSourceFolder = TestData.Get(@"TestData", "WixStdBa");
+                var dataFolder = TestData.Get(@"TestData", ".Data");
                 var intermediateFolder = Path.Combine(baseFolder, "obj");
 
                 var compileResult = WixRunner.Execute(new[]
@@ -191,6 +200,7 @@ namespace WixToolsetTest.BootstrapperApplications
                     Path.Combine(bundleSourceFolder, "Bundle.wxs"),
                     "-ext", TestData.Get(@"WixToolset.BootstrapperApplications.wixext.dll"),
                     "-intermediateFolder", intermediateFolder,
+                    "-bindpath", dataFolder,
                     "-o", bundleFile,
                 });
                 compileResult.AssertSuccess();
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/.Data/test.ico b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/.Data/test.ico
new file mode 100644
index 00000000..906ce324
--- /dev/null
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/.Data/test.ico
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/BootstrapperApplicationId.wxs b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/BootstrapperApplicationId.wxs
index e1945f53..44c15999 100644
--- a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/BootstrapperApplicationId.wxs
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/BootstrapperApplicationId.wxs
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs"
      xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
-    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2">
+    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2"
+            IconSourceFile="test.ico">
         <BootstrapperApplication Id="Custom">
             <bal:WixStandardBootstrapperApplication LicenseUrl="http://wixtoolset.org/about/license/" Theme="hyperlinkLicense" />
         </BootstrapperApplication>
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/Bundle.wxs b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/Bundle.wxs
index c17b53ff..e0346900 100644
--- a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/Bundle.wxs
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/Bundle.wxs
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs"
      xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
-    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2">
+    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2"
+            IconSourceFile="test.ico">
         <BootstrapperApplication>
             <bal:WixStandardBootstrapperApplication LicenseUrl="http://wixtoolset.org/about/license/" Theme="hyperlinkLicense" />
         </BootstrapperApplication>
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayFilesInUseDialogConditionBundle.wxs b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayFilesInUseDialogConditionBundle.wxs
index 1041eb39..098d6f83 100644
--- a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayFilesInUseDialogConditionBundle.wxs
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayFilesInUseDialogConditionBundle.wxs
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs"
      xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
-    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2">
+    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2"
+            IconSourceFile="test.ico">
         <BootstrapperApplication>
             <bal:WixStandardBootstrapperApplication LicenseUrl="http://wixtoolset.org/about/license/" Theme="hyperlinkLicense" />
         </BootstrapperApplication>
diff --git a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayInternalUIConditionBundle.wxs b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayInternalUIConditionBundle.wxs
index f08cfe6a..2dea66fd 100644
--- a/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayInternalUIConditionBundle.wxs
+++ b/src/ext/Bal/test/WixToolsetTest.BootstrapperApplications/TestData/WixStdBa/DisplayInternalUIConditionBundle.wxs
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs"
      xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
-    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2">
+    <Bundle Name="WixStdBa" Version="1.0.0.0" Manufacturer="Example Corporation" UpgradeCode="75D5D534-E177-4689-AAE9-CAC1C39002C2"
+            IconSourceFile="test.ico">
         <BootstrapperApplication>
             <bal:WixStandardBootstrapperApplication LicenseUrl="http://wixtoolset.org/about/license/" Theme="hyperlinkLicense" />
         </BootstrapperApplication>