- Move marshaling state machine into stub manager from ifstub.

dlls/ole32/compobj_private.h

@@ -53,10 +53,12 @@ typedef struct apartment APARTMENT;
 
 typedef enum ifstub_state
 {
-    IFSTUB_STATE_NORMAL_MARSHALED,
-    IFSTUB_STATE_NORMAL_UNMARSHALED,
-    IFSTUB_STATE_TABLE_MARSHALED
-} IFSTUB_STATE;
+    STUBSTATE_NORMAL_MARSHALED,
+    STUBSTATE_NORMAL_UNMARSHALED,
+    STUBSTATE_TABLE_WEAK_MARSHALED,
+    STUBSTATE_TABLE_WEAK_UNMARSHALED,
+    STUBSTATE_TABLE_STRONG,
+} STUB_STATE;
 
 /* an interface stub */
 struct ifstub   
@@ -66,7 +68,6 @@ struct ifstub
     IID               iid;        /* RO */
     IPID              ipid;       /* RO */
     IUnknown         *iface;      /* RO */
-    IFSTUB_STATE      state;      /* CS stub_manager->lock */
 };
 
 
@@ -78,11 +79,12 @@ struct stub_manager
     CRITICAL_SECTION  lock;
     APARTMENT        *apt;        /* owning apt (RO) */
 
-    ULONG             extrefs;    /* number of 'external' references (LOCK) */
+    ULONG             extrefs;    /* number of 'external' references (CS lock) */
     ULONG             refs;       /* internal reference count (CS apt->cs) */
     OID               oid;        /* apartment-scoped unique identifier (RO) */
     IUnknown         *object;     /* the object we are managing the stub for (RO) */
     ULONG             next_ipid;  /* currently unused (LOCK) */
+    STUB_STATE        state;      /* state machine (CS lock) */
 };
 
 /* imported interface proxy */
@@ -164,15 +166,16 @@ HRESULT MARSHAL_GetStandardMarshalCF(LPVOID *ppv);
 
 ULONG stub_manager_int_addref(struct stub_manager *This);
 ULONG stub_manager_int_release(struct stub_manager *This);
-struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object);
+struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object, MSHLFLAGS mshlflags);
 ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs);
 ULONG stub_manager_ext_release(struct stub_manager *m, ULONG refs);
-struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid, BOOL tablemarshal);
+struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid);
 struct stub_manager *get_stub_manager(APARTMENT *apt, OID oid);
 struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object);
 void apartment_disconnect_object(APARTMENT *apt, void *object);
-BOOL stub_manager_notify_unmarshal(struct stub_manager *m, const IPID *ipid);
-BOOL stub_manager_is_table_marshaled(struct stub_manager *m, const IPID *ipid);
+BOOL stub_manager_notify_unmarshal(struct stub_manager *m);
+BOOL stub_manager_is_table_marshaled(struct stub_manager *m);
+void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs);
 HRESULT register_ifstub(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnknown *obj, MSHLFLAGS mshlflags);
 HRESULT ipid_to_stub_manager(const IPID *ipid, APARTMENT **stub_apt, struct stub_manager **stubmgr_ret);
 IRpcStubBuffer *ipid_to_apt_and_stubbuffer(const IPID *ipid, APARTMENT **stub_apt);

dlls/ole32/marshal.c

@@ -117,13 +117,15 @@ HRESULT register_ifstub(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnkn
 
     stdobjref->oxid = apt->oxid;
 
+    /* FIXME: what happens if we register an interface twice with different
+     * marshaling flags? */
     if ((manager = get_stub_manager_from_object(apt, obj)))
         TRACE("registering new ifstub on pre-existing manager\n");
     else
     {
         TRACE("constructing new stub manager\n");
 
-        manager = new_stub_manager(apt, obj);
+        manager = new_stub_manager(apt, obj, mshlflags);
         if (!manager)
             return E_OUTOFMEMORY;
     }
@@ -131,7 +133,7 @@ HRESULT register_ifstub(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnkn
 
     tablemarshal = ((mshlflags & MSHLFLAGS_TABLESTRONG) || (mshlflags & MSHLFLAGS_TABLEWEAK));
 
-    ifstub = stub_manager_new_ifstub(manager, stub, obj, riid, tablemarshal);
+    ifstub = stub_manager_new_ifstub(manager, stub, obj, riid);
     if (!ifstub)
     {
         IRpcStubBuffer_Release(stub);
@@ -909,7 +911,7 @@ StdMarshalImpl_UnmarshalInterface(LPMARSHAL iface, IStream *pStm, REFIID riid, v
       hres = IUnknown_QueryInterface(stubmgr->object, riid, ppv);
       
       /* unref the ifstub. FIXME: only do this on success? */
-      if (!stub_manager_is_table_marshaled(stubmgr, &stdobjref.ipid))
+      if (!stub_manager_is_table_marshaled(stubmgr))
           stub_manager_ext_release(stubmgr, 1);
 
       stub_manager_int_release(stubmgr);
@@ -925,7 +927,7 @@ StdMarshalImpl_UnmarshalInterface(LPMARSHAL iface, IStream *pStm, REFIID riid, v
   {
       if ((stubmgr = get_stub_manager(stub_apt, stdobjref.oid)))
       {
-          if (!stub_manager_notify_unmarshal(stubmgr, &stdobjref.ipid))
+          if (!stub_manager_notify_unmarshal(stubmgr))
               hres = CO_E_OBJNOTCONNECTED;
 
           stub_manager_int_release(stubmgr);
@@ -980,9 +982,7 @@ StdMarshalImpl_ReleaseMarshalData(LPMARSHAL iface, IStream *pStm) {
         return RPC_E_INVALID_OBJREF;
     }
 
-    /* FIXME: don't release if table-weak and already unmarshaled an object */
-    /* FIXME: this should also depend on stdobjref.cPublicRefs */
-    stub_manager_ext_release(stubmgr, 1);
+    stub_manager_release_marshal_data(stubmgr, stdobjref.cPublicRefs);
 
     stub_manager_int_release(stubmgr);
     COM_ApartmentRelease(apt);

dlls/ole32/stubmanager.c

@@ -29,6 +29,7 @@
 
 #include <assert.h>
 #include <stdarg.h>
+#include <limits.h>
 
 #include "windef.h"
 #include "winbase.h"
@@ -48,7 +49,7 @@ static struct ifstub *stub_manager_ipid_to_ifstub(struct stub_manager *m, const
 /* 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)
+struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object, MSHLFLAGS mshlflags)
 {
     struct stub_manager *sm;
 
@@ -74,6 +75,13 @@ struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object)
      * the marshalled ifptr.
      */
     sm->extrefs = 0;
+
+    if (mshlflags & MSHLFLAGS_TABLESTRONG)
+        sm->state = STUBSTATE_TABLE_STRONG;
+    else if (mshlflags & MSHLFLAGS_TABLEWEAK)
+        sm->state = STUBSTATE_TABLE_WEAK_UNMARSHALED;
+    else
+        sm->state = STUBSTATE_NORMAL_MARSHALED;
     
     EnterCriticalSection(&apt->cs);
     sm->oid    = apt->oidc++;
@@ -228,8 +236,16 @@ ULONG stub_manager_int_release(struct stub_manager *This)
 /* add some external references (ie from a client that unmarshaled an ifptr) */
 ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs)
 {
-    ULONG rc = InterlockedExchangeAdd(&m->extrefs, refs) + refs;
+    ULONG rc;
+
+    EnterCriticalSection(&m->lock);
+    
+    /* make sure we don't overflow extrefs */
+    refs = min(refs, (ULONG_MAX-1 - m->extrefs));
+    rc = (m->extrefs += refs);
 
+    LeaveCriticalSection(&m->lock);
+    
     TRACE("added %lu refs to %p (oid %s), rc is now %lu\n", refs, m, wine_dbgstr_longlong(m->oid), rc);
 
     return rc;
@@ -238,7 +254,15 @@ ULONG stub_manager_ext_addref(struct stub_manager *m, ULONG refs)
 /* remove some external references */
 ULONG stub_manager_ext_release(struct stub_manager *m, ULONG refs)
 {
-    ULONG rc = InterlockedExchangeAdd(&m->extrefs, -refs) - refs;
+    ULONG rc;
+
+    EnterCriticalSection(&m->lock);
+
+    /* make sure we don't underflow extrefs */
+    refs = min(refs, m->extrefs);
+    rc = (m->extrefs -= refs);
+
+    LeaveCriticalSection(&m->lock);
     
     TRACE("removed %lu refs from %p (oid %s), rc is now %lu\n", refs, m, wine_dbgstr_longlong(m->oid), rc);
 
@@ -370,12 +394,12 @@ static inline HRESULT generate_ipid(struct stub_manager *m, IPID *ipid)
 }
 
 /* registers a new interface stub COM object with the stub manager and returns registration record */
-struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid, BOOL tablemarshal)
+struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid)
 {
     struct ifstub *stub;
 
-    TRACE("oid=%s, stubbuffer=%p, iptr=%p, iid=%s, tablemarshal=%s\n",
-          wine_dbgstr_longlong(m->oid), sb, iptr, debugstr_guid(iid), tablemarshal ? "TRUE" : "FALSE");
+    TRACE("oid=%s, stubbuffer=%p, iptr=%p, iid=%s\n",
+          wine_dbgstr_longlong(m->oid), sb, iptr, debugstr_guid(iid));
 
     stub = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(struct ifstub));
     if (!stub) return NULL;
@@ -386,11 +410,6 @@ struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *s
     /* no need to ref this, same object as sb */
     stub->iface = iptr;
 
-    if (tablemarshal)
-        stub->state = IFSTUB_STATE_TABLE_MARSHALED;
-    else
-        stub->state = IFSTUB_STATE_NORMAL_MARSHALED;
-
     stub->iid = *iid;
 
     /* FIXME: hack for IRemUnknown because we don't notify SCM of our IPID
@@ -430,33 +449,30 @@ static void stub_manager_delete_ifstub(struct stub_manager *m, struct ifstub *if
 }
 
 /* returns TRUE if it is possible to unmarshal, FALSE otherwise. */
-BOOL stub_manager_notify_unmarshal(struct stub_manager *m, const IPID *ipid)
+BOOL stub_manager_notify_unmarshal(struct stub_manager *m)
 {
-    struct ifstub *ifstub;
     BOOL ret;
 
-    ifstub = stub_manager_ipid_to_ifstub(m, ipid);
-    if (!ifstub)
-    {
-        WARN("Can't find ifstub for OID %s, IPID %s\n",
-            wine_dbgstr_longlong(m->oid), wine_dbgstr_guid(ipid));
-        return FALSE;
-    }
-
     EnterCriticalSection(&m->lock);
 
-    switch (ifstub->state)
+    switch (m->state)
     {
-    case IFSTUB_STATE_TABLE_MARSHALED:
+    case STUBSTATE_TABLE_STRONG:
+    case STUBSTATE_TABLE_WEAK_MARSHALED:
+        /* no transition */
         ret = TRUE;
         break;
-    case IFSTUB_STATE_NORMAL_MARSHALED:
-        ifstub->state = IFSTUB_STATE_NORMAL_UNMARSHALED;
+    case STUBSTATE_TABLE_WEAK_UNMARSHALED:
+        m->state = STUBSTATE_TABLE_WEAK_MARSHALED;
+        ret = TRUE;
+        break;
+    case STUBSTATE_NORMAL_MARSHALED:
+        m->state = STUBSTATE_NORMAL_UNMARSHALED;
         ret = TRUE;
         break;
     default:
-        WARN("object OID %s, IPID %s already unmarshaled\n",
-            wine_dbgstr_longlong(m->oid), wine_dbgstr_guid(ipid));
+        WARN("object OID %s already unmarshaled\n",
+            wine_dbgstr_longlong(m->oid));
         ret = FALSE;
         break;
     }
@@ -466,22 +482,39 @@ BOOL stub_manager_notify_unmarshal(struct stub_manager *m, const IPID *ipid)
     return ret;
 }
 
-/* is an ifstub table marshaled? */
-BOOL stub_manager_is_table_marshaled(struct stub_manager *m, const IPID *ipid)
+void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs)
 {
-    struct ifstub *ifstub;
-    BOOL ret;
+    EnterCriticalSection(&m->lock);
 
-    ifstub = stub_manager_ipid_to_ifstub(m, ipid);
-    if (!ifstub)
+    switch (m->state)
     {
-        WARN("Can't find ifstub for OID %s, IPID %s\n",
-            wine_dbgstr_longlong(m->oid), wine_dbgstr_guid(ipid));
-        return FALSE;
+    case STUBSTATE_NORMAL_MARSHALED:
+    case STUBSTATE_NORMAL_UNMARSHALED: /* FIXME: check this */
+        /* nothing to change */
+        break;
+    case STUBSTATE_TABLE_WEAK_UNMARSHALED:
+    case STUBSTATE_TABLE_STRONG:
+        refs = 1;
+        break;
+    case STUBSTATE_TABLE_WEAK_MARSHALED:
+        refs = 0; /* like native */
+        break;
     }
 
+    stub_manager_ext_release(m, refs);
+
+    LeaveCriticalSection(&m->lock);
+}
+
+/* is an ifstub table marshaled? */
+BOOL stub_manager_is_table_marshaled(struct stub_manager *m)
+{
+    BOOL ret;
+
     EnterCriticalSection(&m->lock);
-    ret = (ifstub->state == IFSTUB_STATE_TABLE_MARSHALED);
+    ret = ((m->state == STUBSTATE_TABLE_STRONG) ||
+           (m->state == STUBSTATE_TABLE_WEAK_MARSHALED) ||
+           (m->state == STUBSTATE_TABLE_WEAK_UNMARSHALED));
     LeaveCriticalSection(&m->lock);
 
     return ret;
@@ -693,7 +726,7 @@ HRESULT start_apartment_remote_unknown()
         {
             STDOBJREF stdobjref; /* dummy - not used */
             /* register it with the stub manager */
-            hr = register_ifstub(COM_CurrentApt(), &stdobjref, &IID_IRemUnknown, (IUnknown *)pRemUnknown, MSHLFLAGS_NORMAL);
+            hr = register_ifstub(apt, &stdobjref, &IID_IRemUnknown, (IUnknown *)pRemUnknown, MSHLFLAGS_NORMAL);
             /* release our reference to the object as the stub manager will manage the life cycle for us */
             IRemUnknown_Release(pRemUnknown);
             if (hr == S_OK)

dlls/ole32/tests/marshal.c

@@ -669,14 +669,14 @@ static void test_tableweak_marshal_releasedata1()
     IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
     release_host_object(tid);
 
-    todo_wine { ok_more_than_one_lock(); }
+    ok_more_than_one_lock();
 
     IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
     hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy2);
-    todo_wine { ok_ole_success(hr, CoUnmarshalInterface); }
+    ok_ole_success(hr, CoUnmarshalInterface);
     IStream_Release(pStream);
 
-    todo_wine { ok_more_than_one_lock(); }
+    ok_more_than_one_lock();
 
     IUnknown_Release(pProxy1);
     if (pProxy2)
@@ -712,20 +712,19 @@ static void test_tableweak_marshal_releasedata2()
     IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
     release_host_object(tid);
 
-    todo_wine
-    {
-
     ok_no_locks();
 
     IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
     hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy);
+    todo_wine
+    {
     ok(hr == CO_E_OBJNOTREG,
        "CoUnmarshalInterface should have failed with CO_E_OBJNOTREG, but returned 0x%08lx instead\n",
        hr);
+    }
     IStream_Release(pStream);
 
     ok_no_locks();
-    }
 
     end_host_object(tid, thread);
 }