[swift-inspect] Inspect all processes and add summary option

With this change swift-inspect can inspect all processes to see if metadata allocation iteration is enabled. We also added summary option that sorts metadata to popularity.

Author: w6sec

tools/swift-inspect/Sources/swift-inspect/DarwinRemoteProcess.swift

@@ -20,6 +20,8 @@ internal final class DarwinRemoteProcess: RemoteProcess {
   public typealias ProcessHandle = task_t
 
   private var task: task_t
+  internal var processIdentifier: ProcessIdentifier
+  internal lazy var processName = getProcessName(processId: processIdentifier) ?? "<unknown process>"
 
   public var process: ProcessHandle { task }
   public private(set) var context: SwiftReflectionContextRef!
@@ -114,20 +116,32 @@ internal final class DarwinRemoteProcess: RemoteProcess {
   }
 
   init?(processId: ProcessIdentifier, forkCorpse: Bool) {
+    processIdentifier = processId
     var task: task_t = task_t()
     let taskResult = task_for_pid(mach_task_self_, processId, &task)
     guard taskResult == KERN_SUCCESS else {
-      print("unable to get task for pid \(processId): \(String(cString: mach_error_string(taskResult))) \(hex: taskResult)")
+      print("unable to get task for pid \(processId): \(String(cString: mach_error_string(taskResult))) \(hex: taskResult)",
+        to: &Std.err)
       return nil
     }
 
     if forkCorpse {
       var corpse = task_t()
-      let corpseResult = task_generate_corpse(task, &corpse)
-      if corpseResult == KERN_SUCCESS {
-        task = corpse
-      } else {
-        print("unable to fork corpse for pid \(processId): \(String(cString: mach_error_string(corpseResult))) \(hex: corpseResult)")
+      let maxRetry = 6
+      for retry in 0..<maxRetry {
+        let corpseResult = task_generate_corpse(task, &corpse)
+        if corpseResult == KERN_SUCCESS {
+          task_stop_peeking(task)
+          mach_port_deallocate(mach_task_self_, task)
+          task = corpse
+          break
+        }
+        if corpseResult != KERN_RESOURCE_SHORTAGE || retry == maxRetry {
+          print("unable to fork corpse for pid \(processId): \(String(cString: mach_error_string(corpseResult))) \(hex: corpseResult)",
+            to: &Std.err)
+          return nil
+        }
+        sleep(UInt32(1 << retry))
       }
     }
 
@@ -161,6 +175,7 @@ internal final class DarwinRemoteProcess: RemoteProcess {
   deinit {
     task_stop_peeking(self.task)
     mach_port_deallocate(mach_task_self_, self.task)
+    mach_port_mod_refs(mach_task_self_, self.task, MACH_PORT_RIGHT_SEND, -1);
   }
 
   func symbolicate(_ address: swift_addr_t) -> (module: String?, symbol: String?) {

tools/swift-inspect/Sources/swift-inspect/Operations/DumpGenericMetadata.swift

@@ -55,14 +55,25 @@ private struct Metadata: Encodable {
   }
 }
 
+private struct ProcessMetadata: Encodable {
+  var name: String
+  var pid: ProcessIdentifier
+  var metadata: [Metadata]
+}
+
+private struct MetadataSummary: Encodable {
+  var totalSize: Int
+  var processes: Set<String>
+}
+
 internal struct Output: TextOutputStream {
   let fileHandle: FileHandle
   init(_ outputFile: String? = nil) throws {
     if let outputFile {
       if FileManager().createFile(atPath: outputFile, contents: nil) {
-        self.fileHandle = try FileHandle(forWritingAtPath: outputFile)!
+        self.fileHandle = FileHandle(forWritingAtPath: outputFile)!
       } else {
-        print("Unable to create file \(outputFile)")
+        print("Unable to create file \(outputFile)", to: &Std.err)
         exit(1)
       }
     } else {
@@ -75,7 +86,6 @@ internal struct Output: TextOutputStream {
       fileHandle.write(encodedString)
     }
   }
-
 }
 
 internal struct DumpGenericMetadata: ParsableCommand {
@@ -92,6 +102,9 @@ internal struct DumpGenericMetadata: ParsableCommand {
   var genericMetadataOptions: GenericMetadataOptions
 
   func run() throws {
+    disableStdErrBuffer()
+    var metadataSummary = [String: MetadataSummary]()
+    var allProcesses = [ProcessMetadata]()
     try inspect(options: options) { process in
       let allocations: [swift_metadata_allocation_t] =
           try process.context.allocations.sorted()
@@ -117,19 +130,47 @@ internal struct DumpGenericMetadata: ParsableCommand {
                         isArrayOfClass: process.context.isArrayOfClass(pointer),
                         garbage: garbage,
                         backtrace: currentBacktrace)
+      } // generics
+
+      // Update summary
+      generics.forEach { metadata in
+        if let allocation = metadata.allocation {
+          let name = metadata.name
+          if metadataSummary.keys.contains(name) {
+            metadataSummary[name]!.totalSize += allocation.size
+            metadataSummary[name]!.processes.insert(process.processName)
+          } else {
+            metadataSummary[name] = MetadataSummary(totalSize: allocation.size,
+                                                    processes: Set([process.processName]))
+            }
+        }
       }
 
       if genericMetadataOptions.json {
-        try dumpJson(process: process, generics: generics)
-      } else {
+        let processMetadata = ProcessMetadata(name: process.processName,
+                                              pid: process.processIdentifier as! ProcessIdentifier,
+                                              metadata: generics)
+        allProcesses.append(processMetadata)
+      } else if !genericMetadataOptions.summary {
         try dumpText(process: process, generics: generics)
-      }
+        }
+    } // inspect
+
+    if genericMetadataOptions.json {
+      if genericMetadataOptions.summary {
+        try dumpJson(of: metadataSummary)
+      } else {
+        try dumpJson(of: allProcesses)
+        }
+    } else if genericMetadataOptions.summary {
+      try dumpTextSummary(of: metadataSummary)
     }
   }
 
   private func dumpText(process: any RemoteProcess, generics: [Metadata]) throws {
     var errorneousMetadata: [(ptr: swift_reflection_ptr_t, name: String)] = []
     var output = try Output(genericMetadataOptions.outputFile)
+    print("\(process.processName)(\(process.processIdentifier)):\n", to: &output)
     print("Address", "Allocation", "Size", "Offset", "isArrayOfClass", "Name", separator: "\t", to: &output)
     generics.forEach {
       print("\(hex: $0.ptr)", terminator: "\t", to: &output)
@@ -149,22 +190,18 @@ internal struct DumpGenericMetadata: ParsableCommand {
     }
 
     if errorneousMetadata.count > 0 {
-      print("Error: The following metadata was not found in any DATA or AUTH segments, may be garbage.", to: &output)
+      print("Warning: The following metadata was not found in any DATA or AUTH segments, may be garbage.", to: &output)
       errorneousMetadata.forEach {
         print("\(hex: $0.ptr)\t\($0.name)", to: &output)
       }
     }
+    print("", to: &output)
   }
 
-  private func dumpJson(process: any RemoteProcess,
-                        generics: [Metadata]) throws {
-    struct AllMetadataEntries: Encodable {
-      var metadata: [Metadata]
-    }
-    let allMetadataEntries = AllMetadataEntries(metadata: generics)
+  private func dumpJson(of: (any Encodable)) throws {
     let encoder = JSONEncoder()
-    encoder.outputFormatting = .prettyPrinted
-    let data = try encoder.encode(allMetadataEntries)
+    encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
+    let data = try encoder.encode(of)
     let jsonOutput = String(data: data, encoding: .utf8)!
     if let outputFile = genericMetadataOptions.outputFile {
       try jsonOutput.write(toFile: outputFile, atomically: true, encoding: .utf8)
@@ -173,4 +210,23 @@ internal struct DumpGenericMetadata: ParsableCommand {
     }
   }
 
+  private func dumpTextSummary(of: [String: MetadataSummary]) throws {
+    var output = try Output(genericMetadataOptions.outputFile)
+    print("Size", "Owners", "Name", separator: "\t", to: &output)
+    var totalSize = 0
+    var unknownSize = 0
+    of.sorted { first, second in
+      first.value.processes.count > second.value.processes.count
+    }
+    .forEach { summary in
+      totalSize += summary.value.totalSize
+      if summary.key == "<unknown>" {
+        unknownSize += summary.value.totalSize
+      }
+      print(summary.value.totalSize, summary.value.processes.count, summary.key,
+        separator: "\t", to: &output)
+    }
+    print("\nTotal size:\t\(totalSize / 1024) KiB", to: &output)
+    print("Unknown size:\t\(unknownSize / 1024) KiB", to: &output)
+  }
 }

tools/swift-inspect/Sources/swift-inspect/Process.swift

@@ -43,6 +43,62 @@ private func refersToSelf(_ str: String) -> Bool {
   // No match.
   return false
 }
+
+internal func getRemoteProcess(processId: ProcessIdentifier,
+                               options: UniversalOptions) -> (any RemoteProcess)? {
+  return DarwinRemoteProcess(processId: processId,
+                             forkCorpse: options.forkCorpse)
+}
+
+internal func getProcessName(processId: ProcessIdentifier) -> String? {
+  var info = proc_bsdinfo()
+  let bsdinfoSize = Int32(MemoryLayout<proc_bsdinfo>.stride)
+  let size = proc_pidinfo(processId, PROC_PIDTBSDINFO, 0, &info, bsdinfoSize)
+  if (size != bsdinfoSize) {
+    return nil
+  }
+ let processName = withUnsafeBytes(of: info.pbi_name) { buffer in
+    let nonnullBuffer = buffer.prefix { $0 != 0 }
+    return String(decoding: nonnullBuffer, as: UTF8.self)
+  }
+  return processName
+}
+
+internal func getAllProcesses(options: UniversalOptions) -> [ProcessIdentifier]? {
+  var ProcessIdentifiers = [ProcessIdentifier]()
+  let kinfo_stride = MemoryLayout<kinfo_proc>.stride
+  var bufferSize: Int = 0
+  var name: [Int32] = [CTL_KERN, KERN_PROC, KERN_PROC_ALL]
+
+  guard sysctl(&name, u_int(name.count), nil, &bufferSize, nil, 0) == 0 else {
+    return nil
+  }
+  let count = bufferSize / kinfo_stride
+  var buffer = Array(repeating: kinfo_proc(), count: count)
+  guard sysctl(&name, u_int(name.count), &buffer, &bufferSize, nil, 0) == 0 else {
+    return nil
+  }
+  let newCount = bufferSize / kinfo_stride
+  if count > newCount {
+    buffer.dropLast(count - newCount)
+  }
+  let sorted = buffer.sorted { first, second in
+    first.kp_proc.p_pid > second.kp_proc.p_pid
+  }
+  let myPid = getpid()
+  for kinfo in sorted {
+    let pid = kinfo.kp_proc.p_pid
+    if pid <= 1 {
+      break
+    }
+    if pid == myPid { // skip self
+      continue
+    }
+    ProcessIdentifiers.append(pid)
+  }
+  return ProcessIdentifiers
+}
+
 #elseif os(Windows)
 import WinSDK
 
@@ -81,6 +137,12 @@ internal func process(matching: String) -> ProcessIdentifier? {
 
   return matches.first?.0
 }
+
+internal func getRemoteProcess(processId: ProcessIdentifier,
+                               options: UniversalOptions) -> (any RemoteProcess)? {
+  return WindowsRemoteProcess(processId: processId)
+}
+
 #else
 #error("Unsupported platform")
 #endif

tools/swift-inspect/Sources/swift-inspect/RemoteProcess.swift

@@ -18,6 +18,8 @@ internal protocol RemoteProcess: AnyObject {
 
   var process: ProcessHandle { get }
   var context: SwiftReflectionContextRef! { get }
+  var processIdentifier: ProcessIdentifier { get }
+  var processName: String { get }
 
   typealias QueryDataLayoutFunction =
       @convention(c) (UnsafeMutableRawPointer?, DataLayoutQueryType,

tools/swift-inspect/Sources/swift-inspect/String+Extensions.swift

@@ -11,10 +11,27 @@
 //===----------------------------------------------------------------------===//
 
 import SwiftRemoteMirror
+import Foundation
 
 extension DefaultStringInterpolation {
   mutating func appendInterpolation<T>(hex: T) where T: BinaryInteger {
     appendInterpolation("0x")
     appendInterpolation(String(hex, radix: 16))
   }
 }
+
+enum Std {
+  struct File: TextOutputStream {
+    var underlying: UnsafeMutablePointer<FILE>
+
+    mutating func write(_ string: String) {
+      fputs(string, underlying)
+    }
+  }
+
+  static var err = File(underlying: stderr)
+}
+
+internal func disableStdErrBuffer() {
+  setvbuf(stderr, nil, Int32(_IONBF), 0)
+}

tools/swift-inspect/Sources/swift-inspect/Symbolication+Extensions.swift

@@ -15,18 +15,6 @@
 import Foundation
 import SymbolicationShims
 
-enum Std {
-  struct File: TextOutputStream {
-    var underlying: UnsafeMutablePointer<FILE>
-
-    mutating func write(_ string: String) {
-      fputs(string, underlying)
-    }
-  }
-
-  static var err = File(underlying: stderr)
-}
-
 private let symbolicationPath =
   "/System/Library/PrivateFrameworks/Symbolication.framework/Symbolication"
 private let symbolicationHandle = dlopen(symbolicationPath, RTLD_LAZY)!
@@ -133,7 +121,7 @@ func CSSymbolOwnerForeachSymbol(
 }
 
 func CSSymbolOwnerGetSymbolWithMangledName(
-  _ owner: CSTypeRef, 
+  _ owner: CSTypeRef,
   _ name: String
 ) -> CSTypeRef {
   Sym.CSSymbolOwnerGetSymbolWithMangledName(owner, name)
@@ -196,7 +184,7 @@ func task_start_peeking(_ task: task_t) -> Bool {
   if result == KERN_SUCCESS {
     return true
   }
-  
+
   print("task_start_peeking failed: \(machErrStr(result))", to: &Std.err)
   return false
 }

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

@@ -23,6 +23,8 @@ internal final class WindowsRemoteProcess: RemoteProcess {
 
   public private(set) var process: ProcessHandle
   public private(set) var context: SwiftReflectionContextRef!
+  public private(set) var processIdentifier: ProcessIdentifier
+  public private(set) var processName: String = "<unknown process>"
 
   private var hSwiftCore: HMODULE = HMODULE(bitPattern: -1)!
 
@@ -136,6 +138,7 @@ internal final class WindowsRemoteProcess: RemoteProcess {
   }
 
   init?(processId: ProcessIdentifier) {
+    self.processIdentifier = processId
     // Get process handle.
     self.process =
       OpenProcess(