- The apartment reference should be held while the stub manager

dlls/ole32/compobj_private.h

@@ -183,8 +183,8 @@ ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs);
 ULONG stub_manager_ext_release(struct stub_manager *m, ULONG refs);
 IRpcStubBuffer *stub_manager_ipid_to_stubbuffer(struct stub_manager *m, const IPID *iid);
 struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid, BOOL tablemarshal);
-struct stub_manager *get_stub_manager(OXID oxid, OID oid);
-struct stub_manager *get_stub_manager_from_object(OXID oxid, void *object);
+struct stub_manager *get_stub_manager(APARTMENT *apt, OID oid);
+struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object);
 
 IRpcStubBuffer *mid_to_stubbuffer(wine_marshal_id *mid);
 

dlls/ole32/marshal.c

@@ -79,9 +79,16 @@ get_facbuf_for_iid(REFIID riid,IPSFactoryBuffer **facbuf) {
 IRpcStubBuffer *mid_to_stubbuffer(wine_marshal_id *mid)
 {
     IRpcStubBuffer *ret;
+    APARTMENT *apt;
     struct stub_manager *m;
 
-    if (!(m = get_stub_manager(mid->oxid, mid->oid)))
+    if (!(apt = COM_ApartmentFromOXID(mid->oxid, TRUE)))
+    {
+        WARN("Could not map OXID %s to apartment object\n", wine_dbgstr_longlong(mid->oxid));
+        return NULL;
+    }
+
+    if (!(m = get_stub_manager(apt, mid->oid)))
     {
         WARN("unknown OID %s\n", wine_dbgstr_longlong(mid->oid));
         return NULL;
@@ -90,6 +97,9 @@ IRpcStubBuffer *mid_to_stubbuffer(wine_marshal_id *mid)
     ret = stub_manager_ipid_to_stubbuffer(m, &mid->ipid);
 
     stub_manager_int_release(m);
+
+    COM_ApartmentRelease(apt);
+
     return ret;
 }
 
@@ -98,23 +108,30 @@ static HRESULT register_ifstub(wine_marshal_id *mid, REFIID riid, IUnknown *obj,
 {
     struct stub_manager *manager = NULL;
     struct ifstub       *ifstub;
+    APARTMENT           *apt;
+
+    if (!(apt = COM_ApartmentFromOXID(mid->oxid, TRUE)))
+    {
+        ERR("Could not map OXID %s to apartment object\n", wine_dbgstr_longlong(mid->oxid));
+        return E_UNEXPECTED;
+    }
 
     /* mid->oid of zero means create a new stub manager */
-    
-    if (mid->oid && (manager = get_stub_manager(mid->oxid, mid->oid)))
+
+    if (mid->oid && (manager = get_stub_manager(apt, mid->oid)))
     {
         TRACE("registering new ifstub on pre-existing manager\n");
     }
     else
     {
-        struct apartment *apt;
-        
         TRACE("constructing new stub manager\n");
 
-        apt = COM_ApartmentFromOXID(mid->oxid, TRUE);
         manager = new_stub_manager(apt, obj);
-        COM_ApartmentRelease(apt);
-        if (!manager) return E_OUTOFMEMORY;
+        if (!manager)
+        {
+            COM_ApartmentRelease(apt);
+            return E_OUTOFMEMORY;
+        }
 
         mid->oid = manager->oid;
     }
@@ -125,6 +142,7 @@ static HRESULT register_ifstub(wine_marshal_id *mid, REFIID riid, IUnknown *obj,
         stub_manager_int_release(manager);
         /* FIXME: should we do another release to completely destroy the
          * stub manager? */
+        COM_ApartmentRelease(apt);
         return E_OUTOFMEMORY;
     }
 
@@ -133,6 +151,7 @@ static HRESULT register_ifstub(wine_marshal_id *mid, REFIID riid, IUnknown *obj,
     mid->ipid = ifstub->ipid;
 
     stub_manager_int_release(manager);
+    COM_ApartmentRelease(apt);
     return S_OK;
 }
 
@@ -489,9 +508,16 @@ StdMarshalImpl_MarshalInterface(
   IPSFactoryBuffer     *psfacbuf;
   BOOL                  tablemarshal;
   struct stub_manager  *manager;
+  APARTMENT            *apt = COM_CurrentApt();
     
   TRACE("(...,%s,...)\n",debugstr_guid(riid));
 
+  if (!apt)
+  {
+      ERR("Apartment not initialized\n");
+      return CO_E_NOTINITIALIZED;
+  }
+
   start_apartment_listener_thread(); /* just to be sure we have one running. */
 
   hres = get_facbuf_for_iid(riid,&psfacbuf);
@@ -508,11 +534,11 @@ StdMarshalImpl_MarshalInterface(
   if (tablemarshal) FIXME("table marshalling unimplemented\n");
 
   /* now fill out the MID */
-  mid.oxid = COM_CurrentApt()->oxid;
+  mid.oxid = apt->oxid;
  
   IUnknown_QueryInterface((LPUNKNOWN)pv, riid, (LPVOID*)&pUnk);
  
-  if ((manager = get_stub_manager_from_object(mid.oxid, pUnk)))
+  if ((manager = get_stub_manager_from_object(apt, pUnk)))
   {
       mid.oid = manager->oid;
       stub_manager_int_release(manager);
@@ -551,14 +577,17 @@ StdMarshalImpl_UnmarshalInterface(LPMARSHAL iface, IStream *pStm, REFIID riid, v
 
   TRACE("(...,%s,....)\n",debugstr_guid(riid));
 
-  if (!apt) return CO_E_NOTINITIALIZED;
+  if (!apt)
+  {
+      ERR("Apartment not initialized\n");
+      return CO_E_NOTINITIALIZED;
+  }
 
   hres = IStream_Read(pStm,&mid,sizeof(mid),&res);
   if (hres) return hres;
   
   /* check if we're marshalling back to ourselves */
-  /* FIXME: commented out until we can get the tests passing with it uncommented. */
-  if (/*(apt->oxid == mid.oxid) &&*/ (stubmgr = get_stub_manager(mid.oxid, mid.oid)))
+  if ((apt->oxid == mid.oxid) && (stubmgr = get_stub_manager(apt, mid.oid)))
   {
       TRACE("Unmarshalling object marshalled in same apartment for iid %s, returning original object %p\n", debugstr_guid(riid), stubmgr->object);
     
@@ -600,13 +629,20 @@ StdMarshalImpl_ReleaseMarshalData(LPMARSHAL iface, IStream *pStm) {
     ULONG                res;
     HRESULT              hres;
     struct stub_manager *stubmgr;
+    APARTMENT           *apt;
 
     TRACE("iface=%p, pStm=%p\n", iface, pStm);
     
     hres = IStream_Read(pStm,&mid,sizeof(mid),&res);
     if (hres) return hres;
 
-    if (!(stubmgr = get_stub_manager(mid.oxid, mid.oid)))
+    if (!(apt = COM_ApartmentFromOXID(mid.oxid, TRUE)))
+    {
+        WARN("Could not map OXID %s to apartment object\n", wine_dbgstr_longlong(mid.oxid));
+        return RPC_E_INVALID_OBJREF;
+    }
+
+    if (!(stubmgr = get_stub_manager(apt, mid.oid)))
     {
         ERR("could not map MID to stub manager, oxid=%s, oid=%s\n",
             wine_dbgstr_longlong(mid.oxid), wine_dbgstr_longlong(mid.oid));
@@ -616,6 +652,7 @@ StdMarshalImpl_ReleaseMarshalData(LPMARSHAL iface, IStream *pStm) {
     stub_manager_ext_release(stubmgr, 1);
 
     stub_manager_int_release(stubmgr);
+    COM_ApartmentRelease(apt);
 
     return S_OK;
 }

dlls/ole32/rpc.c

@@ -789,6 +789,7 @@ COM_RpcReceive(wine_pipe *xpipe) {
         wine_rpc_disconnect_header header;
         struct stub_manager *stubmgr;
         DWORD magic = 0xcafebabe;
+        APARTMENT *apt;
 
         hres = read_pipe(xhPipe, &header, sizeof(header));
         if (hres) {
@@ -798,15 +799,23 @@ COM_RpcReceive(wine_pipe *xpipe) {
 
         TRACE("read disconnect header\n");
 
-        if (!(stubmgr = get_stub_manager(header.mid.oxid, header.mid.oid)))
+        if (!(apt = COM_ApartmentFromOXID(header.mid.oxid, TRUE)))
+        {
+            ERR("Could not map OXID %s to apartment object in disconnect\n", wine_dbgstr_longlong(header.mid.oxid));
+            goto disconnect_end;
+        }
+
+        if (!(stubmgr = get_stub_manager(apt, header.mid.oid)))
         {
             ERR("could not locate stub to disconnect, mid.oid=%s\n", wine_dbgstr_longlong(header.mid.oid));
+            COM_ApartmentRelease(apt);
             goto disconnect_end;
         }
 
         stub_manager_ext_release(stubmgr, 1);
 
         stub_manager_int_release(stubmgr);
+        COM_ApartmentRelease(apt);
 
 disconnect_end:
         write_pipe(xhPipe, &magic, sizeof(magic));

dlls/ole32/stubmanager.c

@@ -107,26 +107,50 @@ static void stub_manager_delete(struct stub_manager *m)
     HeapFree(GetProcessHeap(), 0, m);
 }
 
-/* gets the stub manager associated with an object - caller must call
- * release on the stub manager when it is no longer needed */
-struct stub_manager *get_stub_manager_from_object(OXID oxid, void *object)
+/* 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 *result = NULL;
     struct list         *cursor;
-    APARTMENT           *apt;
 
-    if (!(apt = COM_ApartmentFromOXID(oxid, TRUE)))
+    EnterCriticalSection(&apt->cs);
+    LIST_FOR_EACH( cursor, &apt->stubmgrs )
     {
-        WARN("Could not map OXID %s to apartment object\n", wine_dbgstr_longlong(oxid));
-        return NULL;
+        struct stub_manager *m = LIST_ENTRY( cursor, struct stub_manager, entry );
+
+        if (m->object == object)
+        {
+            result = m;
+            stub_manager_int_addref(result);
+            break;
+        }
     }
+    LeaveCriticalSection(&apt->cs);
+
+    if (result)
+        TRACE("found %p for object %p\n", result, object);
+    else
+        TRACE("not found for object %p\n", object);
+
+    return result;    
+}
+
+/* gets the stub manager associated with an object id - 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(APARTMENT *apt, OID oid)
+{
+    struct stub_manager *result = NULL;
+    struct list         *cursor;
 
     EnterCriticalSection(&apt->cs);
     LIST_FOR_EACH( cursor, &apt->stubmgrs )
     {
         struct stub_manager *m = LIST_ENTRY( cursor, struct stub_manager, entry );
 
-        if (m->object == object)
+        if (m->oid == oid)
         {
             result = m;
             stub_manager_int_addref(result);
@@ -135,11 +159,12 @@ struct stub_manager *get_stub_manager_from_object(OXID oxid, void *object)
     }
     LeaveCriticalSection(&apt->cs);
 
-    COM_ApartmentRelease(apt);
-    
-    TRACE("found %p from object %p\n", result, object);
+    if (result)
+        TRACE("found %p for oid %s\n", result, wine_dbgstr_longlong(oid));
+    else
+        TRACE("not found for oid %s\n", wine_dbgstr_longlong(oid));
 
-    return result;    
+    return result;
 }
 
 /* increments the internal refcount */
@@ -174,41 +199,6 @@ ULONG stub_manager_int_release(struct stub_manager *This)
     return refs;
 }
 
-/* gets the stub manager associated with an object id - caller must call
- * release on the stub manager when it is no longer needed */
-struct stub_manager *get_stub_manager(OXID oxid, OID oid)
-{
-    struct stub_manager *result = NULL;
-    struct list         *cursor;
-    APARTMENT           *apt;
-
-    if (!(apt = COM_ApartmentFromOXID(oxid, TRUE)))
-    {
-        WARN("Could not map OXID %s to apartment object\n", wine_dbgstr_longlong(oxid));
-        return NULL;
-    }
-    
-    EnterCriticalSection(&apt->cs);
-    LIST_FOR_EACH( cursor, &apt->stubmgrs )
-    {
-        struct stub_manager *m = LIST_ENTRY( cursor, struct stub_manager, entry );
-
-        if (m->oid == oid)
-        {
-            result = m;
-            stub_manager_int_addref(result);
-            break;
-        }
-    }
-    LeaveCriticalSection(&apt->cs);
-
-    COM_ApartmentRelease(apt);
-
-    TRACE("found %p from oid %s\n", result, wine_dbgstr_longlong(oid));
-
-    return result;
-}
-
 /* add some external references (ie from a client that unmarshaled an ifptr) */
 ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs)
 {

dlls/ole32/tests/marshal.c

@@ -317,7 +317,7 @@ static void test_marshal_stub_apartment_shutdown()
 
     end_host_object(tid, thread);
 
-    todo_wine { ok_no_locks(); }
+    ok_no_locks();
 
     IUnknown_Release(pProxy);