[PATCH v3 1/1] vccorlib140: Fix potential memory leak in allocation functions.

[Vibhav Pant]

From: Vibhav Pant

---
 dlls/vccorlib140/vccorlib.c | 39 ++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/dlls/vccorlib140/vccorlib.c b/dlls/vccorlib140/vccorlib.c
index 0071c1a4591..fbbc23d7573 100644
--- a/dlls/vccorlib140/vccorlib.c
+++ b/dlls/vccorlib140/vccorlib.c
@@ -85,26 +85,43 @@ HRESULT WINAPI GetIidsFn(unsigned int count, unsigned int *copied, const GUID *s
     return S_OK;
 }

+static BOOL try_Allocate(size_t size, void **addr)
+{
+    *addr = malloc(size);
+    return !!*addr;
+}
+
 void *__cdecl Allocate(size_t size)
 {
     void *addr;

TRACE("(%Iu)\n", size);

-    addr = malloc(size);
-    if (!addr)
+    if (!try_Allocate(size, &addr))
         __abi_WinRTraiseOutOfMemoryException();
     return addr;
 }

-void *__cdecl AllocateException(size_t size)
+static BOOL try_AllocateException(size_t size, void **addr)
 {
     struct exception_alloc *base;

+    if (!try_Allocate(offsetof(struct exception_alloc, data[size]), (void **)&base))
+        return FALSE;
+    *addr = &base->data;
+    return TRUE;
+}
+
+void *__cdecl AllocateException(size_t size)
+{
+    void *addr;
+
     TRACE("(%Iu)\n", size);

-    base = Allocate(offsetof(struct exception_alloc, data[size]));
-    return &base->data;
+    if (!try_AllocateException(size, &addr))
+        __abi_WinRTraiseOutOfMemoryException();
+
+    return addr;
 }

void __cdecl Free(void *addr)
@@ -204,7 +221,11 @@ void *__cdecl AllocateWithWeakRef(ptrdiff_t offset, size_t size)
     if (size > inline_max)
     {
         weakref = Allocate(sizeof(*weakref));
-        object = Allocate(size);
+        if (!try_Allocate(size, &object))
+        {
+            Free(weakref);
+            __abi_WinRTraiseOutOfMemoryException();
+        }
         weakref->is_inline = FALSE;
     }
     else /* Perform an inline allocation */
@@ -233,7 +254,11 @@ void *__cdecl AllocateExceptionWithWeakRef(ptrdiff_t offset, size_t size)

/* AllocateExceptionWithWeakRef does not store the control block inline, regardless of size. */
     weakref = Allocate(sizeof(*weakref));
-    excp = AllocateException(size);
+    if (!try_AllocateException(size, &excp))
+    {
+        Free(weakref);
+        __abi_WinRTraiseOutOfMemoryException();
+    }
     *(struct control_block **)((char *)excp + offset) = weakref;
     weakref->IWeakReference_iface.lpVtbl = &control_block_vtbl;
     weakref->object = excp;