Skip to content
This repository was archived by the owner on Dec 27, 2025. It is now read-only.

Commit 8385c74

Browse files
author
leikun
committed
fix AspectToken remove bug
1 parent ab3a811 commit 8385c74

1 file changed

Lines changed: 78 additions & 40 deletions

File tree

‎Aspects.m‎

Lines changed: 78 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ - (NSArray *)aspects_arguments;
9090
NSString *const AspectErrorDomain = @"AspectErrorDomain";
9191
static NSString *const AspectsSubclassSuffix = @"_Aspects_";
9292
static NSString *const AspectsMessagePrefix = @"aspects_";
93+
void *AspectsAliasSelecterName2ContainerDictionary = &AspectsAliasSelecterName2ContainerDictionary;
9394

9495
@implementation NSObject (Aspects)
9596

@@ -272,6 +273,7 @@ static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSE
272273
Class klass = aspect_hookClass(self, error);
273274
Method targetMethod = class_getInstanceMethod(klass, selector);
274275
IMP targetMethodIMP = method_getImplementation(targetMethod);
276+
NSString *className = NSStringFromClass(klass);
275277
if (!aspect_isMsgForwardIMP(targetMethodIMP)) {
276278
// Make a method alias for the existing method implementation, it not already copied.
277279
const char *typeEncoding = method_getTypeEncoding(targetMethod);
@@ -284,7 +286,16 @@ static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSE
284286
// We use forwardInvocation to hook in.
285287
class_replaceMethod(klass, selector, aspect_getMsgForwardIMP(self, selector), typeEncoding);
286288
AspectLog(@"Aspects: Installed hook for -[%@ %@].", klass, NSStringFromSelector(selector));
289+
} else if ([className hasSuffix:AspectsSubclassSuffix]) {
290+
/**
291+
* global selector hack -> singal instance selector hack -> remove global selector hack
292+
* it will clean up forward method in class, but also in class_Aspects_
293+
* it should have 2 methods in class and class_Aspects_
294+
*/
295+
const char *typeEncoding = method_getTypeEncoding(targetMethod);
296+
class_addMethod(klass, selector, method_getImplementation(targetMethod), typeEncoding);
287297
}
298+
288299
}
289300

290301
// Will undo the runtime changes made.
@@ -297,8 +308,47 @@ static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) {
297308
if (isMetaClass) {
298309
klass = (Class)self;
299310
}
311+
312+
// Deregister global tracked selector
313+
aspect_deregisterTrackedSelector(self, selector);
314+
// Get the aspect container and check if there are any hooks remaining. Clean up if there are not.
315+
AspectsContainer *container = aspect_getContainerForObject(self, selector);
316+
if (!container.hasAspects) {
317+
// Destroy the container
318+
aspect_destroyContainerForObject(self, selector);
319+
}
320+
321+
if (isMetaClass) {
322+
aspect_cleanupHookedForwardSelector(klass, selector);
323+
// if there is no global selector hack
324+
if (aspect_getAliasSelectorName2ContainerDictionaryForObject(self).count == 0) {
325+
// Class is most likely swizzled in place. Undo that.
326+
aspect_undoSwizzleClassInPlace((Class)self);
327+
}
328+
} else {
329+
if (aspect_getContainerForClass(klass, selector) == nil) {
330+
aspect_cleanupHookedForwardSelector(klass, selector);
331+
}
332+
//recover class
333+
if (aspect_getAliasSelectorName2ContainerDictionaryForObject(self).count == 0) {
334+
// Figure out how the class was modified to undo the changes.
335+
NSString *className = NSStringFromClass(klass);
336+
if ([className hasSuffix:AspectsSubclassSuffix]) {
337+
Class originalClass = NSClassFromString([className stringByReplacingOccurrencesOfString:AspectsSubclassSuffix withString:@""]);
338+
NSCAssert(originalClass != nil, @"Original class must exist");
339+
object_setClass(self, originalClass);
340+
AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(originalClass));
341+
342+
// We can only dispose the class pair if we can ensure that no instances exist using our subclass.
343+
// Since we don't globally track this, we can't ensure this - but there's also not much overhead in keeping it around.
344+
//objc_disposeClassPair(object.class);
345+
}
346+
}
347+
}
348+
}
300349

301-
// Check if the method is marked as forwarded and undo that.
350+
// clean up forward method
351+
static void aspect_cleanupHookedForwardSelector(Class klass, SEL selector) {
302352
Method targetMethod = class_getInstanceMethod(klass, selector);
303353
IMP targetMethodIMP = method_getImplementation(targetMethod);
304354
if (aspect_isMsgForwardIMP(targetMethodIMP)) {
@@ -308,40 +358,10 @@ static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) {
308358
Method originalMethod = class_getInstanceMethod(klass, aliasSelector);
309359
IMP originalIMP = method_getImplementation(originalMethod);
310360
NSCAssert(originalMethod, @"Original implementation for %@ not found %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass);
311-
361+
312362
class_replaceMethod(klass, selector, originalIMP, typeEncoding);
313363
AspectLog(@"Aspects: Removed hook for -[%@ %@].", klass, NSStringFromSelector(selector));
314364
}
315-
316-
// Deregister global tracked selector
317-
aspect_deregisterTrackedSelector(self, selector);
318-
319-
// Get the aspect container and check if there are any hooks remaining. Clean up if there are not.
320-
AspectsContainer *container = aspect_getContainerForObject(self, selector);
321-
if (!container.hasAspects) {
322-
// Destroy the container
323-
aspect_destroyContainerForObject(self, selector);
324-
325-
// Figure out how the class was modified to undo the changes.
326-
NSString *className = NSStringFromClass(klass);
327-
if ([className hasSuffix:AspectsSubclassSuffix]) {
328-
Class originalClass = NSClassFromString([className stringByReplacingOccurrencesOfString:AspectsSubclassSuffix withString:@""]);
329-
NSCAssert(originalClass != nil, @"Original class must exist");
330-
object_setClass(self, originalClass);
331-
AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(originalClass));
332-
333-
// We can only dispose the class pair if we can ensure that no instances exist using our subclass.
334-
// Since we don't globally track this, we can't ensure this - but there's also not much overhead in keeping it around.
335-
//objc_disposeClassPair(object.class);
336-
}else {
337-
// Class is most likely swizzled in place. Undo that.
338-
if (isMetaClass) {
339-
aspect_undoSwizzleClassInPlace((Class)self);
340-
}else if (self.class != klass) {
341-
aspect_undoSwizzleClassInPlace(klass);
342-
}
343-
}
344-
}
345365
}
346366

347367
///////////////////////////////////////////////////////////////////////////////////////////
@@ -477,7 +497,7 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
477497
SEL originalSelector = invocation.selector;
478498
SEL aliasSelector = aspect_aliasForSelector(invocation.selector);
479499
invocation.selector = aliasSelector;
480-
AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector);
500+
AspectsContainer *objectContainer = aspect_getContainerForObject(self, aliasSelector);
481501
AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);
482502
AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation];
483503
NSArray *aspectsToRemove = nil;
@@ -524,14 +544,27 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
524544
///////////////////////////////////////////////////////////////////////////////////////////
525545
#pragma mark - Aspect Container Management
526546

547+
static NSMutableDictionary *aspect_getAliasSelectorName2ContainerDictionaryForObject(NSObject *self) {
548+
NSMutableDictionary *aliasSelectorName2ContainerDictionary = objc_getAssociatedObject(self, AspectsAliasSelecterName2ContainerDictionary);
549+
if (!aliasSelectorName2ContainerDictionary) {
550+
aliasSelectorName2ContainerDictionary = [NSMutableDictionary dictionary];
551+
objc_setAssociatedObject(self, AspectsAliasSelecterName2ContainerDictionary, aliasSelectorName2ContainerDictionary, OBJC_ASSOCIATION_RETAIN);
552+
}
553+
return aliasSelectorName2ContainerDictionary;
554+
}
555+
527556
// Loads or creates the aspect container.
528557
static AspectsContainer *aspect_getContainerForObject(NSObject *self, SEL selector) {
529558
NSCParameterAssert(self);
530-
SEL aliasSelector = aspect_aliasForSelector(selector);
531-
AspectsContainer *aspectContainer = objc_getAssociatedObject(self, aliasSelector);
559+
SEL aliasSelector = selector;
560+
if (![NSStringFromSelector(selector) hasPrefix:AspectsMessagePrefix]) {
561+
aliasSelector = aspect_aliasForSelector(selector);
562+
}
563+
NSString *selectorName = NSStringFromSelector(aliasSelector);
564+
AspectsContainer *aspectContainer = aspect_getAliasSelectorName2ContainerDictionaryForObject(self)[selectorName];
532565
if (!aspectContainer) {
533566
aspectContainer = [AspectsContainer new];
534-
objc_setAssociatedObject(self, aliasSelector, aspectContainer, OBJC_ASSOCIATION_RETAIN);
567+
[aspect_getAliasSelectorName2ContainerDictionaryForObject(self) setObject:aspectContainer forKey:selectorName];
535568
}
536569
return aspectContainer;
537570
}
@@ -540,17 +573,22 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
540573
NSCParameterAssert(klass);
541574
AspectsContainer *classContainer = nil;
542575
do {
543-
classContainer = objc_getAssociatedObject(klass, selector);
544-
if (classContainer.hasAspects) break;
576+
AspectsContainer *container = aspect_getContainerForObject((NSObject *)klass, selector);
577+
if (container.hasAspects) {
578+
classContainer = container;
579+
break;
580+
}
545581
}while ((klass = class_getSuperclass(klass)));
546-
582+
547583
return classContainer;
548584
}
549585

550586
static void aspect_destroyContainerForObject(id<NSObject> self, SEL selector) {
551587
NSCParameterAssert(self);
552588
SEL aliasSelector = aspect_aliasForSelector(selector);
553-
objc_setAssociatedObject(self, aliasSelector, nil, OBJC_ASSOCIATION_RETAIN);
589+
NSString *selectorName = NSStringFromSelector(aliasSelector);
590+
NSMutableDictionary *aliasSelectorName2ContainerDictionary = aspect_getAliasSelectorName2ContainerDictionaryForObject(self);
591+
[aliasSelectorName2ContainerDictionary removeObjectForKey:selectorName];
554592
}
555593

556594
///////////////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)