refactor

Author: p-x9

Sources/SwiftHook/SwiftHook+function.swift

@@ -17,50 +17,27 @@ extension SwiftHook {
         _ first: String,
         _ second: String,
         isMangled: Bool
-    ) -> Bool {
-        if isMangled {
-            return exchangeFuncImplementationMangled(first, second)
-        } else {
-            return exchangeFuncImplementationDeMangled(first, second)
-        }
-    }
-}
-
-extension SwiftHook {
-    private static func exchangeFuncImplementationMangled(
-        _ first: String,
-        _ second: String
-    ) -> Bool {
-        var firstSymbol: UnsafeMutableRawPointer?
-        var secondSymbol: UnsafeMutableRawPointer?
-
-        for i in 0..<_dyld_image_count() {
-            let dylib = _dyld_get_image_name(i)
-            guard let handle = dlopen(dylib, RTLD_NOW) else {
-                continue
-            }
-            if firstSymbol == nil {
-                firstSymbol = dlsym(handle, first.cString(using: .utf8))
-            }
-            if secondSymbol == nil {
-                secondSymbol = dlsym(handle, second.cString(using: .utf8))
-            }
-
-            if firstSymbol != nil && secondSymbol != nil { break }
-        }
+    ) throws {
+        var isSucceeded = false
 
-        return exchangeFuncImplementation(
+        isSucceeded = try _exchangeFuncImplementation(
             first,
             second,
-            firstSymbol,
-            secondSymbol
+            isMangled: isMangled
         )
+
+        if isSucceeded { return }
+
+        throw SwiftHookError.failedToExchangeFuncImplementation
     }
+}
 
-    private static func exchangeFuncImplementationDeMangled(
+extension SwiftHook {
+    private static func _exchangeFuncImplementation(
         _ first: String,
-        _ second: String
-    ) -> Bool {
+        _ second: String,
+        isMangled: Bool
+    ) throws -> Bool {
         var first: String = first
         var second: String = second
 
@@ -69,23 +46,39 @@ extension SwiftHook {
 
         for i in 0..<_dyld_image_count() {
             let machO = MachOImage(ptr: _dyld_get_image_header(i))
-            if let symbol = machO.symbol(named: first, mangled: false) {
+            if let symbol = machO.symbol(
+                named: first,
+                mangled: isMangled
+            ) {
                 firstSymbol = .init(
                     mutating: machO.ptr.advanced(by: symbol.offset)
                 )
                 first = String(cString: symbol.nameC + 1)
             }
-            if let symbol = machO.symbol(named: second, mangled: false) {
+            if let symbol = machO.symbol(
+                named: second,
+                mangled: isMangled
+            ) {
                 secondSymbol = .init(
                     mutating: machO.ptr.advanced(by: symbol.offset)
                 )
                 second = String(cString: symbol.nameC + 1)
             }
 
-
             if firstSymbol != nil && secondSymbol != nil { break }
         }
 
+        if firstSymbol == nil && secondSymbol == nil {
+            throw SwiftHookError.firstAndSecondSymbolAreNotFound
+        }
+
+        if firstSymbol == nil {
+            throw SwiftHookError.firstSymbolIsNotFound
+        }
+        if secondSymbol == nil {
+            throw SwiftHookError.secondSymbolIsNotFound
+        }
+
         return exchangeFuncImplementation(
             first,
             second,
@@ -94,6 +87,7 @@ extension SwiftHook {
         )
     }
 
+    @discardableResult
     private static func exchangeFuncImplementation(
         _ first: String,
         _ second: String,
@@ -124,4 +118,3 @@ extension SwiftHook {
         return false
     }
 }
-

Sources/SwiftHook/SwiftHook+method.swift

@@ -3,7 +3,7 @@
 //
 //
 //  Created by p-x9 on 2023/12/22.
-//  
+//
 //
 
 import Foundation
@@ -31,6 +31,8 @@ extension SwiftHook {
             for: `class`
         )
         if isSucceeded { return }
+
+        throw SwiftHookError.failedToExchangeMethodImplementation
     }
 }
 
@@ -103,21 +105,3 @@ extension SwiftHook {
         return true
     }
 }
-
-@discardableResult
-func swizzle(class: AnyClass, orig origSelector: Selector, hooked hookedSelector: Selector) -> Bool {
-    guard let origMethod = class_getInstanceMethod(`class`, origSelector),
-          let hookedMethod = class_getInstanceMethod(`class`, hookedSelector) else {
-        return false
-    }
-
-    let didAddMethod = class_addMethod(`class`, origSelector,
-                                       method_getImplementation(hookedMethod),
-                                       method_getTypeEncoding(hookedMethod))
-    if didAddMethod {
-        class_replaceMethod(`class`, hookedSelector, method_getImplementation(origMethod), method_getTypeEncoding(origMethod))
-        return true
-    }
-    method_exchangeImplementations(origMethod, hookedMethod)
-    return true
-    }

Sources/SwiftHook/SwiftHookError.swift

@@ -0,0 +1,18 @@
+//
+//  SwiftHookError.swift
+//
+//
+//  Created by p-x9 on 2024/01/08.
+//  
+//
+
+import Foundation
+
+public enum SwiftHookError: Error {
+    case failedToExchangeFuncImplementation
+    case firstSymbolIsNotFound
+    case secondSymbolIsNotFound
+    case firstAndSecondSymbolAreNotFound
+
+    case failedToExchangeMethodImplementation
+}

Tests/SwiftHookTests/SwiftHookTests.swift

@@ -9,73 +9,73 @@ final class SwiftHookTests: XCTestCase {
         _ = disableExclusivityChecking
     }
 
-    func testExchangeFunc() {
+    func testExchangeFunc() throws {
         if setjump(&buf) != 0 {
             return
         }
 
         // XXXXhook_assertionFailure ⇔ _assertionFailure
-        print(SwiftHook.exchangeFuncImplementation(
+        try SwiftHook.exchangeFuncImplementation(
             "$s14SwiftHookTests25XXXXhook_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_A2ISus6UInt32VtF",
             "$ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_A2HSus6UInt32VtF",
             isMangled: true
-        ))
+        )
 
         var optional: Int?
         let forceUnwrapped = optional!
     }
 
-    func testExchangeDemangledFunc() {
+    func testExchangeDemangledFunc() throws {
         if setjump(&buf) != 0 {
             return
         }
 
         // XXXXhook_assertionFailure ⇔ _assertionFailure
-        print(SwiftHook.exchangeFuncImplementation(
+        try SwiftHook.exchangeFuncImplementation(
             "SwiftHookTests.XXXXhook_assertionFailure(_: Swift.StaticString, _: Swift.StaticString, file: Swift.StaticString, line: Swift.UInt, flags: Swift.UInt32) -> Swift.Never",
             "Swift._assertionFailure(_: Swift.StaticString, _: Swift.StaticString, file: Swift.StaticString, line: Swift.UInt, flags: Swift.UInt32) -> Swift.Never",
             isMangled: false
-        ))
+        )
 
         var optional: Int?
         let forceUnwrapped = optional!
     }
 }
 
 extension SwiftHookTests {
-    func testExchangeFuncInSelfImage() {
+    func testExchangeFuncInSelfImage() throws {
         if setjump(&buf) != 0 {
             return
         }
 
         // hook_assertionFailure ⇔ XXXXhook_assertionFailure
-        print(SwiftHook.exchangeFuncImplementation(
+        try SwiftHook.exchangeFuncImplementation(
             "$s14SwiftHookTests21hook_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAISus6UInt32VtF",
             "$s14SwiftHookTests25XXXXhook_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAISus6UInt32VtF",
             isMangled: true
-        ))
+        )
 
         hook_assertionFailure("aaa", "bbbb", flags: 0)
     }
 
-    func testExchangeDemangledFuncInSelfImage() {
+    func testExchangeDemangledFuncInSelfImage() throws {
         if setjump(&buf) != 0 {
             return
         }
 
         // hook_assertionFailure ⇔ XXXXhook_assertionFailure
-        print(SwiftHook.exchangeFuncImplementation(
+        try SwiftHook.exchangeFuncImplementation(
             "SwiftHookTests.hook_assertionFailure(_: Swift.StaticString, _: Swift.String, file: Swift.StaticString, line: Swift.UInt, flags: Swift.UInt32) -> Swift.Never",
             "SwiftHookTests.XXXXhook_assertionFailure(_: Swift.StaticString, _: Swift.String, file: Swift.StaticString, line: Swift.UInt, flags: Swift.UInt32) -> Swift.Never",
             isMangled: false
-        ))
+        )
 
         hook_assertionFailure("aaa", "bbbb", flags: 0)
     }
 }
 
 extension SwiftHookTests {
-    func testExchangeStructFunction() {
+    func testExchangeStructFunction() throws {
         let item = StructItem()
 
         print(item.printA())
@@ -84,11 +84,11 @@ extension SwiftHookTests {
         XCTAssertEqual(item.printB(), "B")
 
         // printA ⇔ printB
-        print(SwiftHook.exchangeFuncImplementation(
+        try SwiftHook.exchangeFuncImplementation(
             "SwiftHookTests.StructItem.printA() -> Swift.String",
             "SwiftHookTests.StructItem.printB() -> Swift.String",
             isMangled: false
-        ))
+        )
 
         print(item.printA())
         print(item.printB())