13 files changed, 374 insertions, 48 deletions

diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 43786701..2a34aaad 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -85,6 +85,7 @@ enum BA_FUNCTIONS_MESSAGE
 
     BA_FUNCTIONS_MESSAGE_ONTHEMELOADED = 1024,
     BA_FUNCTIONS_MESSAGE_WNDPROC,
+    BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLLOADING,
 };
 
 typedef HRESULT(WINAPI *PFN_BA_FUNCTIONS_PROC)(
@@ -94,6 +95,10 @@ typedef HRESULT(WINAPI *PFN_BA_FUNCTIONS_PROC)(
     __in_opt LPVOID pvContext
     );
 
+// Should be the same as THEME_FIRST_ASSIGN_CONTROL_ID.
+// BAFunctions must only assign ids in the range [BAFUNCTIONS_FIRST_ASSIGN_CONTROL_ID, 0x8000) to avoid collisions.
+const WORD BAFUNCTIONS_FIRST_ASSIGN_CONTROL_ID = 0x4000;
+
 struct BA_FUNCTIONS_CREATE_ARGS
 {
     DWORD cbSize;
@@ -108,6 +113,19 @@ struct BA_FUNCTIONS_CREATE_RESULTS
     LPVOID pvBAFunctionsProcContext;
 };
 
+struct BA_FUNCTIONS_ONTHEMECONTROLLOADING_ARGS
+{
+    DWORD cbSize;
+    LPCWSTR wzName;
+};
+
+struct BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS
+{
+    DWORD cbSize;
+    BOOL fProcessed;
+    WORD wId;
+};
+
 struct BA_FUNCTIONS_ONTHEMELOADED_ARGS
 {
     DWORD cbSize;
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index a3054709..c2c8a6dc 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -841,6 +841,15 @@ public: // IBAFunctions
         return E_NOTIMPL;
     }
 
+    virtual STDMETHODIMP OnThemeControlLoading(
+        __in LPCWSTR /*wzName*/,
+        __inout BOOL* /*pfProcessed*/,
+        __inout WORD* /*pwId*/
+        )
+    {
+        return S_OK;
+    }
+
 protected:
     CBalBaseBAFunctions(
         __in HMODULE hModule,
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
index 8d1227fc..efe22ddd 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -24,6 +24,15 @@ static HRESULT BalBaseBAFunctionsProcWndProc(
     return pBAFunctions->WndProc(pArgs->pTheme, pArgs->hWnd, pArgs->uMsg, pArgs->wParam, pArgs->lParam, &pResults->lres);
 }
 
+static HRESULT BalBaseBAFunctionsProcOnThemeControlLoading(
+    __in IBAFunctions* pBAFunctions,
+    __in BA_FUNCTIONS_ONTHEMECONTROLLOADING_ARGS* pArgs,
+    __inout BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS* pResults
+    )
+{
+    return pBAFunctions->OnThemeControlLoading(pArgs->wzName, &pResults->fProcessed, &pResults->wId);
+}
+
 /*******************************************************************
 BalBaseBAFunctionsProc - requires pvContext to be of type IBAFunctions.
 Provides a default mapping between the message based BAFunctions interface and
@@ -125,6 +134,9 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_WNDPROC:
             hr = BalBaseBAFunctionsProcWndProc(pBAFunctions, reinterpret_cast<BA_FUNCTIONS_WNDPROC_ARGS*>(pvArgs), reinterpret_cast<BA_FUNCTIONS_WNDPROC_RESULTS*>(pvResults));
             break;
+        case BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLLOADING:
+            hr = BalBaseBAFunctionsProcOnThemeControlLoading(pBAFunctions, reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLLOADING_ARGS*>(pvArgs), reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS*>(pvResults));
+            break;
         }
     }
 
diff --git a/src/api/burn/balutil/inc/IBAFunctions.h b/src/api/burn/balutil/inc/IBAFunctions.h
index 7d8a07fa..d41b7c9b 100644
--- a/src/api/burn/balutil/inc/IBAFunctions.h
+++ b/src/api/burn/balutil/inc/IBAFunctions.h
@@ -31,4 +31,12 @@ DECLARE_INTERFACE_IID_(IBAFunctions, IBootstrapperApplication, "0FB445ED-17BD-49
         __inout LPVOID pvResults,
         __in_opt LPVOID pvContext
         ) = 0;
+
+    // OnThemeControlLoading - Called while creating a control for the theme.
+    //
+    STDMETHOD(OnThemeControlLoading)(
+        __in LPCWSTR wzName,
+        __inout BOOL* pfProcessed,
+        __inout WORD* pwId
+        ) = 0;
 };
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 74e0b7d3..8cdd31ce 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -77,10 +77,15 @@ static LPCWSTR vrgwzPageNames[] = {
     L"Failure",
 };
 
+// The range [0, 100) is unused to avoid collisions with system ids,
+// the range [100, 0x4000) is unused to avoid collisions with thmutil,
+// the range [0x4000, 0x8000) is unused to avoid collisions with BAFunctions.
+const WORD WIXSTDBA_FIRST_ASSIGN_CONTROL_ID = 0x8000; 
+
 enum WIXSTDBA_CONTROL
 {
     // Welcome page
-    WIXSTDBA_CONTROL_INSTALL_BUTTON = THEME_FIRST_ASSIGN_CONTROL_ID,
+    WIXSTDBA_CONTROL_INSTALL_BUTTON = WIXSTDBA_FIRST_ASSIGN_CONTROL_ID,
     WIXSTDBA_CONTROL_EULA_RICHEDIT,
     WIXSTDBA_CONTROL_EULA_LINK,
     WIXSTDBA_CONTROL_EULA_ACCEPT_CHECKBOX,
@@ -2833,6 +2838,9 @@ private:
             }
             break;
 
+        case WM_THMUTIL_LOADING_CONTROL:
+            return pBA->OnThemeLoadingControl(reinterpret_cast<THEME_LOADINGCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADINGCONTROL_RESULTS*>(lParam));
+
         case WM_QUERYENDSESSION:
             fCancel = true;
             pBA->OnSystemShutdown(static_cast<DWORD>(lParam), &fCancel);
@@ -2956,7 +2964,7 @@ private:
         BA_FUNCTIONS_ONTHEMELOADED_ARGS themeLoadedArgs = { };
         BA_FUNCTIONS_ONTHEMELOADED_RESULTS themeLoadedResults = { };
 
-        hr = ThemeLoadControls(m_pTheme, vrgInitControls, countof(vrgInitControls));
+        hr = ThemeLoadControls(m_pTheme);
         BalExitOnFailure(hr, "Failed to load theme controls.");
 
         C_ASSERT(COUNT_WIXSTDBA_PAGE == countof(vrgwzPageNames));
@@ -3029,6 +3037,57 @@ private:
         return SUCCEEDED(hr);
     }
 
+    BOOL OnThemeLoadingControl(
+        __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+        __in THEME_LOADINGCONTROL_RESULTS* pResults
+        )
+    {
+        HRESULT hr = S_OK;
+        BOOL fProcessed = FALSE;
+        BA_FUNCTIONS_ONTHEMECONTROLLOADING_ARGS themeControlLoadingArgs = { };
+        BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS themeControlLoadingResults = { };
+
+        for (DWORD iAssignControl = 0; iAssignControl < countof(vrgInitControls); ++iAssignControl)
+        {
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, pArgs->pThemeControl->sczName, -1, vrgInitControls[iAssignControl].wzName, -1))
+            {
+                fProcessed = TRUE;
+                pResults->wId = vrgInitControls[iAssignControl].wId;
+                ExitFunction();
+            }
+        }
+
+        if (m_pfnBAFunctionsProc)
+        {
+            themeControlLoadingArgs.cbSize = sizeof(themeControlLoadingArgs);
+            themeControlLoadingArgs.wzName = pArgs->pThemeControl->sczName;
+
+            themeControlLoadingResults.cbSize = sizeof(themeControlLoadingResults);
+            themeControlLoadingResults.wId = pResults->wId;
+
+            hr = m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLLOADING, &themeControlLoadingArgs, &themeControlLoadingResults, m_pvBAFunctionsProcContext);
+
+            if (E_NOTIMPL == hr)
+            {
+                hr = S_OK;
+            }
+            else
+            {
+                BalExitOnFailure(hr, "BAFunctions OnThemeControlLoading failed.");
+
+                if (themeControlLoadingResults.fProcessed)
+                {
+                    fProcessed = TRUE;
+                    pResults->wId = themeControlLoadingResults.wId;
+                }
+            }
+        }
+
+    LExit:
+        pResults->hr = hr;
+        return fProcessed || FAILED(hr);
+    }
+
 
     //
     // OnShowFailure - display the failure page.
@@ -3857,7 +3916,7 @@ private:
         BalExitOnNullWithLastError(pfnBAFunctionsCreate, hr, "Failed to get BAFunctionsCreate entry-point from: %ls", sczBafPath);
 
         bafCreateArgs.cbSize = sizeof(bafCreateArgs);
-        bafCreateArgs.qwBAFunctionsAPIVersion = MAKEQWORDVERSION(0, 0, 0, 2); // TODO: need to decide whether to keep this, and if so when to update it.
+        bafCreateArgs.qwBAFunctionsAPIVersion = MAKEQWORDVERSION(2021, 9, 20, 0);
         bafCreateArgs.pBootstrapperCreateArgs = &m_createArgs;
 
         bafCreateResults.cbSize = sizeof(bafCreateResults);
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h b/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
index 6ac3711f..dc554219 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
@@ -95,6 +95,15 @@ typedef enum THEME_WINDOW_INITIAL_POSITION
     THEME_WINDOW_INITIAL_POSITION_CENTER_MONITOR_FROM_COORDINATES,
 } THEME_WINDOW_INITIAL_POSITION;
 
+// These messages are sent by thmutil to the parent window created in ThemeCreateParentWindow.
+// thmutil reserves the last values of WM_USER's range.
+typedef enum _WM_THMUTIL
+{
+    // Sent while creating a control.
+    // wparam is THEME_LOADINGCONTROL_ARGS* and lparam is THEME_LOADINGCONTROL_RESULTS*.
+    // Return code is TRUE if it was processed.
+    WM_THMUTIL_LOADING_CONTROL = WM_APP - 1,
+} WM_THMUTIL;
 
 struct THEME_COLUMN
 {
@@ -171,7 +180,7 @@ struct THEME_ASSIGN_CONTROL_ID
     LPCWSTR wzName; // name of control to match
 };
 
-const DWORD THEME_FIRST_ASSIGN_CONTROL_ID = 1024; // Recommended first control id to be assigned.
+const WORD THEME_FIRST_ASSIGN_CONTROL_ID = 0x4000; // Recommended first control id to be assigned.
 
 struct THEME_CONTROL
 {
@@ -325,7 +334,7 @@ struct THEME_FONT
 
 struct THEME
 {
-    WORD wId;
+    WORD wNextControlId;
 
     BOOL fAutoResize;
     BOOL fForceResize;
@@ -388,6 +397,24 @@ struct THEME
     LPVOID pvVariableContext;
 };
 
+typedef struct _THEME_LOADINGCONTROL_ARGS
+{
+    DWORD cbSize;
+    const THEME_CONTROL* pThemeControl;
+} THEME_LOADINGCONTROL_ARGS;
+
+typedef struct _THEME_LOADINGCONTROL_RESULTS
+{
+    DWORD cbSize;
+    HRESULT hr;
+
+    // Used to apply a specific id to the control (usually used for WM_COMMAND).
+    // If assigning an id, it is recommended to start with THEME_FIRST_ASSIGN_CONTROL_ID to avoid collisions with well known ids such as IDOK and IDCANCEL.
+    // The values [100, THEME_FIRST_ASSIGN_CONTROL_ID) are reserved for thmutil.
+    // Due to this value being packed into 16 bits for many system window messages, this is restricted to a WORD.
+    WORD wId;
+} THEME_LOADINGCONTROL_RESULTS;
+
 
 /********************************************************************
  ThemeInitialize - initialized theme management.
@@ -472,9 +499,7 @@ HRESULT DAPI ThemeCreateParentWindow(
 
 *******************************************************************/
 HRESULT DAPI ThemeLoadControls(
-    __in THEME* pTheme,
-    __in_ecount_opt(cAssignControlIds) const THEME_ASSIGN_CONTROL_ID* rgAssignControlIds,
-    __in DWORD cAssignControlIds
+    __in THEME* pTheme
     );
 
 /********************************************************************
diff --git a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
index 53257d8e..d796bbaf 100644
--- a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
@@ -39,6 +39,7 @@
 #define LWS_NOPREFIX        0x0004
 #endif
 
+const WORD THEME_FIRST_AUTO_ASSIGN_CONTROL_ID = 100;
 const DWORD THEME_INVALID_ID = 0xFFFFFFFF;
 const COLORREF THEME_INVISIBLE_COLORREF = 0xFFFFFFFF;
 const DWORD GROW_FONT_INSTANCES = 3;
@@ -272,11 +273,14 @@ static HRESULT FindImageList(
     __in_z LPCWSTR wzImageListName,
     __out HIMAGELIST *phImageList
     );
+static HRESULT OnLoadingControl(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pControl,
+    __inout WORD* pwId
+    );
 static HRESULT LoadControls(
     __in THEME* pTheme,
-    __in_opt THEME_CONTROL* pParentControl,
-    __in_ecount_opt(cAssignControlIds) const THEME_ASSIGN_CONTROL_ID* rgAssignControlIds,
-    __in DWORD cAssignControlIds
+    __in_opt THEME_CONTROL* pParentControl
     );
 static HRESULT ShowControl(
     __in THEME* pTheme,
@@ -871,9 +875,7 @@ LExit:
 
 
 DAPI_(HRESULT) ThemeLoadControls(
-    __in THEME* pTheme,
-    __in_ecount_opt(cAssignControlIds) const THEME_ASSIGN_CONTROL_ID* rgAssignControlIds,
-    __in DWORD cAssignControlIds
+    __in THEME* pTheme
     )
 {
     HRESULT hr = S_OK;
@@ -883,7 +885,7 @@ DAPI_(HRESULT) ThemeLoadControls(
         ThmExitOnFailure(hr = E_INVALIDSTATE, "ThemeLoadControls called before theme parent window created.");
     }
 
-    hr = LoadControls(pTheme, NULL, rgAssignControlIds, cAssignControlIds);
+    hr = LoadControls(pTheme, NULL);
 
 LExit:
     return hr;
@@ -1844,8 +1846,6 @@ static HRESULT ParseTheme(
     __out THEME** ppTheme
     )
 {
-    static WORD wThemeId = 0;
-
     HRESULT hr = S_OK;
     THEME* pTheme = NULL;
     IXMLDOMElement *pThemeElement = NULL;
@@ -1858,8 +1858,8 @@ static HRESULT ParseTheme(
     pTheme = static_cast<THEME*>(MemAlloc(sizeof(THEME), TRUE));
     ThmExitOnNull(pTheme, hr, E_OUTOFMEMORY, "Failed to allocate memory for theme.");
 
-    pTheme->wId = ++wThemeId;
     pTheme->nDpi = USER_DEFAULT_SCREEN_DPI;
+    pTheme->wNextControlId = THEME_FIRST_AUTO_ASSIGN_CONTROL_ID;
 
     // Parse the optional background resource image.
     hr = GetAttributeImageFileOrResource(hModule, wzRelativePath, pThemeElement, &pBitmap);
@@ -5248,7 +5248,7 @@ static BOOL OnButtonClicked(
 LExit:
     return fHandled;
 }
-    
+
 static BOOL OnDpiChanged(
     __in THEME* pTheme,
     __in WPARAM wParam,
@@ -5934,11 +5934,38 @@ static LRESULT CALLBACK StaticOwnerDrawWndProc(
     }
 }
 
+static HRESULT OnLoadingControl(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pControl,
+    __inout WORD* pwId
+    )
+{
+    HRESULT hr = S_OK;
+    THEME_LOADINGCONTROL_ARGS loadingControlArgs = { };
+    THEME_LOADINGCONTROL_RESULTS loadingControlResults = { };
+
+    loadingControlArgs.cbSize = sizeof(loadingControlArgs);
+    loadingControlArgs.pThemeControl = pControl;
+
+    loadingControlResults.cbSize = sizeof(loadingControlResults);
+    loadingControlResults.hr = E_NOTIMPL;
+    loadingControlResults.wId = *pwId;
+
+    if (::SendMessageW(pTheme->hwndParent, WM_THMUTIL_LOADING_CONTROL, reinterpret_cast<WPARAM>(&loadingControlArgs), reinterpret_cast<LPARAM>(&loadingControlResults)))
+    {
+        hr = loadingControlResults.hr;
+        if (SUCCEEDED(hr))
+        {
+            *pwId = loadingControlResults.wId;
+        }
+    }
+
+    return hr;
+}
+
 static HRESULT LoadControls(
     __in THEME* pTheme,
-    __in_opt THEME_CONTROL* pParentControl,
-    __in_ecount_opt(cAssignControlIds) const THEME_ASSIGN_CONTROL_ID* rgAssignControlIds,
-    __in DWORD cAssignControlIds
+    __in_opt THEME_CONTROL* pParentControl
     )
 {
     HRESULT hr = S_OK;
@@ -6101,16 +6128,16 @@ static HRESULT LoadControls(
         }
         ThmExitOnNull(wzWindowClass, hr, E_INVALIDDATA, "Failed to configure control %u because of unknown type: %u", i, pControl->type);
 
-        // Default control ids to the theme id and its index in the control array, unless there
-        // is a specific id to assign to a named control.
-        WORD wControlId = MAKEWORD(i, pTheme->wId);
-        for (DWORD iAssignControl = 0; pControl->sczName && iAssignControl < cAssignControlIds; ++iAssignControl)
+        // Default control ids to the next id, unless there is a specific id to assign to a control.
+        WORD wControlId = THEME_FIRST_AUTO_ASSIGN_CONTROL_ID;
+        hr = OnLoadingControl(pTheme, pControl, &wControlId);
+        ThmExitOnFailure(hr, "ThmLoadingControl failed.");
+
+        // This range is reserved for thmutil. The process will run out of available window handles before reaching the end of the range.
+        if (THEME_FIRST_AUTO_ASSIGN_CONTROL_ID <= wControlId && THEME_FIRST_ASSIGN_CONTROL_ID > wControlId)
         {
-            if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, pControl->sczName, -1, rgAssignControlIds[iAssignControl].wzName, -1))
-            {
-                wControlId = rgAssignControlIds[iAssignControl].wId;
-                break;
-            }
+            wControlId = pTheme->wNextControlId;
+            pTheme->wNextControlId += 1;
         }
 
         pControl->wId = wControlId;
@@ -6296,7 +6323,7 @@ static HRESULT LoadControls(
 
         if (pControl->cControls)
         {
-            hr = LoadControls(pTheme, pControl, rgAssignControlIds, cAssignControlIds);
+            hr = LoadControls(pTheme, pControl);
             ThmExitOnFailure(hr, "Failed to load child controls.");
         }
     }
diff --git a/src/samples/thmviewer/display.cpp b/src/samples/thmviewer/display.cpp
index 52fa3cf8..c0e6c7e1 100644
--- a/src/samples/thmviewer/display.cpp
+++ b/src/samples/thmviewer/display.cpp
@@ -329,7 +329,7 @@ static BOOL DisplayOnCreate(
 {
     HRESULT hr = S_OK;
 
-    hr = ThemeLoadControls(pTheme, NULL, 0);
+    hr = ThemeLoadControls(pTheme);
     ExitOnFailure(hr, "Failed to load theme controls");
 
     // Pre-populate some control types with data.
diff --git a/src/samples/thmviewer/thmviewer.cpp b/src/samples/thmviewer/thmviewer.cpp
index f83182d3..cffa3851 100644
--- a/src/samples/thmviewer/thmviewer.cpp
+++ b/src/samples/thmviewer/thmviewer.cpp
@@ -14,10 +14,6 @@ enum THMVWR_CONTROL
     THMVWR_CONTROL_TREE = THEME_FIRST_ASSIGN_CONTROL_ID,
 };
 
-static THEME_ASSIGN_CONTROL_ID vrgInitControls[] = {
-    { THMVWR_CONTROL_TREE, L"Tree" },
-};
-
 // Internal functions
 
 static HRESULT ProcessCommandLine(
@@ -52,6 +48,10 @@ static void OnNewTheme(
     __in HWND hWnd,
     __in HANDLE_THEME* pHandle
     );
+static BOOL OnThemeLoadingControl(
+    __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+    __in THEME_LOADINGCONTROL_RESULTS* pResults
+    );
 static void CALLBACK ThmviewerTraceError(
     __in_z LPCSTR szFile,
     __in int iLine,
@@ -353,7 +353,7 @@ static LRESULT CALLBACK MainWndProc(
 
     case WM_CREATE:
         {
-            HRESULT hr = ThemeLoadControls(vpTheme, vrgInitControls, countof(vrgInitControls));
+            HRESULT hr = ThemeLoadControls(vpTheme);
             if (FAILED(hr))
             {
                 return -1;
@@ -400,6 +400,9 @@ static LRESULT CALLBACK MainWndProc(
         }
         }
         break;
+
+    case WM_THMUTIL_LOADING_CONTROL:
+        return OnThemeLoadingControl(reinterpret_cast<THEME_LOADINGCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADINGCONTROL_RESULTS*>(lParam));
     }
 
     return ThemeDefWindowProc(vpTheme, hWnd, uMsg, wParam, lParam);
@@ -541,3 +544,17 @@ static void OnNewTheme(
         ThemeSendControlMessage(pTheme, THMVWR_CONTROL_TREE, TVM_SELECTITEM, TVGN_CARET, reinterpret_cast<LPARAM>(htiSelected));
     }
 }
+
+static BOOL OnThemeLoadingControl(
+    __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+    __in THEME_LOADINGCONTROL_RESULTS* pResults
+    )
+{
+    if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, pArgs->pThemeControl->sczName, -1, L"Tree", -1))
+    {
+        pResults->wId = THMVWR_CONTROL_TREE;
+    }
+
+    pResults->hr = S_OK;
+    return TRUE;
+}
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index 8b49cab6..b35b4e02 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -6,6 +6,11 @@
 
 static const LPCWSTR BAFTHMUTILTESTING_WINDOW_CLASS = L"BafThmUtilTesting";
 
+enum BAF_CONTROL
+{
+    BAF_CONTROL_INSTALL_TEST_BUTTON = BAFUNCTIONS_FIRST_ASSIGN_CONTROL_ID,
+};
+
 enum BAFTHMUTILTESTING_CONTROL
 {
     BAFTHMUTILTESTING_CONTROL_LISTVIEW_TOP_LEFT = THEME_FIRST_ASSIGN_CONTROL_ID,
@@ -59,24 +64,41 @@ public: // IBAFunctions
         __inout LRESULT* plRes
         )
     {
-        HRESULT hr = S_OK;
-
-        __super::WndProc(pTheme, hWnd, uMsg, wParam, lParam, plRes);
-
-        // Show our window when any button is clicked.
         switch (uMsg)
         {
         case WM_COMMAND:
             switch (HIWORD(wParam))
             {
             case BN_CLICKED:
-                OnShowTheme();
+                switch (LOWORD(wParam))
+                {
+                case BAF_CONTROL_INSTALL_TEST_BUTTON:
+                    OnShowTheme();
+                    *plRes = 0;
+                    return S_OK;
+                }
+
                 break;
             }
             break;
         }
 
-        return hr;
+        return __super::WndProc(pTheme, hWnd, uMsg, wParam, lParam, plRes);
+    }
+
+    virtual STDMETHODIMP OnThemeControlLoading(
+        __in LPCWSTR wzName,
+        __inout BOOL* pfProcessed,
+        __inout WORD* pwId
+        )
+    {
+        if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzName, -1, L"InstallTestButton", -1))
+        {
+            *pfProcessed = TRUE;
+            *pwId = BAF_CONTROL_INSTALL_TEST_BUTTON;
+        }
+
+        return S_OK;
     }
 
 private:
@@ -229,6 +251,9 @@ private:
             }
             break;
 
+        case WM_THMUTIL_LOADING_CONTROL:
+            return pBaf->OnThemeLoadingControl(reinterpret_cast<THEME_LOADINGCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADINGCONTROL_RESULTS*>(lParam));
+
         case WM_TIMER:
             if (!lParam && pBaf)
             {
@@ -255,7 +280,7 @@ private:
         HWND hwndBottomLeft = NULL;
         HWND hwndBottomRight = NULL;
 
-        hr = ThemeLoadControls(m_pBafTheme, vrgInitControls, countof(vrgInitControls));
+        hr = ThemeLoadControls(m_pBafTheme);
         BalExitOnFailure(hr, "Failed to load theme controls.");
 
         hwndTopLeft = ::GetDlgItem(m_pBafTheme->hwndParent, BAFTHMUTILTESTING_CONTROL_LISTVIEW_TOP_LEFT);
@@ -333,6 +358,24 @@ private:
         return SUCCEEDED(hr);
     }
 
+    BOOL OnThemeLoadingControl(
+        __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+        __in THEME_LOADINGCONTROL_RESULTS* pResults
+        )
+    {
+        for (DWORD iAssignControl = 0; iAssignControl < countof(vrgInitControls); ++iAssignControl)
+        {
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, pArgs->pThemeControl->sczName, -1, vrgInitControls[iAssignControl].wzName, -1))
+            {
+                pResults->wId = vrgInitControls[iAssignControl].wId;
+                break;
+            }
+        }
+
+        pResults->hr = S_OK;
+        return TRUE;
+    }
+
     void UpdateProgressBarProgress()
     {
         static DWORD dwProgress = 0;
diff --git a/src/test/burn/TestData/Manual/BundleA/BundleA.wixproj b/src/test/burn/TestData/Manual/BundleA/BundleA.wixproj
index dcfd3b7e..907b85c9 100644
--- a/src/test/burn/TestData/Manual/BundleA/BundleA.wixproj
+++ b/src/test/burn/TestData/Manual/BundleA/BundleA.wixproj
@@ -2,7 +2,7 @@
 <Project Sdk="WixToolset.Sdk">
   <PropertyGroup>
     <OutputType>Bundle</OutputType>
-    <BA>hyperlinkLicense</BA>
+    <BA>customHyperlinkLicense</BA>
     <UpgradeCode>{98ACBCF6-B54A-46AF-8990-DFB8795B965B}</UpgradeCode>
   </PropertyGroup>
   <ItemGroup>
diff --git a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
index 1706f4e8..20706b6a 100644
--- a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
+++ b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
@@ -3,7 +3,8 @@
 
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
   <Fragment>
-    <BootstrapperApplication>
+    <BootstrapperApplication Id="customHyperlinkLicense">
+      <bal:WixStandardBootstrapperApplication LicenseUrl="" Theme="hyperlinkLicense" ThemeFile="CustomHyperlinkTheme.xml" />
       <Payload SourceFile="$(var.BafThmUtilTesting.TargetPath)" bal:BAFunctions="yes" />
     </BootstrapperApplication>
     <PackageGroup Id="BundlePackages">
diff --git a/src/test/burn/TestData/Manual/BundleA/CustomHyperlinkTheme.xml b/src/test/burn/TestData/Manual/BundleA/CustomHyperlinkTheme.xml
new file mode 100644
index 00000000..b8157193
--- /dev/null
+++ b/src/test/burn/TestData/Manual/BundleA/CustomHyperlinkTheme.xml
@@ -0,0 +1,107 @@
+<?xml version="1.0" encoding="utf-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. -->
+
+
+<Theme xmlns="http://wixtoolset.org/schemas/v4/thmutil">
+    <Font Id="0" Height="-12" Weight="500" Foreground="windowtext" Background="window">Segoe UI</Font>
+    <Font Id="1" Height="-24" Weight="500" Foreground="windowtext">Segoe UI</Font>
+    <Font Id="2" Height="-22" Weight="500" Foreground="graytext">Segoe UI</Font>
+    <Font Id="3" Height="-12" Weight="500" Foreground="windowtext" Background="window">Segoe UI</Font>
+
+    <Window Width="485" Height="300" HexStyle="100a0000" FontId="0" Caption="#(loc.Caption)">
+        <ImageControl X="11" Y="11" Width="64" Height="64" ImageFile="logo.png" Visible="yes"/>
+        <Label X="80" Y="11" Width="-11" Height="64" FontId="1" Visible="yes" DisablePrefix="yes">#(loc.Title)</Label>
+
+        <Page Name="Help">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.HelpHeader)</Label>
+            <Label X="11" Y="112" Width="-11" Height="-35" FontId="3" DisablePrefix="yes">#(loc.HelpText)</Label>
+            <Button Name="HelpCloseButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.HelpCloseButton)</Text>
+                <CloseWindowAction />
+            </Button>
+        </Page>
+        <Page Name="Install">
+            <Hypertext Name="EulaHyperlink" X="11" Y="121" Width="-11" Height="51" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallLicenseLinkText)</Hypertext>
+            <Checkbox Name="EulaAcceptCheckbox" X="-11" Y="-41" Width="260" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallTestButton" X="-251" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">Test</Button>
+            <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
+                <Text>#(loc.InstallOptionsButton)</Text>
+                <ChangePageAction Page="Options" />
+            </Button>
+            <Button Name="InstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.InstallInstallButton)</Button>
+            <Button Name="InstallCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.InstallCancelButton)</Text>
+                <CloseWindowAction />
+            </Button>
+        </Page>
+        <Page Name="Options">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.OptionsHeader)</Label>
+            <Label X="11" Y="121" Width="-11" Height="17" FontId="3" DisablePrefix="yes">#(loc.OptionsLocationLabel)</Label>
+            <Editbox Name="InstallFolder" X="11" Y="143" Width="-91" Height="21" TabStop="yes" FontId="3" FileSystemAutoComplete="yes" />
+            <Button Name="BrowseButton" X="-11" Y="142" Width="75" Height="23" TabStop="yes" FontId="3">
+                <Text>#(loc.OptionsBrowseButton)</Text>
+                <BrowseDirectoryAction VariableName="InstallFolder" />
+            </Button>
+            <Button Name="OptionsOkButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.OptionsOkButton)</Text>
+                <ChangePageAction Page="Install" />
+            </Button>
+            <Button Name="OptionsCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.OptionsCancelButton)</Text>
+                <ChangePageAction Page="Install" Cancel="yes" />
+            </Button>
+        </Page>
+        <Page Name="Progress">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ProgressHeader)</Label>
+            <Label X="11" Y="121" Width="70" Height="17" FontId="3" DisablePrefix="yes">#(loc.ProgressLabel)</Label>
+            <Label Name="OverallProgressPackageText" X="85" Y="121" Width="-11" Height="17" FontId="3" DisablePrefix="yes">#(loc.OverallProgressPackageText)</Label>
+            <Progressbar Name="OverallCalculatedProgressbar" X="11" Y="143" Width="-11" Height="15" />
+            <Button Name="ProgressCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ProgressCancelButton)</Button>
+        </Page>
+        <Page Name="Modify">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
+            <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
+            <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.ModifyCancelButton)</Text>
+                <CloseWindowAction />
+            </Button>
+        </Page>
+        <Page Name="Success">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">
+                <Text>#(loc.SuccessHeader)</Text>
+                <Text Condition="WixBundleAction = 2">#(loc.SuccessLayoutHeader)</Text>
+                <Text Condition="WixBundleAction = 3">#(loc.SuccessUninstallHeader)</Text>
+                <Text Condition="WixBundleAction = 5">#(loc.SuccessInstallHeader)</Text>
+                <Text Condition="WixBundleAction = 7">#(loc.SuccessRepairHeader)</Text>
+            </Label>
+            <Button Name="LaunchButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.SuccessLaunchButton)</Button>
+            <Label X="-11" Y="-51" Width="400" Height="34" FontId="3" DisablePrefix="yes" VisibleCondition="WixStdBARestartRequired">
+                <Text>#(loc.SuccessRestartText)</Text>
+                <Text Condition="WixBundleAction = 3">#(loc.SuccessUninstallRestartText)</Text>
+            </Label>
+            <Button Name="SuccessRestartButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.SuccessRestartButton)</Button>
+            <Button Name="SuccessCloseButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.SuccessCloseButton)</Text>
+                <CloseWindowAction />
+            </Button>
+        </Page>
+        <Page Name="Failure">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">
+                <Text>#(loc.FailureHeader)</Text>
+                <Text Condition="WixBundleAction = 2">#(loc.FailureLayoutHeader)</Text>
+                <Text Condition="WixBundleAction = 3">#(loc.FailureUninstallHeader)</Text>
+                <Text Condition="WixBundleAction = 5">#(loc.FailureInstallHeader)</Text>
+                <Text Condition="WixBundleAction = 7">#(loc.FailureRepairHeader)</Text>
+            </Label>
+            <Hypertext Name="FailureLogFileLink" X="11" Y="121" Width="-11" Height="42" FontId="3" TabStop="yes" HideWhenDisabled="yes">#(loc.FailureHyperlinkLogText)</Hypertext>
+            <Hypertext Name="FailureMessageText" X="22" Y="163" Width="-11" Height="51" FontId="3" TabStop="yes" HideWhenDisabled="yes" />
+            <Label X="-11" Y="-51" Width="400" Height="34" FontId="3" DisablePrefix="yes" VisibleCondition="WixStdBARestartRequired">#(loc.FailureRestartText)</Label>
+            <Button Name="FailureRestartButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.FailureRestartButton)</Button>
+            <Button Name="FailureCloseButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
+                <Text>#(loc.FailureCloseButton)</Text>
+                <CloseWindowAction />
+            </Button>
+        </Page>
+    </Window>
+</Theme>