From 8385c74ddc3194e135bd87e04e26b8a0dc7f66fa Mon Sep 17 00:00:00 2001 From: leikun Date: Mon, 6 Jun 2016 18:56:26 +0800 Subject: [PATCH 1/3] fix AspectToken remove bug --- Aspects.m | 118 ++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 78 insertions(+), 40 deletions(-) diff --git a/Aspects.m b/Aspects.m index c907066..93d7d17 100644 --- a/Aspects.m +++ b/Aspects.m @@ -90,6 +90,7 @@ - (NSArray *)aspects_arguments; NSString *const AspectErrorDomain = @"AspectErrorDomain"; static NSString *const AspectsSubclassSuffix = @"_Aspects_"; static NSString *const AspectsMessagePrefix = @"aspects_"; +void *AspectsAliasSelecterName2ContainerDictionary = &AspectsAliasSelecterName2ContainerDictionary; @implementation NSObject (Aspects) @@ -272,6 +273,7 @@ static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSE Class klass = aspect_hookClass(self, error); Method targetMethod = class_getInstanceMethod(klass, selector); IMP targetMethodIMP = method_getImplementation(targetMethod); + NSString *className = NSStringFromClass(klass); if (!aspect_isMsgForwardIMP(targetMethodIMP)) { // Make a method alias for the existing method implementation, it not already copied. const char *typeEncoding = method_getTypeEncoding(targetMethod); @@ -284,7 +286,16 @@ static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSE // We use forwardInvocation to hook in. class_replaceMethod(klass, selector, aspect_getMsgForwardIMP(self, selector), typeEncoding); AspectLog(@"Aspects: Installed hook for -[%@ %@].", klass, NSStringFromSelector(selector)); + } else if ([className hasSuffix:AspectsSubclassSuffix]) { + /** + * global selector hack -> singal instance selector hack -> remove global selector hack + * it will clean up forward method in class, but also in class_Aspects_ + * it should have 2 methods in class and class_Aspects_ + */ + const char *typeEncoding = method_getTypeEncoding(targetMethod); + class_addMethod(klass, selector, method_getImplementation(targetMethod), typeEncoding); } + } // Will undo the runtime changes made. @@ -297,8 +308,47 @@ static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) { if (isMetaClass) { klass = (Class)self; } + + // Deregister global tracked selector + aspect_deregisterTrackedSelector(self, selector); + // Get the aspect container and check if there are any hooks remaining. Clean up if there are not. + AspectsContainer *container = aspect_getContainerForObject(self, selector); + if (!container.hasAspects) { + // Destroy the container + aspect_destroyContainerForObject(self, selector); + } + + if (isMetaClass) { + aspect_cleanupHookedForwardSelector(klass, selector); + // if there is no global selector hack + if (aspect_getAliasSelectorName2ContainerDictionaryForObject(self).count == 0) { + // Class is most likely swizzled in place. Undo that. + aspect_undoSwizzleClassInPlace((Class)self); + } + } else { + if (aspect_getContainerForClass(klass, selector) == nil) { + aspect_cleanupHookedForwardSelector(klass, selector); + } + //recover class + if (aspect_getAliasSelectorName2ContainerDictionaryForObject(self).count == 0) { + // Figure out how the class was modified to undo the changes. + NSString *className = NSStringFromClass(klass); + if ([className hasSuffix:AspectsSubclassSuffix]) { + Class originalClass = NSClassFromString([className stringByReplacingOccurrencesOfString:AspectsSubclassSuffix withString:@""]); + NSCAssert(originalClass != nil, @"Original class must exist"); + object_setClass(self, originalClass); + AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(originalClass)); + + // We can only dispose the class pair if we can ensure that no instances exist using our subclass. + // Since we don't globally track this, we can't ensure this - but there's also not much overhead in keeping it around. + //objc_disposeClassPair(object.class); + } + } + } +} - // Check if the method is marked as forwarded and undo that. +// clean up forward method +static void aspect_cleanupHookedForwardSelector(Class klass, SEL selector) { Method targetMethod = class_getInstanceMethod(klass, selector); IMP targetMethodIMP = method_getImplementation(targetMethod); if (aspect_isMsgForwardIMP(targetMethodIMP)) { @@ -308,40 +358,10 @@ static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) { Method originalMethod = class_getInstanceMethod(klass, aliasSelector); IMP originalIMP = method_getImplementation(originalMethod); NSCAssert(originalMethod, @"Original implementation for %@ not found %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass); - + class_replaceMethod(klass, selector, originalIMP, typeEncoding); AspectLog(@"Aspects: Removed hook for -[%@ %@].", klass, NSStringFromSelector(selector)); } - - // Deregister global tracked selector - aspect_deregisterTrackedSelector(self, selector); - - // Get the aspect container and check if there are any hooks remaining. Clean up if there are not. - AspectsContainer *container = aspect_getContainerForObject(self, selector); - if (!container.hasAspects) { - // Destroy the container - aspect_destroyContainerForObject(self, selector); - - // Figure out how the class was modified to undo the changes. - NSString *className = NSStringFromClass(klass); - if ([className hasSuffix:AspectsSubclassSuffix]) { - Class originalClass = NSClassFromString([className stringByReplacingOccurrencesOfString:AspectsSubclassSuffix withString:@""]); - NSCAssert(originalClass != nil, @"Original class must exist"); - object_setClass(self, originalClass); - AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(originalClass)); - - // We can only dispose the class pair if we can ensure that no instances exist using our subclass. - // Since we don't globally track this, we can't ensure this - but there's also not much overhead in keeping it around. - //objc_disposeClassPair(object.class); - }else { - // Class is most likely swizzled in place. Undo that. - if (isMetaClass) { - aspect_undoSwizzleClassInPlace((Class)self); - }else if (self.class != klass) { - aspect_undoSwizzleClassInPlace(klass); - } - } - } } /////////////////////////////////////////////////////////////////////////////////////////// @@ -477,7 +497,7 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL SEL originalSelector = invocation.selector; SEL aliasSelector = aspect_aliasForSelector(invocation.selector); invocation.selector = aliasSelector; - AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector); + AspectsContainer *objectContainer = aspect_getContainerForObject(self, aliasSelector); AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector); AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation]; NSArray *aspectsToRemove = nil; @@ -524,14 +544,27 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL /////////////////////////////////////////////////////////////////////////////////////////// #pragma mark - Aspect Container Management +static NSMutableDictionary *aspect_getAliasSelectorName2ContainerDictionaryForObject(NSObject *self) { + NSMutableDictionary *aliasSelectorName2ContainerDictionary = objc_getAssociatedObject(self, AspectsAliasSelecterName2ContainerDictionary); + if (!aliasSelectorName2ContainerDictionary) { + aliasSelectorName2ContainerDictionary = [NSMutableDictionary dictionary]; + objc_setAssociatedObject(self, AspectsAliasSelecterName2ContainerDictionary, aliasSelectorName2ContainerDictionary, OBJC_ASSOCIATION_RETAIN); + } + return aliasSelectorName2ContainerDictionary; +} + // Loads or creates the aspect container. static AspectsContainer *aspect_getContainerForObject(NSObject *self, SEL selector) { NSCParameterAssert(self); - SEL aliasSelector = aspect_aliasForSelector(selector); - AspectsContainer *aspectContainer = objc_getAssociatedObject(self, aliasSelector); + SEL aliasSelector = selector; + if (![NSStringFromSelector(selector) hasPrefix:AspectsMessagePrefix]) { + aliasSelector = aspect_aliasForSelector(selector); + } + NSString *selectorName = NSStringFromSelector(aliasSelector); + AspectsContainer *aspectContainer = aspect_getAliasSelectorName2ContainerDictionaryForObject(self)[selectorName]; if (!aspectContainer) { aspectContainer = [AspectsContainer new]; - objc_setAssociatedObject(self, aliasSelector, aspectContainer, OBJC_ASSOCIATION_RETAIN); + [aspect_getAliasSelectorName2ContainerDictionaryForObject(self) setObject:aspectContainer forKey:selectorName]; } return aspectContainer; } @@ -540,17 +573,22 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL NSCParameterAssert(klass); AspectsContainer *classContainer = nil; do { - classContainer = objc_getAssociatedObject(klass, selector); - if (classContainer.hasAspects) break; + AspectsContainer *container = aspect_getContainerForObject((NSObject *)klass, selector); + if (container.hasAspects) { + classContainer = container; + break; + } }while ((klass = class_getSuperclass(klass))); - + return classContainer; } static void aspect_destroyContainerForObject(id self, SEL selector) { NSCParameterAssert(self); SEL aliasSelector = aspect_aliasForSelector(selector); - objc_setAssociatedObject(self, aliasSelector, nil, OBJC_ASSOCIATION_RETAIN); + NSString *selectorName = NSStringFromSelector(aliasSelector); + NSMutableDictionary *aliasSelectorName2ContainerDictionary = aspect_getAliasSelectorName2ContainerDictionaryForObject(self); + [aliasSelectorName2ContainerDictionary removeObjectForKey:selectorName]; } /////////////////////////////////////////////////////////////////////////////////////////// From cbd833bccb266ddfcda616497a1d179dbb6f992f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=90=A8?= Date: Sat, 11 Jun 2016 15:46:12 +0800 Subject: [PATCH 2/3] modify repeat create AspectsContainer --- Aspects.m | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/Aspects.m b/Aspects.m index 93d7d17..cd7d9dd 100644 --- a/Aspects.m +++ b/Aspects.m @@ -312,7 +312,7 @@ static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) { // Deregister global tracked selector aspect_deregisterTrackedSelector(self, selector); // Get the aspect container and check if there are any hooks remaining. Clean up if there are not. - AspectsContainer *container = aspect_getContainerForObject(self, selector); + AspectsContainer *container = aspect_getContainerForObjectWithoutCreate(self, selector); if (!container.hasAspects) { // Destroy the container aspect_destroyContainerForObject(self, selector); @@ -497,7 +497,7 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL SEL originalSelector = invocation.selector; SEL aliasSelector = aspect_aliasForSelector(invocation.selector); invocation.selector = aliasSelector; - AspectsContainer *objectContainer = aspect_getContainerForObject(self, aliasSelector); + AspectsContainer *objectContainer = aspect_getContainerForObjectWithoutCreate(self, aliasSelector); AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector); AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation]; NSArray *aspectsToRemove = nil; @@ -569,11 +569,22 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL return aspectContainer; } +static AspectsContainer *aspect_getContainerForObjectWithoutCreate(NSObject *self, SEL selector) { + NSCParameterAssert(self); + SEL aliasSelector = selector; + if (![NSStringFromSelector(selector) hasPrefix:AspectsMessagePrefix]) { + aliasSelector = aspect_aliasForSelector(selector); + } + NSString *selectorName = NSStringFromSelector(aliasSelector); + AspectsContainer *aspectContainer = aspect_getAliasSelectorName2ContainerDictionaryForObject(self)[selectorName]; + return aspectContainer; +} + static AspectsContainer *aspect_getContainerForClass(Class klass, SEL selector) { NSCParameterAssert(klass); AspectsContainer *classContainer = nil; do { - AspectsContainer *container = aspect_getContainerForObject((NSObject *)klass, selector); + AspectsContainer *container = aspect_getContainerForObjectWithoutCreate((id)klass, selector); if (container.hasAspects) { classContainer = container; break; From c0a0f49cc4e65e8a9d50b8d5d27290594ebe93d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=90=A8?= Date: Sat, 11 Jun 2016 15:46:26 +0800 Subject: [PATCH 3/3] add test --- .../AspectsDemoTests/AspectsDemoTests.m | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m b/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m index 2bc32fc..ef0d317 100644 --- a/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m +++ b/AspectsDemo/AspectsDemoTests/AspectsDemoTests.m @@ -534,6 +534,98 @@ - (void)testAutoDeregistration { XCTAssertFalse([aspectToken remove], @"Must not able to deregister again"); } +- (void)testMultiInstanceTokenDeregistration { + TestClass *testClass = [TestClass new]; + + __block BOOL testCallCalled = NO; + id aspectToken = [testClass aspect_hookSelector:@selector(testCall) withOptions:AspectPositionInstead usingBlock:^(id info) { + testCallCalled = YES; + } error:NULL]; + XCTAssertNotNil(aspectToken, @"Must return a token."); + + __block BOOL called = NO; + id blockAspectToken = [testClass aspect_hookSelector:@selector(testCallAndExecuteBlock:) withOptions:AspectPositionAfter usingBlock:^{ + called = YES; + } error:NULL]; + XCTAssertNotNil(blockAspectToken, @"Must return a token."); + + [testClass testCallAndExecuteBlock:NULL]; + XCTAssertTrue(called, @"Flag must have been set."); + + [testClass testCall]; + XCTAssertTrue(testCallCalled, @"Hook must work."); + + XCTAssertNotEqualObjects(testClass.class, object_getClass(testClass), @"Object must have a custom subclass."); + XCTAssertTrue([aspectToken remove], @"Deregistration must work"); + testCallCalled = NO; + [testClass testCall]; + XCTAssertFalse(testCallCalled, @"Hook must no longer work."); + + called = NO; + [testClass testCallAndExecuteBlock:NULL]; + XCTAssertTrue(called, @"Flag must have been set."); + + XCTAssertTrue([blockAspectToken remove], @"Deregistration must work"); + called = NO; + [testClass testCallAndExecuteBlock:NULL]; + XCTAssertFalse(testCallCalled, @"Hook must no longer work."); + + XCTAssertEqualObjects(testClass.class, object_getClass(testClass), @"Object must not have a custom subclass."); + + XCTAssertFalse([aspectToken remove], @"Deregistration must not work twice"); + XCTAssertFalse([blockAspectToken remove], @"Deregistration must not work twice"); +} + +- (void)testInstanceAndGlobalTokenDeregistration { + TestClass *testClass = [TestClass new]; + + __block BOOL instanceCallCalled = NO; + id aspectToken = [testClass aspect_hookSelector:@selector(testCall) withOptions:AspectPositionInstead usingBlock:^(id info) { + instanceCallCalled = YES; + } error:NULL]; + XCTAssertNotNil(aspectToken, @"Must return a token."); + + __block BOOL globalCallCalled = NO; + id globalAspectToken = [TestClass aspect_hookSelector:@selector(testCall) withOptions:AspectPositionInstead usingBlock:^(id info) { + globalCallCalled = YES; + } error:NULL]; + XCTAssertNotNil(globalAspectToken, @"Must return a token."); + + [testClass testCall]; + XCTAssertTrue(globalAspectToken, @"Hook must work."); + XCTAssertTrue(instanceCallCalled, @"Hook must work."); + + XCTAssertTrue([aspectToken remove], @"Deregistration must work"); + instanceCallCalled = NO; + globalCallCalled = NO; + [testClass testCall]; + XCTAssertFalse(instanceCallCalled, @"Hook must no longer work."); + XCTAssertTrue(globalCallCalled, @"Hook must work."); + + aspectToken = [testClass aspect_hookSelector:@selector(testCall) withOptions:AspectPositionInstead usingBlock:^(id info) { + instanceCallCalled = YES; + } error:NULL]; + XCTAssertNotNil(aspectToken, @"Must return a token."); + + XCTAssertTrue([globalAspectToken remove], @"Deregistration must work"); + globalCallCalled = NO; + instanceCallCalled = NO; + [testClass testCall]; + XCTAssertTrue(instanceCallCalled, @"Hook must work."); + XCTAssertFalse(globalCallCalled, @"Hook must no longer work."); + + XCTAssertTrue([aspectToken remove], @"Deregistration must work"); + instanceCallCalled = NO; + globalCallCalled = NO; + [testClass testCall]; + XCTAssertFalse(instanceCallCalled, @"Hook must work."); + XCTAssertFalse(globalCallCalled, @"Hook must no longer work."); + + Method forwardInvocationMethod = class_getInstanceMethod(testClass.class, @selector(forwardInvocation:)); + Method objectMethod = class_getInstanceMethod(NSObject.class, @selector(forwardInvocation:)); + XCTAssertEqual(method_getImplementation(forwardInvocationMethod), method_getImplementation(objectMethod), @"Implementations must not be equal"); +} + /////////////////////////////////////////////////////////////////////////////////////////// #pragma mark - Test KVO