Kevin Puetz : rpcrt4: Add a refcount owned by MIDL_STUB_MESSAGE.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: a7be4256f6081684e8fb0258564f9358eb36cfcf
URL:    https://gitlab.winehq.org/wine/wine/-/commit/a7be4256f6081684e8fb0258564f935...

Author: Kevin Puetz 
Date:   Mon Oct 10 23:31:08 2022 +0000

rpcrt4: Add a refcount owned by MIDL_STUB_MESSAGE.

In cases where this could have been use-after-free, exceptions were
caught/hidden by RpcTryFinally, but still lead to leaks since
NdrProxyFreeBuffer wasn't able to call IRPCChannelBuffer::FreeBuffer.

StdProxy_GetChannel() now AddRef() on its return value (used to set
__proxy_frame::_StubMsg::pRpcChannelBuffer), and NdrProxyFreeBuffer()
calls the corresponding Release() and clears the now-weak pointer.
This makes the behavior of these function match the observed test
results, and fixes the crash/leak when a proxy is released mid-Invoke.

---

dlls/rpcrt4/cproxy.c      | 5 +++++
 dlls/rpcrt4/tests/cstub.c | 8 ++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/dlls/rpcrt4/cproxy.c b/dlls/rpcrt4/cproxy.c
index da3e19e50f3..426b1425d3a 100644
--- a/dlls/rpcrt4/cproxy.c
+++ b/dlls/rpcrt4/cproxy.c
@@ -454,6 +454,9 @@ static void StdProxy_GetChannel(LPVOID iface,
   StdProxyImpl *This = impl_from_proxy_obj( iface );
   TRACE("(%p)->GetChannel(%p) %s\n",This,ppChannel,This->name);

+  if(This->pChannel)
+    IRpcChannelBuffer_AddRef(This->pChannel);
+
   *ppChannel = This->pChannel;
 }

@@ -585,6 +588,8 @@ void WINAPI NdrProxyFreeBuffer(void *This,
                                  (RPCOLEMESSAGE*)pStubMsg->RpcMsg);
     pStubMsg->fBufferValid = TRUE;
   }
+  IRpcChannelBuffer_Release(pStubMsg->pRpcChannelBuffer);
+  pStubMsg->pRpcChannelBuffer = NULL;
 }

/***********************************************************************
diff --git a/dlls/rpcrt4/tests/cstub.c b/dlls/rpcrt4/tests/cstub.c
index 57e69389ebd..d44ee36088c 100644
--- a/dlls/rpcrt4/tests/cstub.c
+++ b/dlls/rpcrt4/tests/cstub.c
@@ -1527,7 +1527,7 @@ static void test_ChannelBufferRefCount(IPSFactoryBuffer *ppsf)

NdrProxyInitialize(proxy_if1, &rpcMessage, &stubMessage, &stubDesc, 0);
     /* stubMessage should add its own refcount on test_chanbuf */
-    todo_wine ok(test_chanbuf.RefCount == 2, "got %ld\n", test_chanbuf.RefCount);
+    ok(test_chanbuf.RefCount == 2, "got %ld\n", test_chanbuf.RefCount);
     ok(stubMessage.pRpcChannelBuffer != NULL, "NULL pRocChannelBuffer\n");

/* stubMessage doesn't add its own refcounts on proxy_if1 or proxy_buffer,
@@ -1536,17 +1536,17 @@ static void test_ChannelBufferRefCount(IPSFactoryBuffer *ppsf)
      * this unadvise could be reentrant to Invoke because SendReceive pumps STA messages.
      * The source would then erase that conection point entry and Release the proxy. */
     IRpcProxyBuffer_Disconnect(proxy_buffer);
-    todo_wine ok(test_chanbuf.RefCount == 1, "got %ld\n", test_chanbuf.RefCount);
+    ok(test_chanbuf.RefCount == 1, "got %ld\n", test_chanbuf.RefCount);
     IRpcProxyBuffer_Release(proxy_buffer);
     refs = IUnknown_Release(proxy_if1);
     ok(refs == 0, "got %ld\n", refs);
-    todo_wine ok(test_chanbuf.RefCount == 1, "got %ld\n", test_chanbuf.RefCount);
+    ok(test_chanbuf.RefCount == 1, "got %ld\n", test_chanbuf.RefCount);

/* NdrProxyFreeBuffer must not dereference the the now-freed proxy_if1,
      * yet should still free the remaining reference on test_chanbuf */
     NdrProxyFreeBuffer(proxy_if1, &stubMessage);
     ok(test_chanbuf.RefCount == 0, "got %ld\n", test_chanbuf.RefCount);
-    todo_wine ok(!stubMessage.pRpcChannelBuffer, "dangling pRpcChannelBuffer = %p\n", stubMessage.pRpcChannelBuffer);
+    ok(!stubMessage.pRpcChannelBuffer, "dangling pRpcChannelBuffer = %p\n", stubMessage.pRpcChannelBuffer);
 }

START_TEST( cstub )