Added support for cleanup of the dynamic subclass

Author: lukaskubanek

Sources/InterposeKit/Hooks/HookStrategy/ObjectHookStrategy/ObjectHookRegistry.swift

@@ -13,7 +13,7 @@ internal enum ObjectHookRegistry {
 
         objc_setAssociatedObject(
             block,
-            &self.associatedKey,
+            &ObjectHookRegistryKey,
             WeakReference(handle),
             .OBJC_ASSOCIATION_RETAIN
         )
@@ -26,16 +26,16 @@ internal enum ObjectHookRegistry {
 
         guard let reference = objc_getAssociatedObject(
             block,
-            &self.associatedKey
+            &ObjectHookRegistryKey
         ) as? WeakReference<ObjectHookHandle> else { return nil }
 
         return reference.object
     }
 
-    private static var associatedKey: UInt8 = 0
-    
 }
 
+fileprivate var ObjectHookRegistryKey: UInt8 = 0
+
 fileprivate class WeakReference<T: AnyObject>: NSObject {
 
     fileprivate init(_ object: T?) {

Sources/InterposeKit/Hooks/HookStrategy/ObjectHookStrategy/ObjectHookStrategy.swift

@@ -143,6 +143,8 @@ internal final class ObjectHookStrategy: HookStrategy {
 
         self.appliedHookIMP = hookIMP
         self.storedOriginalIMP = originalIMP
+        
+        self.object.incrementHookCount()
         ObjectHookRegistry.register(self.handle, for: hookIMP)
 
         Interpose.log("Replaced implementation for -[\(self.class) \(self.selector)] IMP: \(originalIMP) -> \(hookIMP)")
@@ -200,6 +202,11 @@ internal final class ObjectHookStrategy: HookStrategy {
         }
 
         Interpose.log("Restored implementation for -[\(self.class) \(self.selector)] IMP: \(originalIMP)")
+        
+        // Decrement the hook count and if this was the last hook, uninstall the dynamic subclass.
+        if self.object.decrementHookCount() {
+            ObjectSubclassManager.uninstallSubclass(for: object)
+        }
     }
 
     // ============================================================================ //
@@ -233,3 +240,54 @@ internal final class ObjectHookStrategy: HookStrategy {
     }
 
 }
+
+extension NSObject {
+    
+    /// Increments the number of active object-based hooks on this instance by one.
+    fileprivate func incrementHookCount() {
+        self.hookCount += 1
+    }
+    
+    /// Decrements the number of active object-based hooks on this instance by one and returns
+    /// `true` if this was the last hook, or `false` otherwise.
+    fileprivate func decrementHookCount() -> Bool {
+        guard self.hookCount > 0 else { return false }
+        self.hookCount -= 1
+        return self.hookCount == 0
+    }
+    
+    /// The current number of active object-based hooks on this instance.
+    ///
+    /// Internally stored using associated objects. Always returns a non-negative value.
+    private var hookCount: Int {
+        get {
+            guard let count = objc_getAssociatedObject(
+                self,
+                &ObjectHookCountKey
+            ) as? NSNumber else { return 0 }
+            
+            return count.intValue
+        }
+        set {
+            let newCount = max(0, newValue)
+            if newCount == 0 {
+                objc_setAssociatedObject(
+                    self,
+                    &ObjectHookCountKey,
+                    nil,
+                    .OBJC_ASSOCIATION_ASSIGN
+                )
+            } else {
+                objc_setAssociatedObject(
+                    self,
+                    &ObjectHookCountKey,
+                    NSNumber(value: newCount),
+                    .OBJC_ASSOCIATION_RETAIN_NONATOMIC
+                )
+            }
+        }
+    }
+    
+}
+
+private var ObjectHookCountKey: UInt8 = 0

Sources/InterposeKit/Hooks/HookStrategy/ObjectHookStrategy/ObjectSubclassManager.swift

@@ -65,7 +65,29 @@ internal enum ObjectSubclassManager {
     internal static func uninstallSubclass(
         for object: NSObject
     ) {
-        fatalError("Not yet implemented")
+        // Get the InterposeKit-managed dynamic subclass installed on the object.
+        guard let dynamicSubclass = self.installedSubclass(for: object) else { return }
+        
+        // Retrieve the original class (superclass of the dynamic subclass) we want to restore
+        // the object to.
+        guard let originalClass = class_getSuperclass(dynamicSubclass) else { return }
+        
+        // Restore the object’s class to its original class.
+        object_setClass(object, originalClass)
+        
+        Interpose.log({
+            let subclassName = NSStringFromClass(dynamicSubclass)
+            let originalClassName = NSStringFromClass(originalClass)
+            let objectAddress = String(format: "%p", object)
+            return "Removed subclass: \(subclassName), restored \(originalClassName) on object \(objectAddress)"
+        }())
+        
+        // Dispose of the dynamic subclass.
+        //
+        // This is safe to call here because all hooks have been reverted. Unfortunately, we can’t
+        // validate this explicitly, as `objc_disposeClassPair(...)` offers no feedback mechanism
+        // and will silently fail if the subclass is still in use.
+        objc_disposeClassPair(dynamicSubclass)
     }
 
     // ============================================================================ //

Tests/InterposeKitTests/ObjectHookTests.swift

@@ -4,6 +4,7 @@ import XCTest
 fileprivate class ExampleClass: NSObject {
     @objc dynamic var intValue = 1
     @objc dynamic func doSomething() {}
+    @objc dynamic func doSomethingElse() {}
     @objc dynamic var arrayValue: [String] { ["base"] }
 }
 
@@ -48,7 +49,7 @@ final class ObjectHookTests: XCTestCase {
         )
     }
 
-    func testMultipleHooks() throws {
+    func testMultipleHooksOnSingleObject() throws {
         let object = ExampleClass()
         XCTAssertEqual(object.arrayValue, ["base"])
 
@@ -82,7 +83,7 @@ final class ObjectHookTests: XCTestCase {
         XCTAssertEqual(object.arrayValue, ["base"])
     }
 
-    func testHookOnMultipleObjects() throws {
+    func testHooksOnMultipleObjects() throws {
         let object1 = ExampleClass()
         let object2 = ExampleClass()
 
@@ -201,7 +202,7 @@ final class ObjectHookTests: XCTestCase {
         _ = token
     }
 
-    func testCleanUp_implementationPreserved() throws {
+    func testCleanUp_hook_implementationPreserved() throws {
         let object = ExampleClass()
         var deallocated = false
 
@@ -223,7 +224,7 @@ final class ObjectHookTests: XCTestCase {
         XCTAssertFalse(deallocated)
     }
 
-    func testCleanUp_implementationDeallocated() throws {
+    func testCleanUp_hook_implementationDeallocated() throws {
         let object = ExampleClass()
         var deallocated = false
 
@@ -246,5 +247,56 @@ final class ObjectHookTests: XCTestCase {
 
         XCTAssertTrue(deallocated)
     }
+    
+    func testCleanUp_dynamicSubclass() throws {
+        let object = ExampleClass()
+        
+        // Original class
+        XCTAssertTrue(object_getClass(object) == ExampleClass.self)
+        
+        let hook1 = try object.applyHook(
+            for: #selector(ExampleClass.doSomething),
+            methodSignature: (@convention(c) (NSObject, Selector) -> Void).self,
+            hookSignature: (@convention(block) (NSObject) -> Void).self
+        ) { hook in
+            return { `self` in hook.original(self, hook.selector) }
+        }
+        
+        let hook2 = try object.applyHook(
+            for: #selector(ExampleClass.doSomething),
+            methodSignature: (@convention(c) (NSObject, Selector) -> Void).self,
+            hookSignature: (@convention(block) (NSObject) -> Void).self
+        ) { hook in
+            return { `self` in hook.original(self, hook.selector) }
+        }
+        
+        let hook3 = try object.applyHook(
+            for: #selector(ExampleClass.doSomethingElse),
+            methodSignature: (@convention(c) (NSObject, Selector) -> Void).self,
+            hookSignature: (@convention(block) (NSObject) -> Void).self
+        ) { hook in
+            return { `self` in hook.original(self, hook.selector) }
+        }
+        
+        // Dynamic subclass
+        XCTAssertTrue(object_getClass(object) != ExampleClass.self)
+        
+        try hook1.revert()
+        try hook2.revert()
+        
+        // Dynamic subclass
+        XCTAssertTrue(object_getClass(object) != ExampleClass.self)
+        
+        // Back to original subclass after reverting the last hook
+        try hook3.revert()
+        XCTAssertTrue(object_getClass(object) == ExampleClass.self)
+        
+        try hook2.apply()
+        XCTAssertTrue(object_getClass(object) != ExampleClass.self)
+        
+        // Back to original subclass after reverting the last hook
+        try hook2.revert()
+        XCTAssertTrue(object_getClass(object) == ExampleClass.self)
+    }
 
 }