Skip to content

Commit

Permalink
fix AspectToken remove bug
Browse files Browse the repository at this point in the history
  • Loading branch information
leikun committed Jun 6, 2016
1 parent ab3a811 commit 8385c74
Showing 1 changed file with 78 additions and 40 deletions.
118 changes: 78 additions & 40 deletions Aspects.m
View file
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_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)) {
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_getContainerForObject(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,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;
}
Expand All @@ -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<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

3 comments on commit 8385c74

@l1002k
Copy link

@l1002k l1002k commented on 8385c74 Jun 6, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fix bug:#91

@xindizhiyin2014
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not I can't see the code when I'm use pod 'AspectsV1.4.2','1.4.2'

@xindizhiyin2014
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks the authors, but I want to ask why not merge the code into master? so I'm not sure, I should use this code in my project.

Please sign in to comment.