ole32: Correctly track the number of inits when auto-joining the MTA.

Signed-off-by: Huw Davies <huw@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/ole32/compobj.c

@@ -1724,12 +1724,6 @@ HWND apartment_getwindow(const struct apartment *apt)
     return apt->win;
 }
 
-void apartment_joinmta(void)
-{
-    apartment_addref(MTA);
-    COM_CurrentInfo()->apt = MTA;
-}
-
 static void COM_TlsDestroy(void)
 {
     struct oletls *info = NtCurrentTeb()->ReservedForOle;
@@ -1834,6 +1828,40 @@ HRESULT WINAPI CoRevokeInitializeSpy(ULARGE_INTEGER cookie)
     return S_OK;
 }
 
+HRESULT enter_apartment( struct oletls *info, DWORD model )
+{
+    HRESULT hr = S_OK;
+
+    if (!info->apt)
+    {
+        if (!apartment_get_or_create( model ))
+            return E_OUTOFMEMORY;
+    }
+    else if (!apartment_is_model( info->apt, model ))
+    {
+        WARN( "Attempt to change threading model of this apartment from %s to %s\n",
+              info->apt->multi_threaded ? "multi-threaded" : "apartment threaded",
+              model & COINIT_APARTMENTTHREADED ? "apartment threaded" : "multi-threaded" );
+        return RPC_E_CHANGED_MODE;
+    }
+    else
+        hr = S_FALSE;
+
+    info->inits++;
+
+    return hr;
+}
+
+void leave_apartment( struct oletls *info )
+{
+    if (!--info->inits)
+    {
+        if (info->ole_inits)
+            WARN( "Uninitializing apartment while Ole is still initialized\n" );
+        apartment_release( info->apt );
+        info->apt = NULL;
+    }
+}
 
 /******************************************************************************
  *		CoInitialize	[OLE32.@]
@@ -1892,8 +1920,7 @@ HRESULT WINAPI CoInitialize(LPVOID lpReserved)
 HRESULT WINAPI DECLSPEC_HOTPATCH CoInitializeEx(LPVOID lpReserved, DWORD dwCoInit)
 {
   struct oletls *info = COM_CurrentInfo();
-  HRESULT hr = S_OK;
-  APARTMENT *apt;
+  HRESULT hr;
 
   TRACE("(%p, %x)\n", lpReserved, (int)dwCoInit);
 
@@ -1922,24 +1949,7 @@ HRESULT WINAPI DECLSPEC_HOTPATCH CoInitializeEx(LPVOID lpReserved, DWORD dwCoIni
   if (info->spy)
       IInitializeSpy_PreInitialize(info->spy, dwCoInit, info->inits);
 
-  if (!(apt = info->apt))
-  {
-    apt = apartment_get_or_create(dwCoInit);
-    if (!apt) return E_OUTOFMEMORY;
-  }
-  else if (!apartment_is_model(apt, dwCoInit))
-  {
-    /* Changing the threading model after it's been set is illegal. If this warning is triggered by Wine
-       code then we are probably using the wrong threading model to implement that API. */
-    ERR("Attempt to change threading model of this apartment from %s to %s\n",
-        apt->multi_threaded ? "multi-threaded" : "apartment threaded",
-        dwCoInit & COINIT_APARTMENTTHREADED ? "apartment threaded" : "multi-threaded");
-    return RPC_E_CHANGED_MODE;
-  }
-  else
-    hr = S_FALSE;
-
-  info->inits++;
+  hr = enter_apartment( info, dwCoInit );
 
   if (info->spy)
       IInitializeSpy_PostInitialize(info->spy, hr, dwCoInit, info->inits);
@@ -1986,13 +1996,7 @@ void WINAPI DECLSPEC_HOTPATCH CoUninitialize(void)
     return;
   }
 
-  if (!--info->inits)
-  {
-    if (info->ole_inits)
-      WARN("uninitializing apartment while Ole is still initialized\n");
-    apartment_release(info->apt);
-    info->apt = NULL;
-  }
+  leave_apartment( info );
 
   /*
    * Decrease the reference count.

dlls/ole32/compobj_private.h

@@ -244,8 +244,8 @@ static inline HRESULT apartment_getoxid(const struct apartment *apt, OXID *oxid)
 }
 HRESULT apartment_createwindowifneeded(struct apartment *apt) DECLSPEC_HIDDEN;
 HWND apartment_getwindow(const struct apartment *apt) DECLSPEC_HIDDEN;
-void apartment_joinmta(void) DECLSPEC_HIDDEN;
-
+HRESULT enter_apartment(struct oletls *info, DWORD model) DECLSPEC_HIDDEN;
+void leave_apartment(struct oletls *info) DECLSPEC_HIDDEN;
 
 /* DCOM messages used by the apartment window (not compatible with native) */
 #define DM_EXECUTERPC   (WM_USER + 0) /* WPARAM = 0, LPARAM = (struct dispatch_params *) */

dlls/ole32/rpc.c

@@ -1495,16 +1495,17 @@ static void __RPC_STUB dispatch_rpc(RPC_MESSAGE *msg)
     else
     {
         BOOL joined = FALSE;
-        if (!COM_CurrentInfo()->apt)
+        struct oletls *info = COM_CurrentInfo();
+
+        if (!info->apt)
         {
-            apartment_joinmta();
+            enter_apartment(info, COINIT_MULTITHREADED);
             joined = TRUE;
         }
         RPC_ExecuteCall(params);
         if (joined)
         {
-            apartment_release(COM_CurrentInfo()->apt);
-            COM_CurrentInfo()->apt = NULL;
+            leave_apartment(info);
         }
     }
 

dlls/ole32/tests/marshal.c

@@ -2968,6 +2968,12 @@ static HRESULT WINAPI local_server_GetClassID(IPersist *iface, CLSID *clsid)
     hr = CoDisconnectObject((IUnknown *)iface, 0);
     ok(hr == S_OK, "got %08x\n", hr);
 
+    /* Initialize and uninitialize the apartment to show that we
+     * remain in the autojoined mta */
+    hr = pCoInitializeEx( NULL, COINIT_MULTITHREADED );
+    ok( hr == S_FALSE, "got %08x\n", hr );
+    CoUninitialize();
+
     return S_OK;
 }
 
@@ -3749,7 +3755,7 @@ START_TEST(marshal)
     argc = winetest_get_mainargs( &argv );
     if (argc > 2 && (!strcmp(argv[2], "-Embedding")))
     {
-        pCoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
+        pCoInitializeEx(NULL, COINIT_MULTITHREADED);
         test_register_local_server();
         CoUninitialize();