Add more thmutil window messages to simplify handling control events.

9 files changed, 564 insertions, 213 deletions

diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 2a34aaad..ab6ea4d7 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -86,6 +86,8 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONTHEMELOADED = 1024,
     BA_FUNCTIONS_MESSAGE_WNDPROC,
     BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLLOADING,
+    BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMCOMMAND,
+    BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMNOTIFY,
 };
 
 typedef HRESULT(WINAPI *PFN_BA_FUNCTIONS_PROC)(
@@ -126,6 +128,38 @@ struct BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS
     WORD wId;
 };
 
+struct BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_ARGS
+{
+    DWORD cbSize;
+    WPARAM wParam;
+    LPCWSTR wzName;
+    WORD wId;
+    HWND hWnd;
+};
+
+struct BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_RESULTS
+{
+    DWORD cbSize;
+    BOOL fProcessed;
+    LRESULT lResult;
+};
+
+struct BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_ARGS
+{
+    DWORD cbSize;
+    LPNMHDR lParam;
+    LPCWSTR wzName;
+    WORD wId;
+    HWND hWnd;
+};
+
+struct BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_RESULTS
+{
+    DWORD cbSize;
+    BOOL fProcessed;
+    LRESULT lResult;
+};
+
 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 c2c8a6dc..f6c33f58 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -850,6 +850,30 @@ public: // IBAFunctions
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnThemeControlWmCommand(
+        __in WPARAM /*wParam*/,
+        __in LPCWSTR /*wzName*/,
+        __in WORD /*wId*/,
+        __in HWND /*hWnd*/,
+        __inout BOOL* /*pfProcessed*/,
+        __inout LRESULT* /*plResult*/
+        )
+    {
+        return S_OK;
+    }
+
+    virtual STDMETHODIMP OnThemeControlWmNotify(
+        __in LPNMHDR /*lParam*/,
+        __in LPCWSTR /*wzName*/,
+        __in WORD /*wId*/,
+        __in HWND /*hWnd*/,
+        __inout BOOL* /*pfProcessed*/,
+        __inout LRESULT* /*plResult*/
+        )
+    {
+        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 efe22ddd..1d51c5b6 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -33,6 +33,24 @@ static HRESULT BalBaseBAFunctionsProcOnThemeControlLoading(
     return pBAFunctions->OnThemeControlLoading(pArgs->wzName, &pResults->fProcessed, &pResults->wId);
 }
 
+static HRESULT BalBaseBAFunctionsProcOnThemeControlWmCommand(
+    __in IBAFunctions* pBAFunctions,
+    __in BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_ARGS* pArgs,
+    __inout BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_RESULTS* pResults
+    )
+{
+    return pBAFunctions->OnThemeControlWmCommand(pArgs->wParam, pArgs->wzName, pArgs->wId, pArgs->hWnd, &pResults->fProcessed, &pResults->lResult);
+}
+
+static HRESULT BalBaseBAFunctionsProcOnThemeControlWmNotify(
+    __in IBAFunctions* pBAFunctions,
+    __in BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_ARGS* pArgs,
+    __inout BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_RESULTS* pResults
+    )
+{
+    return pBAFunctions->OnThemeControlWmNotify(pArgs->lParam, pArgs->wzName, pArgs->wId, pArgs->hWnd, &pResults->fProcessed, &pResults->lResult);
+}
+
 /*******************************************************************
 BalBaseBAFunctionsProc - requires pvContext to be of type IBAFunctions.
 Provides a default mapping between the message based BAFunctions interface and
@@ -137,6 +155,12 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLLOADING:
             hr = BalBaseBAFunctionsProcOnThemeControlLoading(pBAFunctions, reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLLOADING_ARGS*>(pvArgs), reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLLOADING_RESULTS*>(pvResults));
             break;
+        case BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMCOMMAND:
+            hr = BalBaseBAFunctionsProcOnThemeControlWmCommand(pBAFunctions, reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_ARGS*>(pvArgs), reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_RESULTS*>(pvResults));
+            break;
+        case BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMNOTIFY:
+            hr = BalBaseBAFunctionsProcOnThemeControlWmNotify(pBAFunctions, reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_ARGS*>(pvArgs), reinterpret_cast<BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_RESULTS*>(pvResults));
+            break;
         }
     }
 
diff --git a/src/api/burn/balutil/inc/IBAFunctions.h b/src/api/burn/balutil/inc/IBAFunctions.h
index d41b7c9b..63395e1e 100644
--- a/src/api/burn/balutil/inc/IBAFunctions.h
+++ b/src/api/burn/balutil/inc/IBAFunctions.h
@@ -39,4 +39,26 @@ DECLARE_INTERFACE_IID_(IBAFunctions, IBootstrapperApplication, "0FB445ED-17BD-49
         __inout BOOL* pfProcessed,
         __inout WORD* pwId
         ) = 0;
+
+    // OnThemeControlWmCommand - Called when WM_COMMAND is received for a control.
+    //
+    STDMETHOD(OnThemeControlWmCommand)(
+        __in WPARAM wParam,
+        __in LPCWSTR wzName,
+        __in WORD wId,
+        __in HWND hWnd,
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
+        ) = 0;
+
+    // OnThemeControlWmNotify - Called when WM_NOTIFY is received for a control.
+    //
+    STDMETHOD(OnThemeControlWmNotify)(
+        __in LPNMHDR lParam,
+        __in LPCWSTR wzName,
+        __in WORD wId,
+        __in HWND hWnd,
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
+        ) = 0;
 };
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 8cdd31ce..2283880c 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -2878,65 +2878,11 @@ private:
             pBA->OnShowFailure();
             return 0;
 
-        case WM_COMMAND:
-            switch (HIWORD(wParam))
-            {
-            case BN_CLICKED:
-                switch (LOWORD(wParam))
-                {
-                case WIXSTDBA_CONTROL_EULA_ACCEPT_CHECKBOX:
-                    pBA->OnClickAcceptCheckbox();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_INSTALL_BUTTON:
-                    pBA->OnClickInstallButton();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_REPAIR_BUTTON:
-                    pBA->OnClickRepairButton();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_UNINSTALL_BUTTON:
-                    pBA->OnClickUninstallButton();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_LAUNCH_BUTTON:
-                    pBA->OnClickLaunchButton();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_SUCCESS_RESTART_BUTTON: __fallthrough;
-                case WIXSTDBA_CONTROL_FAILURE_RESTART_BUTTON:
-                    pBA->OnClickRestartButton();
-                    return 0;
-
-                case WIXSTDBA_CONTROL_PROGRESS_CANCEL_BUTTON:
-                    pBA->OnClickCloseButton();
-                    return 0;
-                }
-                break;
-            }
-            break;
+        case WM_THMUTIL_CONTROL_WM_COMMAND:
+            return pBA->OnThemeControlWmCommand(reinterpret_cast<THEME_CONTROLWMCOMMAND_ARGS*>(wParam), reinterpret_cast<THEME_CONTROLWMCOMMAND_RESULTS*>(lParam));
 
-        case WM_NOTIFY:
-            if (lParam)
-            {
-                LPNMHDR pnmhdr = reinterpret_cast<LPNMHDR>(lParam);
-                switch (pnmhdr->code)
-                {
-                case NM_CLICK: __fallthrough;
-                case NM_RETURN:
-                    switch (static_cast<DWORD>(pnmhdr->idFrom))
-                    {
-                    case WIXSTDBA_CONTROL_EULA_LINK:
-                        pBA->OnClickEulaLink();
-                        return 1;
-                    case WIXSTDBA_CONTROL_FAILURE_LOGFILE_LINK:
-                        pBA->OnClickLogFileLink();
-                        return 1;
-                    }
-                }
-            }
-            break;
+        case WM_THMUTIL_CONTROL_WM_NOTIFY:
+            return pBA->OnThemeControlWmNotify(reinterpret_cast<THEME_CONTROLWMNOTIFY_ARGS*>(wParam), reinterpret_cast<THEME_CONTROLWMNOTIFY_RESULTS*>(lParam));
         }
 
         if (pBA && pBA->m_pTaskbarList && uMsg == pBA->m_uTaskbarButtonCreatedMessage)
@@ -3666,6 +3612,154 @@ private:
         ReleaseStr(sczLogFile);
     }
 
+    BOOL OnThemeControlWmCommand(
+        __in const THEME_CONTROLWMCOMMAND_ARGS* pArgs,
+        __in THEME_CONTROLWMCOMMAND_RESULTS* pResults
+        )
+    {
+        HRESULT hr = S_OK;
+        BOOL fProcessed = FALSE;
+        BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_ARGS themeControlWmCommandArgs = { };
+        BA_FUNCTIONS_ONTHEMECONTROLWMCOMMAND_RESULTS themeControlWmCommandResults = { };
+
+        switch (HIWORD(pArgs->wParam))
+        {
+        case BN_CLICKED:
+            switch (pArgs->pThemeControl->wId)
+            {
+            case WIXSTDBA_CONTROL_EULA_ACCEPT_CHECKBOX:
+                OnClickAcceptCheckbox();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_INSTALL_BUTTON:
+                OnClickInstallButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_REPAIR_BUTTON:
+                OnClickRepairButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_UNINSTALL_BUTTON:
+                OnClickUninstallButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_LAUNCH_BUTTON:
+                OnClickLaunchButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_SUCCESS_RESTART_BUTTON: __fallthrough;
+            case WIXSTDBA_CONTROL_FAILURE_RESTART_BUTTON:
+                OnClickRestartButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
+            case WIXSTDBA_CONTROL_PROGRESS_CANCEL_BUTTON:
+                OnClickCloseButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+            }
+            break;
+        }
+
+        if (m_pfnBAFunctionsProc)
+        {
+            themeControlWmCommandArgs.cbSize = sizeof(themeControlWmCommandArgs);
+            themeControlWmCommandArgs.wParam = pArgs->wParam;
+            themeControlWmCommandArgs.wzName = pArgs->pThemeControl->sczName;
+            themeControlWmCommandArgs.wId = pArgs->pThemeControl->wId;
+            themeControlWmCommandArgs.hWnd = pArgs->pThemeControl->hWnd;
+
+            themeControlWmCommandResults.cbSize = sizeof(themeControlWmCommandResults);
+            themeControlWmCommandResults.lResult = pResults->lResult;
+
+            hr = m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMCOMMAND, &themeControlWmCommandArgs, &themeControlWmCommandResults, m_pvBAFunctionsProcContext);
+            if (E_NOTIMPL != hr)
+            {
+                BalExitOnFailure(hr, "BAFunctions OnThemeControlWmCommand failed.");
+
+                if (themeControlWmCommandResults.fProcessed)
+                {
+                    fProcessed = TRUE;
+                    pResults->lResult = themeControlWmCommandResults.lResult;
+                    ExitFunction();
+                }
+            }
+        }
+
+LExit:
+        return fProcessed;
+    }
+
+    BOOL OnThemeControlWmNotify(
+        __in const THEME_CONTROLWMNOTIFY_ARGS* pArgs,
+        __in THEME_CONTROLWMNOTIFY_RESULTS* pResults
+        )
+    {
+        HRESULT hr = S_OK;
+        BOOL fProcessed = FALSE;
+        BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_ARGS themeControlWmNotifyArgs = { };
+        BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_RESULTS themeControlWmNotifyResults = { };
+
+        switch (pArgs->lParam->code)
+        {
+        case NM_CLICK: __fallthrough;
+        case NM_RETURN:
+            switch (pArgs->pThemeControl->wId)
+            {
+            case WIXSTDBA_CONTROL_EULA_LINK:
+                OnClickEulaLink();
+                fProcessed = TRUE;
+                pResults->lResult = 1;
+                ExitFunction();
+            case WIXSTDBA_CONTROL_FAILURE_LOGFILE_LINK:
+                OnClickLogFileLink();
+                fProcessed = TRUE;
+                pResults->lResult = 1;
+                ExitFunction();
+            }
+        }
+
+        if (m_pfnBAFunctionsProc)
+        {
+            themeControlWmNotifyArgs.cbSize = sizeof(themeControlWmNotifyArgs);
+            themeControlWmNotifyArgs.lParam = pArgs->lParam;
+            themeControlWmNotifyArgs.wzName = pArgs->pThemeControl->sczName;
+            themeControlWmNotifyArgs.wId = pArgs->pThemeControl->wId;
+            themeControlWmNotifyArgs.hWnd = pArgs->pThemeControl->hWnd;
+
+            themeControlWmNotifyResults.cbSize = sizeof(themeControlWmNotifyResults);
+            themeControlWmNotifyResults.lResult = pResults->lResult;
+
+            hr = m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONTHEMECONTROLWMCOMMAND, &themeControlWmNotifyArgs, &themeControlWmNotifyResults, m_pvBAFunctionsProcContext);
+            if (E_NOTIMPL != hr)
+            {
+                BalExitOnFailure(hr, "BAFunctions OnThemeControlWmNotify failed.");
+
+                if (themeControlWmNotifyResults.fProcessed)
+                {
+                    fProcessed = TRUE;
+                    pResults->lResult = themeControlWmNotifyResults.lResult;
+                    ExitFunction();
+                }
+            }
+        }
+
+LExit:
+        return fProcessed;
+    }
+
 
     //
     // SetState
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h b/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
index dc554219..9557c11c 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/thmutil.h
@@ -103,6 +103,14 @@ typedef enum _WM_THMUTIL
     // 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,
+    // Sent when WM_COMMAND is received for a control.
+    // wparam is THEME_CONTROLWMCOMMAND_ARGS* and lparam is THEME_CONTROLWMCOMMAND_RESULTS*.
+    // Return code is TRUE if it was processed.
+    WM_THMUTIL_CONTROL_WM_COMMAND = WM_APP - 2,
+    // Sent when WM_NOTIFY is received for a control.
+    // wparam is THEME_CONTROLWMNOTIFY_ARGS* and lparam is THEME_CONTROLWMNOTIFY_RESULTS*.
+    // Return code is TRUE to prevent further processing of the message.
+    WM_THMUTIL_CONTROL_WM_NOTIFY = WM_APP - 3,
 } WM_THMUTIL;
 
 struct THEME_COLUMN
@@ -397,6 +405,32 @@ struct THEME
     LPVOID pvVariableContext;
 };
 
+typedef struct _THEME_CONTROLWMCOMMAND_ARGS
+{
+    DWORD cbSize;
+    WPARAM wParam;
+    const THEME_CONTROL* pThemeControl;
+} THEME_CONTROLWMCOMMAND_ARGS;
+
+typedef struct _THEME_CONTROLWMCOMMAND_RESULTS
+{
+    DWORD cbSize;
+    LRESULT lResult;
+} THEME_CONTROLWMCOMMAND_RESULTS;
+
+typedef struct _THEME_CONTROLWMNOTIFY_ARGS
+{
+    DWORD cbSize;
+    LPNMHDR lParam;
+    const THEME_CONTROL* pThemeControl;
+} THEME_CONTROLWMNOTIFY_ARGS;
+
+typedef struct _THEME_CONTROLWMNOTIFY_RESULTS
+{
+    DWORD cbSize;
+    LRESULT lResult;
+} THEME_CONTROLWMNOTIFY_RESULTS;
+
 typedef struct _THEME_LOADINGCONTROL_ARGS
 {
     DWORD cbSize;
diff --git a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
index d796bbaf..a3c5d80c 100644
--- a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
@@ -421,12 +421,10 @@ static void CALLBACK OnBillboardTimer(
     );
 static void OnBrowseDirectory(
     __in THEME* pTheme,
-    __in HWND hWnd,
     __in const THEME_ACTION* pAction
     );
 static BOOL OnButtonClicked(
     __in THEME* pTheme,
-    __in HWND hWnd,
     __in const THEME_CONTROL* pControl
     );
 static BOOL OnDpiChanged(
@@ -434,20 +432,37 @@ static BOOL OnDpiChanged(
     __in WPARAM wParam,
     __in LPARAM lParam
     );
+static BOOL OnHypertextClicked(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pThemeControl,
+    __in PNMLINK pnmlink
+    );
 static void OnNcCreate(
     __in THEME* pTheme,
     __in HWND hWnd,
     __in LPARAM lParam
     );
-static HRESULT OnRichEditEnLink(
-    __in LPARAM lParam,
-    __in HWND hWndRichEdit,
-    __in HWND hWnd
+static BOOL OnNotifyEnLink(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pThemeControl,
+    __in ENLINK* link
     );
-static BOOL ControlIsType(
-    __in const THEME* pTheme,
-    __in DWORD dwControl,
-    __in THEME_CONTROL_TYPE type
+static BOOL OnNotifyEnMsgFilter(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pThemeControl,
+    __in MSGFILTER* msgFilter
+    );
+static BOOL OnWmCommand(
+    __in THEME* pTheme,
+    __in WPARAM wParam,
+    __in const THEME_CONTROL* pThemeControl,
+    __inout LRESULT* plResult
+    );
+static BOOL OnWmNotify(
+    __in THEME* pTheme,
+    __in LPNMHDR lParam,
+    __in const THEME_CONTROL* pThemeControl,
+    __inout LRESULT* plResult
     );
 static const THEME_CONTROL* FindControlFromHWnd(
     __in const THEME* pTheme,
@@ -5089,7 +5104,6 @@ static void CALLBACK OnBillboardTimer(
 
 static void OnBrowseDirectory(
     __in THEME* pTheme,
-    __in HWND hWnd,
     __in const THEME_ACTION* pAction
     )
 {
@@ -5098,7 +5112,7 @@ static void OnBrowseDirectory(
     BROWSEINFOW browseInfo = { };
     PIDLIST_ABSOLUTE pidl = NULL;
 
-    browseInfo.hwndOwner = hWnd;
+    browseInfo.hwndOwner = pTheme->hwndParent;
     browseInfo.pszDisplayName = wzPath;
     browseInfo.lpszTitle = pTheme->sczCaption;
     browseInfo.ulFlags = BIF_RETURNONLYFSDIRS | BIF_USENEWUI;
@@ -5148,7 +5162,6 @@ LExit:
 
 static BOOL OnButtonClicked(
     __in THEME* pTheme,
-    __in HWND hWnd,
     __in const THEME_CONTROL* pControl
     )
 {
@@ -5191,11 +5204,11 @@ static BOOL OnButtonClicked(
                 switch (pChosenAction->type)
                 {
                 case THEME_ACTION_TYPE_BROWSE_DIRECTORY:
-                    OnBrowseDirectory(pTheme, hWnd, pChosenAction);
+                    OnBrowseDirectory(pTheme, pChosenAction);
                     break;
 
                 case THEME_ACTION_TYPE_CLOSE_WINDOW:
-                    ::SendMessageW(hWnd, WM_CLOSE, 0, 0);
+                    ::SendMessageW(pTheme->hwndParent, WM_CLOSE, 0, 0);
                     break;
 
                 case THEME_ACTION_TYPE_CHANGE_PAGE:
@@ -5272,6 +5285,25 @@ LExit:
     return !fIgnored;
 }
 
+static BOOL OnHypertextClicked(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* /*pThemeControl*/,
+    __in PNMLINK pnmlink
+    )
+{
+    BOOL fProcessed = FALSE;
+    HRESULT hr = S_OK;
+    LITEM litem = pnmlink->item;
+
+    hr = ShelExec(litem.szUrl, NULL, L"open", NULL, SW_SHOWDEFAULT, pTheme->hwndParent, NULL);
+    ThmExitOnFailure(hr, "Failed to launch hypertext link: %ls", litem.szUrl);
+
+    fProcessed = TRUE;
+
+LExit:
+    return fProcessed;
+}
+
 static void OnNcCreate(
     __in THEME* pTheme,
     __in HWND hWnd,
@@ -5291,63 +5323,185 @@ static void OnNcCreate(
     }
 }
 
-static HRESULT OnRichEditEnLink(
-    __in LPARAM lParam,
-    __in HWND hWndRichEdit,
-    __in HWND hWnd
+static BOOL OnNotifyEnLink(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pThemeControl,
+    __in ENLINK* link
     )
 {
+    BOOL fProcessed = FALSE;
     HRESULT hr = S_OK;
     LPWSTR sczLink = NULL;
-    ENLINK* link = reinterpret_cast<ENLINK*>(lParam);
+    TEXTRANGEW tr = { };
 
-    switch (link->msg)
+    // Hyperlink clicks from rich-edit control.
+    if (THEME_CONTROL_TYPE_RICHEDIT == pThemeControl->type)
     {
-    case WM_LBUTTONDOWN:
+        switch (link->msg)
         {
-        hr = StrAlloc(&sczLink, link->chrg.cpMax - link->chrg.cpMin + 2);
-        ThmExitOnFailure(hr, "Failed to allocate string for link.");
+        case WM_LBUTTONDOWN:
+            hr = StrAlloc(&sczLink, (SIZE_T)2 + link->chrg.cpMax - link->chrg.cpMin);
+            ThmExitOnFailure(hr, "Failed to allocate string for link.");
 
-        TEXTRANGEW tr;
-        tr.chrg.cpMin = link->chrg.cpMin;
-        tr.chrg.cpMax = link->chrg.cpMax;
-        tr.lpstrText = sczLink;
+            tr.chrg.cpMin = link->chrg.cpMin;
+            tr.chrg.cpMax = link->chrg.cpMax;
+            tr.lpstrText = sczLink;
 
-        if (0 < ::SendMessageW(hWndRichEdit, EM_GETTEXTRANGE, 0, reinterpret_cast<LPARAM>(&tr)))
-        {
-            hr = ShelExec(sczLink, NULL, L"open", NULL, SW_SHOWDEFAULT, hWnd, NULL);
-            ThmExitOnFailure(hr, "Failed to launch link: %ls", sczLink);
+            if (0 < ::SendMessageW(pThemeControl->hWnd, EM_GETTEXTRANGE, 0, reinterpret_cast<LPARAM>(&tr)))
+            {
+                hr = ShelExec(sczLink, NULL, L"open", NULL, SW_SHOWDEFAULT, pTheme->hwndParent, NULL);
+                ThmExitOnFailure(hr, "Failed to launch link: %ls", sczLink);
+
+                fProcessed = TRUE;
+            }
+
+            break;
+
+        case WM_SETCURSOR:
+            ::SetCursor(vhCursorHand);
+            fProcessed = TRUE;
+
+            break;
         }
-        
-        break;
+    }
+
+LExit:
+    ReleaseStr(sczLink);
+
+    return fProcessed;
+}
+
+static BOOL OnNotifyEnMsgFilter(
+    __in THEME* pTheme,
+    __in const THEME_CONTROL* pThemeControl,
+    __in MSGFILTER* msgFilter
+    )
+{
+    BOOL fProcessed = FALSE;
+
+    // Tab/Shift+Tab support for rich-edit control.
+    if (THEME_CONTROL_TYPE_RICHEDIT == pThemeControl->type)
+    {
+        switch (msgFilter->msg)
+        {
+        case WM_KEYDOWN:
+            if (VK_TAB == msgFilter->wParam)
+            {
+                BOOL fShift = 0x8000 & ::GetKeyState(VK_SHIFT);
+                HWND hwndFocus = ::GetNextDlgTabItem(pTheme->hwndParent, pThemeControl->hWnd, fShift);
+                ::SetFocus(hwndFocus);
+
+                fProcessed = TRUE;
+            }
+            break;
         }
+    }
+
+    return fProcessed;
+}
 
-    case WM_SETCURSOR:
-        ::SetCursor(vhCursorHand);
+static BOOL OnWmCommand(
+    __in THEME* pTheme,
+    __in WPARAM wParam,
+    __in const THEME_CONTROL* pThemeControl,
+    __inout LRESULT* plResult
+    )
+{
+    BOOL fProcessed = FALSE;
+    THEME_CONTROLWMCOMMAND_ARGS args = { };
+    THEME_CONTROLWMCOMMAND_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.wParam = wParam;
+    args.pThemeControl = pThemeControl;
+
+    results.cbSize = sizeof(results);
+    results.lResult = *plResult;
+
+    if (::SendMessageW(pTheme->hwndParent, WM_THMUTIL_CONTROL_WM_COMMAND, reinterpret_cast<WPARAM>(&args), reinterpret_cast<LPARAM>(&results)))
+    {
+        fProcessed = TRUE;
+        *plResult = results.lResult;
+        ExitFunction();
+    }
+
+    switch (HIWORD(wParam))
+    {
+    case BN_CLICKED:
+        if (OnButtonClicked(pTheme, pThemeControl))
+        {
+            fProcessed = TRUE;
+            *plResult = 0;
+            ExitFunction();
+        }
         break;
     }
 
 LExit:
-    ReleaseStr(sczLink);
-
-    return hr;
+    return fProcessed;
 }
 
-static BOOL ControlIsType(
-    __in const THEME* pTheme,
-    __in DWORD dwControl,
-    __in const THEME_CONTROL_TYPE type
+static BOOL OnWmNotify(
+    __in THEME* pTheme,
+    __in LPNMHDR lParam,
+    __in const THEME_CONTROL* pThemeControl,
+    __inout LRESULT* plResult
     )
 {
-    BOOL fIsType = FALSE;
-    HWND hWnd = ::GetDlgItem(pTheme->hwndParent, dwControl);
-    if (hWnd)
+    BOOL fProcessed = FALSE;
+    THEME_CONTROLWMNOTIFY_ARGS args = { };
+    THEME_CONTROLWMNOTIFY_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.lParam = lParam;
+    args.pThemeControl = pThemeControl;
+
+    results.cbSize = sizeof(results);
+    results.lResult = *plResult;
+
+    if (::SendMessageW(pTheme->hwndParent, WM_THMUTIL_CONTROL_WM_NOTIFY, reinterpret_cast<WPARAM>(&args), reinterpret_cast<LPARAM>(&results)))
     {
-        const THEME_CONTROL* pControl = FindControlFromHWnd(pTheme, hWnd);
-        fIsType = (pControl && type == pControl->type);
+        fProcessed = TRUE;
+        *plResult = results.lResult;
+        ExitFunction();
     }
 
-    return fIsType;
+    switch (lParam->code)
+    {
+    case EN_MSGFILTER:
+        if (OnNotifyEnMsgFilter(pTheme, pThemeControl, reinterpret_cast<MSGFILTER*>(lParam)))
+        {
+            fProcessed = TRUE;
+            *plResult = 1;
+            ExitFunction();
+        }
+        break;
+
+    case EN_LINK:
+        if (OnNotifyEnLink(pTheme, pThemeControl, reinterpret_cast<ENLINK*>(lParam)))
+        {
+            fProcessed = TRUE;
+            *plResult = 1;
+            ExitFunction();
+        }
+
+    case NM_CLICK: __fallthrough;
+    case NM_RETURN:
+        switch (pThemeControl->type)
+        {
+        case THEME_CONTROL_TYPE_HYPERTEXT:
+            // Clicks on a hypertext/syslink control.
+            if (OnHypertextClicked(pTheme, pThemeControl, reinterpret_cast<PNMLINK>(lParam)))
+            {
+                fProcessed = TRUE;
+                *plResult = 1;
+                ExitFunction();
+            }
+        }
+    }
+
+LExit:
+    return fProcessed;
 }
 
 static const THEME_CONTROL* FindControlFromHWnd(
@@ -5780,6 +5934,9 @@ static LRESULT CALLBACK ControlGroupDefWindowProc(
     __in LPARAM lParam
     )
 {
+    LRESULT lres = 0;
+    const THEME_CONTROL* pThemeControl = NULL;
+
     if (pTheme)
     {
         switch (uMsg)
@@ -5827,55 +5984,22 @@ static LRESULT CALLBACK ControlGroupDefWindowProc(
             if (lParam)
             {
                 LPNMHDR pnmhdr = reinterpret_cast<LPNMHDR>(lParam);
-                switch (pnmhdr->code)
+                pThemeControl = FindControlFromHWnd(pTheme, pnmhdr->hwndFrom);
+                if (pThemeControl && OnWmNotify(pTheme, pnmhdr, pThemeControl, &lres))
                 {
-                    // Tab/Shift+Tab support for rich-edit control.
-                case EN_MSGFILTER:
-                {
-                    MSGFILTER* msgFilter = reinterpret_cast<MSGFILTER*>(lParam);
-                    if (WM_KEYDOWN == msgFilter->msg && VK_TAB == msgFilter->wParam)
-                    {
-                        BOOL fShift = 0x8000 & ::GetKeyState(VK_SHIFT);
-                        HWND hwndFocus = ::GetNextDlgTabItem(hWnd, msgFilter->nmhdr.hwndFrom, fShift);
-                        ::SetFocus(hwndFocus);
-                        return 1;
-                    }
-                    break;
-                }
-
-                // Hyperlink clicks from rich-edit control.
-                case EN_LINK:
-                    return SUCCEEDED(OnRichEditEnLink(lParam, pnmhdr->hwndFrom, hWnd));
-
-                    // Clicks on a hypertext/syslink control.
-                case NM_CLICK: __fallthrough;
-                case NM_RETURN:
-                    if (ControlIsType(pTheme, static_cast<DWORD>(pnmhdr->idFrom), THEME_CONTROL_TYPE_HYPERTEXT))
-                    {
-                        PNMLINK pnmlink = reinterpret_cast<PNMLINK>(lParam);
-                        LITEM litem = pnmlink->item;
-                        ShelExec(litem.szUrl, NULL, L"open", NULL, SW_SHOWDEFAULT, hWnd, NULL);
-                        return 1;
-                    }
-
-                    return 0;
+                    return lres;
                 }
             }
             break;
 
         case WM_COMMAND:
-            switch (HIWORD(wParam))
+            if (lParam)
             {
-            case BN_CLICKED:
-                if (lParam)
+                pThemeControl = FindControlFromHWnd(pTheme, (HWND)lParam);
+                if (pThemeControl && OnWmCommand(pTheme, wParam, pThemeControl, &lres))
                 {
-                    const THEME_CONTROL* pControl = FindControlFromHWnd(pTheme, (HWND)lParam);
-                    if (pControl && OnButtonClicked(pTheme, hWnd, pControl))
-                    {
-                        return 0;
-                    }
+                    return lres;
                 }
-                break;
             }
             break;
         }
diff --git a/src/samples/thmviewer/thmviewer.cpp b/src/samples/thmviewer/thmviewer.cpp
index cffa3851..e593d6ad 100644
--- a/src/samples/thmviewer/thmviewer.cpp
+++ b/src/samples/thmviewer/thmviewer.cpp
@@ -52,6 +52,10 @@ static BOOL OnThemeLoadingControl(
     __in const THEME_LOADINGCONTROL_ARGS* pArgs,
     __in THEME_LOADINGCONTROL_RESULTS* pResults
     );
+static BOOL OnThemeControlWmNotify(
+    __in const THEME_CONTROLWMNOTIFY_ARGS* pArgs,
+    __in THEME_CONTROLWMNOTIFY_RESULTS* pResults
+    );
 static void CALLBACK ThmviewerTraceError(
     __in_z LPCSTR szFile,
     __in int iLine,
@@ -377,32 +381,11 @@ static LRESULT CALLBACK MainWndProc(
         ::PostQuitMessage(0);
         break;
 
-    case WM_NOTIFY:
-        {
-        NMHDR* pnmhdr = reinterpret_cast<NMHDR*>(lParam);
-        switch (pnmhdr->code)
-        {
-        case TVN_SELCHANGEDW:
-            {
-            NMTREEVIEWW* ptv = reinterpret_cast<NMTREEVIEWW*>(lParam);
-            ::PostThreadMessageW(vdwDisplayThreadId, WM_THMVWR_SHOWPAGE, SW_HIDE, ptv->itemOld.lParam);
-            ::PostThreadMessageW(vdwDisplayThreadId, WM_THMVWR_SHOWPAGE, SW_SHOW, ptv->itemNew.lParam);
-            }
-            break;
-
-        //case NM_DBLCLK:
-        //    TVITEM item = { };
-        //    item.mask = TVIF_PARAM;
-        //    item.hItem = TreeView_GetSelection(pnmhdr->hwndFrom);
-        //    TreeView_GetItem(pnmhdr->hwndFrom, &item);
-        //    ::PostThreadMessageW(vdwDisplayThreadId, WM_THMVWR_SHOWPAGE, SW_SHOW, item.lParam);
-        //    return 1;
-        }
-        }
-        break;
-
     case WM_THMUTIL_LOADING_CONTROL:
         return OnThemeLoadingControl(reinterpret_cast<THEME_LOADINGCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADINGCONTROL_RESULTS*>(lParam));
+
+    case WM_THMUTIL_CONTROL_WM_NOTIFY:
+        return OnThemeControlWmNotify(reinterpret_cast<THEME_CONTROLWMNOTIFY_ARGS*>(wParam), reinterpret_cast<THEME_CONTROLWMNOTIFY_RESULTS*>(lParam));
     }
 
     return ThemeDefWindowProc(vpTheme, hWnd, uMsg, wParam, lParam);
@@ -558,3 +541,29 @@ static BOOL OnThemeLoadingControl(
     pResults->hr = S_OK;
     return TRUE;
 }
+
+static BOOL OnThemeControlWmNotify(
+    __in const THEME_CONTROLWMNOTIFY_ARGS* pArgs,
+    __in THEME_CONTROLWMNOTIFY_RESULTS* /*pResults*/
+    )
+{
+    BOOL fProcessed = FALSE;
+
+    switch (pArgs->lParam->code)
+    {
+    case TVN_SELCHANGEDW:
+        switch (pArgs->pThemeControl->wId)
+        {
+        case THMVWR_CONTROL_TREE:
+            NMTREEVIEWW* ptv = reinterpret_cast<NMTREEVIEWW*>(pArgs->lParam);
+            ::PostThreadMessageW(vdwDisplayThreadId, WM_THMVWR_SHOWPAGE, SW_HIDE, ptv->itemOld.lParam);
+            ::PostThreadMessageW(vdwDisplayThreadId, WM_THMVWR_SHOWPAGE, SW_SHOW, ptv->itemNew.lParam);
+
+            fProcessed = TRUE;
+            break;
+        }
+        break;
+    }
+
+    return fProcessed;
+}
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index b35b4e02..a5bcba3e 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -43,62 +43,48 @@ static void CALLBACK BafThmUtilTestingTraceError(
 class CBafThmUtilTesting : public CBalBaseBAFunctions
 {
 public: // IBAFunctions
-    /*virtual STDMETHODIMP OnThemeLoaded(
-        THEME* pTheme,
-        WIX_LOCALIZATION* pWixLoc
+    virtual STDMETHODIMP OnThemeControlLoading(
+        __in LPCWSTR wzName,
+        __inout BOOL* pfProcessed,
+        __inout WORD* pwId
         )
     {
-        HRESULT hr = S_OK;
-
-        hr = __super::OnThemeLoaded(pTheme, pWixLoc);
+        if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzName, -1, L"InstallTestButton", -1))
+        {
+            *pfProcessed = TRUE;
+            *pwId = BAF_CONTROL_INSTALL_TEST_BUTTON;
+        }
 
-        return hr;
-    }*/
+        return S_OK;
+    }
 
-    virtual STDMETHODIMP WndProc(
-        __in THEME* pTheme,
-        __in HWND hWnd,
-        __in UINT uMsg,
+    virtual STDMETHODIMP OnThemeControlWmCommand(
         __in WPARAM wParam,
-        __in LPARAM lParam,
-        __inout LRESULT* plRes
+        __in LPCWSTR /*wzName*/,
+        __in WORD wId,
+        __in HWND /*hWnd*/,
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
         )
     {
-        switch (uMsg)
+        HRESULT hr = S_OK;
+
+        switch (HIWORD(wParam))
         {
-        case WM_COMMAND:
-            switch (HIWORD(wParam))
+        case BN_CLICKED:
+            switch (wId)
             {
-            case BN_CLICKED:
-                switch (LOWORD(wParam))
-                {
-                case BAF_CONTROL_INSTALL_TEST_BUTTON:
-                    OnShowTheme();
-                    *plRes = 0;
-                    return S_OK;
-                }
-
+            case BAF_CONTROL_INSTALL_TEST_BUTTON:
+                OnShowTheme();
+                *pfProcessed = TRUE;
+                *plResult = 0;
                 break;
             }
-            break;
-        }
-
-        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;
+            break;
         }
 
-        return S_OK;
+        return hr;
     }
 
 private: