PR Feedback: improve PTrace implementation

Author: andrurogerz

tools/swift-inspect/Sources/SwiftInspectLinux/PTrace.swift

@@ -15,25 +15,31 @@ import LinuxSystemHeaders
 
 public class PTrace {
   enum PTraceError: Error {
-    case ptraceFailure(_ command: Int32, pid: pid_t, errno: Int32 = get_errno())
-    case waitFailure(pid: pid_t, errno: Int32 = get_errno())
-    case unexpectedWaitStatus(pid: pid_t, status: Int32, sigInfo: siginfo_t? = nil)
+    case operationFailure(_ command: CInt, pid: pid_t, errno: CInt = get_errno())
+    case waitFailure(pid: pid_t, errno: CInt = get_errno())
+    case unexpectedWaitStatus(pid: pid_t, status: CInt, sigInfo: siginfo_t? = nil)
   }
 
   let pid: pid_t
 
+  // Provides scoped access to a PTrace object.
+  public static func withAttachedProcess(pid: pid_t, _ closure: (PTrace) throws -> Void) throws {
+    let ptrace = try PTrace(pid);
+    try closure(ptrace)
+  }
+
   // Initializing a PTrace instance attaches to the target process, waits for
   // it to stop, and leaves it in a stopped state. The caller may resume the
   // process by calling cont().
-  public init(process pid: pid_t) throws {
+  init(_ pid: pid_t) throws {
     guard ptrace_attach(pid) != -1 else {
-      throw PTraceError.ptraceFailure(PTRACE_ATTACH, pid: pid)
+      throw PTraceError.operationFailure(PTRACE_ATTACH, pid: pid)
     }
 
     while true {
-      var status: Int32 = 0
+      var status: CInt = 0
       let result = waitpid(pid, &status, 0)
-      guard result != -1 else {
+      if result == -1 {
         if get_errno() == EINTR { continue }
         throw PTraceError.waitFailure(pid: pid)
       }
@@ -49,38 +55,38 @@ public class PTrace {
   deinit { _ = ptrace_detach(self.pid) }
 
   public func cont() throws {
-    guard ptrace_cont(self.pid) != -1 else {
-      throw PTraceError.ptraceFailure(PTRACE_CONT, pid: self.pid)
+    if ptrace_cont(self.pid) == -1 {
+      throw PTraceError.operationFailure(PTRACE_CONT, pid: self.pid)
     }
   }
 
   public func interrupt() throws {
-    guard ptrace_interrupt(self.pid) != -1 else {
-      throw PTraceError.ptraceFailure(PTRACE_INTERRUPT, pid: self.pid)
+    if ptrace_interrupt(self.pid) == -1 {
+      throw PTraceError.operationFailure(PTRACE_INTERRUPT, pid: self.pid)
     }
   }
 
   public func getSigInfo() throws -> siginfo_t {
     var sigInfo = siginfo_t()
-    guard ptrace_getsiginfo(self.pid, &sigInfo) != -1 else {
-      throw PTraceError.ptraceFailure(PTRACE_GETSIGINFO, pid: self.pid)
+    if ptrace_getsiginfo(self.pid, &sigInfo) == -1 {
+      throw PTraceError.operationFailure(PTRACE_GETSIGINFO, pid: self.pid)
     }
 
     return sigInfo
   }
 
   public func pokeData(addr: UInt64, value: UInt64) throws {
-    guard ptrace_pokedata(self.pid, UInt(addr), UInt(value)) != -1 else {
-      throw PTraceError.ptraceFailure(PTRACE_POKEDATA, pid: self.pid)
+    if ptrace_pokedata(self.pid, UInt(addr), UInt(value)) == -1 {
+      throw PTraceError.operationFailure(PTRACE_POKEDATA, pid: self.pid)
     }
   }
 
   public func getRegSet() throws -> RegisterSet {
     var regSet = RegisterSet()
     try withUnsafeMutableBytes(of: &regSet) {
       var vec = iovec(iov_base: $0.baseAddress!, iov_len: MemoryLayout<RegisterSet>.size)
-      guard ptrace_getregset(self.pid, NT_PRSTATUS, &vec) != -1 else {
-        throw PTraceError.ptraceFailure(PTRACE_GETREGSET, pid: self.pid)
+      if ptrace_getregset(self.pid, NT_PRSTATUS, &vec) == -1 {
+        throw PTraceError.operationFailure(PTRACE_GETREGSET, pid: self.pid)
       }
     }
 
@@ -91,21 +97,19 @@ public class PTrace {
     var regSetCopy = regSet
     try withUnsafeMutableBytes(of: &regSetCopy) {
       var vec = iovec(iov_base: $0.baseAddress!, iov_len: MemoryLayout<RegisterSet>.size)
-      guard ptrace_setregset(self.pid, NT_PRSTATUS, &vec) != -1 else {
-        throw PTraceError.ptraceFailure(PTRACE_SETREGSET, pid: self.pid)
+      if ptrace_setregset(self.pid, NT_PRSTATUS, &vec) == -1 {
+        throw PTraceError.operationFailure(PTRACE_SETREGSET, pid: self.pid)
       }
     }
   }
 
-  // Call the function at the specified address in the attached process. Caller
-  // may pass up to six 8-byte arguments. The optional callback is invoked when
-  // the process is stopped with a SIGTRAP signal. In this case, the caller is
-  // responsible for taking action on the signal.
-  public func callRemoteFunction(
-    at address: UInt64, with args: [UInt64] = [], onTrap callback: (() throws -> Void)? = nil
+  // Call the function at the specified address in the attached process with the
+  // provided argument array. The optional callback is invoked when the process
+  // is stopped with a SIGTRAP signal. In this case, the caller is responsible
+  // for taking action on the signal.
+  public func jump(
+    to address: UInt64, with args: [UInt64] = [], _ callback: (() throws -> Void)? = nil
   ) throws -> UInt64 {
-    precondition(args.count <= 6, "callRemoteFunction supports max of 6 arguments")
-
     let origRegs = try self.getRegSet()
     defer { try? self.setRegSet(regSet: origRegs) }
 
@@ -119,7 +123,7 @@ public class PTrace {
     try self.setRegSet(regSet: newRegs)
     try self.cont()
 
-    var status: Int32 = 0
+    var status: CInt = 0
     while true {
       let result = waitpid(self.pid, &status, 0)
       guard result != -1 else {

tools/swift-inspect/Sources/SwiftInspectLinux/RegisterSet.swift

@@ -19,17 +19,20 @@ public typealias RegisterSet = user_pt_regs
 extension RegisterSet {
   public static var trapInstructionSize: UInt { return 4 }  // brk #0x0
 
-  public func setupCall(
-    _ ptrace: PTrace, to funcAddr: UInt64, with args: [UInt64], returnTo returnAddr: UInt64
-  ) throws -> RegisterSet {
-    precondition(args.count <= 6)
+  public func setupCall(_ ptrace: PTrace, to funcAddr: UInt64, with args: [UInt64], returnTo returnAddr: UInt64) throws -> RegisterSet {
+    // The first 8 arguments are passed in regsters. Any additional arguments
+    // must be pushed on the stack, which is not implemented.
+    precondition(args.count <= 8)
+
     var registers = self
     registers.regs.0 = args.count > 0 ? args[0] : 0
     registers.regs.1 = args.count > 1 ? args[1] : 0
     registers.regs.2 = args.count > 2 ? args[2] : 0
     registers.regs.3 = args.count > 3 ? args[3] : 0
     registers.regs.4 = args.count > 4 ? args[4] : 0
     registers.regs.5 = args.count > 5 ? args[5] : 0
+    registers.regs.6 = args.count > 6 ? args[6] : 0
+    registers.regs.7 = args.count > 7 ? args[7] : 0
     registers.pc = funcAddr
     registers.regs.30 = returnAddr  // link register (x30)
     return registers
@@ -50,10 +53,11 @@ public typealias RegisterSet = pt_regs
 extension RegisterSet {
   public static var trapInstructionSize: UInt { return 1 }  // int3
 
-  public func setupCall(
-    _ ptrace: PTrace, to funcAddr: UInt64, with args: [UInt64], returnTo returnAddr: UInt64
-  ) throws -> RegisterSet {
+  public func setupCall(_ ptrace: PTrace, to funcAddr: UInt64, with args: [UInt64], returnTo returnAddr: UInt64) throws -> RegisterSet {
+    // The first 6 arguments are passed in registers. Any additional arguments
+    // must be pushed on the stack, which is not implemented.
     precondition(args.count <= 6)
+
     var registers = self
     registers.rdi = UInt(args.count > 0 ? args[0] : 0)
     registers.rsi = UInt(args.count > 1 ? args[1] : 0)

tools/swift-inspect/Sources/swift-inspect/AndroidRemoteProcess.swift

@@ -49,8 +49,6 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
     }
   }
 
-  let ptrace: SwiftInspectLinux.PTrace
-
   // We call mmap/munmap in the remote process to alloc/free memory for our own
   // use without impacting existing allocations in the remote process.
   lazy var mmapSymbol: RemoteSymbol = RemoteSymbol("mmap", self.symbolCache)
@@ -63,16 +61,6 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
   lazy var mallocEnableSymbol: RemoteSymbol = RemoteSymbol("malloc_enable", self.symbolCache)
   lazy var mallocIterateSymbol: RemoteSymbol = RemoteSymbol("malloc_iterate", self.symbolCache)
 
-  override init?(processId: ProcessIdentifier) {
-    do {
-      self.ptrace = try SwiftInspectLinux.PTrace(process: processId)
-    } catch {
-      print("failed initialization: \(error)")
-      return nil
-    }
-    super.init(processId: processId)
-  }
-
   // Linux and Android have no supported method to enumerate allocations in the
   // heap of a remote process. Android does, however, support the malloc_iterate
   // API, which enumerates allocations in the current process. We leverage this
@@ -81,26 +69,27 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
   override internal func iterateHeap(_ body: (swift_addr_t, UInt64) -> Void) {
     var regionCount = 0
     var allocCount = 0
-    for entry in self.memoryMap.entries {
-      // Limiting malloc_iterate calls to only memory regions that are known
-      // to contain heap allocations is not strictly necessary but it does
-      // significantly improve the speed of heap iteration.
-      guard entry.isHeapRegion() else { continue }
-
-      // collect all of the allocations in this heap region
-      let allocations: [(base: swift_addr_t, len: UInt64)]
-      do {
-        allocations = try self.iterateHeapRegion(region: entry)
-        regionCount += 1
-      } catch {
-        print("failed iterating remote heap: \(error)")
-        return
+    do {
+      try PTrace.withAttachedProcess(pid: self.processIdentifier) { ptrace in
+        for entry in self.memoryMap.entries {
+          // Limiting malloc_iterate calls to only memory regions that are known
+          // to contain heap allocations is not strictly necessary but it does
+          // significantly improve the speed of heap iteration.
+          guard entry.isHeapRegion() else { continue }
+
+          // collect all of the allocations in this heap region
+          let allocations: [(base: swift_addr_t, len: UInt64)]
+          allocations = try self.iterateHeapRegion(ptrace, region: entry)
+          regionCount += 1
+          allocCount += allocations.count
+
+          // process all of the collected allocations
+          for alloc in allocations { body(alloc.base, alloc.len) }
+        }
       }
-
-      allocCount += allocations.count
-
-      // process all of the collected allocations
-      for alloc in allocations { body(alloc.base, alloc.len) }
+    } catch {
+      print("failed iterating remote heap: \(error)")
+      return
     }
 
     if regionCount == 0 {
@@ -116,16 +105,16 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
 
   // Iterate a single heap region in the remote process and return an array
   // of (base, len) pairs describing each heap allocation in the region.
-  internal func iterateHeapRegion(region: MemoryMap.Entry) throws -> [(
+  internal func iterateHeapRegion(_ ptrace: PTrace, region: MemoryMap.Entry) throws -> [(
     base: swift_addr_t, len: UInt64
   )] {
     // Allocate a page-sized buffer in the remote process that malloc_iterate
     // will populaate with metadata describing each heap entry it enumerates.
     let dataLen = sysconf(Int32(_SC_PAGESIZE))
     let remoteDataAddr = try self.mmapRemote(
-      len: dataLen, prot: PROT_READ | PROT_WRITE, flags: MAP_ANON | MAP_PRIVATE)
+      ptrace, len: dataLen, prot: PROT_READ | PROT_WRITE, flags: MAP_ANON | MAP_PRIVATE)
     defer {
-      _ = try? self.munmapRemote(addr: remoteDataAddr, len: dataLen)
+      _ = try? self.munmapRemote(ptrace, addr: remoteDataAddr, len: dataLen)
     }
 
     // Allocate and inialize a local buffer that will be used to copy metadata
@@ -142,9 +131,9 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
     // executed in the remote process.
     let codeLen = heap_iterate_callback_len()
     let remoteCodeAddr = try mmapRemote(
-      len: codeLen, prot: PROT_READ | PROT_WRITE | PROT_EXEC, flags: MAP_ANON | MAP_PRIVATE)
+      ptrace, len: codeLen, prot: PROT_READ | PROT_WRITE | PROT_EXEC, flags: MAP_ANON | MAP_PRIVATE)
     defer {
-      _ = try? self.munmapRemote(addr: remoteCodeAddr, len: codeLen)
+      _ = try? self.munmapRemote(ptrace, addr: remoteCodeAddr, len: codeLen)
     }
 
     // Copy the malloc_iterate callback implementation to the remote process.
@@ -158,9 +147,9 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
 
     // Disable malloc/free while enumerating the region to get a consistent
     // snapshot of existing allocations.
-    try self.mallocDisableRemote()
+    try self.mallocDisableRemote(ptrace)
     defer {
-      _ = try? self.mallocEnableRemote()
+      _ = try? self.mallocEnableRemote(ptrace)
     }
 
     // Collects (base, len) pairs describing each heap allocation in the remote
@@ -169,7 +158,7 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
 
     let regionLen = region.endAddr - region.startAddr
     let args = [region.startAddr, regionLen, remoteCodeAddr, remoteDataAddr]
-    _ = try self.ptrace.callRemoteFunction(at: mallocIterateAddr, with: args) {
+    _ = try ptrace.jump(to: mallocIterateAddr, with: args) {
       // This callback is invoked when a SIGTRAP is encountered in the remote
       // process. In this context, this signal indicates there is no more room
       // in the allocated metadata region (see AndroidCLib/heap.c).
@@ -183,11 +172,11 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
       }
       try self.process.writeMem(remoteAddr: remoteDataAddr, localAddr: buffer, len: UInt(dataLen))
 
-      var regs = try self.ptrace.getRegSet()
+      var regs = try ptrace.getRegSet()
       regs.step(RegisterSet.trapInstructionSize)
 
-      try self.ptrace.setRegSet(regSet: regs)
-      try self.ptrace.cont()
+      try ptrace.setRegSet(regSet: regs)
+      try ptrace.cont()
     }
 
     try self.process.readMem(remoteAddr: remoteDataAddr, localAddr: buffer, len: UInt(dataLen))
@@ -219,37 +208,37 @@ internal final class AndroidRemoteProcess: LinuxRemoteProcess {
   }
 
   // call mmap in the remote process with the provided arguments
-  internal func mmapRemote(len: Int, prot: Int32, flags: Int32) throws -> UInt64 {
+  internal func mmapRemote(_ ptrace: PTrace, len: Int, prot: Int32, flags: Int32) throws -> UInt64 {
     guard let sym = self.mmapSymbol.addr else {
       throw RemoteProcessError.missingSymbol(self.mmapSymbol.name)
     }
     let args = [0, UInt64(len), UInt64(prot), UInt64(flags)]
-    return try self.ptrace.callRemoteFunction(at: sym, with: args)
+    return try ptrace.jump(to: sym, with: args)
   }
 
   // call munmap in the remote process with the provdied arguments
-  internal func munmapRemote(addr: UInt64, len: Int) throws -> UInt64 {
+  internal func munmapRemote(_ ptrace: PTrace, addr: UInt64, len: Int) throws -> UInt64 {
     guard let sym = self.munmapSymbol.addr else {
       throw RemoteProcessError.missingSymbol(self.munmapSymbol.name)
     }
     let args: [UInt64] = [addr, UInt64(len)]
-    return try self.ptrace.callRemoteFunction(at: sym, with: args)
+    return try ptrace.jump(to: sym, with: args)
   }
 
   // call malloc_disable in the remote process
-  internal func mallocDisableRemote() throws {
+  internal func mallocDisableRemote(_ ptrace: PTrace) throws {
     guard let sym = self.mallocDisableSymbol.addr else {
       throw RemoteProcessError.missingSymbol(self.mallocDisableSymbol.name)
     }
-    _ = try self.ptrace.callRemoteFunction(at: sym)
+    _ = try ptrace.jump(to: sym)
   }
 
   // call malloc_enable in the remote process
-  internal func mallocEnableRemote() throws {
+  internal func mallocEnableRemote(_ ptrace: PTrace) throws {
     guard let sym = self.mallocEnableSymbol.addr else {
       throw RemoteProcessError.missingSymbol(self.mallocEnableSymbol.name)
     }
-    _ = try self.ptrace.callRemoteFunction(at: sym)
+    _ = try ptrace.jump(to: sym)
   }
 }