[PATCH] combase: Check all apartment types when looking for an apartment belonging to a thread.

[Dmitry Timoshkov]

When an application has both apartment-threaded and multi-threaded apartments
then apartment_findfromtid() might return just the first one, and later stub manager
assiciated with that apartment may not have an appropriate connection. In order
to find the appropriate connection it's needed to check all apartment types
belonging to a thread.

Signed-off-by: Dmitry Timoshkov <dmitry at baikal.ru>
---
 dlls/combase/apartment.c       |  4 ++--
 dlls/combase/combase_private.h |  2 +-
 dlls/combase/stubmanager.c     | 34 +++++++++++++++++++++++++++++-----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/dlls/combase/apartment.c b/dlls/combase/apartment.c
index 108d6a71c5c..23dfaa5da2c 100644
--- a/dlls/combase/apartment.c
+++ b/dlls/combase/apartment.c
@@ -656,14 +656,14 @@ struct apartment * apartment_findfromoxid(OXID oxid)
 /* gets the apartment which has a given creator thread ID. The caller must
  * release the reference from the apartment as soon as the apartment pointer
  * is no longer required. */
-struct apartment * apartment_findfromtid(DWORD tid)
+struct apartment * apartment_findfromtid(DWORD tid, BOOL multi_threaded)
 {
     struct apartment *result = NULL, *apt;
 
     EnterCriticalSection(&apt_cs);
     LIST_FOR_EACH_ENTRY(apt, &apts, struct apartment, entry)
     {
-        if (apt->tid == tid)
+        if (apt->tid == tid && apt->multi_threaded == multi_threaded)
         {
             result = apt;
             apartment_addref(result);
diff --git a/dlls/combase/combase_private.h b/dlls/combase/combase_private.h
index 9247af4ebb6..7d67e6f439a 100644
--- a/dlls/combase/combase_private.h
+++ b/dlls/combase/combase_private.h
@@ -170,7 +170,7 @@ IUnknown *com_get_registered_class_object(const struct apartment *apartment, REF
         DWORD clscontext) DECLSPEC_HIDDEN;
 void apartment_revoke_all_classes(const struct apartment *apt) DECLSPEC_HIDDEN;
 struct apartment * apartment_findfromoxid(OXID oxid) DECLSPEC_HIDDEN;
-struct apartment * apartment_findfromtid(DWORD tid) DECLSPEC_HIDDEN;
+struct apartment * apartment_findfromtid(DWORD tid, BOOL multi_threaded) DECLSPEC_HIDDEN;
 
 HRESULT marshal_object(struct apartment *apt, STDOBJREF *stdobjref, REFIID riid, IUnknown *object,
         DWORD dest_context, void *dest_context_data, MSHLFLAGS mshlflags) DECLSPEC_HIDDEN;
diff --git a/dlls/combase/stubmanager.c b/dlls/combase/stubmanager.c
index e9f20962e3b..8a910f63190 100644
--- a/dlls/combase/stubmanager.c
+++ b/dlls/combase/stubmanager.c
@@ -499,21 +499,45 @@ 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
-        *stub_apt = apartment_findfromtid(ipid->Data2);
+    else /* Look for multi-threaded apartment */
+        *stub_apt = apartment_findfromtid(ipid->Data2, TRUE);
+
+    TRACE("ipid %s (tid %08x) => stub_apt %p\n", debugstr_guid(ipid), ipid->Data2, *stub_apt);
+
     if (!*stub_apt)
     {
         TRACE("Couldn't find apartment corresponding to TID 0x%04x\n", ipid->Data2);
         return RPC_E_INVALID_OBJECT;
     }
     *stubmgr_ret = get_stub_manager_from_ipid(*stub_apt, ipid, ifstub);
-    if (!*stubmgr_ret)
+    TRACE("stub_apt %p, ipid %s (tid %08x) => stubmgr_ret %p\n", *stub_apt, debugstr_guid(ipid), ipid->Data2, *stubmgr_ret);
+
+    if (*stubmgr_ret) return S_OK;
+
+    apartment_release(*stub_apt);
+
+    if (ipid->Data2 == 0xffff) /* IRemUnknown */
     {
-        apartment_release(*stub_apt);
         *stub_apt = NULL;
         return RPC_E_INVALID_OBJECT;
     }
-    return S_OK;
+
+    /* Look for apartment-threaded apartment */
+    *stub_apt = apartment_findfromtid(ipid->Data2, FALSE);
+    if (!*stub_apt)
+    {
+        TRACE("Couldn't find apartment corresponding to TID 0x%04x\n", ipid->Data2);
+        return RPC_E_INVALID_OBJECT;
+    }
+
+    *stubmgr_ret = get_stub_manager_from_ipid(*stub_apt, ipid, ifstub);
+    TRACE("stub_apt %p, ipid %s (tid %08x) => stubmgr_ret %p\n", *stub_apt, debugstr_guid(ipid), ipid->Data2, *stubmgr_ret);
+
+    if (*stubmgr_ret) return S_OK;
+
+    apartment_release(*stub_apt);
+    *stub_apt = NULL;
+    return RPC_E_INVALID_OBJECT;
 }
 
 static HRESULT ipid_to_stub_manager(const IPID *ipid, struct apartment **stub_apt, struct stub_manager **stub)
-- 
2.31.1