rpcrt4: Handle memory allocation error when creating OLE stream instance.

dlls/rpcrt4/ndr_ole.c

@@ -95,14 +95,15 @@ static HRESULT WINAPI RpcStream_QueryInterface(LPSTREAM iface,
                                               REFIID riid,
                                               LPVOID *obj)
 {
-  RpcStreamImpl *This = impl_from_IStream(iface);
   if (IsEqualGUID(&IID_IUnknown, riid) ||
       IsEqualGUID(&IID_ISequentialStream, riid) ||
       IsEqualGUID(&IID_IStream, riid)) {
-    *obj = This;
-    InterlockedIncrement( &This->RefCount );
+    *obj = iface;
+    IStream_AddRef(iface);
     return S_OK;
   }
+
+  *obj = NULL;
   return E_NOINTERFACE;
 }
 
@@ -120,7 +121,6 @@ static ULONG WINAPI RpcStream_Release(LPSTREAM iface)
     TRACE("size=%d\n", *This->size);
     This->pMsg->Buffer = This->data + *This->size;
     HeapFree(GetProcessHeap(),0,This);
-    return 0;
   }
   return ref;
 }
@@ -212,11 +212,13 @@ static const IStreamVtbl RpcStream_Vtbl =
   NULL  /* Clone */
 };
 
-static LPSTREAM RpcStream_Create(PMIDL_STUB_MESSAGE pStubMsg, BOOL init)
+static HRESULT RpcStream_Create(PMIDL_STUB_MESSAGE pStubMsg, BOOL init, IStream **stream)
 {
   RpcStreamImpl *This;
+
+  *stream = NULL;
   This = HeapAlloc(GetProcessHeap(),HEAP_ZERO_MEMORY,sizeof(RpcStreamImpl));
-  if (!This) return NULL;
+  if (!This) return E_OUTOFMEMORY;
   This->IStream_iface.lpVtbl = &RpcStream_Vtbl;
   This->RefCount = 1;
   This->pMsg = pStubMsg;
@@ -225,7 +227,8 @@ static LPSTREAM RpcStream_Create(PMIDL_STUB_MESSAGE pStubMsg, BOOL init)
   This->pos = 0;
   if (init) *This->size = 0;
   TRACE("init size=%d\n", *This->size);
-  return (LPSTREAM)This;
+  *stream = &This->IStream_iface;
+  return S_OK;
 }
 
 static const IID* get_ip_iid(PMIDL_STUB_MESSAGE pStubMsg, unsigned char *pMemory, PFORMAT_STRING pFormat)
@@ -260,19 +263,17 @@ unsigned char * WINAPI NdrInterfacePointerMarshall(PMIDL_STUB_MESSAGE pStubMsg,
   pStubMsg->MaxCount = 0;
   if (!LoadCOM()) return NULL;
   if (pStubMsg->Buffer + sizeof(DWORD) <= (unsigned char *)pStubMsg->RpcMsg->Buffer + pStubMsg->BufferLength) {
-    stream = RpcStream_Create(pStubMsg, TRUE);
-    if (stream) {
+    hr = RpcStream_Create(pStubMsg, TRUE, &stream);
+    if (hr == S_OK) {
       if (pMemory)
         hr = COM_MarshalInterface(stream, riid, (LPUNKNOWN)pMemory,
                                   pStubMsg->dwDestContext, pStubMsg->pvDestContext,
                                   MSHLFLAGS_NORMAL);
-      else
-        hr = S_OK;
-
       IStream_Release(stream);
-      if (FAILED(hr))
-        RpcRaiseException(hr);
     }
+
+    if (FAILED(hr))
+      RpcRaiseException(hr);
   }
   return NULL;
 }
@@ -292,13 +293,14 @@ unsigned char * WINAPI NdrInterfacePointerUnmarshall(PMIDL_STUB_MESSAGE pStubMsg
   if (!LoadCOM()) return NULL;
   *(LPVOID*)ppMemory = NULL;
   if (pStubMsg->Buffer + sizeof(DWORD) < (unsigned char *)pStubMsg->RpcMsg->Buffer + pStubMsg->BufferLength) {
-    stream = RpcStream_Create(pStubMsg, FALSE);
-    if (!stream) RpcRaiseException(E_OUTOFMEMORY);
-    if (*((RpcStreamImpl *)stream)->size != 0)
-      hr = COM_UnmarshalInterface(stream, &IID_NULL, (LPVOID*)ppMemory);
-    else
-      hr = S_OK;
-    IStream_Release(stream);
+    hr = RpcStream_Create(pStubMsg, FALSE, &stream);
+    if (hr == S_OK) {
+      if (*((RpcStreamImpl *)stream)->size != 0)
+        hr = COM_UnmarshalInterface(stream, &IID_NULL, (LPVOID*)ppMemory);
+
+      IStream_Release(stream);
+    }
+
     if (FAILED(hr))
         RpcRaiseException(hr);
   }