[PATCH v2 2/6] uiautomationcore: Increment module reference count when starting provider thread.

[Connor McAdams]

From: Connor McAdams

Signed-off-by: Connor McAdams 
---
 dlls/uiautomationcore/uia_main.c     | 23 +++++++++++++++++++++++
 dlls/uiautomationcore/uia_private.h  |  2 ++
 dlls/uiautomationcore/uia_provider.c | 10 ++++++++--
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/dlls/uiautomationcore/uia_main.c b/dlls/uiautomationcore/uia_main.c
index 37657032b53..f141b4ca716 100644
--- a/dlls/uiautomationcore/uia_main.c
+++ b/dlls/uiautomationcore/uia_main.c
@@ -27,6 +27,8 @@

WINE_DEFAULT_DEBUG_CHANNEL(uiautomation);

+HMODULE huia_module;
+
 struct uia_object_wrapper
 {
     IUnknown IUnknown_iface;
@@ -349,3 +351,24 @@ HRESULT WINAPI UiaDisconnectProvider(IRawElementProviderSimple *provider)
     FIXME("(%p): stub\n", provider);
     return E_NOTIMPL;
 }
+
+/***********************************************************************
+ *          DllMain (uiautomationcore.@)
+ */
+BOOL WINAPI DllMain(HINSTANCE hinst, DWORD reason, void *reserved)
+{
+    TRACE("(%p, %ld, %p)\n", hinst, reason, reserved);
+
+    switch (reason)
+    {
+    case DLL_PROCESS_ATTACH:
+        DisableThreadLibraryCalls(hinst);
+        huia_module = hinst;
+        break;
+
+    default:
+        break;
+    }
+
+    return TRUE;
+}
diff --git a/dlls/uiautomationcore/uia_private.h b/dlls/uiautomationcore/uia_private.h
index ed53da58471..f92d4bd3721 100644
--- a/dlls/uiautomationcore/uia_private.h
+++ b/dlls/uiautomationcore/uia_private.h
@@ -21,6 +21,8 @@
 #include "uiautomation.h"
 #include "uia_classes.h"

+extern HMODULE huia_module DECLSPEC_HIDDEN;
+
 enum uia_prop_type {
     PROP_TYPE_UNKNOWN,
     PROP_TYPE_ELEM_PROP,
diff --git a/dlls/uiautomationcore/uia_provider.c b/dlls/uiautomationcore/uia_provider.c
index 9ba74a4271f..f05a3e1ff8a 100644
--- a/dlls/uiautomationcore/uia_provider.c
+++ b/dlls/uiautomationcore/uia_provider.c
@@ -1193,7 +1193,7 @@ static DWORD WINAPI uia_provider_thread_proc(void *arg)
     {
         WARN("CreateWindow failed: %ld\n", GetLastError());
         CoUninitialize();
-        ExitThread(1);
+        FreeLibraryAndExitThread(huia_module, 1);
     }

SetWindowLongPtrW(hwnd, GWLP_WNDPROC, (LONG_PTR)uia_provider_thread_msg_proc);
@@ -1212,7 +1212,7 @@ static DWORD WINAPI uia_provider_thread_proc(void *arg)

DestroyWindow(hwnd);
     CoUninitialize();
-    ExitThread(0);
+    FreeLibraryAndExitThread(huia_module, 0);
 }

static BOOL uia_start_provider_thread(void)
@@ -1224,12 +1224,18 @@ static BOOL uia_start_provider_thread(void)
     {
         HANDLE ready_event;
         HANDLE events[2];
+        HMODULE hmodule;
         DWORD wait_obj;

+        /* Increment DLL reference count. */
+        GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+                (const WCHAR *)uia_start_provider_thread, &hmodule);
+
         events[0] = ready_event = CreateEventW(NULL, FALSE, FALSE, NULL);
         if (!(provider_thread.hthread = CreateThread(NULL, 0, uia_provider_thread_proc,
                 ready_event, 0, NULL)))
         {
+            FreeLibrary(hmodule);
             started = FALSE;
             goto exit;
         }