Authors: Mike Hearn <mh@codeweavers.com>, Robert Shearman <rob@codeweavers.com>

dlls/ole32/compobj_private.h

@@ -71,6 +71,7 @@ struct ifstub
     IID               iid;        /* RO */
     IPID              ipid;       /* RO */
     IUnknown         *iface;      /* RO */
+    MSHLFLAGS         flags;      /* so we can enforce process-local marshalling rules (RO) */
 };
 
 
@@ -87,7 +88,14 @@ struct stub_manager
     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) */
+
+    /* We need to keep a count of the outstanding marshals, so we can enforce the
+     * marshalling rules (ie, you can only unmarshal normal marshals once). Note
+     * that these counts do NOT include unmarshalled interfaces, once a stream is
+     * unmarshalled and a proxy set up, this count is decremented.
+     */
+
+    ULONG             norm_refs;  /* refcount of normal marshals (CS lock) */
 };
 
 /* imported interface proxy */
@@ -170,15 +178,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, MSHLFLAGS mshlflags);
+struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object);
 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);
+struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid, MSHLFLAGS flags);
+struct ifstub *stub_manager_find_ifstub(struct stub_manager *m, REFIID iid, MSHLFLAGS flags);
 struct stub_manager *get_stub_manager(APARTMENT *apt, OID oid);
 struct stub_manager *get_stub_manager_from_object(APARTMENT *apt, void *object);
-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);
+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);
+void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs, const IPID *ipid);
 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);
 HRESULT start_apartment_remote_unknown(void);

dlls/ole32/marshal.c

@@ -84,7 +84,7 @@ inline static HRESULT get_facbuf_for_iid(REFIID riid, IPSFactoryBuffer **facbuf)
         &IID_IPSFactoryBuffer, (LPVOID*)facbuf);
 }
 
-/* creates a new stub manager */
+/* marshals an object into a STDOBJREF structure */
 HRESULT marshal_object(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnknown *object, MSHLFLAGS mshlflags)
 {
     struct stub_manager *manager;
@@ -135,15 +135,13 @@ HRESULT marshal_object(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnkno
     else
         stdobjref->flags = SORF_NULL;
 
-    /* FIXME: what happens if we register an interface twice with different
-     * marshaling flags? */
     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, mshlflags);
+        manager = new_stub_manager(apt, object);
         if (!manager)
         {
             if (stub) IRpcStubBuffer_Release(stub);
@@ -155,15 +153,20 @@ HRESULT marshal_object(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnkno
 
     tablemarshal = ((mshlflags & MSHLFLAGS_TABLESTRONG) || (mshlflags & MSHLFLAGS_TABLEWEAK));
 
-    ifstub = stub_manager_new_ifstub(manager, stub, iobject, riid);
-    IUnknown_Release(iobject);
-    if (stub) IRpcStubBuffer_Release(stub);
+    /* make sure ifstub that we are creating is unique */
+    ifstub = stub_manager_find_ifstub(manager, riid, mshlflags);
     if (!ifstub)
     {
-        stub_manager_int_release(manager);
-        /* FIXME: should we do another release to completely destroy the
-         * stub manager? */
-        return E_OUTOFMEMORY;
+        ifstub = stub_manager_new_ifstub(manager, stub, iobject, riid, mshlflags);
+        IUnknown_Release(iobject);
+        if (stub) IRpcStubBuffer_Release(stub);
+        if (!ifstub)
+        {
+            stub_manager_int_release(manager);
+            /* FIXME: should we do another release to completely destroy the
+             * stub manager? */
+            return E_OUTOFMEMORY;
+        }
     }
 
     if (!tablemarshal)
@@ -1034,7 +1037,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))
+        if (!stub_manager_is_table_marshaled(stubmgr, &stdobjref.ipid))
             stub_manager_ext_release(stubmgr, stdobjref.cPublicRefs);
 
         stub_manager_int_release(stubmgr);
@@ -1050,7 +1053,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))
+            if (!stub_manager_notify_unmarshal(stubmgr, &stdobjref.ipid))
                 hres = CO_E_OBJNOTCONNECTED;
 
             stub_manager_int_release(stubmgr);
@@ -1111,7 +1114,7 @@ StdMarshalImpl_ReleaseMarshalData(LPMARSHAL iface, IStream *pStm)
         return RPC_E_INVALID_OBJREF;
     }
 
-    stub_manager_release_marshal_data(stubmgr, stdobjref.cPublicRefs);
+    stub_manager_release_marshal_data(stubmgr, stdobjref.cPublicRefs, &stdobjref.ipid);
 
     stub_manager_int_release(stubmgr);
     apartment_release(apt);
@@ -1343,6 +1346,18 @@ HRESULT WINAPI CoGetMarshalSizeMax(ULONG *pulSize, REFIID riid, IUnknown *pUnk,
 }
 
 
+static void dump_MSHLFLAGS(MSHLFLAGS flags)
+{
+    if (flags & MSHLFLAGS_TABLESTRONG)
+        TRACE(" MSHLFLAGS_TABLESTRONG");
+    if (flags & MSHLFLAGS_TABLEWEAK)
+        TRACE(" MSHLFLAGS_TABLEWEAK");
+    if (!(flags & (MSHLFLAGS_TABLESTRONG|MSHLFLAGS_TABLEWEAK)))
+        TRACE(" MSHLFLAGS_NORMAL");
+    if (flags & MSHLFLAGS_NOPING)
+        TRACE(" MSHLFLAGS_NOPING");
+}
+
 /***********************************************************************
  *		CoMarshalInterface	[OLE32.@]
  *
@@ -1384,8 +1399,10 @@ HRESULT WINAPI CoMarshalInterface(IStream *pStream, REFIID riid, IUnknown *pUnk,
     OBJREF objref;
     LPMARSHAL pMarshal;
 
-    TRACE("(%p, %s, %p, %lx, %p, %lx)\n", pStream, debugstr_guid(riid), pUnk,
-        dwDestContext, pvDestContext, mshlFlags);
+    TRACE("(%p, %s, %p, %lx, %p,", pStream, debugstr_guid(riid), pUnk,
+        dwDestContext, pvDestContext);
+    dump_MSHLFLAGS(mshlFlags);
+    TRACE(")\n");
 
     if (pUnk == NULL)
         return E_INVALIDARG;

dlls/ole32/stubmanager.c

@@ -49,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, MSHLFLAGS mshlflags)
+struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object)
 {
     struct stub_manager *sm;
 
@@ -79,15 +79,8 @@ struct stub_manager *new_stub_manager(APARTMENT *apt, IUnknown *object, MSHLFLAG
      */
     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++;
+    sm->oid = apt->oidc++;
     list_add_head(&apt->stubmgrs, &sm->entry);
     LeaveCriticalSection(&apt->cs);
 
@@ -301,6 +294,25 @@ static struct ifstub *stub_manager_ipid_to_ifstub(struct stub_manager *m, const
     return result;
 }
 
+struct ifstub *stub_manager_find_ifstub(struct stub_manager *m, REFIID iid, MSHLFLAGS flags)
+{
+    struct ifstub  *result = NULL;
+    struct ifstub  *ifstub;
+
+    EnterCriticalSection(&m->lock);
+    LIST_FOR_EACH_ENTRY( ifstub, &m->ifstubs, struct ifstub, entry )
+    {
+        if (IsEqualIID(iid, &ifstub->iid) && (ifstub->flags == flags))
+        {
+            result = ifstub;
+            break;
+        }
+    }
+    LeaveCriticalSection(&m->lock);
+
+    return result;
+}
+
 /* gets the stub manager associated with an ipid - 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 */
@@ -402,7 +414,7 @@ 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)
+struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *sb, IUnknown *iptr, REFIID iid, MSHLFLAGS flags)
 {
     struct ifstub *stub;
 
@@ -417,11 +429,16 @@ struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *s
 
     IUnknown_AddRef(iptr);
     stub->iface = iptr;
-
+    stub->flags = flags;
     stub->iid = *iid;
 
-    /* FIXME: hack for IRemUnknown because we don't notify SCM of our IPID
-     * yet, so we need to use a well-known one */
+    /* 
+     * FIXME: this is a hack for marshalling IRemUnknown. In real
+     * DCOM, the IPID of the IRemUnknown interface is generated like
+     * any other and passed to the OXID resolver which then returns it
+     * when queried. We don't have an OXID resolver yet so instead we
+     * use a magic IPID reserved for IRemUnknown.
+     */
     if (IsEqualIID(iid, &IID_IRemUnknown))
     {
         stub->ipid.Data1 = 0xffffffff;
@@ -435,6 +452,8 @@ struct ifstub *stub_manager_new_ifstub(struct stub_manager *m, IRpcStubBuffer *s
 
     EnterCriticalSection(&m->lock);
     list_add_head(&m->ifstubs, &stub->entry);
+    /* every normal marshal is counted so we don't allow more than we should */
+    if (flags & MSHLFLAGS_NORMAL) m->norm_refs++;
     LeaveCriticalSection(&m->lock);
 
     TRACE("ifstub %p created with ipid %s\n", stub, debugstr_guid(&stub->ipid));
@@ -457,34 +476,29 @@ 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)
+BOOL stub_manager_notify_unmarshal(struct stub_manager *m, const IPID *ipid)
 {
-    BOOL ret;
+    BOOL ret = TRUE;
+    struct ifstub *ifstub;
+
+    if (!(ifstub = stub_manager_ipid_to_ifstub(m, ipid)))
+    {
+        ERR("attempted unmarshal of unknown IPID %s\n", debugstr_guid(ipid));
+        return FALSE;
+    }
 
     EnterCriticalSection(&m->lock);
 
-    switch (m->state)
+    /* track normal marshals so we can enforce rules whilst in-process */
+    if (ifstub->flags & MSHLFLAGS_NORMAL)
     {
-    case STUBSTATE_TABLE_STRONG:
-    case STUBSTATE_TABLE_WEAK_MARSHALED:
-        /* no transition */
-        ret = TRUE;
-        break;
-    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 already unmarshaled\n",
-            wine_dbgstr_longlong(m->oid));
-        ret = TRUE; /* FIXME: the state management should be per-ifstub, so
-                     * it is disabled at the moment so that InstallShield
-                     * works again */
-        break;
+        if (m->norm_refs)
+            m->norm_refs--;
+        else
+        {
+            ERR("attempted invalid normal unmarshal, norm_refs is zero\n");
+            ret = FALSE;
+        }
     }
 
     LeaveCriticalSection(&m->lock);
@@ -492,42 +506,30 @@ BOOL stub_manager_notify_unmarshal(struct stub_manager *m)
     return ret;
 }
 
-void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs)
+/* handles refcounting for CoReleaseMarshalData */
+void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs, const IPID *ipid)
 {
-    EnterCriticalSection(&m->lock);
-
-    switch (m->state)
-    {
-    case STUBSTATE_NORMAL_MARSHALED:
-    case STUBSTATE_NORMAL_UNMARSHALED: /* FIXME: check this */
-        /* nothing to change */
-        break;
-    case STUBSTATE_TABLE_WEAK_UNMARSHALED:
-    case STUBSTATE_TABLE_STRONG:
+    struct ifstub *ifstub;
+ 
+    if (!(ifstub = stub_manager_ipid_to_ifstub(m, ipid)))
+        return;
+ 
+    if (ifstub->flags & MSHLFLAGS_TABLEWEAK)
+        refs = 0;
+    else
         refs = 1;
-        break;
-    case STUBSTATE_TABLE_WEAK_MARSHALED:
-        refs = 0; /* like native */
-        break;
-    }
-
-    LeaveCriticalSection(&m->lock);
 
     stub_manager_ext_release(m, refs);
 }
 
 /* is an ifstub table marshaled? */
-BOOL stub_manager_is_table_marshaled(struct stub_manager *m)
+BOOL stub_manager_is_table_marshaled(struct stub_manager *m, const IPID *ipid)
 {
-    BOOL ret;
-
-    EnterCriticalSection(&m->lock);
-    ret = ((m->state == STUBSTATE_TABLE_STRONG) ||
-           (m->state == STUBSTATE_TABLE_WEAK_MARSHALED) ||
-           (m->state == STUBSTATE_TABLE_WEAK_UNMARSHALED));
-    LeaveCriticalSection(&m->lock);
-
-    return ret;
+    struct ifstub *ifstub = stub_manager_ipid_to_ifstub(m, ipid);
+ 
+    assert( ifstub );
+    
+    return ifstub->flags & (MSHLFLAGS_TABLESTRONG | MSHLFLAGS_TABLEWEAK);
 }