Merge pull request #11 from p-x9/feature/improve-read-class-protocol

Return protocols and classes paired with their defining Mach-O files

Author: p-x9

Sources/MachOObjCSection/Extension/DyldCache+.swift

@@ -115,3 +115,14 @@ extension DyldCache {
         return nil
     }
 }
+
+extension DyldCache {
+    func machO(containing unslidAddress: UInt64) -> MachOFile? {
+        for machO in self.machOFiles() {
+            if machO.contains(unslidAddress: unslidAddress) {
+                return machO
+            }
+        }
+        return nil
+    }
+}

Sources/MachOObjCSection/Extension/DyldCacheLoaded+.swift

@@ -99,3 +99,23 @@ extension DyldCacheLoaded {
         return nil
     }
 }
+
+extension DyldCacheLoaded {
+    func machO(containing ptr: UnsafeRawPointer) -> MachOImage? {
+        for machO in machOImages() {
+            if machO.contains(ptr: ptr) {
+                return machO
+            }
+        }
+        return nil
+    }
+
+    func machO(containing unslidAddress: UInt64) -> MachOImage? {
+        for machO in self.machOImages() {
+            if machO.contains(unslidAddress: unslidAddress) {
+                return machO
+            }
+        }
+        return nil
+    }
+}

Sources/MachOObjCSection/Model/Protocol/ObjCProtocolList32.swift

@@ -26,13 +26,13 @@ public struct ObjCProtocolList32: ObjCProtocolListProtocol {
 extension ObjCProtocolList32 {
     public func protocols(
         in machO: MachOImage
-    ) -> [ObjCProtocol]? {
+    ) -> [(MachOImage, ObjCProtocol)]? {
         _readProtocols(in: machO, pointerType: UInt32.self)
     }
 
     public func protocols(
         in machO: MachOFile
-    ) -> [ObjCProtocol]? {
+    ) -> [(MachOFile, ObjCProtocol)]? {
         guard !isListOfLists else {
             assertionFailure()
             return nil
@@ -63,23 +63,32 @@ extension ObjCProtocolList32 {
                 let offset = $0 + numericCast(headerStartOffset)
                 var resolvedOffset = offset
 
+                var targetMachO = machO
+
                 var fileHandle = machO.fileHandle
 
                 if let (_cache, _offset) = machO.cacheAndFileOffset(
                     fromStart: offset
                 ) {
                     resolvedOffset = _offset
                     fileHandle = _cache.fileHandle
+
+                    let unslidAddress = offset + _cache.mainCacheHeader.sharedRegionStart
+                    if !targetMachO.contains(unslidAddress: unslidAddress),
+                       let machO = _cache.machO(containing: unslidAddress) {
+                        targetMachO = machO
+                    }
                 }
 
                 let layout: ObjCProtocol32.Layout = fileHandle.read(
                     offset: numericCast(resolvedOffset),
                     swapHandler: { _ in }
                 )
-                return .init(
+                let `protocol`: ObjCProtocol = .init(
                     layout: layout,
                     offset: numericCast(offset) - machO.headerStartOffset
                 )
+                return (targetMachO, `protocol`)
             }
     }
 }

Sources/MachOObjCSection/Model/Protocol/ObjCProtocolList64.swift

@@ -26,13 +26,13 @@ public struct ObjCProtocolList64: ObjCProtocolListProtocol {
 extension ObjCProtocolList64 {
     public func protocols(
         in machO: MachOImage
-    ) -> [ObjCProtocol]? {
+    ) -> [(MachOImage, ObjCProtocol)]? {
         _readProtocols(in: machO, pointerType: UInt64.self)
     }
 
     public func protocols(
         in machO: MachOFile
-    ) -> [ObjCProtocol]? {
+    ) -> [(MachOFile, ObjCProtocol)]? {
         guard !isListOfLists else {
             assertionFailure()
             return nil
@@ -63,23 +63,32 @@ extension ObjCProtocolList64 {
                 let offset = machO.fileOffset(of: $0) + numericCast(headerStartOffset)
                 var resolvedOffset = offset
 
+                var targetMachO = machO
+
                 var fileHandle = machO.fileHandle
 
                 if let (_cache, _offset) = machO.cacheAndFileOffset(
                     fromStart: offset
                 ) {
                     resolvedOffset = _offset
                     fileHandle = _cache.fileHandle
+
+                    let unslidAddress = offset + _cache.mainCacheHeader.sharedRegionStart
+                    if !targetMachO.contains(unslidAddress: unslidAddress),
+                       let machO = _cache.machO(containing: unslidAddress) {
+                        targetMachO = machO
+                    }
                 }
 
                 let layout: ObjCProtocol64.Layout = fileHandle.read(
                     offset:  numericCast(resolvedOffset),
                     swapHandler: { _ in }
                 )
-                return .init(
+                let `protocol`: ObjCProtocol = .init(
                     layout: layout,
                     offset: numericCast(offset) - machO.headerStartOffset
                 )
+                return (targetMachO, `protocol`)
             }
     }
 }