Fix thread freezing when there is only one thread

1 files changed, 39 insertions, 22 deletions

diff --git a/src/hook.c b/src/hook.c
index aba9e35..2eb33c0 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -260,8 +260,10 @@ static VOID ProcessThreadIPs(HANDLE hThread, UINT pos, UINT action)
 }
 
 //-------------------------------------------------------------------------
-static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
+static BOOL EnumerateThreads(PFROZEN_THREADS pThreads)
 {
+    BOOL succeeded = FALSE;
+
     HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
     if (hSnapshot != INVALID_HANDLE_VALUE)
     {
@@ -269,6 +271,7 @@ static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
         te.dwSize = sizeof(THREADENTRY32);
         if (Thread32First(hSnapshot, &te))
         {
+            succeeded = TRUE;
             do
             {
                 if (te.dwSize >= (FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(DWORD))
@@ -281,20 +284,22 @@ static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
                         pThreads->pItems
                             = (LPDWORD)HeapAlloc(g_hHeap, 0, pThreads->capacity * sizeof(DWORD));
                         if (pThreads->pItems == NULL)
+                        {
+                            succeeded = FALSE;
                             break;
+                        }
                     }
                     else if (pThreads->size >= pThreads->capacity)
                     {
+                        pThreads->capacity *= 2;
                         LPDWORD p = (LPDWORD)HeapReAlloc(
-                            g_hHeap, 0, pThreads->pItems, (pThreads->capacity * 2) * sizeof(DWORD));
+                            g_hHeap, 0, pThreads->pItems, pThreads->capacity * sizeof(DWORD));
                         if (p == NULL)
                         {
-                            HeapFree(g_hHeap, 0, pThreads->pItems);
-                            pThreads->pItems = NULL;
+                            succeeded = FALSE;
                             break;
                         }
 
-                        pThreads->capacity *= 2;
                         pThreads->pItems = p;
                     }
                     pThreads->pItems[pThreads->size++] = te.th32ThreadID;
@@ -302,22 +307,35 @@ static VOID EnumerateThreads(PFROZEN_THREADS pThreads)
 
                 te.dwSize = sizeof(THREADENTRY32);
             } while (Thread32Next(hSnapshot, &te));
+
+            if (succeeded && GetLastError() != ERROR_NO_MORE_FILES)
+                succeeded = FALSE;
+
+            if (!succeeded && pThreads->pItems != NULL)
+            {
+                HeapFree(g_hHeap, 0, pThreads->pItems);
+                pThreads->pItems = NULL;
+            }
         }
         CloseHandle(hSnapshot);
     }
+
+    return succeeded;
 }
 
 //-------------------------------------------------------------------------
 static MH_STATUS Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
 {
+    MH_STATUS status = MH_OK;
+
     pThreads->pItems   = NULL;
     pThreads->capacity = 0;
     pThreads->size     = 0;
-    EnumerateThreads(pThreads);
-
-    MH_STATUS status = MH_OK;
-
-    if (pThreads->pItems != NULL)
+    if (!EnumerateThreads(pThreads))
+    {
+        status = MH_ERROR_MEMORY_ALLOC;
+    }
+    else if (pThreads->pItems != NULL)
     {
         UINT i;
         for (i = 0; i < pThreads->size; ++i)
@@ -331,10 +349,6 @@ static MH_STATUS Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
             }
         }
     }
-    else
-    {
-        status = MH_ERROR_MEMORY_ALLOC;
-    }
 
     return status;
 }
@@ -342,18 +356,21 @@ static MH_STATUS Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
 //-------------------------------------------------------------------------
 static VOID Unfreeze(PFROZEN_THREADS pThreads)
 {
-    UINT i;
-    for (i = 0; i < pThreads->size; ++i)
+    if (pThreads->pItems != NULL)
     {
-        HANDLE hThread = OpenThread(THREAD_ACCESS, FALSE, pThreads->pItems[i]);
-        if (hThread != NULL)
+        UINT i;
+        for (i = 0; i < pThreads->size; ++i)
         {
-            ResumeThread(hThread);
-            CloseHandle(hThread);
+            HANDLE hThread = OpenThread(THREAD_ACCESS, FALSE, pThreads->pItems[i]);
+            if (hThread != NULL)
+            {
+                ResumeThread(hThread);
+                CloseHandle(hThread);
+            }
         }
-    }
 
-    HeapFree(g_hHeap, 0, pThreads->pItems);
+        HeapFree(g_hHeap, 0, pThreads->pItems);
+    }
 }
 
 //-------------------------------------------------------------------------