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

Fixes 8104

5 files changed, 104 insertions, 13 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