[PATCH 2/3] ole32: Improve CoRegisterMallocSpy/CoRevokeMallocSpy state and parameter validation

[Nikolay Sivov]

Signed-off-by: Nikolay Sivov <nsivov at codeweavers.com>
---
 dlls/ole32/ifs.c           |  7 +++++--
 dlls/ole32/tests/compobj.c | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/dlls/ole32/ifs.c b/dlls/ole32/ifs.c
index 0ca9b71..7469fb2 100644
--- a/dlls/ole32/ifs.c
+++ b/dlls/ole32/ifs.c
@@ -464,8 +464,9 @@ HRESULT WINAPI CoRegisterMallocSpy(LPMALLOCSPY pMallocSpy)
 	IMallocSpy* pSpy;
         HRESULT hres = E_INVALIDARG;
 
-	TRACE("\n");
+	TRACE("%p\n", pMallocSpy);
 
+	if(!pMallocSpy) return E_INVALIDARG;
 	if(Malloc32.pSpy) return CO_E_OBJISREG;
 
         EnterCriticalSection(&IMalloc32_SpyCS);
@@ -504,7 +505,9 @@ HRESULT WINAPI CoRevokeMallocSpy(void)
 
         EnterCriticalSection(&IMalloc32_SpyCS);
 
-	if (Malloc32.SpyedAllocationsLeft) {
+	if (!Malloc32.pSpy)
+	    hres = CO_E_OBJNOTREG;
+	else if (Malloc32.SpyedAllocationsLeft) {
             TRACE("SpyReleasePending with %u allocations left\n", Malloc32.SpyedAllocationsLeft);
 	    Malloc32.SpyReleasePending = TRUE;
 	    hres = E_ACCESSDENIED;
diff --git a/dlls/ole32/tests/compobj.c b/dlls/ole32/tests/compobj.c
index 39beee4..b88029f 100644
--- a/dlls/ole32/tests/compobj.c
+++ b/dlls/ole32/tests/compobj.c
@@ -3051,9 +3051,21 @@ static void test_IMallocSpy(void)
     IMalloc *imalloc;
     HRESULT hr;
 
+    hr = CoRegisterMallocSpy(NULL);
+    ok(hr == E_INVALIDARG, "got 0x%08x\n", hr);
+
+    hr = CoRevokeMallocSpy();
+    ok(hr == CO_E_OBJNOTREG, "got 0x%08x\n", hr);
+
     hr = CoRegisterMallocSpy(&testspy);
     ok(hr == S_OK, "got 0x%08x\n", hr);
 
+    hr = CoRegisterMallocSpy(NULL);
+    ok(hr == E_INVALIDARG, "got 0x%08x\n", hr);
+
+    hr = CoRegisterMallocSpy(&testspy);
+    ok(hr == CO_E_OBJISREG, "got 0x%08x\n", hr);
+
     imalloc = NULL;
     hr = CoGetMalloc(MEMCTX_TASK, &imalloc);
     ok(hr == S_OK, "got 0x%08x\n", hr);
@@ -3065,6 +3077,9 @@ static void test_IMallocSpy(void)
 
     hr = CoRevokeMallocSpy();
     ok(hr == S_OK, "got 0x%08x\n", hr);
+
+    hr = CoRevokeMallocSpy();
+    ok(hr == CO_E_OBJNOTREG, "got 0x%08x\n", hr);
 }
 
 static void init_funcs(void)
-- 
2.7.0