Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix AspectToken remove bug #93

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

fix AspectToken remove bug #93

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8385c74
fix AspectToken remove bug
Jun 6, 2016
cbd833b
modify repeat create AspectsContainer
l1002k Jun 11, 2016
c0a0f49
add test
l1002k Jun 11, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
129 changes: 89 additions & 40 deletions Aspects.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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);
Expand All @@ -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.
Expand All @@ -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_getContainerForObjectWithoutCreate(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)) {
Expand All @@ -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);
}
}
}
}

///////////////////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -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_getContainerForObjectWithoutCreate(self, aliasSelector);
AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);
AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation];
NSArray *aspectsToRemove = nil;
Expand Down Expand Up @@ -524,33 +544,62 @@ 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;
}

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 {
classContainer = objc_getAssociatedObject(klass, selector);
if (classContainer.hasAspects) break;
AspectsContainer *container = aspect_getContainerForObjectWithoutCreate((id)klass, selector);
if (container.hasAspects) {
classContainer = container;
break;
}
}while ((klass = class_getSuperclass(klass)));

return classContainer;
}

static void aspect_destroyContainerForObject(id<NSObject> 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];
}

///////////////////////////////////////////////////////////////////////////////////////////
Expand Down
92 changes: 92 additions & 0 deletions AspectsDemo/AspectsDemoTests/AspectsDemoTests.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<AspectInfo> info) {
testCallCalled = YES;
} error:NULL];
XCTAssertNotNil(aspectToken, @"Must return a token.");

__block BOOL called = NO;
id<AspectToken> 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<AspectInfo> 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<AspectInfo> 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<AspectInfo> 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

Expand Down