steipete · doggy · Feb 15, 2017 · Mar 3, 2017 · Mar 3, 2017 · Oct 31, 2017
diff --git a/Aspects.m b/Aspects.m
@@ -279,6 +279,13 @@ static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSE
         if (![klass instancesRespondToSelector:aliasSelector]) {
             __unused BOOL addedAlias = class_addMethod(klass, aliasSelector, method_getImplementation(targetMethod), typeEncoding);
             NSCAssert(addedAlias, @"Original implementation for %@ is already copied to %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass);
+        } else {
+            // update exist alias method
+            Method aliasMethod = class_getInstanceMethod(klass, aliasSelector);
+            IMP aliasMethodIMP = method_getImplementation(aliasMethod);
+            if (aliasMethodIMP != targetMethodIMP) {
+                class_replaceMethod(klass, aliasSelector, targetMethodIMP, typeEncoding);
+            }
         }
 
         // We use forwardInvocation to hook in.
@@ -458,6 +465,79 @@ static void aspect_undoSwizzleClassInPlace(Class klass) {
     });
 }
 
+static BOOL aspect_invokeAlias(NSInvocation *invocation, SEL originalSelector)
+{
+    BOOL respondsToAlias = YES;
+    Class klass = object_getClass(invocation.target);
+    SEL aliasSelector = invocation.selector;
+
+    do {
+        if ((respondsToAlias = [klass instancesRespondToSelector:aliasSelector])) {
+            // invoke alias method
+            Method aliasInvocationMethod = class_getInstanceMethod(klass, aliasSelector);
+            IMP aliasInvocation = method_getImplementation(aliasInvocationMethod);
+
+            const char *typeEncoding = method_getTypeEncoding(aliasInvocationMethod);
+            IMP msgForwardInvocation = class_replaceMethod(klass, originalSelector, aliasInvocation, typeEncoding);
+
+            invocation.selector = originalSelector;
+            [invocation invoke];
+
+            class_replaceMethod(klass, originalSelector, msgForwardInvocation, typeEncoding);
+            invocation.selector = aliasSelector;
+            break;
+        }
+    }while (!respondsToAlias && (klass = class_getSuperclass(klass)));
+
+    return respondsToAlias;
+}
+
+static BOOL aspect_invokeOriginalForwarder(__unsafe_unretained NSObject *self, NSInvocation *invocation)
+{
+    SEL forwardInvocationSEL = @selector(forwardInvocation:);
+    SEL originalForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
+
+    // grab original ForwardInvocation saved previously
+    BOOL respondsToParent = [self respondsToSelector:originalForwardInvocationSEL];
+
+    // trying to find implementation on parent(s)
+    if (! respondsToParent) {
+        Method dummyObjectMethod = class_getInstanceMethod(NSObject.class, forwardInvocationSEL);
+        IMP dummyImplementation = method_getImplementation(dummyObjectMethod);
+
+        Class klass = object_getClass(invocation.target);
+
+        BOOL aspectsFound = NO;
+
+        do {
+            if ([klass instancesRespondToSelector:forwardInvocationSEL]) {
+                // skip Aspects' forwardInvocation method(s)
+                Method parentInvocationMethod = class_getInstanceMethod(klass, forwardInvocationSEL);
+                IMP parentForwarder = method_getImplementation(parentInvocationMethod);
+
+                // skip until found the aspectForwarder
+                if (aspectsFound) {
+                    if (dummyImplementation != parentForwarder) {
+                        // setup forwarder
+                        const char *typeEncoding = method_getTypeEncoding(parentInvocationMethod);
+                        class_replaceMethod(klass, originalForwardInvocationSEL, parentForwarder, typeEncoding);
+                        respondsToParent = YES;
+                        break;
+                    }
+                }else {
+                    aspectsFound = ((IMP)__ASPECTS_ARE_BEING_CALLED__ == parentForwarder);
+                }
+            }
+        }while ((klass = class_getSuperclass(klass)));
+    }
+
+    // ... then call it
+    if (respondsToParent) {
+        ((void( *)(id, SEL, NSInvocation *))objc_msgSend)(self, originalForwardInvocationSEL, invocation);
+    }
+    return respondsToParent;
+}
+
 ///////////////////////////////////////////////////////////////////////////////////////////
 #pragma mark - Aspect Invoke Point
 
@@ -492,28 +572,21 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
         aspect_invoke(classContainer.insteadAspects, info);
         aspect_invoke(objectContainer.insteadAspects, info);
     }else {
-        Class klass = object_getClass(invocation.target);
-        do {
-            if ((respondsToAlias = [klass instancesRespondToSelector:aliasSelector])) {
-                [invocation invoke];
-                break;
-            }
-        }while (!respondsToAlias && (klass = class_getSuperclass(klass)));
+        respondsToAlias = aspect_invokeAlias(invocation, originalSelector);
+        // If no hooks are installed, try to call original implementation (usually to throw an exception)
+        if (!respondsToAlias) {
+            invocation.selector = originalSelector;
+            respondsToAlias = aspect_invokeOriginalForwarder(self, invocation);
+        }
     }
 
     // After hooks.
     aspect_invoke(classContainer.afterAspects, info);
     aspect_invoke(objectContainer.afterAspects, info);
 
-    // If no hooks are installed, call original implementation (usually to throw an exception)
     if (!respondsToAlias) {
-        invocation.selector = originalSelector;
-        SEL originalForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
-        if ([self respondsToSelector:originalForwardInvocationSEL]) {
-            ((void( *)(id, SEL, NSInvocation *))objc_msgSend)(self, originalForwardInvocationSEL, invocation);
-        }else {
-            [self doesNotRecognizeSelector:invocation.selector];
-        }
+        // Throw an exception
+        [self doesNotRecognizeSelector:invocation.selector];
     }
 
     // Remove any hooks that are queued for deregistration.

diff --git a/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m b/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m
@@ -12,7 +12,6 @@
 
 @interface TestClass : NSObject
 @property (nonatomic, copy) NSString *string;
-@property (nonatomic, assign) BOOL kvoTestCalled;
 - (void)testCall;
 - (void)testCallAndExecuteBlock:(dispatch_block_t)block;
 - (double)callReturnsDouble;
@@ -551,10 +550,37 @@ - (void)testAutoDeregistration {
     XCTAssertFalse([aspectToken remove], @"Must not able to deregister again");
 }
 
+@end
+
 ///////////////////////////////////////////////////////////////////////////////////////////
 #pragma mark - Test KVO
 
+@protocol TestKVOClassProtocol <NSObject>
+@property (nonatomic, assign) BOOL kvoTestCalled;
+@end
+
+#define TestKVOClass(suffix) \
+    Test##suffix
+
+#define declareTestKVOClass(suffix) \
+    @interface TestKVOClass(suffix) : NSObject <TestKVOClassProtocol> \
+    @property (nonatomic, copy) NSString *string; \
+    @property (nonatomic, assign) BOOL kvoTestCalled; \
+    @end \
+    @implementation TestKVOClass(suffix) @end
+
+declareTestKVOClass(InstanceHookThenKVO)
+declareTestKVOClass(KVOThenInstanceHook)
+declareTestKVOClass(ClassHookThenKVO)
+declareTestKVOClass(KVOThenClassHook)
+
+@interface AspectsKVOTests : XCTestCase @end
+@implementation AspectsKVOTests
+
 - (void)testKVOCoexistance {
+#pragma push_macro( "TestClass" )
+#define TestClass TestKVOClass(InstanceHookThenKVO)
+
     TestClass *testClass = [TestClass new];
 
     __block BOOL hookCalled = NO;
@@ -576,37 +602,140 @@ - (void)testKVOCoexistance {
     XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
 
     XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
+
+#pragma pop_macro( "TestClass" )
 }
 
-// TODO: Pre-registeded KVO is currently not working.
-//- (void)testKVOCoexistanceWithPreregisteredKVO {
-//    TestClass *testClass = [TestClass new];
-//    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must be not set");
-//    [testClass addObserver:self forKeyPath:NSStringFromSelector(@selector(string)) options:0 context:_cmd];
-//    testClass.string = @"test";
-//    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
-//
-//    __block BOOL hookCalled = NO;
-//    [testClass aspect_hookSelector:@selector(setString:) withOptions:AspectPositionAfter usingBlock:^(id instance, NSArray *arguments) {
-//        NSLog(@"Aspect hook!");
-//        hookCalled = YES;
-//    }];
-//
-//    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must be not set");
-//    testClass.string = @"test";
-//    XCTAssertTrue(hookCalled, @"Hook must be called");
-//    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
-//    [testClass removeObserver:self forKeyPath:NSStringFromSelector(@selector(string)) context:_cmd];
-//    hookCalled = NO;
-//    testClass.kvoTestCalled = NO;
-//    testClass.string = @"test2";
-//    XCTAssertTrue(hookCalled, @"Hook must be called");
-//    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
-//}
+// Pre-registeded KVO
+- (void)testKVOCoexistanceWithPreregisteredKVO {
+#pragma push_macro( "TestClass" )
+#define TestClass TestKVOClass(KVOThenInstanceHook)
+
+    TestClass *testClass = [TestClass new];
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must be not set");
+
+    // Step 1: kvo dynamic-subclassing
+    [testClass addObserver:self forKeyPath:NSStringFromSelector(@selector(string)) options:0 context:_cmd];
+    [testClass addObserver:self forKeyPath:NSStringFromSelector(@selector(kvoTestCalled)) options:0 context:_cmd];
+
+    // Step 2: kvo validation
+    testClass.string = @"test";
+    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
+
+    // Step 3: Instance-hooking
+    __block BOOL hookCalled = NO;
+    id aspectToken = [testClass aspect_hookSelector:@selector(setString:) withOptions:AspectPositionBefore usingBlock:^(id instance, NSArray *arguments) {
+        NSLog(@"Aspect hook!");
+        hookCalled = YES;
+    } error:nil];
+
+    // Step 4: call w/ Observer
+    testClass.kvoTestCalled = NO;
+    testClass.string = @"test";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
+
+    // Step 5: call w/o Observer
+    [testClass removeObserver:self forKeyPath:NSStringFromSelector(@selector(string)) context:_cmd];
+    hookCalled = NO;
+    testClass.kvoTestCalled = NO;
+    testClass.string = @"test2";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
+
+    // Step 6: aspect restore
+    XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
+
+    // Step 7: KVO isa-swizzing restored
+    [testClass removeObserver:self forKeyPath:NSStringFromSelector(@selector(kvoTestCalled)) context:_cmd];
+    hookCalled = NO;
+    testClass.kvoTestCalled = NO;
+    XCTAssertFalse(hookCalled, @"Hook must be not called");
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
+
+#pragma pop_macro( "TestClass" )
+}
+
+#pragma mark - Test KVO with class hook
+- (void)testKVOClassHookCoexistence {
+#pragma push_macro( "TestClass" )
+#define TestClass TestKVOClass(ClassHookThenKVO)
+
+    // Step 1: Class-hooking
+    __block BOOL hookCalled = NO;
+    id aspectToken = [TestClass aspect_hookSelector:@selector(setString:) withOptions:AspectPositionAfter usingBlock:^(id<AspectInfo> info, NSString *string) {
+        NSLog(@"Aspect hook!");
+        hookCalled = YES;
+    } error:NULL];
+
+    // Step 2: kvo dynamic-subclassing
+    TestClass *testClass = [TestClass new];
+    [testClass addObserver:self forKeyPath:NSStringFromSelector(@selector(string)) options:0 context:_cmd];
+
+    // Step 3: validation
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must be not set");
+
+    // Step 3.1: call w/ Observer
+    testClass.string = @"test";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
+
+    // Step 3.2: call w/o Observer
+    [testClass removeObserver:self forKeyPath:NSStringFromSelector(@selector(string)) context:_cmd];
+    hookCalled = NO;
+    testClass.kvoTestCalled = NO;
+    testClass.string = @"test2";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
+
+    XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
+
+#pragma pop_macro( "TestClass" )
+}
 
-- (void)observeValueForKeyPath:(NSString *)keyPath ofObject:(id)object change:(NSDictionary *)change context:(void *)context {
-    NSLog(@"KVO!");
-    ((TestClass *)object).kvoTestCalled = YES;
+// Pre-registeded KVO
+- (void)testKVOClassHookCoexistenceWithPreregisteredKVO {
+#pragma push_macro( "TestClass" )
+#define TestClass TestKVOClass(KVOThenClassHook)
+
+    // Step 1: kvo dynamic-subclassing
+    TestClass *testClass = [TestClass new];
+    [testClass addObserver:self forKeyPath:NSStringFromSelector(@selector(string)) options:0 context:_cmd];
+
+    // Step 2: Class-hooking
+    __block BOOL hookCalled = NO;
+    id aspectToken = [TestClass aspect_hookSelector:@selector(setString:) withOptions:AspectPositionAfter usingBlock:^(id<AspectInfo> info, NSString *string) {
+        NSLog(@"Aspect hook!");
+        hookCalled = YES;
+    } error:NULL];
+
+    // Step 3: validation
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must be not set");
+
+    // Step 3.1: call w/ Observer
+    testClass.string = @"test";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertTrue(testClass.kvoTestCalled, @"KVO must work");
+
+    // Step 3.2: call w/o Observer
+    [testClass removeObserver:self forKeyPath:NSStringFromSelector(@selector(string)) context:_cmd];
+    hookCalled = NO;
+    testClass.kvoTestCalled = NO;
+    testClass.string = @"test2";
+    XCTAssertTrue(hookCalled, @"Hook must be called");
+    XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");
+
+    XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
+
+#pragma pop_macro( "TestClass" )
+}
+
+
+- (void)observeValueForKeyPath:(NSString *)keyPath ofObject:(id<TestKVOClassProtocol>)object change:(NSDictionary *)change context:(void *)context {
+    if (! [keyPath isEqualToString:NSStringFromSelector(@selector(kvoTestCalled))]) {
+        NSLog(@"KVO!");
+        object.kvoTestCalled = YES;
+    }
 }
 
 @end

diff --git a/README.md b/README.md
@@ -160,8 +160,6 @@ Aspects uses quite some runtime trickery to achieve what it does. You can mostly
 An important limitation is that for class-based hooking, a method can only be hooked once within the subclass hierarchy. [See #2](https://github.com/steipete/Aspects/issues/2)
 This does not apply for objects that are hooked. Aspects creates a dynamic subclass here and has full control.
 
-KVO works if observers are created after your calls `aspect_hookSelector:` It most likely will crash the other way around. Still looking for workarounds here - any help appreciated.
-
 Because of ugly implementation details on the ObjC runtime, methods that return unions that also contain structs might not work correctly unless this code runs on the arm64 runtime.
 
 Credits
@@ -183,6 +181,10 @@ MIT licensed, Copyright (c) 2014 Peter Steinberger, [email protected], [@steipe
 Release Notes
 -----------------
 
+Version 1.4.3
+
+- Works with KVO. Only with one limitation: Aspects works on a pre-registeded KVO'd object until all observers has been removed. More detailed info at `TestCase:testKVOCoexistanceWithPreregisteredKVO` and [PR #115](https://github.com/steipete/Aspects/pull/115)
+
 Version 1.4.2
 
 - Allow to hook different subclasses.