From 6c965f4ab78e180b0c291b1013249010302216f0 Mon Sep 17 00:00:00 2001
From: Jonathan Grynspan <jgrynspan@apple.com>
Date: Mon, 6 Jan 2025 15:24:16 -0500
Subject: [PATCH] Use the Toolhelp32 API to enumerate loaded Win32 modules.

This PR replaces our call to `EnumProcessModules()` with one to
`CreateToolhelp32Snapshot(TH32CS_SNAPMODULE)`. Three reasons:

1. `EnumProcessModules()` requires us to specify a large, fixed-size buffer to
   contain all the `HMODULE` handles;
2. `EnumProcessModules()` does not own any references to the handles it returns,
   meaning that a module can be unloaded while we are iterating over them (while
   `CreateToolhelp32Snapshot()` temporarily bumps the refcounts of the handles
   it produces); and
3. `CreateToolhelp32Snapshot()` lets us produce a lazy sequence of `HMODULE`
   values rather than an array, letting us write somewhat Swiftier code that
   uses it.

The overhead of using `CreateToolhelp32Snapshot()` was negligible (below the
noise level when measuring).
---
 .../Support/Additions/WinSDKAdditions.swift   | 52 +++++++++++++++++++
 Sources/Testing/Support/GetSymbol.swift       | 23 ++------
 2 files changed, 56 insertions(+), 19 deletions(-)

diff --git a/Sources/Testing/Support/Additions/WinSDKAdditions.swift b/Sources/Testing/Support/Additions/WinSDKAdditions.swift
index 488d52dd6..9b902c5d1 100644
--- a/Sources/Testing/Support/Additions/WinSDKAdditions.swift
+++ b/Sources/Testing/Support/Additions/WinSDKAdditions.swift
@@ -50,4 +50,56 @@ let STATUS_SIGNAL_CAUGHT_BITS = {
 
   return result
 }()
+
+// MARK: - HMODULE members
+
+extension HMODULE {
+  /// A helper type that manages state for ``HMODULE/all``.
+  private final class _AllState {
+    /// The toolhelp snapshot.
+    var snapshot: HANDLE?
+
+    /// The module iterator.
+    var me = MODULEENTRY32W()
+
+    deinit {
+      if let snapshot {
+        CloseHandle(snapshot)
+      }
+    }
+  }
+
+  /// All modules loaded in the current process.
+  ///
+  /// - Warning: It is possible for one or more modules in this sequence to be
+  ///   unloaded while you are iterating over it. To minimize the risk, do not
+  ///   discard the sequence until iteration is complete. Modules containing
+  ///   Swift code can never be safely unloaded.
+  static var all: some Sequence<Self> {
+    sequence(state: _AllState()) { state in
+      if let snapshot = state.snapshot {
+        // We have already iterated over the first module. Return the next one.
+        if Module32NextW(snapshot, &state.me) {
+          return state.me.hModule
+        }
+      } else {
+        // Create a toolhelp snapshot that lists modules.
+        guard let snapshot = CreateToolhelp32Snapshot(DWORD(TH32CS_SNAPMODULE), 0) else {
+          return nil
+        }
+        state.snapshot = snapshot
+
+        // Initialize the iterator for use by the resulting sequence and return
+        // the first module.
+        state.me.dwSize = DWORD(MemoryLayout.stride(ofValue: state.me))
+        if Module32FirstW(snapshot, &state.me) {
+          return state.me.hModule
+        }
+      }
+
+      // Reached the end of the iteration.
+      return nil
+    }
+  }
+}
 #endif
diff --git a/Sources/Testing/Support/GetSymbol.swift b/Sources/Testing/Support/GetSymbol.swift
index 264bc0daa..b0f057088 100644
--- a/Sources/Testing/Support/GetSymbol.swift
+++ b/Sources/Testing/Support/GetSymbol.swift
@@ -70,25 +70,10 @@ func symbol(in handle: ImageAddress? = nil, named symbolName: String) -> UnsafeR
       }
     }
 
-    // Find all the modules loaded in the current process. We assume there
-    // aren't more than 1024 loaded modules (as does Microsoft sample code.)
-    return withUnsafeTemporaryAllocation(of: HMODULE?.self, capacity: 1024) { hModules in
-      let byteCount = DWORD(hModules.count * MemoryLayout<HMODULE?>.stride)
-      var byteCountNeeded: DWORD = 0
-      guard K32EnumProcessModules(GetCurrentProcess(), hModules.baseAddress!, byteCount, &byteCountNeeded) else {
-        return nil
-      }
-
-      // Enumerate all modules looking for one containing the given symbol.
-      let hModuleCount = min(hModules.count, Int(byteCountNeeded) / MemoryLayout<HMODULE?>.stride)
-      let hModulesEnd = hModules.index(hModules.startIndex, offsetBy: hModuleCount)
-      for hModule in hModules[..<hModulesEnd] {
-        if let hModule, let result = GetProcAddress(hModule, symbolName) {
-          return unsafeBitCast(result, to: UnsafeRawPointer.self)
-        }
-      }
-      return nil
-    }
+    return HMODULE.all.lazy
+      .compactMap { GetProcAddress($0, symbolName) }
+      .map { unsafeBitCast($0, to: UnsafeRawPointer.self) }
+      .first
   }
 #else
 #warning("Platform-specific implementation missing: Dynamic loading unavailable")