ole32: Implement MTA pinning with CoIncrementMTAUsage()/CoDecrementMTAUsage().

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

dlls/api-ms-win-core-com-l1-1-0/api-ms-win-core-com-l1-1-0.spec

@@ -10,7 +10,7 @@
 @ stdcall CoCreateInstanceEx(ptr ptr long ptr long ptr) ole32.CoCreateInstanceEx
 @ stub CoCreateInstanceFromApp
 @ stub CoDecodeProxy
-@ stub CoDecrementMTAUsage
+@ stdcall CoDecrementMTAUsage(ptr) ole32.CoDecrementMTAUsage
 @ stdcall CoDisableCallCancellation(ptr) ole32.CoDisableCallCancellation
 @ stub CoDisconnectContext
 @ stdcall CoDisconnectObject(ptr long) ole32.CoDisconnectObject
@@ -35,7 +35,7 @@
 @ stub CoGetStdMarshalEx
 @ stdcall CoGetTreatAsClass(ptr ptr) ole32.CoGetTreatAsClass
 @ stdcall CoImpersonateClient() ole32.CoImpersonateClient
-@ stub CoIncrementMTAUsage
+@ stdcall CoIncrementMTAUsage(ptr) ole32.CoIncrementMTAUsage
 @ stdcall CoInitializeEx(ptr long) ole32.CoInitializeEx
 @ stdcall CoInitializeSecurity(ptr long ptr ptr long long ptr long ptr) ole32.CoInitializeSecurity
 @ stub CoInvalidateRemoteMachineBindings

dlls/api-ms-win-core-com-l1-1-1/api-ms-win-core-com-l1-1-1.spec

@@ -10,7 +10,7 @@
 @ stdcall CoCreateInstanceEx(ptr ptr long ptr long ptr) ole32.CoCreateInstanceEx
 @ stub CoCreateInstanceFromApp
 @ stub CoDecodeProxy
-@ stub CoDecrementMTAUsage
+@ stdcall CoDecrementMTAUsage(ptr) ole32.CoDecrementMTAUsage
 @ stdcall CoDisableCallCancellation(ptr) ole32.CoDisableCallCancellation
 @ stub CoDisconnectContext
 @ stdcall CoDisconnectObject(ptr long) ole32.CoDisconnectObject
@@ -35,7 +35,7 @@
 @ stub CoGetStdMarshalEx
 @ stdcall CoGetTreatAsClass(ptr ptr) ole32.CoGetTreatAsClass
 @ stdcall CoImpersonateClient() ole32.CoImpersonateClient
-@ stub CoIncrementMTAUsage
+@ stdcall CoIncrementMTAUsage(ptr) ole32.CoIncrementMTAUsage
 @ stdcall CoInitializeEx(ptr long) ole32.CoInitializeEx
 @ stdcall CoInitializeSecurity(ptr long ptr ptr long long ptr long ptr) ole32.CoInitializeSecurity
 @ stub CoInvalidateRemoteMachineBindings

dlls/combase/combase.spec

@@ -81,7 +81,7 @@
 @ stub CoCreateObjectInContext
 @ stub CoDeactivateObject
 @ stub CoDecodeProxy
-@ stub CoDecrementMTAUsage
+@ stdcall CoDecrementMTAUsage(ptr) ole32.CoDecrementMTAUsage
 @ stdcall CoDisableCallCancellation(ptr) ole32.CoDisableCallCancellation
 @ stub CoDisconnectContext
 @ stdcall CoDisconnectObject(ptr long) ole32.CoDisconnectObject
@@ -116,7 +116,7 @@
 @ stub CoGetSystemSecurityPermissions
 @ stdcall CoGetTreatAsClass(ptr ptr) ole32.CoGetTreatAsClass
 @ stdcall CoImpersonateClient() ole32.CoImpersonateClient
-@ stub CoIncrementMTAUsage
+@ stdcall CoIncrementMTAUsage(ptr) ole32.CoIncrementMTAUsage
 @ stdcall CoInitializeEx(ptr long) ole32.CoInitializeEx
 @ stdcall CoInitializeSecurity(ptr long ptr ptr long long ptr long ptr) ole32.CoInitializeSecurity
 @ stdcall CoInitializeWOW(long long) ole32.CoInitializeWOW

dlls/ole32/compobj.c

@@ -637,6 +637,7 @@ static APARTMENT *apartment_construct(DWORD model)
     list_init(&apt->proxies);
     list_init(&apt->stubmgrs);
     list_init(&apt->loaded_dlls);
+    list_init(&apt->usage_cookies);
     apt->ipidc = 0;
     apt->refs = 1;
     apt->remunk_exported = FALSE;
@@ -5268,6 +5269,72 @@ HRESULT WINAPI CoGetApartmentType(APTTYPE *type, APTTYPEQUALIFIER *qualifier)
     return info->apt ? S_OK : CO_E_NOTINITIALIZED;
 }
 
+struct mta_cookie
+{
+    struct list entry;
+};
+
+/***********************************************************************
+ *           CoIncrementMTAUsage [OLE32.@]
+ */
+HRESULT WINAPI CoIncrementMTAUsage(CO_MTA_USAGE_COOKIE *cookie)
+{
+    struct mta_cookie *mta_cookie;
+
+    TRACE("%p\n", cookie);
+
+    *cookie = NULL;
+
+    if (!(mta_cookie = heap_alloc(sizeof(*mta_cookie))))
+        return E_OUTOFMEMORY;
+
+    EnterCriticalSection(&csApartment);
+
+    if (MTA)
+        apartment_addref(MTA);
+    else
+        MTA = apartment_construct(COINIT_MULTITHREADED);
+    list_add_head(&MTA->usage_cookies, &mta_cookie->entry);
+
+    LeaveCriticalSection(&csApartment);
+
+    *cookie = (CO_MTA_USAGE_COOKIE)mta_cookie;
+
+    return S_OK;
+}
+
+/***********************************************************************
+ *           CoDecrementMTAUsage [OLE32.@]
+ */
+HRESULT WINAPI CoDecrementMTAUsage(CO_MTA_USAGE_COOKIE cookie)
+{
+    struct mta_cookie *mta_cookie = (struct mta_cookie *)cookie;
+
+    TRACE("%p\n", cookie);
+
+    EnterCriticalSection(&csApartment);
+
+    if (MTA)
+    {
+        struct mta_cookie *cur;
+
+        LIST_FOR_EACH_ENTRY(cur, &MTA->usage_cookies, struct mta_cookie, entry)
+        {
+            if (mta_cookie == cur)
+            {
+                list_remove(&cur->entry);
+                heap_free(cur);
+                apartment_release(MTA);
+                break;
+            }
+        }
+    }
+
+    LeaveCriticalSection(&csApartment);
+
+    return S_OK;
+}
+
 /***********************************************************************
  *           CoDisableCallCancellation [OLE32.@]
  */

dlls/ole32/compobj_private.h

@@ -151,6 +151,8 @@ struct apartment
   HWND win;                /* message window (LOCK) */
   LPMESSAGEFILTER filter;  /* message filter (CS cs) */
   BOOL main;               /* is this a main-threaded-apartment? (RO) */
+  /* MTA-only */
+  struct list usage_cookies; /* Used for refcount control with CoIncrementMTAUsage()/CoDecrementMTAUsage(). */
 };
 
 struct init_spy

dlls/ole32/ole32.spec

@@ -14,6 +14,7 @@
 @ stdcall CoCreateGuid(ptr)
 @ stdcall CoCreateInstance(ptr ptr long ptr ptr)
 @ stdcall CoCreateInstanceEx(ptr ptr long ptr long ptr)
+@ stdcall CoDecrementMTAUsage(ptr)
 @ stdcall CoDisableCallCancellation(ptr)
 @ stdcall CoDisconnectObject(ptr long)
 @ stdcall CoDosDateTimeToFileTime(long long ptr) kernel32.DosDateTimeToFileTime
@@ -47,6 +48,7 @@
 @ stub CoGetTIDFromIPID
 @ stdcall CoGetTreatAsClass(ptr ptr)
 @ stdcall CoImpersonateClient()
+@ stdcall CoIncrementMTAUsage(ptr)
 @ stdcall CoInitialize(ptr)
 @ stdcall CoInitializeEx(ptr long)
 @ stdcall CoInitializeSecurity(ptr long ptr ptr long long ptr long ptr)

dlls/ole32/tests/compobj.c

@@ -76,6 +76,8 @@ static HRESULT (WINAPI * pCoGetTreatAsClass)(REFCLSID clsidOld, LPCLSID pClsidNe
 static HRESULT (WINAPI * pCoTreatAsClass)(REFCLSID clsidOld, REFCLSID pClsidNew);
 static HRESULT (WINAPI * pCoGetContextToken)(ULONG_PTR *token);
 static HRESULT (WINAPI * pCoGetApartmentType)(APTTYPE *type, APTTYPEQUALIFIER *qualifier);
+static HRESULT (WINAPI * pCoIncrementMTAUsage)(CO_MTA_USAGE_COOKIE *cookie);
+static HRESULT (WINAPI * pCoDecrementMTAUsage)(CO_MTA_USAGE_COOKIE cookie);
 static LONG (WINAPI * pRegDeleteKeyExA)(HKEY, LPCSTR, REGSAM, DWORD);
 static LONG (WINAPI * pRegOverridePredefKey)(HKEY key, HKEY override);
 
@@ -3818,6 +3820,8 @@ static void init_funcs(void)
     pCoTreatAsClass = (void*)GetProcAddress(hOle32,"CoTreatAsClass");
     pCoGetContextToken = (void*)GetProcAddress(hOle32, "CoGetContextToken");
     pCoGetApartmentType = (void*)GetProcAddress(hOle32, "CoGetApartmentType");
+    pCoIncrementMTAUsage = (void*)GetProcAddress(hOle32, "CoIncrementMTAUsage");
+    pCoDecrementMTAUsage = (void*)GetProcAddress(hOle32, "CoDecrementMTAUsage");
     pRegDeleteKeyExA = (void*)GetProcAddress(hAdvapi32, "RegDeleteKeyExA");
     pRegOverridePredefKey = (void*)GetProcAddress(hAdvapi32, "RegOverridePredefKey");
 
@@ -3916,6 +3920,42 @@ static void test_CoGetCurrentProcess(void)
     ok(id2 && id2 != id, "Unexpected id from another thread.\n");
 }
 
+static void test_mta_usage(void)
+{
+    CO_MTA_USAGE_COOKIE cookie, cookie2;
+    HRESULT hr;
+
+    if (!pCoIncrementMTAUsage)
+    {
+        win_skip("CoIncrementMTAUsage() is not available.\n");
+        return;
+    }
+
+    test_apt_type(APTTYPE_CURRENT, APTTYPEQUALIFIER_NONE);
+
+    cookie = 0;
+    hr = pCoIncrementMTAUsage(&cookie);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    ok(cookie != NULL, "Unexpected cookie %p.\n", cookie);
+
+    cookie2 = 0;
+    hr = pCoIncrementMTAUsage(&cookie2);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    ok(cookie2 != NULL && cookie2 != cookie, "Unexpected cookie %p.\n", cookie2);
+
+    test_apt_type(APTTYPE_MTA, APTTYPEQUALIFIER_IMPLICIT_MTA);
+
+    hr = pCoDecrementMTAUsage(cookie);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    test_apt_type(APTTYPE_MTA, APTTYPEQUALIFIER_IMPLICIT_MTA);
+
+    hr = pCoDecrementMTAUsage(cookie2);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    test_apt_type(APTTYPE_CURRENT, APTTYPEQUALIFIER_NONE);
+}
+
 START_TEST(compobj)
 {
     init_funcs();
@@ -3967,6 +4007,7 @@ START_TEST(compobj)
     test_GlobalOptions();
     test_implicit_mta();
     test_CoGetCurrentProcess();
+    test_mta_usage();
 
     DeleteFileA( testlib );
 }

include/objbase.h

@@ -287,12 +287,16 @@ typedef enum tagCOINIT
     COINIT_SPEED_OVER_MEMORY  = 0x8  /* Trade memory for speed */
 } COINIT;
 
+DECLARE_HANDLE(CO_MTA_USAGE_COOKIE);
+
 HRESULT WINAPI CoInitialize(LPVOID lpReserved);
 HRESULT WINAPI CoInitializeEx(LPVOID lpReserved, DWORD dwCoInit);
 void WINAPI CoUninitialize(void);
 DWORD WINAPI CoGetCurrentProcess(void);
 HRESULT WINAPI CoGetCurrentLogicalThreadId(GUID *id);
 HRESULT WINAPI CoGetApartmentType(APTTYPE *type, APTTYPEQUALIFIER *qualifier);
+HRESULT WINAPI CoIncrementMTAUsage(CO_MTA_USAGE_COOKIE *cookie);
+HRESULT WINAPI CoDecrementMTAUsage(CO_MTA_USAGE_COOKIE cookie);
 
 HINSTANCE WINAPI CoLoadLibrary(LPOLESTR lpszLibName, BOOL bAutoFree);
 void WINAPI CoFreeAllLibraries(void);