ole32: Use proper IUnknown implementation when finding and creating stub manager.

dlls/ole32/compobj.c

@@ -3605,32 +3605,8 @@ HRESULT WINAPI CoLockObjectExternal(
     apt = COM_CurrentApt();
     if (!apt) return CO_E_NOTINITIALIZED;
 
-    stubmgr = get_stub_manager_from_object(apt, pUnk);
+    stubmgr = get_stub_manager_from_object(apt, pUnk, fLock);
-    
+    if (!stubmgr)
-    if (stubmgr)
-    {
-        if (fLock)
-            stub_manager_ext_addref(stubmgr, 1, FALSE);
-        else
-            stub_manager_ext_release(stubmgr, 1, FALSE, fLastUnlockReleases);
-        
-        stub_manager_int_release(stubmgr);
-
-        return S_OK;
-    }
-    else if (fLock)
-    {
-        stubmgr = new_stub_manager(apt, pUnk);
-
-        if (stubmgr)
-        {
-            stub_manager_ext_addref(stubmgr, 1, FALSE);
-            stub_manager_int_release(stubmgr);
-        }
-
-        return S_OK;
-    }
-    else
     {
         WARN("stub object not found %p\n", pUnk);
         /* Note: native is pretty broken here because it just silently
@@ -3638,6 +3614,14 @@ HRESULT WINAPI CoLockObjectExternal(
          * think that the object was disconnected, when it actually wasn't */
         return S_OK;
     }
+
+    if (fLock)
+        stub_manager_ext_addref(stubmgr, 1, FALSE);
+    else
+        stub_manager_ext_release(stubmgr, 1, FALSE, fLastUnlockReleases);
+
+    stub_manager_int_release(stubmgr);
+    return S_OK;
 }
 
 /***********************************************************************

dlls/ole32/compobj_private.h

@@ -185,14 +185,13 @@ HRESULT FTMarshalCF_Create(REFIID riid, LPVOID *ppv) DECLSPEC_HIDDEN;
 /* Stub Manager */
 
 ULONG stub_manager_int_release(struct stub_manager *This) DECLSPEC_HIDDEN;
-struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object) DECLSPEC_HIDDEN;
 ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs, BOOL tableweak) DECLSPEC_HIDDEN;
 ULONG stub_manager_ext_release(struct stub_manager *m, ULONG refs, BOOL tableweak, BOOL last_unlock_releases) DECLSPEC_HIDDEN;
 struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid,
      DWORD dest_context, void *dest_context_data, MSHLFLAGS flags) DECLSPEC_HIDDEN;
 struct ifstub *stub_manager_find_ifstub(struct stub_manager *m, REFIID iid, MSHLFLAGS flags) DECLSPEC_HIDDEN;
 struct stub_manager *get_stub_manager(APARTMENT *apt, OID oid) DECLSPEC_HIDDEN;
-struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object) DECLSPEC_HIDDEN;
+struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, IUnknown *object, BOOL alloc) DECLSPEC_HIDDEN;
 BOOL stub_manager_notify_unmarshal(struct stub_manager *m, const IPID *ipid) DECLSPEC_HIDDEN;
 BOOL stub_manager_is_table_marshaled(struct stub_manager *m, const IPID *ipid) DECLSPEC_HIDDEN;
 void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs, const IPID *ipid, BOOL tableweak) DECLSPEC_HIDDEN;

dlls/ole32/marshal.c

@@ -132,22 +132,14 @@ HRESULT marshal_object(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnkno
     if (hr != S_OK)
         return hr;
 
+    if (!(manager = get_stub_manager_from_object(apt, object, TRUE)))
+        return E_OUTOFMEMORY;
+
     stdobjref->flags = SORF_NULL;
     if (mshlflags & MSHLFLAGS_TABLEWEAK)
         stdobjref->flags |= SORFP_TABLEWEAK;
     if (mshlflags & MSHLFLAGS_NOPING)
         stdobjref->flags |= SORF_NOPING;
-
-    if ((manager = get_stub_manager_from_object(apt, object)))
-        TRACE("registering new ifstub on pre-existing manager\n");
-    else
-    {
-        TRACE("constructing new stub manager\n");
-
-        manager = new_stub_manager(apt, object);
-        if (!manager)
-            return E_OUTOFMEMORY;
-    }
     stdobjref->oid = manager->oid;
 
     tablemarshal = ((mshlflags & MSHLFLAGS_TABLESTRONG) || (mshlflags & MSHLFLAGS_TABLEWEAK));

dlls/ole32/stubmanager.c

@@ -168,7 +168,7 @@ struct ifstub *stub_manager_find_ifstub(struct stub_manager *m, REFIID iid, MSHL
 /* creates a new stub manager and adds it into the apartment. caller must
  * release stub manager when it is no longer required. the apartment and
  * external refs together take one implicit ref */
-struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object)
+static struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object)
 {
     struct stub_manager *sm;
     HRESULT hres;
@@ -298,10 +298,18 @@ ULONG stub_manager_int_release(struct stub_manager *This)
 /* gets the stub manager associated with an object - caller must have
  * a reference to the apartment while a reference to the stub manager is held.
  * it must also call release on the stub manager when it is no longer needed */
-struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object)
+struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, IUnknown *obj, BOOL alloc)
 {
     struct stub_manager *result = NULL;
     struct list         *cursor;
+    IUnknown *object;
+    HRESULT hres;
+
+    hres = IUnknown_QueryInterface(obj, &IID_IUnknown, (void**)&object);
+    if (FAILED(hres)) {
+        ERR("QueryInterface(IID_IUnknown failed): %08x\n", hres);
+        return NULL;
+    }
 
     EnterCriticalSection(&apt->cs);
     LIST_FOR_EACH( cursor, &apt->stubmgrs )
@@ -317,11 +325,16 @@ struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object)
     }
     LeaveCriticalSection(&apt->cs);
 
-    if (result)
+    if (result) {
         TRACE("found %p for object %p\n", result, object);
-    else
+    }else if (alloc) {
+        TRACE("not found, creating new stub manager...\n");
+        result = new_stub_manager(apt, object);
+    }else {
         TRACE("not found for object %p\n", object);
+    }
 
+    IUnknown_Release(object);
     return result;    
 }