swift-inspect: ensure that we eject any injected code

We would previously fail to eject the injected code on a failure.  This
would prevent a future introspection into the process as well as leave
the file open with an incremented retain count in the kernel space which
would prevent the file from being deleted.

In the future, when the application is able to treat the injected code
as a resource, this resource would be temporarily extracted, but would
no longer be possible to delete until a reboot (with a registration of
the deletion) due to the retained code.

Take the opportunity to rename some functions to take advantage of
labelled parameters and trailing function syntax.  This makes the code a
small amount easier to read.

Author: compnerd

tools/swift-inspect/Sources/swift-inspect/WindowsRemoteProcess.swift

@@ -161,16 +161,13 @@ internal final class WindowsRemoteProcess: RemoteProcess {
     self.context = context
 
     // Locate swiftCore.dll in the target process and load modules.
-    iterateRemoteModules(
-      dwProcessId: processId,
-      closure: { (entry, module) in
-        // FIXME(compnerd) support static linking at some point
-        if module == "swiftCore.dll" {
-          self.hSwiftCore = entry.hModule
-        }
-        _ = swift_reflection_addImage(
-          context, unsafeBitCast(entry.modBaseAddr, to: swift_addr_t.self))
-      })
+    modules(of: processId) { (entry, module) in
+      // FIXME(compnerd) support static linking at some point
+      if module == "swiftCore.dll" { self.hSwiftCore = entry.hModule }
+      _ = swift_reflection_addImage(context,
+                                    unsafeBitCast(entry.modBaseAddr,
+                                                  to: swift_addr_t.self))
+    }
     if self.hSwiftCore == HMODULE(bitPattern: -1) {
       // FIXME(compnerd) log error
       return nil
@@ -321,7 +318,7 @@ internal final class WindowsRemoteProcess: RemoteProcess {
       print("Failed to find remote LoadLibraryW/FreeLibrary addresses")
       return
     }
-    let (loadLibraryAddr, freeLibraryAddr) = (remoteAddrs[0], remoteAddrs[1])
+    let (loadLibraryAddr, pfnFreeLibrary) = (remoteAddrs[0], remoteAddrs[1])
     let hThread: HANDLE = CreateRemoteThread(
       self.process, nil, 0, loadLibraryAddr,
       dllPathRemote, 0, nil)
@@ -331,6 +328,18 @@ internal final class WindowsRemoteProcess: RemoteProcess {
     }
     defer { CloseHandle(hThread) }
 
+    defer {
+      // Always perform the code ejection process even if the heap walk fails.
+      // The module cannot re-execute the heap walk and will leave a retain
+      // count behind, preventing the module from being unlinked on the file
+      // system as well as leave code in the inspected process.  This will
+      // eventually be an issue for treating the injected code as a resource
+      // which is extracted temporarily.
+      if !eject(module: dllPathRemote, from: dwProcessId, pfnFreeLibrary) {
+        print("Failed to unload the remote dll")
+      }
+    }
+
     // The main heap iteration loop.
     outer: while true {
       let wait = WaitForSingleObject(hReadEvent, waitTimeoutMs)
@@ -381,14 +390,6 @@ internal final class WindowsRemoteProcess: RemoteProcess {
       print("LoadLibraryW failed \(threadExitCode)")
       return
     }
-
-    // Unload the dll and deallocate the dll path from the remote process
-    if !unloadDllAndPathRemote(
-      dwProcessId: dwProcessId, dllPathRemote: dllPathRemote, freeLibraryAddr: freeLibraryAddr)
-    {
-      print("Failed to unload the remote dll")
-      return
-    }
   }
 
   private func allocateDllPathRemote() -> UnsafeMutableRawPointer? {
@@ -430,26 +431,27 @@ internal final class WindowsRemoteProcess: RemoteProcess {
       }
   }
 
-  private func unloadDllAndPathRemote(
-    dwProcessId: DWORD, dllPathRemote: UnsafeMutableRawPointer,
-    freeLibraryAddr: LPTHREAD_START_ROUTINE
-  ) -> Bool {
+  /// Eject the injected code from the instrumented process.
+  ///
+  /// Performs the necessary clean up to remove the injected code from the
+  /// instrumented process once the heap walk is complete.
+  private func eject(module dllPathRemote: UnsafeMutableRawPointer,
+                     from dwProcessId: DWORD,
+                     _ freeLibraryAddr: LPTHREAD_START_ROUTINE) -> Bool {
     // Get the dll module handle in the remote process to use it to
     // unload it below.
+
     // GetExitCodeThread returns a DWORD (32-bit) but the HMODULE
     // returned from LoadLibraryW is a 64-bit pointer and may be truncated.
     // So, search for it using the snapshot instead.
-    guard
-      let hDllModule = findRemoteModule(
-        dwProcessId: dwProcessId, moduleName: "SwiftInspectClient.dll")
-    else {
+    guard let hModule = find(module: "SwiftInspectClient.dll", in: dwProcessId) else {
       print("Failed to find the client dll")
       return false
     }
     // Unload the dll from the remote process
     let hUnloadThread = CreateRemoteThread(
       self.process, nil, 0, freeLibraryAddr,
-      UnsafeMutableRawPointer(hDllModule), 0, nil)
+      UnsafeMutableRawPointer(hModule), 0, nil)
     if hUnloadThread == HANDLE(bitPattern: 0) {
       print("CreateRemoteThread for unload failed \(GetLastError())")
       return false
@@ -476,7 +478,7 @@ internal final class WindowsRemoteProcess: RemoteProcess {
     return true
   }
 
-  private func iterateRemoteModules(dwProcessId: DWORD, closure: (MODULEENTRY32W, String) -> Void) {
+  private func modules(of dwProcessId: DWORD, _ closure: (MODULEENTRY32W, String) -> Void) {
     let hModuleSnapshot: HANDLE =
       CreateToolhelp32Snapshot(DWORD(TH32CS_SNAPMODULE), dwProcessId)
     if hModuleSnapshot == INVALID_HANDLE_VALUE {
@@ -503,22 +505,18 @@ internal final class WindowsRemoteProcess: RemoteProcess {
     } while Module32NextW(hModuleSnapshot, &entry)
   }
 
-  private func findRemoteModule(dwProcessId: DWORD, moduleName: String) -> HMODULE? {
-    var hDllModule: HMODULE? = nil
-    iterateRemoteModules(
-      dwProcessId: dwProcessId,
-      closure: { (entry, module) in
-        if module == moduleName {
-          hDllModule = entry.hModule
-        }
-      })
-    return hDllModule
+  private func find(module named: String, in dwProcessId: DWORD) -> HMODULE? {
+    var hModule: HMODULE?
+    modules(of: dwProcessId) { (entry, module) in
+      if module == named { hModule = entry.hModule }
+    }
+    return hModule
   }
 
   private func findRemoteAddresses(dwProcessId: DWORD, moduleName: String, symbols: [String])
     -> [LPTHREAD_START_ROUTINE]?
   {
-    guard let hDllModule = findRemoteModule(dwProcessId: dwProcessId, moduleName: moduleName) else {
+    guard let hDllModule = find(module: moduleName, in: dwProcessId) else {
       print("Failed to find remote module \(moduleName)")
       return nil
     }