Support weak references for object interposers

CHANGELOG.md

@@ -6,7 +6,7 @@
 
 ##### Enhancements
 
-* None.
+* Support weakly referenced object interposers, avoiding retain cycles when an interposer is associated with its target. Thanks to [@ezoushen](https://github.com/ezoushen).
 
 ##### Bug Fixes
 

README.md

@@ -87,6 +87,22 @@ InterposeKit can hook classes and object. Class hooking is similar to swizzling,
 
 Caveat: Hooking will fail with an error if the object uses KVO or is backed by Core Foundation, such as `NSURL`. These objects rely on runtime behavior that is incompatible with InterposeKit's dynamic subclass. Using KVO after a hook was created is supported and will not cause issues.
 
+Object interposers retain their target by default. Use a weak reference when associating an interposer back to its target, avoiding a retain cycle:
+
+```swift
+let interposer = try Interpose(.weak(testObj)) {
+    try $0.prepareHook(
+        #selector(TestClass.sayHi),
+        methodSignature: (@convention(c) (AnyObject, Selector) -> String).self,
+        hookSignature: (@convention(block) (AnyObject) -> String).self
+    ) { store in
+        { object in store.original(object, store.selector) + "weakly held" }
+    }
+}
+```
+
+`interposer.object` becomes `nil` after a weakly referenced target deallocates. Operations that require the target then throw `InterposeError.objectDeallocated`.
+
 ## Various ways to define the signature
 
 Next to using  `methodSignature` and `hookSignature`, following variants to define the signature are also possible:

Sources/InterposeKit/InterposeError.swift

@@ -30,6 +30,9 @@ public enum InterposeError: LocalizedError {
     /// Core Foundation-backed objects do not support isa-swizzling from Swift.
     case coreFoundationObjectDetected(AnyObject)
 
+    /// A weakly referenced object deallocated before an operation that requires it.
+    case objectDeallocated
+
     /// Object is lying about it's actual class metadata.
     /// This usually happens when other swizzling libraries (like Aspects) also interfere with a class.
     /// While this might just work, it's not worth risking a crash, so similar to KVO this case is rejected.
@@ -73,6 +76,8 @@ extension InterposeError: Equatable {
             return "Unable to hook object that uses Key Value Observing: \(obj)"
         case .coreFoundationObjectDetected(let obj):
             return "Unable to hook Core Foundation-backed object: \(obj)"
+        case .objectDeallocated:
+            return "Unable to hook an object that has been deallocated"
         case .objectPosingAsDifferentClass(let obj, let actualClass):
             return "Unable to hook \(type(of: obj)) posing as \(NSStringFromClass(actualClass))/"
         case .invalidState(let expectedState):

Sources/InterposeKit/InterposeKit.swift

@@ -3,6 +3,25 @@ import CoreFoundation
 #endif
 import Foundation
 
+private final class InterposeObjectReferenceStorage {
+    private let strongObject: AnyObject?
+    private weak var weakObject: AnyObject?
+
+    init(strong object: AnyObject) {
+        strongObject = object
+        weakObject = nil
+    }
+
+    init(weak object: AnyObject) {
+        strongObject = nil
+        weakObject = object
+    }
+
+    var object: AnyObject? {
+        strongObject ?? weakObject
+    }
+}
+
 extension NSObject {
     /// Hook an `@objc dynamic` instance method via selector  on the current object or class..
     @discardableResult public func hook<MethodSignature, HookSignature> (
@@ -34,13 +53,45 @@ extension NSObject {
 /// Methods are hooked via replacing the implementation, instead of the usual exchange.
 /// Supports both swizzling classes and individual objects.
 final public class Interpose {
+    /// Controls whether an object-based interposer retains its target.
+    public struct ObjectReference {
+        private let storage: InterposeObjectReferenceStorage
+
+        init(strong object: AnyObject) {
+            storage = InterposeObjectReferenceStorage(strong: object)
+        }
+
+        init(weak object: AnyObject) {
+            storage = InterposeObjectReferenceStorage(weak: object)
+        }
+
+        /// Creates a reference that retains the target object.
+        public static func strong(_ object: NSObject) -> ObjectReference {
+            ObjectReference(strong: object)
+        }
+
+        /// Creates a reference that allows the target object to deallocate.
+        public static func weak(_ object: NSObject) -> ObjectReference {
+            ObjectReference(weak: object)
+        }
+
+        /// The target object, or `nil` after a weakly referenced target deallocates.
+        public var object: AnyObject? {
+            storage.object
+        }
+    }
+
     /// Stores swizzle hooks and executes them at once.
     public let `class`: AnyClass
     /// Lists all hooks for the current interpose class object.
     public private(set) var hooks: [AnyHook] = []
 
     /// If Interposing is object-based, this is set.
-    public let object: AnyObject?
+    public var object: AnyObject? {
+        objectReference?.object
+    }
+
+    private let objectReference: ObjectReference?
 
     // Checks if a object is posing as a different class
     // via implementing 'class' and returning something else.
@@ -87,7 +138,7 @@ final public class Interpose {
     /// If `builder` is present, `apply()` is automatically called.
     public init(_ `class`: AnyClass, builder: ((Interpose) throws -> Void)? = nil) throws {
         self.class = `class`
-        self.object = nil
+        self.objectReference = nil
 
         // Only apply if a builder is present
         if let builder = builder {
@@ -96,10 +147,18 @@ final public class Interpose {
     }
 
     /// Initialize with a single object to interpose.
-    public init(_ object: NSObject, builder: ((Interpose) throws -> Void)? = nil) throws {
-        self.object = object
-        self.class = type(of: object)
+    public convenience init(_ object: NSObject, builder: ((Interpose) throws -> Void)? = nil) throws {
+        try self.init(.strong(object), builder: builder)
+    }
 
+    /// Initialize with a strong or weak reference to a single object.
+    public init(_ objectReference: ObjectReference, builder: ((Interpose) throws -> Void)? = nil) throws {
+        guard let object = objectReference.object else {
+            throw InterposeError.objectDeallocated
+        }
+
+        self.objectReference = objectReference
+        self.class = type(of: object)
         try Self.validateObjectForHooking(object)
 
         // Only apply if a builder is present
@@ -145,8 +204,9 @@ final public class Interpose {
         _ implementation: (TypedHook<MethodSignature, HookSignature>) -> HookSignature?)
         throws -> TypedHook<MethodSignature, HookSignature> {
             var hook: TypedHook<MethodSignature, HookSignature>
-            if let object = self.object {
-                hook = try ObjectHook(object: object, selector: selector, implementation: implementation)
+            if let objectReference = self.objectReference {
+                hook = try ObjectHook(
+                    objectReference: objectReference, selector: selector, implementation: implementation)
             } else {
                 hook = try ClassHook(class: `class`, selector: selector, implementation: implementation)
             }

Sources/InterposeKit/InterposeSubclass.swift

@@ -18,9 +18,6 @@ class InterposeSubclass {
         }
     }
 
-    /// The object that is being hooked.
-    let object: AnyObject
-
     /// Subclass that we create on the fly
     private(set) var dynamicClass: AnyClass
 
@@ -30,12 +27,11 @@ class InterposeSubclass {
     /// Making KVO and Object-based hooking work at the same time is difficult.
     /// If we make a dynamic subclass over KVO, invalidating the token crashes in cache_getImp.
     init(object: AnyObject) throws {
-        self.object = object
         dynamicClass = type(of: object) // satisfy set to something
-        dynamicClass = try getExistingSubclass() ?? createSubclass()
+        dynamicClass = try Self.getExistingSubclass(for: object) ?? Self.createSubclass(for: object)
     }
 
-    private func createSubclass() throws -> AnyClass {
+    private static func createSubclass(for object: AnyObject) throws -> AnyClass {
         let perceivedClass: AnyClass = type(of: object)
         let actualClass: AnyClass = object_getClass(object)!
 
@@ -50,7 +46,7 @@ class InterposeSubclass {
                 return existingClass
             } else {
                 guard let subclass: AnyClass = objc_allocateClassPair(actualClass, cString, 0) else { return nil }
-                replaceGetClass(in: subclass, decoy: perceivedClass)
+                Self.replaceGetClass(in: subclass, decoy: perceivedClass)
                 objc_registerClassPair(subclass)
                 return subclass
             }
@@ -67,7 +63,7 @@ class InterposeSubclass {
     }
 
     /// We need to reuse a dynamic subclass if the object already has one.
-    private func getExistingSubclass() -> AnyClass? {
+    private static func getExistingSubclass(for object: AnyObject) -> AnyClass? {
         let actualClass: AnyClass = object_getClass(object)!
         if Self.isInterposeSubclass(actualClass) {
             return actualClass
@@ -80,7 +76,7 @@ class InterposeSubclass {
     }
 
     #if !os(Linux)
-    private func replaceGetClass(in class: AnyClass, decoy perceivedClass: AnyClass) {
+    private static func replaceGetClass(in class: AnyClass, decoy perceivedClass: AnyClass) {
         // crashes on linux
         let getClass: @convention(block) (AnyObject) -> AnyClass = { _ in
             perceivedClass
@@ -112,6 +108,6 @@ class InterposeSubclass {
     #else
     func addSuperTrampoline(selector: Selector) { }
     class var supportsSuperTrampolines: Bool { return false }
-    private func replaceGetClass(in class: AnyClass, decoy perceivedClass: AnyClass) {}
+    private static func replaceGetClass(in class: AnyClass, decoy perceivedClass: AnyClass) {}
     #endif
 }

Sources/InterposeKit/ObjectHook.swift

@@ -7,7 +7,14 @@ extension Interpose {
     final public class ObjectHook<MethodSignature, HookSignature>: TypedHook<MethodSignature, HookSignature> {
 
         /// The object that is being hooked.
-        public let object: AnyObject
+        public var object: AnyObject {
+            guard let object = objectReference.object else {
+                preconditionFailure("The interposed object has been deallocated")
+            }
+            return object
+        }
+
+        private let objectReference: Interpose.ObjectReference
 
         /// Subclass that we create on the fly
         var interposeSubclass: InterposeSubclass?
@@ -16,10 +23,21 @@ extension Interpose {
         let generatesSuperIMP = InterposeSubclass.supportsSuperTrampolines
 
         /// Initialize a new hook to interpose an instance method.
-        public init(object: AnyObject, selector: Selector,
-                    implementation: (ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
+        public convenience init(object: AnyObject, selector: Selector,
+                                implementation: (ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
+            try self.init(
+                objectReference: Interpose.ObjectReference(strong: object),
+                selector: selector,
+                implementation: implementation)
+        }
+
+        init(objectReference: Interpose.ObjectReference, selector: Selector,
+             implementation: (ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
+            guard let object = objectReference.object else {
+                throw InterposeError.objectDeallocated
+            }
             try Interpose.validateObjectForHooking(object)
-            self.object = object
+            self.objectReference = objectReference
             try super.init(class: type(of: object), selector: selector)
             let block = implementation(self) as AnyObject
             try validateImplementationBlock(block)
@@ -85,6 +103,9 @@ extension Interpose {
 
         override func replaceImplementation() throws {
             let method = try validate()
+            guard let object = objectReference.object else {
+                throw InterposeError.objectDeallocated
+            }
 
             // Check if there's an existing subclass we can reuse.
             // Create one at runtime if there is none.
@@ -180,7 +201,8 @@ extension Interpose {
 #if DEBUG
 extension Interpose.ObjectHook: CustomDebugStringConvertible {
     public var debugDescription: String {
-        return "\(selector) of \(object) -> \(String(describing: original))"
+        let objectDescription = objectReference.object.map(String.init(describing:)) ?? "<deallocated>"
+        return "\(selector) of \(objectDescription) -> \(String(describing: original))"
     }
 }
 #endif

Tests/InterposeKitTests/ObjectInterposeTests.swift

@@ -4,6 +4,54 @@
 
 final class ObjectInterposeTests: InterposeKitTestCase {
 
+    #if !os(Linux)
+    func testWeakInterposeBreaksAssociatedObjectCycle() throws {
+        weak var weakObject: TestClass?
+        var associationKey: UInt8 = 0
+
+        try autoreleasepool {
+            let object = TestClass()
+            weakObject = object
+            let interpose = try Interpose(.weak(object)) {
+                try $0.prepareHook(
+                    #selector(TestClass.sayHi),
+                    methodSignature: (@convention(c) (AnyObject, Selector) -> String).self,
+                    hookSignature: (@convention(block) (AnyObject) -> String).self
+                ) { store in
+                    { object in store.original(object, store.selector) }
+                }
+            }
+            objc_setAssociatedObject(object, &associationKey, interpose, .OBJC_ASSOCIATION_RETAIN)
+
+            XCTAssertNotNil(interpose.object)
+            XCTAssertEqual(object.sayHi(), testClassHi)
+        }
+
+        XCTAssertNil(weakObject)
+    }
+
+    func testWeakInterposeReportsDeallocatedObject() throws {
+        var interpose: Interpose?
+
+        autoreleasepool {
+            let object = TestClass()
+            interpose = try? Interpose(.weak(object))
+            XCTAssertNotNil(interpose?.object)
+        }
+
+        XCTAssertNil(interpose?.object)
+        XCTAssertThrowsError(try interpose?.prepareHook(
+            #selector(TestClass.sayHi),
+            methodSignature: (@convention(c) (AnyObject, Selector) -> String).self,
+            hookSignature: (@convention(block) (AnyObject) -> String).self
+        ) { _ in
+            { _ in testClassHi }
+        }) { error in
+            XCTAssertEqual(error as? InterposeError, .objectDeallocated)
+        }
+    }
+    #endif
+
     func testRejectsMismatchedHookReturnType() {
         let testObj = TestClass()
 
@@ -12,7 +60,7 @@
             methodSignature: (@convention(c) (AnyObject, Selector) -> Int).self,
             hookSignature: (@convention(block) (AnyObject) -> Double).self
         ) { _ in
             { _ in 1.0 }
         }) { error in
             guard case InterposeError.incompatibleHookSignature = error else {
                 return XCTFail("Unexpected error: \(error)")
@@ -28,7 +76,7 @@
             methodSignature: (@convention(c) (AnyObject, Selector, String) -> String).self,
             hookSignature: (@convention(block) (AnyObject, Int) -> String).self
         ) { _ in
             { _, _ in "" }
         }) { error in
             guard case InterposeError.incompatibleHookSignature = error else {
                 return XCTFail("Unexpected error: \(error)")
@@ -48,7 +96,7 @@
             methodSignature: (@convention(c) (AnyObject, Selector) -> String).self,
             hookSignature: (@convention(block) (AnyObject) -> String).self
         ) { _ in
             { _ in "www.facebook.com" }
         }) { error in
             XCTAssertEqual(error as? InterposeError, expectedError)
         }