ole32: Implement chained IInitializeSpy support.

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/ole32/compobj.c

@@ -1738,11 +1738,21 @@ static void COM_TlsDestroy(void)
     struct oletls *info = NtCurrentTeb()->ReservedForOle;
     if (info)
     {
+        struct init_spy *cursor, *cursor2;
+
         if (info->apt) apartment_release(info->apt);
         if (info->errorinfo) IErrorInfo_Release(info->errorinfo);
         if (info->state) IUnknown_Release(info->state);
-        if (info->spy) IInitializeSpy_Release(info->spy);
+
+        LIST_FOR_EACH_ENTRY_SAFE(cursor, cursor2, &info->spies, struct init_spy, entry)
+        {
+            list_remove(&cursor->entry);
+            IInitializeSpy_Release(cursor->spy);
+            heap_free(cursor);
+        }
+
         if (info->context_token) IObjContext_Release(info->context_token);
+
         HeapFree(GetProcessHeap(), 0, info);
         NtCurrentTeb()->ReservedForOle = NULL;
     }
@@ -1764,6 +1774,19 @@ DWORD WINAPI CoBuildVersion(void)
     return (rmm<<16)+rup;
 }
 
+static struct init_spy *get_spy_entry(struct oletls *info, unsigned int id)
+{
+    struct init_spy *spy;
+
+    LIST_FOR_EACH_ENTRY(spy, &info->spies, struct init_spy, entry)
+    {
+        if (id == spy->id)
+            return spy;
+    }
+
+    return NULL;
+}
+
 /******************************************************************************
  *              CoRegisterInitializeSpy [OLE32.@]
  *
@@ -1783,6 +1806,8 @@ DWORD WINAPI CoBuildVersion(void)
 HRESULT WINAPI CoRegisterInitializeSpy(IInitializeSpy *spy, ULARGE_INTEGER *cookie)
 {
     struct oletls *info = COM_CurrentInfo();
+    struct init_spy *entry;
+    unsigned int id;
     HRESULT hr;
 
     TRACE("(%p, %p)\n", spy, cookie);
@@ -1794,19 +1819,32 @@ HRESULT WINAPI CoRegisterInitializeSpy(IInitializeSpy *spy, ULARGE_INTEGER *cook
         return E_INVALIDARG;
     }
 
-    if (info->spy)
+    hr = IInitializeSpy_QueryInterface(spy, &IID_IInitializeSpy, (void **)&spy);
+    if (FAILED(hr))
+        return hr;
+
+    entry = heap_alloc(sizeof(*entry));
+    if (!entry)
     {
-        FIXME("Already registered?\n");
+        IInitializeSpy_Release(spy);
-        return E_UNEXPECTED;
+        return E_OUTOFMEMORY;
     }
 
-    hr = IInitializeSpy_QueryInterface(spy, &IID_IInitializeSpy, (void **) &info->spy);
+    entry->spy = spy;
-    if (SUCCEEDED(hr))
+
+    id = 0;
+    while (get_spy_entry(info, id) != NULL)
     {
-        cookie->QuadPart = (DWORD_PTR)spy;
+        id++;
-        return S_OK;
     }
-    return hr;
+
+    entry->id = id;
+    list_add_head(&info->spies, &entry->entry);
+
+    cookie->HighPart = GetCurrentThreadId();
+    cookie->LowPart = entry->id;
+
+    return S_OK;
 }
 
 /******************************************************************************
@@ -1827,16 +1865,25 @@ HRESULT WINAPI CoRegisterInitializeSpy(IInitializeSpy *spy, ULARGE_INTEGER *cook
 HRESULT WINAPI CoRevokeInitializeSpy(ULARGE_INTEGER cookie)
 {
     struct oletls *info = COM_CurrentInfo();
+    struct init_spy *spy;
+
     TRACE("(%s)\n", wine_dbgstr_longlong(cookie.QuadPart));
 
-    if (!info || !info->spy || cookie.QuadPart != (DWORD_PTR)info->spy)
+    if (!info || cookie.HighPart != GetCurrentThreadId())
         return E_INVALIDARG;
 
-    IInitializeSpy_Release(info->spy);
+    if ((spy = get_spy_entry(info, cookie.LowPart)))
-    info->spy = NULL;
+    {
+        IInitializeSpy_Release(spy->spy);
+        list_remove(&spy->entry);
+        heap_free(spy);
+
         return S_OK;
     }
 
+    return E_INVALIDARG;
+}
+
 HRESULT enter_apartment( struct oletls *info, DWORD model )
 {
     HRESULT hr = S_OK;
@@ -1929,6 +1976,7 @@ HRESULT WINAPI CoInitialize(LPVOID lpReserved)
 HRESULT WINAPI DECLSPEC_HOTPATCH CoInitializeEx(LPVOID lpReserved, DWORD dwCoInit)
 {
   struct oletls *info = COM_CurrentInfo();
+  struct init_spy *cursor;
   HRESULT hr;
 
   TRACE("(%p, %x)\n", lpReserved, (int)dwCoInit);
@@ -1955,13 +2003,17 @@ HRESULT WINAPI DECLSPEC_HOTPATCH CoInitializeEx(LPVOID lpReserved, DWORD dwCoIni
     RunningObjectTableImpl_Initialize();
   }
 
-  if (info->spy)
+  LIST_FOR_EACH_ENTRY(cursor, &info->spies, struct init_spy, entry)
-      IInitializeSpy_PreInitialize(info->spy, dwCoInit, info->inits);
+  {
+      IInitializeSpy_PreInitialize(cursor->spy, dwCoInit, info->inits);
+  }
 
   hr = enter_apartment( info, dwCoInit );
 
-  if (info->spy)
+  LIST_FOR_EACH_ENTRY(cursor, &info->spies, struct init_spy, entry)
-      IInitializeSpy_PostInitialize(info->spy, hr, dwCoInit, info->inits);
+  {
+      hr = IInitializeSpy_PostInitialize(cursor->spy, hr, dwCoInit, info->inits);
+  }
 
   return hr;
 }
@@ -1985,6 +2037,7 @@ HRESULT WINAPI DECLSPEC_HOTPATCH CoInitializeEx(LPVOID lpReserved, DWORD dwCoIni
 void WINAPI DECLSPEC_HOTPATCH CoUninitialize(void)
 {
   struct oletls * info = COM_CurrentInfo();
+  struct init_spy *cursor;
   LONG lCOMRefCnt;
 
   TRACE("()\n");
@@ -1992,16 +2045,21 @@ void WINAPI DECLSPEC_HOTPATCH CoUninitialize(void)
   /* will only happen on OOM */
   if (!info) return;
 
-  if (info->spy)
+  LIST_FOR_EACH_ENTRY(cursor, &info->spies, struct init_spy, entry)
-      IInitializeSpy_PreUninitialize(info->spy, info->inits);
+  {
+      IInitializeSpy_PreUninitialize(cursor->spy, info->inits);
+  }
 
   /* sanity check */
   if (!info->inits)
   {
       ERR("Mismatched CoUninitialize\n");
 
-    if (info->spy)
+      LIST_FOR_EACH_ENTRY(cursor, &info->spies, struct init_spy, entry)
-        IInitializeSpy_PostUninitialize(info->spy, info->inits);
+      {
+          IInitializeSpy_PostUninitialize(cursor->spy, info->inits);
+      }
+
       return;
   }
 
@@ -2024,8 +2082,11 @@ void WINAPI DECLSPEC_HOTPATCH CoUninitialize(void)
     ERR( "CoUninitialize() - not CoInitialized.\n" );
     InterlockedExchangeAdd(&s_COMLockCount,1); /* restore the lock count. */
   }
-  if (info->spy)
+
-      IInitializeSpy_PostUninitialize(info->spy, info->inits);
+  LIST_FOR_EACH_ENTRY(cursor, &info->spies, struct init_spy, entry)
+  {
+      IInitializeSpy_PostUninitialize(cursor->spy, info->inits);
+  }
 }
 
 /******************************************************************************

dlls/ole32/compobj_private.h

@@ -153,6 +153,13 @@ struct apartment
   BOOL main;               /* is this a main-threaded-apartment? (RO) */
 };
 
+struct init_spy
+{
+    struct list entry;
+    IInitializeSpy *spy;
+    unsigned int id;
+};
+
 /* this is what is stored in TEB->ReservedForOle */
 struct oletls
 {
@@ -160,7 +167,7 @@ struct oletls
     IErrorInfo       *errorinfo;   /* see errorinfo.c */
     IUnknown         *state;       /* see CoSetState */
     DWORD             apt_mask;    /* apartment mask (+0Ch on x86) */
-    IInitializeSpy   *spy;         /* The "SPY" from CoInitializeSpy */
+    void            *unknown0;
     DWORD            inits;        /* number of times CoInitializeEx called */
     DWORD            ole_inits;    /* number of times OleInitialize called */
     GUID             causality_id; /* unique identifier for each COM call */
@@ -171,6 +178,7 @@ struct oletls
     IUnknown        *call_state;    /* current call context (+3Ch on x86) */
     DWORD            unknown2[46];
     IUnknown        *cancel_object; /* cancel object set by CoSetCancelObject (+F8h on x86) */
+    struct list      spies;         /* Spies installed with CoRegisterInitializeSpy */
 };
 
 
@@ -262,7 +270,12 @@ APARTMENT *apartment_get_current_or_mta(void) DECLSPEC_HIDDEN;
 static inline struct oletls *COM_CurrentInfo(void)
 {
     if (!NtCurrentTeb()->ReservedForOle)
-        NtCurrentTeb()->ReservedForOle = heap_alloc_zero(sizeof(struct oletls));
+    {
+        struct oletls *oletls = heap_alloc_zero(sizeof(*oletls));
+        if (oletls)
+            list_init(&oletls->spies);
+        NtCurrentTeb()->ReservedForOle = oletls;
+    }
 
     return NtCurrentTeb()->ReservedForOle;
 }

dlls/ole32/tests/compobj.c

@@ -3363,32 +3363,28 @@ static void test_IInitializeSpy(void)
     cookie.LowPart = 1;
     hr = CoRegisterInitializeSpy(&testinitialize, &cookie);
     ok(hr == S_OK, "got 0x%08x\n", hr);
-todo_wine {
     ok(cookie.HighPart == GetCurrentThreadId(), "got high part 0x%08x, expected 0x%08x\n", cookie.HighPart,
         GetCurrentThreadId());
     ok(cookie.LowPart == 0, "got wrong low part 0x%x\n", cookie.LowPart);
-}
+
     /* register same instance one more time */
     cookie1.HighPart = 0;
     cookie1.LowPart = 0;
     hr = CoRegisterInitializeSpy(&testinitialize, &cookie1);
-todo_wine {
     ok(hr == S_OK, "got 0x%08x\n", hr);
     ok(cookie1.HighPart == GetCurrentThreadId(), "got high part 0x%08x, expected 0x%08x\n", cookie1.HighPart,
         GetCurrentThreadId());
     ok(cookie1.LowPart == 1, "got wrong low part 0x%x\n", cookie1.LowPart);
-}
+
     cookie2.HighPart = 0;
     cookie2.LowPart = 0;
     hr = CoRegisterInitializeSpy(&testinitialize, &cookie2);
-todo_wine {
     ok(hr == S_OK, "got 0x%08x\n", hr);
     ok(cookie2.HighPart == GetCurrentThreadId(), "got high part 0x%08x, expected 0x%08x\n", cookie2.HighPart,
         GetCurrentThreadId());
     ok(cookie2.LowPart == 2, "got wrong low part 0x%x\n", cookie2.LowPart);
-}
+
     hr = CoRevokeInitializeSpy(cookie1);
-todo_wine
     ok(hr == S_OK, "got 0x%08x\n", hr);
 
     hr = CoRevokeInitializeSpy(cookie1);
@@ -3397,21 +3393,18 @@ todo_wine
     cookie1.HighPart = 0;
     cookie1.LowPart = 0;
     hr = CoRegisterInitializeSpy(&testinitialize, &cookie1);
-todo_wine {
     ok(hr == S_OK, "got 0x%08x\n", hr);
     ok(cookie1.HighPart == GetCurrentThreadId(), "got high part 0x%08x, expected 0x%08x\n", cookie1.HighPart,
         GetCurrentThreadId());
     ok(cookie1.LowPart == 1, "got wrong low part 0x%x\n", cookie1.LowPart);
-}
+
     hr = CoRevokeInitializeSpy(cookie);
     ok(hr == S_OK, "got 0x%08x\n", hr);
 
     hr = CoRevokeInitializeSpy(cookie1);
-todo_wine
     ok(hr == S_OK, "got 0x%08x\n", hr);
 
     hr = CoRevokeInitializeSpy(cookie2);
-todo_wine
     ok(hr == S_OK, "got 0x%08x\n", hr);
 }