[PATCH 6/6] ole32: Allow marshalling objects from an implicit MTA.

[Zebediah Figura]

Signed-off-by: Zebediah Figura <z.figura12 at gmail.com>
---
 dlls/ole32/compobj_private.h |  2 +-
 dlls/ole32/marshal.c         |  5 +++--
 dlls/ole32/rpc.c             |  2 +-
 dlls/ole32/stubmanager.c     |  3 +--
 dlls/ole32/tests/marshal.c   | 50 ++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/dlls/ole32/compobj_private.h b/dlls/ole32/compobj_private.h
index dc09d20..212d328 100644
--- a/dlls/ole32/compobj_private.h
+++ b/dlls/ole32/compobj_private.h
@@ -200,7 +200,7 @@ void stub_manager_release_marshal_data(struct stub_manager *m, ULONG refs, const
 void stub_manager_disconnect(struct stub_manager *m) DECLSPEC_HIDDEN;
 HRESULT ipid_get_dispatch_params(const IPID *ipid, APARTMENT **stub_apt, struct stub_manager **manager, IRpcStubBuffer **stub,
                                  IRpcChannelBuffer **chan, IID *iid, IUnknown **iface) DECLSPEC_HIDDEN;
-HRESULT start_apartment_remote_unknown(void) DECLSPEC_HIDDEN;
+HRESULT start_apartment_remote_unknown(APARTMENT *apt) DECLSPEC_HIDDEN;
 
 HRESULT marshal_object(APARTMENT *apt, STDOBJREF *stdobjref, REFIID riid, IUnknown *obj, DWORD dest_context, void *dest_context_data, MSHLFLAGS mshlflags) DECLSPEC_HIDDEN;
 
diff --git a/dlls/ole32/marshal.c b/dlls/ole32/marshal.c
index 822d06e..7c0f541 100644
--- a/dlls/ole32/marshal.c
+++ b/dlls/ole32/marshal.c
@@ -1225,11 +1225,11 @@ StdMarshalImpl_MarshalInterface(
     STDOBJREF             stdobjref;
     ULONG                 res;
     HRESULT               hres;
-    APARTMENT            *apt = COM_CurrentApt();
+    APARTMENT *apt;
 
     TRACE("(...,%s,...)\n", debugstr_guid(riid));
 
-    if (!apt)
+    if (!(apt = apartment_get_current_or_mta()))
     {
         ERR("Apartment not initialized\n");
         return CO_E_NOTINITIALIZED;
@@ -1239,6 +1239,7 @@ StdMarshalImpl_MarshalInterface(
     RPC_StartRemoting(apt);
 
     hres = marshal_object(apt, &stdobjref, riid, pv, dest_context, dest_context_data, mshlflags);
+    apartment_release(apt);
     if (hres != S_OK)
     {
         ERR("Failed to create ifstub, hres=0x%x\n", hres);
diff --git a/dlls/ole32/rpc.c b/dlls/ole32/rpc.c
index 5a7626b..a73d23c 100644
--- a/dlls/ole32/rpc.c
+++ b/dlls/ole32/rpc.c
@@ -1648,7 +1648,7 @@ void RPC_StartRemoting(struct apartment *apt)
 
         /* FIXME: move remote unknown exporting into this function */
     }
-    start_apartment_remote_unknown();
+    start_apartment_remote_unknown(apt);
 }
 
 
diff --git a/dlls/ole32/stubmanager.c b/dlls/ole32/stubmanager.c
index 57048c6..75c4b04 100644
--- a/dlls/ole32/stubmanager.c
+++ b/dlls/ole32/stubmanager.c
@@ -812,11 +812,10 @@ static const IRemUnknownVtbl RemUnknown_Vtbl =
 };
 
 /* starts the IRemUnknown listener for the current apartment */
-HRESULT start_apartment_remote_unknown(void)
+HRESULT start_apartment_remote_unknown(APARTMENT *apt)
 {
     IRemUnknown *pRemUnknown;
     HRESULT hr = S_OK;
-    APARTMENT *apt = COM_CurrentApt();
 
     EnterCriticalSection(&apt->cs);
     if (!apt->remunk_exported)
diff --git a/dlls/ole32/tests/marshal.c b/dlls/ole32/tests/marshal.c
index c5c69f0..e8558dd 100644
--- a/dlls/ole32/tests/marshal.c
+++ b/dlls/ole32/tests/marshal.c
@@ -3475,10 +3475,34 @@ static DWORD CALLBACK implicit_mta_use_proc(void *param)
     return 0;
 }
 
+struct implicit_mta_marshal_data
+{
+    IStream *stream;
+    HANDLE start;
+    HANDLE stop;
+};
+
+static DWORD CALLBACK implicit_mta_marshal_proc(void *param)
+{
+    struct implicit_mta_marshal_data *data = param;
+    HRESULT hr;
+
+    hr = CoMarshalInterface(data->stream, &IID_IClassFactory,
+        (IUnknown *)&Test_ClassFactory, MSHCTX_INPROC, NULL, MSHLFLAGS_NORMAL);
+    ok_ole_success(hr, CoMarshalInterface);
+
+    SetEvent(data->start);
+
+    ok(!WaitForSingleObject(data->stop, 1000), "wait failed\n");
+    return 0;
+}
+
 static void test_implicit_mta(void)
 {
+    struct implicit_mta_marshal_data data;
     HANDLE host_thread, thread;
     IClassFactory *cf;
+    IUnknown *proxy;
     IStream *stream;
     HRESULT hr;
     DWORD tid;
@@ -3529,6 +3553,32 @@ static void test_implicit_mta(void)
 
     end_host_object(tid, host_thread);
 
+    /* Thirdly: we can marshal an object from the implicit MTA and then
+     * unmarshal it into the real one. */
+    data.start = CreateEventA(NULL, FALSE, FALSE, NULL);
+    data.stop  = CreateEventA(NULL, FALSE, FALSE, NULL);
+
+    hr = CreateStreamOnHGlobal(NULL, TRUE, &data.stream);
+    ok_ole_success(hr, CreateStreamOnHGlobal);
+
+    thread = CreateThread(NULL, 0, implicit_mta_marshal_proc, &data, 0, NULL);
+    ok(!WaitForSingleObject(data.start, 1000), "wait failed\n");
+
+    IStream_Seek(data.stream, ullZero, STREAM_SEEK_SET, NULL);
+    hr = CoUnmarshalInterface(data.stream, &IID_IClassFactory, (void **)&cf);
+    ok_ole_success(hr, CoUnmarshalInterface);
+
+    hr = IClassFactory_CreateInstance(cf, NULL, &IID_IUnknown, (void **)&proxy);
+    ok_ole_success(hr, IClassFactory_CreateInstance);
+
+    IUnknown_Release(proxy);
+
+    SetEvent(data.stop);
+    ok(!WaitForSingleObject(thread, 1000), "wait failed\n");
+    CloseHandle(thread);
+
+    IStream_Release(data.stream);
+
     CoUninitialize();
 }
 
-- 
2.7.4