ole32: Handle MSHLFLAGS_TABLEWEAK and MSHLFLAGS_TABLESTRONG when marshaling a proxy.

dlls/ole32/marshal.c

@@ -375,22 +375,29 @@ static HRESULT WINAPI Proxy_MarshalInterface(
     if (SUCCEEDED(hr))
     {
         STDOBJREF stdobjref = ifproxy->stdobjref;
-        ULONG cPublicRefs = ifproxy->refs;
-        ULONG cPublicRefsOld;
 
-        /* optimization - share out proxy's public references if possible
-         * instead of making new proxy do a roundtrip through the server */
-        do
+        stdobjref.cPublicRefs = 0;
+
+        if ((mshlflags != MSHLFLAGS_TABLEWEAK) &&
+            (mshlflags != MSHLFLAGS_TABLESTRONG))
         {
-            ULONG cPublicRefsNew;
-            cPublicRefsOld = cPublicRefs;
-            stdobjref.cPublicRefs = cPublicRefs / 2;
-            cPublicRefsNew = cPublicRefs - stdobjref.cPublicRefs;
-            cPublicRefs = InterlockedCompareExchange(
-                (LONG *)&ifproxy->refs, cPublicRefsNew, cPublicRefsOld);
-        } while (cPublicRefs != cPublicRefsOld);
-
-        if (!stdobjref.cPublicRefs)
+            ULONG cPublicRefs = ifproxy->refs;
+            ULONG cPublicRefsOld;
+            /* optimization - share out proxy's public references if possible
+             * instead of making new proxy do a roundtrip through the server */
+            do
+            {
+                ULONG cPublicRefsNew;
+                cPublicRefsOld = cPublicRefs;
+                stdobjref.cPublicRefs = cPublicRefs / 2;
+                cPublicRefsNew = cPublicRefs - stdobjref.cPublicRefs;
+                cPublicRefs = InterlockedCompareExchange(
+                    (LONG *)&ifproxy->refs, cPublicRefsNew, cPublicRefsOld);
+            } while (cPublicRefs != cPublicRefsOld);
+        }
+
+        /* normal and table-strong marshaling need at least one reference */
+        if (!stdobjref.cPublicRefs && (mshlflags != MSHLFLAGS_TABLEWEAK))
         {
             IRemUnknown *remunk;
             hr = proxy_manager_get_remunknown(This, &remunk);
@@ -399,11 +406,16 @@ static HRESULT WINAPI Proxy_MarshalInterface(
                 HRESULT hrref = S_OK;
                 REMINTERFACEREF rif;
                 rif.ipid = ifproxy->stdobjref.ipid;
-                rif.cPublicRefs = NORMALEXTREFS;
+                rif.cPublicRefs = (mshlflags == MSHLFLAGS_TABLESTRONG) ? 1 : NORMALEXTREFS;
                 rif.cPrivateRefs = 0;
                 hr = IRemUnknown_RemAddRef(remunk, 1, &rif, &hrref);
                 if (hr == S_OK && hrref == S_OK)
-                    stdobjref.cPublicRefs = rif.cPublicRefs;
+                {
+                    /* table-strong marshaling doesn't give the refs to the
+                     * client that unmarshals the STDOBJREF */
+                    if (mshlflags != MSHLFLAGS_TABLESTRONG)
+                        stdobjref.cPublicRefs = rif.cPublicRefs;
+                }
                 else
                     ERR("IRemUnknown_RemAddRef returned with 0x%08x, hrref = 0x%08x\n", hr, hrref);
             }

dlls/ole32/tests/marshal.c

@@ -609,6 +609,112 @@ static void test_proxy_marshal_and_unmarshal2(void)
     end_host_object(tid, thread);
 }
 
+/* tests success case of an interthread marshal and then table-weak-marshaling the proxy */
+static void test_proxy_marshal_and_unmarshal_weak(void)
+{
+    HRESULT hr;
+    IStream *pStream = NULL;
+    IUnknown *pProxy = NULL;
+    IUnknown *pProxy2 = NULL;
+    DWORD tid;
+    HANDLE thread;
+
+    cLocks = 0;
+
+    hr = CreateStreamOnHGlobal(NULL, TRUE, &pStream);
+    ok_ole_success(hr, CreateStreamOnHGlobal);
+    tid = start_host_object(pStream, &IID_IClassFactory, (IUnknown*)&Test_ClassFactory, MSHLFLAGS_NORMAL, &thread);
+
+    ok_more_than_one_lock();
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy);
+    ok_ole_success(hr, CoUnmarshalInterface);
+
+    ok_more_than_one_lock();
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    /* marshal the proxy */
+    hr = CoMarshalInterface(pStream, &IID_IClassFactory, pProxy, MSHCTX_INPROC, NULL, MSHLFLAGS_TABLEWEAK);
+    ok_ole_success(hr, CoMarshalInterface);
+
+    ok_more_than_one_lock();
+
+    /* release the original proxy to test that we successfully keep the
+     * original object alive */
+    IUnknown_Release(pProxy);
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy2);
+    todo_wine
+    ok(hr == CO_E_OBJNOTREG, "CoUnmarshalInterface should return CO_E_OBJNOTREG instead of 0x%08x\n", hr);
+
+    ok_no_locks();
+
+    IStream_Release(pStream);
+
+    end_host_object(tid, thread);
+}
+
+/* tests success case of an interthread marshal and then table-strong-marshaling the proxy */
+static void test_proxy_marshal_and_unmarshal_strong(void)
+{
+    HRESULT hr;
+    IStream *pStream = NULL;
+    IUnknown *pProxy = NULL;
+    IUnknown *pProxy2 = NULL;
+    DWORD tid;
+    HANDLE thread;
+
+    cLocks = 0;
+
+    hr = CreateStreamOnHGlobal(NULL, TRUE, &pStream);
+    ok_ole_success(hr, CreateStreamOnHGlobal);
+    tid = start_host_object(pStream, &IID_IClassFactory, (IUnknown*)&Test_ClassFactory, MSHLFLAGS_NORMAL, &thread);
+
+    ok_more_than_one_lock();
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy);
+    ok_ole_success(hr, CoUnmarshalInterface);
+
+    ok_more_than_one_lock();
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    /* marshal the proxy */
+    hr = CoMarshalInterface(pStream, &IID_IClassFactory, pProxy, MSHCTX_INPROC, NULL, MSHLFLAGS_TABLESTRONG);
+    ok(hr == S_OK /* WinNT */ || hr == E_INVALIDARG /* Win9x */,
+        "CoMarshalInterface should have return S_OK or E_INVALIDARG instead of 0x%08x\n", hr);
+    if (FAILED(hr))
+    {
+        IUnknown_Release(pProxy);
+        goto end;
+    }
+
+    ok_more_than_one_lock();
+
+    /* release the original proxy to test that we successfully keep the
+     * original object alive */
+    IUnknown_Release(pProxy);
+
+    IStream_Seek(pStream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(pStream, &IID_IClassFactory, (void **)&pProxy2);
+    ok_ole_success(hr, CoUnmarshalInterface);
+
+    ok_more_than_one_lock();
+
+    IUnknown_Release(pProxy2);
+
+    ok_more_than_one_lock();
+
+end:
+    IStream_Release(pStream);
+
+    end_host_object(tid, thread);
+
+    ok_no_locks();
+}
+
 /* tests that stubs are released when the containing apartment is destroyed */
 static void test_marshal_stub_apartment_shutdown(void)
 {
@@ -2750,6 +2856,8 @@ START_TEST(marshal)
     test_interthread_marshal_and_unmarshal();
     test_proxy_marshal_and_unmarshal();
     test_proxy_marshal_and_unmarshal2();
+    test_proxy_marshal_and_unmarshal_weak();
+    test_proxy_marshal_and_unmarshal_strong();
     test_marshal_stub_apartment_shutdown();
     test_marshal_proxy_apartment_shutdown();
     test_marshal_proxy_mta_apartment_shutdown();