Nikolay Sivov : ole32: Fix IMalloc::DidAlloc() return value to reflect block validity.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 08f4b6ee0a05fe27e56cb9777dce845dcf1072f8
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=08f4b6ee0a05fe27e56cb9777dce845dcf1072f8

Author: Nikolay Sivov <nsivov at codeweavers.com>
Date:   Thu Apr 16 22:01:48 2020 +0300

ole32: Fix IMalloc::DidAlloc() return value to reflect block validity.

Signed-off-by: Nikolay Sivov <nsivov at codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard at winehq.org>

---

 dlls/ole32/ifs.c           | 59 +++++++++++++++++++++++++++++-----------------
 dlls/ole32/tests/compobj.c | 21 +++++++++++++++++
 2 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/dlls/ole32/ifs.c b/dlls/ole32/ifs.c
index 9c7cb0977f..e747cfc0fe 100644
--- a/dlls/ole32/ifs.c
+++ b/dlls/ole32/ifs.c
@@ -114,6 +114,20 @@ static BOOL AddMemoryLocation(LPVOID * pMem)
         return TRUE;
 }
 
+static void** mallocspy_is_allocation_spyed(const void *mem)
+{
+    void **current = Malloc32.SpyedBlocks;
+
+    while (*current != mem)
+    {
+        current++;
+        if (current >= Malloc32.SpyedBlocks + Malloc32.SpyedBlockTableLength)
+            return NULL;
+    }
+
+    return current;
+}
+
 static BOOL RemoveMemoryLocation(LPCVOID pMem)
 {
         LPVOID * Current;
@@ -122,14 +136,8 @@ static BOOL RemoveMemoryLocation(LPCVOID pMem)
         if (!Malloc32.SpyedBlockTableLength && !SetSpyedBlockTableLength(0x1000))
             return FALSE;
 
-	Current = Malloc32.SpyedBlocks;
-
-	/* find the location */
-	while (*Current != pMem) {
-            Current++;
-            if (Current >= Malloc32.SpyedBlocks + Malloc32.SpyedBlockTableLength)
-                return FALSE; /* not found  */
-	}
+        if (!(Current = mallocspy_is_allocation_spyed(pMem)))
+            return FALSE;
 
 	/* location found */
         Malloc32.SpyedAllocationsLeft--;
@@ -313,25 +321,32 @@ static SIZE_T WINAPI IMalloc_fnGetSize(IMalloc *iface, void *pv)
 /******************************************************************************
  * IMalloc32_DidAlloc [VTABLE]
  */
-static INT WINAPI IMalloc_fnDidAlloc(IMalloc *iface, void *pv)
+static INT WINAPI IMalloc_fnDidAlloc(IMalloc *iface, void *mem)
 {
-        BOOL fSpyed = FALSE;
-	int didAlloc;
+    BOOL spyed_block = FALSE;
+    int did_alloc;
 
-	TRACE("(%p)\n",pv);
+    TRACE("(%p)\n", mem);
 
-	if(Malloc32.pSpy) {
-            EnterCriticalSection(&IMalloc32_SpyCS);
-	    pv = IMallocSpy_PreDidAlloc(Malloc32.pSpy, pv, fSpyed);
-	}
+    if (!mem)
+        return -1;
+
+    if (Malloc32.pSpy)
+    {
+        EnterCriticalSection(&IMalloc32_SpyCS);
+        spyed_block = !!mallocspy_is_allocation_spyed(mem);
+        mem = IMallocSpy_PreDidAlloc(Malloc32.pSpy, mem, spyed_block);
+    }
 
-	didAlloc = -1;
+    did_alloc = HeapValidate(GetProcessHeap(), 0, mem);
 
-	if(Malloc32.pSpy) {
-	    didAlloc = IMallocSpy_PostDidAlloc(Malloc32.pSpy, pv, fSpyed, didAlloc);
-            LeaveCriticalSection(&IMalloc32_SpyCS);
-	}
-	return didAlloc;
+    if (Malloc32.pSpy)
+    {
+        did_alloc = IMallocSpy_PostDidAlloc(Malloc32.pSpy, mem, spyed_block, did_alloc);
+        LeaveCriticalSection(&IMalloc32_SpyCS);
+    }
+
+    return did_alloc;
 }
 
 /******************************************************************************
diff --git a/dlls/ole32/tests/compobj.c b/dlls/ole32/tests/compobj.c
index cc2737bdd1..b0a6ba1ea7 100644
--- a/dlls/ole32/tests/compobj.c
+++ b/dlls/ole32/tests/compobj.c
@@ -3069,6 +3069,8 @@ static void test_CoGetMalloc(void)
 {
     IMalloc *imalloc;
     HRESULT hr;
+    char *ptr;
+    int ret;
 
     if (0) /* crashes on native */
         hr = CoGetMalloc(0, NULL);
@@ -3102,6 +3104,25 @@ static void test_CoGetMalloc(void)
     hr = CoGetMalloc(MEMCTX_TASK, &imalloc);
     ok(hr == S_OK, "got 0x%08x\n", hr);
     ok(imalloc != NULL, "got %p\n", imalloc);
+
+    /* DidAlloc() */
+    ptr = IMalloc_Alloc(imalloc, 16);
+    ok(!!ptr, "Failed to allocate block.\n");
+
+    ret = IMalloc_DidAlloc(imalloc, ptr);
+    ok(ret == 1, "Unexpected return value %d.\n", ret);
+
+    ret = IMalloc_DidAlloc(imalloc, NULL);
+    ok(ret == -1, "Unexpected return value %d.\n", ret);
+
+    ret = IMalloc_DidAlloc(imalloc, (void *)0x1);
+    ok(ret == 0, "Unexpected return value %d.\n", ret);
+
+    ret = IMalloc_DidAlloc(imalloc, ptr + 4);
+    ok(ret == 0, "Unexpected return value %d.\n", ret);
+
+    IMalloc_Free(imalloc, ptr);
+
     IMalloc_Release(imalloc);
 }