Finish implementation

microsoft · Feb 20, 2024 · 6ccd95b · 6ccd95b
1 parent c982bbe
commit 6ccd95b
Show file tree

Hide file tree

Showing 3 changed files with 138 additions and 41 deletions.
diff --git a/CMakeSettings.json b/CMakeSettings.json
diff --git a/include/wil/toolhelp32.h b/include/wil/toolhelp32.h
@@ -1,36 +1,133 @@
+#include "resource.h"
 #ifndef __WIL_TOOLHELP32_INCLUDED
 #define __WIL_TOOLHELP32_INCLUDED
 #include <TlHelp32.h>
+#include <processthreadsapi.h>
 namespace wil
 {
-	namespace details
-	{
+namespace details
+{
+    template <typename TEntry, typename TEnumApi, typename TCallback>
+    void do_enum_snapshot(HANDLE handle, TEntry& entry, TEnumApi&& enumApiFirst, TEnumApi&& enumApiNext, TCallback&& callback)
+    {
+        using result_t = decltype(callback(TEntry{}));
+        bool enumResult = enumApiFirst(handle, &entry);
+        if (!enumResult)
+            return;
 
-	}
+        do
+        {
+            if constexpr (wistd::is_void_v<result_t>)
+            {
+                callback(entry);
+            }
+            else if constexpr (wistd::is_same_v<result_t, bool>)
+            {
+                if (callback(entry))
+                    return;
+            }
+            else
+            {
+                static_assert(
+                    [] {
+                        return false;
+                    }(),
+                    "Callback must return void or bool");
+            }
+            enumResult = enumApiNext(handle, &entry);
+        } while (enumResult);
+    }
+} // namespace details
 
-	template<typename TCallback>
-    void for_each_process(TCallback&& /*callback*/)
-	{
+template <typename TCallback>
+void for_each_process(TCallback&& callback)
+{
+    PROCESSENTRY32 entry{};
+    entry.dwSize = sizeof(entry);
+    details::do_enum_snapshot(
+        unique_handle{CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)}.get(),
+        entry,
+        &Process32First,
+        &Process32Next,
+        wistd::forward<TCallback>(callback));
+}
 
-	}
+template <typename TCallback>
+void for_each_thread(TCallback&& callback)
+{
+    THREADENTRY32 entry{};
+    entry.dwSize = sizeof(entry);
+    details::do_enum_snapshot(
+        unique_handle{CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)}.get(),
+        entry,
+        &Thread32First,
+        &Thread32Next,
+        wistd::forward<TCallback>(callback));
+}
 
-	template<typename TCallback>
-    void for_each_thread(TCallback&& /*callback*/)
-	{
+template <typename TCallback>
+void for_each_module(TCallback&& callback, bool include32For64Bit = false)
+{
+    MODULEENTRY32 entry{};
+    entry.dwSize = sizeof(entry);
+    details::do_enum_snapshot(
+        unique_handle{CreateToolhelp32Snapshot(include32For64Bit ? TH32CS_SNAPMODULE32 : TH32CS_SNAPMODULE, 0)}.get(),
+        entry,
+        &Module32First,
+        &Module32Next,
+        wistd::forward<TCallback>(callback));
+}
 
-	}
+template <typename TCallback>
+void for_each_heap_list(TCallback&& callback)
+{
+    HEAPLIST32 entry{};
+    entry.dwSize = sizeof(entry);
+    details::do_enum_snapshot(
+        unique_handle{CreateToolhelp32Snapshot(TH32CS_SNAPHEAPLIST, 0)}.get(),
+        entry,
+        &Heap32ListFirst,
+        &Heap32ListNext,
+        wistd::forward<TCallback>(callback));
+}
 
-	template<typename TCallback>
-    void for_each_module(TCallback&& /*callback*/)
-	{
+template <typename TCallback>
+void for_each_heap(TCallback&& callback, ULONG_PTR heapId, DWORD pid = GetCurrentProcessId())
+{
+    using result_t = decltype(callback(HEAPENTRY32{}));
 
-	}
+    HEAPENTRY32 entry{};
+    entry.dwSize = sizeof(entry);
 
-	template<typename TCallback>
-    void for_each_heap(TCallback&& /*callback*/)
-	{
+    bool enumResult = Heap32First(&entry, pid, heapId);
+    do
+    {
+        if constexpr (wistd::is_void_v<result_t>)
+        {
+            callback(entry);
+        }
+        else if constexpr (wistd::is_same_v<result_t, bool>)
+        {
+            if (callback(entry))
+                return;
+        }
+        else
+        {
+            static_assert(
+                [] {
+                    return false;
+                }(),
+                "Callback must return void or bool");
+        }
+        enumResult = Heap32Next(&entry);
+    } while (enumResult);
+}
 
-	}
+template <typename TCallback>
+void for_each_heap(TCallback&& callback, HEAPLIST32 const& heapList, DWORD pid = GetCurrentProcessId())
+{
+    for_each_heap(wistd::forward<TCallback>(callback), heapList.th32HeapID, pid);
 }
+} // namespace wil
 
 #endif
diff --git a/tests/Toolhelp32Tests.cpp b/tests/Toolhelp32Tests.cpp
@@ -5,28 +5,44 @@
 
 TEST_CASE("Toolhelp32", "[EnumProcesses]")
 {
-    wil::for_each_process([](PROCESSENTRY32 entry) {
+    wil::for_each_process([](PROCESSENTRY32 const& entry) {
         REQUIRE_FALSE(std::strlen(entry.szExeFile) == 0);
     });
 }
 
 TEST_CASE("Toolhelp32", "[EnumModules]")
 {
-    wil::for_each_module([](MODULEENTRY32 entry) {
+    wil::for_each_module([](MODULEENTRY32 const& entry) {
         REQUIRE_FALSE(std::strlen(entry.szExePath) == 0);
     });
 }
 
 TEST_CASE("Toolhelp32", "[EnumThreads]")
 {
-    wil::for_each_thread([](THREADENTRY32 entry) {
-        REQUIRE_FALSE(entry.th32ThreadID == 0);
+    wil::for_each_thread([pid = GetCurrentProcessId()](THREADENTRY32 const& entry) {
+        if (entry.th32OwnerProcessID == pid)
+        {
+            REQUIRE_FALSE(entry.th32ThreadID == 0);
+        }
     });
 }
 
-TEST_CASE("Toolhelp32", "[EnumHeaps]")
+TEST_CASE("Toolhelp32", "[EnumHeapLists]")
 {
-    wil::for_each_heap([](HEAPLIST32 entry) {
+    wil::for_each_heap_list([](HEAPLIST32 const& entry) {
         REQUIRE_FALSE(entry.th32HeapID == 0);
     });
+}
+
+TEST_CASE("Toolhelp32", "[EnumHeap]")
+{
+    wil::for_each_heap_list([](HEAPLIST32 const& heapListEntry) {
+        REQUIRE_FALSE(heapListEntry.th32HeapID == 0);
+        wil::for_each_heap(
+            [](HEAPENTRY32 const& heapEntry) {
+                REQUIRE_FALSE(heapEntry.dwAddress == 0);
+            },
+            heapListEntry);
+        return false;
+    });
 }