8 files changed, 40 insertions, 47 deletions

diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 0eda95a1..149a12f8 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -178,8 +178,7 @@ struct BA_FUNCTIONS_ONTHEMECONTROLWMNOTIFY_RESULTS
 struct BA_FUNCTIONS_ONTHEMELOADED_ARGS
 {
     DWORD cbSize;
-    THEME* pTheme;
-    WIX_LOCALIZATION* pWixLoc;
+    HWND hWnd;
 };
 
 struct BA_FUNCTIONS_ONTHEMELOADED_RESULTS
@@ -190,7 +189,6 @@ struct BA_FUNCTIONS_ONTHEMELOADED_RESULTS
 struct BA_FUNCTIONS_WNDPROC_ARGS
 {
     DWORD cbSize;
-    THEME* pTheme;
     HWND hWnd;
     UINT uMsg;
     WPARAM wParam;
@@ -200,7 +198,8 @@ struct BA_FUNCTIONS_WNDPROC_ARGS
 struct BA_FUNCTIONS_WNDPROC_RESULTS
 {
     DWORD cbSize;
-    LRESULT lres;
+    BOOL fProcessed;
+    LRESULT lResult;
 };
 
 typedef HRESULT(WINAPI *PFN_BA_FUNCTIONS_CREATE)(
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index d97df350..1493253b 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -5,9 +5,6 @@
 #include <windows.h>
 #include <msiquery.h>
 
-#include "dutil.h"
-#include "locutil.h"
-#include "thmutil.h"
 #include "BAFunctions.h"
 #include "IBAFunctions.h"
 #include "BootstrapperEngine.h"
@@ -803,32 +800,26 @@ public: // IBAFunctions
     }
 
     virtual STDMETHODIMP OnThemeLoaded(
-        THEME* pTheme,
-        WIX_LOCALIZATION* pWixLoc
+        __in HWND hWnd
         )
     {
         HRESULT hr = S_OK;
 
-        m_pTheme = pTheme;
-        m_pWixLoc = pWixLoc;
+        m_hwndParent = hWnd;
 
         return hr;
     }
 
     virtual STDMETHODIMP WndProc(
-        __in THEME* pTheme,
-        __in HWND hWnd,
-        __in UINT uMsg,
-        __in WPARAM wParam,
-        __in LPARAM lParam,
-        __inout LRESULT* plRes
+        __in HWND /*hWnd*/,
+        __in UINT /*uMsg*/,
+        __in WPARAM /*wParam*/,
+        __in LPARAM /*lParam*/,
+        __inout BOOL* /*pfProcessed*/,
+        __inout LRESULT* /*plResult*/
         )
     {
-        HRESULT hr = S_OK;
-
-        *plRes = ThemeDefWindowProc(pTheme, hWnd, uMsg, wParam, lParam);
-
-        return hr;
+        return S_OK;
     }
 
     virtual STDMETHODIMP BAFunctionsProc(
@@ -917,6 +908,5 @@ protected:
     BA_FUNCTIONS_CREATE_ARGS m_bafCreateArgs;
     BOOTSTRAPPER_CREATE_ARGS m_baCreateArgs;
     BOOTSTRAPPER_COMMAND m_command;
-    THEME* m_pTheme;
-    WIX_LOCALIZATION* m_pWixLoc;
+    HWND m_hwndParent;
 };
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
index 2af231e8..135f7c76 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -12,7 +12,7 @@ static HRESULT BalBaseBAFunctionsProcOnThemeLoaded(
     __inout BA_FUNCTIONS_ONTHEMELOADED_RESULTS* /*pResults*/
     )
 {
-    return pBAFunctions->OnThemeLoaded(pArgs->pTheme, pArgs->pWixLoc);
+    return pBAFunctions->OnThemeLoaded(pArgs->hWnd);
 }
 
 static HRESULT BalBaseBAFunctionsProcWndProc(
@@ -21,7 +21,7 @@ static HRESULT BalBaseBAFunctionsProcWndProc(
     __inout BA_FUNCTIONS_WNDPROC_RESULTS* pResults
     )
 {
-    return pBAFunctions->WndProc(pArgs->pTheme, pArgs->hWnd, pArgs->uMsg, pArgs->wParam, pArgs->lParam, &pResults->lres);
+    return pBAFunctions->WndProc(pArgs->hWnd, pArgs->uMsg, pArgs->wParam, pArgs->lParam, &pResults->fProcessed, &pResults->lResult);
 }
 
 static HRESULT BalBaseBAFunctionsProcOnThemeControlLoading(
diff --git a/src/api/burn/balutil/inc/IBAFunctions.h b/src/api/burn/balutil/inc/IBAFunctions.h
index 2e71608d..50126ba1 100644
--- a/src/api/burn/balutil/inc/IBAFunctions.h
+++ b/src/api/burn/balutil/inc/IBAFunctions.h
@@ -7,20 +7,18 @@ DECLARE_INTERFACE_IID_(IBAFunctions, IBootstrapperApplication, "0FB445ED-17BD-49
     // OnThemeLoaded - Called after the BA finished loading all the controls for the theme.
     //
     STDMETHOD(OnThemeLoaded)(
-        THEME* pTheme,
-        WIX_LOCALIZATION* pWixLoc
+        __in HWND hWnd
         ) = 0;
 
     // WndProc - Called if the BA hasn't handled the message.
-    //           The implementation must either return E_NOTIMPL or call ThemeDefWindowProc for unhandled messages.
     //
     STDMETHOD(WndProc)(
-        __in THEME* pTheme,
         __in HWND hWnd,
         __in UINT uMsg,
         __in WPARAM wParam,
         __in LPARAM lParam,
-        __inout LRESULT* plRes
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
         ) = 0;
 
     // BAFunctionsProc - The PFN_BA_FUNCTIONS_PROC can call this method to give the BAFunctions raw access to the callback from WixStdBA.
diff --git a/src/api/burn/test/BalUtilUnitTest/precomp.h b/src/api/burn/test/BalUtilUnitTest/precomp.h
index f4a35e20..ce893906 100644
--- a/src/api/burn/test/BalUtilUnitTest/precomp.h
+++ b/src/api/burn/test/BalUtilUnitTest/precomp.h
@@ -14,8 +14,6 @@
 
 #include <dutil.h>
 #include <dictutil.h>
-#include <locutil.h>
-#include <thmutil.h>
 
 #include <BootstrapperEngine.h>
 #include <BootstrapperApplication.h>
diff --git a/src/ext/Bal/Samples/bafunctions/precomp.h b/src/ext/Bal/Samples/bafunctions/precomp.h
index 832d2250..cda11a40 100644
--- a/src/ext/Bal/Samples/bafunctions/precomp.h
+++ b/src/ext/Bal/Samples/bafunctions/precomp.h
@@ -23,13 +23,13 @@
 #include "dictutil.h"
 //#include "dirutil.h"
 #include "fileutil.h"
-#include "locutil.h"
+//#include "locutil.h"
 //#include "logutil.h"
 #include "pathutil.h"
 //#include "resrutil.h"
 //#include "shelutil.h"
 #include "strutil.h"
-#include "thmutil.h"
+//#include "thmutil.h"
 //#include "uriutil.h"
 //#include "xmlutil.h"
 #include "regutil.h"
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 5546b74f..cdb9d017 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -2783,7 +2783,6 @@ private:
             if (pBA->m_pfnBAFunctionsProc)
             {
                 wndProcArgs.cbSize = sizeof(wndProcArgs);
-                wndProcArgs.pTheme = pTheme;
                 wndProcArgs.hWnd = hWnd;
                 wndProcArgs.uMsg = uMsg;
                 wndProcArgs.wParam = wParam;
@@ -2791,10 +2790,20 @@ private:
                 wndProcResults.cbSize = sizeof(wndProcResults);
 
                 hr = pBA->m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_WNDPROC, &wndProcArgs, &wndProcResults, pBA->m_pvBAFunctionsProcContext);
-                if (E_NOTIMPL != hr)
+
+                if (E_NOTIMPL == hr)
                 {
-                    lres = wndProcResults.lres;
-                    ExitFunction();
+                    hr = S_OK;
+                }
+                else
+                {
+                    BalExitOnFailure(hr, "BAFunctions WndProc failed.");
+
+                    if (wndProcResults.fProcessed)
+                    {
+                        lres = wndProcResults.lResult;
+                        ExitFunction();
+                    }
                 }
             }
         }
@@ -2910,8 +2919,7 @@ private:
         if (m_pfnBAFunctionsProc)
         {
             themeLoadedArgs.cbSize = sizeof(themeLoadedArgs);
-            themeLoadedArgs.pTheme = m_pTheme;
-            themeLoadedArgs.pWixLoc = m_pWixLoc;
+            themeLoadedArgs.hWnd = m_pTheme->hwndParent;
             themeLoadedResults.cbSize = sizeof(themeLoadedResults);
             hr = m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONTHEMELOADED, &themeLoadedArgs, &themeLoadedResults, m_pvBAFunctionsProcContext);
             BalExitOnFailure(hr, "BAFunctions OnThemeLoaded failed.");
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index 8304403a..e5ff9131 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -100,7 +100,7 @@ private:
         hr = CreateTestingWindow();
         BalExitOnFailure(hr, "Failed to create BafThmUtilTesting window.");
 
-        ::EnableWindow(m_pTheme->hwndParent, FALSE);
+        ::EnableWindow(m_hwndParent, FALSE);
 
         // message pump
         while (0 != (fRet = ::GetMessageW(&msg, NULL, 0, 0)))
@@ -118,7 +118,7 @@ private:
         }
 
     LExit:
-        ::EnableWindow(m_pTheme->hwndParent, TRUE);
+        ::EnableWindow(m_hwndParent, TRUE);
 
         DestroyTestingWindow();
 
@@ -130,7 +130,7 @@ private:
     HRESULT CreateTestingWindow()
     {
         HRESULT hr = S_OK;
-        HICON hIcon = reinterpret_cast<HICON>(m_pTheme->hIcon);
+        HICON hIcon = reinterpret_cast<HICON>(m_pBafTheme->hIcon);
         WNDCLASSW wc = { };
         int x = CW_USEDEFAULT;
         int y = CW_USEDEFAULT;
@@ -151,7 +151,7 @@ private:
         wc.hInstance = m_hModule;
         wc.hIcon = hIcon;
         wc.hCursor = ::LoadCursorW(NULL, (LPCWSTR)IDC_ARROW);
-        wc.hbrBackground = m_pTheme->rgFonts[m_pBafTheme->dwFontId].hBackground;
+        wc.hbrBackground = m_pBafTheme->rgFonts[m_pBafTheme->dwFontId].hBackground;
         wc.lpszMenuName = NULL;
         wc.lpszClassName = BAFTHMUTILTESTING_WINDOW_CLASS;
         if (!::RegisterClassW(&wc))
@@ -168,7 +168,7 @@ private:
             y = ptCursor.y;
         }
 
-        hr = ThemeCreateParentWindow(m_pBafTheme, 0, wc.lpszClassName, m_pBafTheme->sczCaption, m_pBafTheme->dwStyle, x, y, m_pTheme->hwndParent, m_hModule, this, THEME_WINDOW_INITIAL_POSITION_CENTER_MONITOR_FROM_COORDINATES, &m_hWndBaf);
+        hr = ThemeCreateParentWindow(m_pBafTheme, 0, wc.lpszClassName, m_pBafTheme->sczCaption, m_pBafTheme->dwStyle, x, y, m_hwndParent, m_hModule, this, THEME_WINDOW_INITIAL_POSITION_CENTER_MONITOR_FROM_COORDINATES, &m_hWndBaf);
         ExitOnFailure(hr, "Failed to create window.");
 
         hr = S_OK;
@@ -216,7 +216,7 @@ private:
         case WM_CLOSE:
             if (pBaf)
             {
-                ::EnableWindow(pBaf->m_pTheme->hwndParent, TRUE);
+                ::EnableWindow(pBaf->m_hwndParent, TRUE);
             }
 
             break;