From 00c8117a7277b5c798ca6078a4f50badeec756ad Mon Sep 17 00:00:00 2001
From: Connor McAdams <cmcadams@codeweavers.com>
Date: Wed, 24 Aug 2022 12:29:46 -0500
Subject: [PATCH] combase: Omit thread ID from the stub manager ipid for MTA
 objects.

It is possible for a thread that creates an MTA to call
CoUninitialize() and not destroy the MTA if another thread has
entered the MTA in the meantime. If the original thread then creates
an STA, subsequent attempts to find the MTA with 'apartment_findfromtid'
will get the original thread's STA. To avoid this, don't set a TID value
in the stub manager IPID value to indicate that the stub resides in the
MTA.

Signed-off-by: Connor McAdams <cmcadams@codeweavers.com>
---
 dlls/combase/stubmanager.c |  6 ++--
 dlls/ole32/tests/marshal.c | 58 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/dlls/combase/stubmanager.c b/dlls/combase/stubmanager.c
index 59782b8c623..5bb87c60c1d 100644
--- a/dlls/combase/stubmanager.c
+++ b/dlls/combase/stubmanager.c
@@ -46,7 +46,7 @@ WINE_DEFAULT_DEBUG_CHANNEL(ole);
 
 /* generates an ipid in the following format (similar to native version):
  * Data1 = apartment-local ipid counter
- * Data2 = apartment creator thread ID
+ * Data2 = apartment creator thread ID, or 0 for an MTA.
  * Data3 = process ID
  * Data4 = random value
  */
@@ -62,7 +62,7 @@ static inline HRESULT generate_ipid(struct stub_manager *m, IPID *ipid)
     }
 
     ipid->Data1 = InterlockedIncrement(&m->apt->ipidc);
-    ipid->Data2 = (USHORT)m->apt->tid;
+    ipid->Data2 = !m->apt->multi_threaded ? (USHORT)m->apt->tid : 0;
     ipid->Data3 = (USHORT)GetCurrentProcessId();
     return S_OK;
 }
@@ -499,6 +499,8 @@ static HRESULT ipid_to_ifstub(const IPID *ipid, struct apartment **stub_apt,
     /* FIXME: hack for IRemUnknown */
     if (ipid->Data2 == 0xffff)
         *stub_apt = apartment_findfromoxid(*(const OXID *)ipid->Data4);
+    else if (!ipid->Data2)
+        *stub_apt = apartment_get_mta();
     else
         *stub_apt = apartment_findfromtid(ipid->Data2);
     if (!*stub_apt)
diff --git a/dlls/ole32/tests/marshal.c b/dlls/ole32/tests/marshal.c
index 6e21ed1889b..7a69ad9ed15 100644
--- a/dlls/ole32/tests/marshal.c
+++ b/dlls/ole32/tests/marshal.c
@@ -4489,6 +4489,63 @@ static void test_channel_hook(void)
     ok_ole_success(hr, CoRegisterMessageFilter);
 }
 
+static DWORD CALLBACK second_mta_thread_proc(void *param)
+{
+    struct implicit_mta_marshal_data *data = param;
+    HRESULT hr;
+
+    /* Second thread now keeps MTA created on first thread alive. */
+    CoInitializeEx(NULL, COINIT_MULTITHREADED);
+    hr = CoMarshalInterface(data->stream, &IID_IClassFactory,
+        (IUnknown *)&Test_ClassFactory, MSHCTX_INPROC, NULL, MSHLFLAGS_NORMAL);
+    ok_ole_success(hr, CoMarshalInterface);
+
+    SetEvent(data->start);
+
+    ok(!WaitForSingleObject(data->stop, 1000), "wait failed\n");
+    CoUninitialize();
+    return 0;
+}
+
+static void test_mta_creation_thread_change_apartment(void)
+{
+    struct implicit_mta_marshal_data data;
+    IClassFactory *cf;
+    IUnknown *proxy;
+    HANDLE thread;
+    HRESULT hr;
+
+    CoInitializeEx(NULL, COINIT_MULTITHREADED);
+
+    hr = CreateStreamOnHGlobal(NULL, TRUE, &data.stream);
+    ok_ole_success(hr, CreateStreamOnHGlobal);
+
+    data.start = CreateEventA(NULL, FALSE, FALSE, NULL);
+    data.stop  = CreateEventA(NULL, FALSE, FALSE, NULL);
+
+    thread = CreateThread(NULL, 0, second_mta_thread_proc, &data, 0, NULL);
+    ok(!WaitForSingleObject(data.start, 1000), "wait failed\n");
+    CoUninitialize();
+
+    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
+
+    IStream_Seek(data.stream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(data.stream, &IID_IClassFactory, (void **)&cf);
+    ok_ole_success(hr, CoUnmarshalInterface);
+
+    hr = IClassFactory_CreateInstance(cf, NULL, &IID_IUnknown, (void **)&proxy);
+    ok_ole_success(hr, IClassFactory_CreateInstance);
+
+    IUnknown_Release(proxy);
+    IStream_Release(data.stream);
+
+    SetEvent(data.stop);
+    ok(!WaitForSingleObject(thread, 1000), "wait failed\n");
+    CloseHandle(thread);
+
+    CoUninitialize();
+}
+
 START_TEST(marshal)
 {
     HMODULE hOle32 = GetModuleHandleA("ole32");
@@ -4511,6 +4568,7 @@ START_TEST(marshal)
 
     test_cocreateinstance_proxy();
     test_implicit_mta();
+    test_mta_creation_thread_change_apartment();
 
     CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);