Report allocation error on thread freezing

2 files changed, 58 insertions, 41 deletions

diff --git a/src/hook.c b/src/hook.c
index ce65e57..a197589 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -288,7 +288,11 @@ static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
                         LPDWORD p = (LPDWORD)HeapReAlloc(
                             g_hHeap, 0, pThreads->pItems, (pThreads->capacity * 2) * sizeof(DWORD));
                         if (p == NULL)
+                        {
+                            HeapFree(g_hHeap, 0, pThreads->pItems);
+                            pThreads->pItems = NULL;
                             break;
+                        }
 
                         pThreads->capacity *= 2;
                         pThreads->pItems = p;
@@ -304,13 +308,15 @@ static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
 }
 
 //-------------------------------------------------------------------------
-static VOID Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
+static MH_STATUS Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
 {
     pThreads->pItems   = NULL;
     pThreads->capacity = 0;
     pThreads->size     = 0;
     EnumerateThreads(pThreads);
 
+    MH_STATUS status = MH_OK;
+
     if (pThreads->pItems != NULL)
     {
         UINT i;
@@ -325,26 +331,29 @@ static VOID Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
             }
         }
     }
+    else
+    {
+        status = MH_ERROR_MEMORY_ALLOC;
+    }
+
+    return status;
 }
 
 //-------------------------------------------------------------------------
 static VOID Unfreeze(PFROZEN_THREADS pThreads)
 {
-    if (pThreads->pItems != NULL)
+    UINT i;
+    for (i = 0; i < pThreads->size; ++i)
     {
-        UINT i;
-        for (i = 0; i < pThreads->size; ++i)
+        HANDLE hThread = OpenThread(THREAD_ACCESS, FALSE, pThreads->pItems[i]);
+        if (hThread != NULL)
         {
-            HANDLE hThread = OpenThread(THREAD_ACCESS, FALSE, pThreads->pItems[i]);
-            if (hThread != NULL)
-            {
-                ResumeThread(hThread);
-                CloseHandle(hThread);
-            }
+            ResumeThread(hThread);
+            CloseHandle(hThread);
         }
-
-        HeapFree(g_hHeap, 0, pThreads->pItems);
     }
+
+    HeapFree(g_hHeap, 0, pThreads->pItems);
 }
 
 //-------------------------------------------------------------------------
@@ -414,19 +423,21 @@ static MH_STATUS EnableAllHooksLL(BOOL enable)
     if (first != INVALID_HOOK_POS)
     {
         FROZEN_THREADS threads;
-        Freeze(&threads, ALL_HOOKS_POS, enable ? ACTION_ENABLE : ACTION_DISABLE);
-
-        for (i = first; i < g_hooks.size; ++i)
+        status = Freeze(&threads, ALL_HOOKS_POS, enable ? ACTION_ENABLE : ACTION_DISABLE);
+        if (status == MH_OK)
         {
-            if (g_hooks.pItems[i].isEnabled != enable)
+            for (i = first; i < g_hooks.size; ++i)
             {
-                status = EnableHookLL(i, enable);
-                if (status != MH_OK)
-                    break;
+                if (g_hooks.pItems[i].isEnabled != enable)
+                {
+                    status = EnableHookLL(i, enable);
+                    if (status != MH_OK)
+                        break;
+                }
             }
-        }
 
-        Unfreeze(&threads);
+            Unfreeze(&threads);
+        }
     }
 
     return status;
@@ -644,11 +655,13 @@ MH_STATUS WINAPI MH_RemoveHook(LPVOID pTarget)
             if (g_hooks.pItems[pos].isEnabled)
             {
                 FROZEN_THREADS threads;
-                Freeze(&threads, pos, ACTION_DISABLE);
-
-                status = EnableHookLL(pos, FALSE);
+                status = Freeze(&threads, pos, ACTION_DISABLE);
+                if (status == MH_OK)
+                {
+                    status = EnableHookLL(pos, FALSE);
 
-                Unfreeze(&threads);
+                    Unfreeze(&threads);
+                }
             }
 
             if (status == MH_OK)
@@ -687,17 +700,19 @@ static MH_STATUS EnableHook(LPVOID pTarget, BOOL enable)
         }
         else
         {
-            FROZEN_THREADS threads;
             UINT pos = FindHookEntry(pTarget);
             if (pos != INVALID_HOOK_POS)
             {
                 if (g_hooks.pItems[pos].isEnabled != enable)
                 {
-                    Freeze(&threads, pos, ACTION_ENABLE);
-
-                    status = EnableHookLL(pos, enable);
+                    FROZEN_THREADS threads;
+                    status = Freeze(&threads, pos, ACTION_ENABLE);
+                    if (status == MH_OK)
+                    {
+                        status = EnableHookLL(pos, enable);
 
-                    Unfreeze(&threads);
+                        Unfreeze(&threads);
+                    }
                 }
                 else
                 {
@@ -804,20 +819,22 @@ MH_STATUS WINAPI MH_ApplyQueued(VOID)
         if (first != INVALID_HOOK_POS)
         {
             FROZEN_THREADS threads;
-            Freeze(&threads, ALL_HOOKS_POS, ACTION_APPLY_QUEUED);
-
-            for (i = first; i < g_hooks.size; ++i)
+            status = Freeze(&threads, ALL_HOOKS_POS, ACTION_APPLY_QUEUED);
+            if (status == MH_OK)
             {
-                PHOOK_ENTRY pHook = &g_hooks.pItems[i];
-                if (pHook->isEnabled != pHook->queueEnable)
+                for (i = first; i < g_hooks.size; ++i)
                 {
-                    status = EnableHookLL(i, pHook->queueEnable);
-                    if (status != MH_OK)
-                        break;
+                    PHOOK_ENTRY pHook = &g_hooks.pItems[i];
+                    if (pHook->isEnabled != pHook->queueEnable)
+                    {
+                        status = EnableHookLL(i, pHook->queueEnable);
+                        if (status != MH_OK)
+                            break;
+                    }
                 }
-            }
 
-            Unfreeze(&threads);
+                Unfreeze(&threads);
+            }
         }
     }
     else
diff --git a/src/trampoline.c b/src/trampoline.c
index c267088..a28c5f3 100644
--- a/src/trampoline.c
+++ b/src/trampoline.c
@@ -208,7 +208,7 @@ BOOL CreateTrampolineFunction(PTRAMPOLINE ct)
                 pCopySrc = &jmp;
                 copySize = sizeof(jmp);
 
-                // Exit the function If it is not in the branch
+                // Exit the function if it is not in the branch.
                 finished = (pOldInst >= jmpDest);
             }
         }