Fix thread freezing when there is only one thread

src/hook.c

@@ -260,15 +260,18 @@ 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)
  {
  THREADENTRY32 te;
  te.dwSize = sizeof(THREADENTRY32);
  if (Thread32First(hSnapshot, &te))
  {
+ succeeded = TRUE;
  do
  {
  if (te.dwSize >= (FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(DWORD))
@@ -281,43 +284,58 @@ 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;
  }
 
  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,29 +349,28 @@ static MH_STATUS Freeze(PFROZEN_THREADS pThreads, UINT pos, UINT action)
  }
  }
  }
- else
- {
- status = MH_ERROR_MEMORY_ALLOC;
- }
 
  return status;
 }
 
 //-------------------------------------------------------------------------
 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);
+ }
 }
 
 //-------------------------------------------------------------------------