Permalink
Browse files

Fix inheritance bugs

Add inherited method to class before injecting to it.
Use per class selectors for pre/postprocesses.
Restore original methods on item manager clear and
remove any implementation blocks.
Update inheritance test.
  • Loading branch information...
1 parent 9124c6d commit a1efd7c55efa4e901333f2622a51d6a3715140a2 @mvasilak mvasilak committed Nov 10, 2013
View
@@ -28,4 +28,6 @@
- (void)skipAfterProcessesWithReturnValue:(void*)pReturnValue;
+- (void)restoreOriginal;
+
@end
View
@@ -30,6 +30,31 @@ - (id)init
return self;
}
+- (void)restoreOriginal
+{
+ Class class = self.isClassMethod ? object_getClass(self.targetClass) : self.targetClass;
+ IMP imp = method_setImplementation(class_getInstanceMethod(class, self.targetSel), class_getMethodImplementation(class, self.originalSel));
+ if (imp) {
+ imp_removeBlock(imp);
+ }
+
+ for (NSValue *value in self.preprocesses) {
+ IMP imp = class_getMethodImplementation(class, [value pointerValue]);
+ if (imp) {
+ imp_removeBlock(imp);
+ }
+ }
+ self.preprocesses = [NSMutableArray array];
+
+ for (NSValue *value in self.postprocesses) {
+ IMP imp = class_getMethodImplementation(class, [value pointerValue]);
+ if (imp) {
+ imp_removeBlock(imp);
+ }
+ }
+ self.postprocesses = [NSMutableArray array];
+}
+
- (void)dealloc
{
if (self.pResult) {
@@ -15,6 +15,7 @@
- (BIItem*)itemForMethodName:(NSString*)methodName forClass:(Class)class;
- (void)setItem:(BIItem*)item forMethodName:(NSString*)methodName forClass:(Class)class;
+- (void)removeItemForMethodName:(NSString*)methodName forClass:(Class)class;
- (void)clear;
@@ -5,6 +5,7 @@
//
#import "BIItemManager.h"
+#import "BIItem.h"
static BIItemManager* sharedInstance = nil;
@@ -26,9 +27,15 @@ - (void)setItem:(BIItem*)item forMethodName:(NSString*)methodName forClass:(Clas
[self.items setObject:item forKey:[self keyForMethodName:methodName forClass:class]];
}
+- (void)removeItemForMethodName:(NSString*)methodName forClass:(Class)class {
+ [self removeItemForKey:[self keyForMethodName:methodName forClass:class]];
+}
+
- (void)clear
{
- self.items = [NSMutableDictionary dictionary];
+ for (NSString *key in [self.items allKeys]) {
+ [self removeItemForKey:key];
+ }
}
#pragma mark - Memory Management
@@ -48,6 +55,14 @@ - (NSString*)keyForMethodName:(NSString*)methodName forClass:(Class)class
return [NSString stringWithFormat:@"%@::%@", NSStringFromClass(class), methodName];
}
+- (void)removeItemForKey:(NSString*)key {
+ BIItem *item = [self.items objectForKey:key];
+ if (!item)
+ return;
+ [item restoreOriginal];
+ [self.items removeObjectForKey:key];
+}
+
#pragma mark - Singleton
+ (BIItemManager*)sharedInstance
View
@@ -180,21 +180,6 @@ + (BOOL)injectToSelector:(SEL)sel forClass:(Class)class preprocess:(id)preproces
{
@try {
NSString* methodName = NSStringFromSelector(sel);
- SEL saveSel = sel_registerName([[BILibUtils saveNameForMethodName:methodName] UTF8String]);
- BOOL isClassMethod = NO;
- Method originalMethod = [BILibUtils getMethodInClass:class selector:sel isClassMethod:&isClassMethod];
- Method savedMethod = [BILibUtils getMethodInClass:class selector:saveSel];
-
- if (!savedMethod) {
- // Save original method
- [BILibUtils addMethodToClass:class selector:saveSel imp:method_getImplementation(originalMethod) typeEncoding:method_getTypeEncoding(originalMethod) isClassMethod:isClassMethod];
- }
-
- if (!originalMethod) {
- NSLog(@"BILib: [%@ %@] is not found.", NSStringFromClass(class), methodName);
- return NO;
- }
-
if ([methodName hasPrefix:@"__mi_"]) {
return NO;
}
@@ -203,27 +188,36 @@ + (BOOL)injectToSelector:(SEL)sel forClass:(Class)class preprocess:(id)preproces
return NO;
}
- @try {
- [NSMethodSignature signatureWithObjCTypes:method_getTypeEncoding(originalMethod)];
- } @catch (NSException* exception) {
- @throw exception;
+ BOOL isClassMethod = NO;
+ Method originalMethod = [BILibUtils getMethodInClass:class selector:sel isClassMethod:&isClassMethod];
+ if (!originalMethod) {
+ NSLog(@"BILib: [%@ %@] is not found.", NSStringFromClass(class), methodName);
+ return NO;
}
// Replace implementation
BIItem* item = [[BIItemManager sharedInstance] itemForMethodName:methodName forClass:class];
if (!item) {
+ @try {
+ [NSMethodSignature signatureWithObjCTypes:method_getTypeEncoding(originalMethod)];
+ } @catch (NSException* exception) {
+ @throw exception;
+ }
+
item = [BIItem new];
item.targetClass = class;
item.targetSel = sel;
- item.originalSel = saveSel;
+ item.originalSel = sel_registerName([[BILibUtils saveNameForMethodName:methodName] UTF8String]);
item.originalMethod = originalMethod;
item.numberOfArguments = method_getNumberOfArguments(originalMethod) - 2;
item.signature = [NSMethodSignature signatureWithObjCTypes:method_getTypeEncoding(originalMethod)];
item.isClassMethod = isClassMethod;
+
+ [BILib replaceImplementationWithItem:item];
+
[[BIItemManager sharedInstance] setItem:item forMethodName:methodName forClass:class];
}
[BILib savePreprocess:preprocess andPostprocess:postprocess withItem:item forMethodName:methodName];
- [BILib replaceImplementationWithItem:item];
} @catch (NSException* exception) {
NSLog(@"BILib handled a exception: %@", exception);
return NO;
@@ -244,7 +238,9 @@ + (void)savePreprocess:(id)preprocess andPostprocess:(id)postprocess withItem:(B
{
if (preprocess) {
// Save preprocess
- SEL preprocessSel = sel_registerName([[BILibUtils preprocessNameForMethodName:methodName index:[item numberOfPreprocess]] UTF8String]);
+ SEL preprocessSel = sel_registerName([[BILibUtils preprocessNameForClassName:NSStringFromClass(item.targetClass)
+ methodName:methodName
+ index:[item numberOfPreprocess]] UTF8String]);
[BILibUtils addMethodToClass:item.targetClass
selector:preprocessSel
imp:imp_implementationWithBlock(preprocess)
@@ -254,7 +250,9 @@ + (void)savePreprocess:(id)preprocess andPostprocess:(id)postprocess withItem:(B
}
if (postprocess) {
// Save postprocess
- SEL postprocessSel = sel_registerName([[BILibUtils postprocessNameForMethodName:methodName index:[item numberOfPostprocess]] UTF8String]);
+ SEL postprocessSel = sel_registerName([[BILibUtils postprocessNameForClassName:NSStringFromClass(item.targetClass)
+ methodName:methodName
+ index:[item numberOfPostprocess]] UTF8String]);
[BILibUtils addMethodToClass:item.targetClass
selector:postprocessSel
imp:imp_implementationWithBlock(postprocess)
@@ -302,7 +300,16 @@ + (void)replaceImplementationWithItem:(BIItem*)item
default: { replaceBlock = REPLACEBLOCK_FOR(int); } break;
}
}
- method_setImplementation(item.originalMethod, imp_implementationWithBlock(replaceBlock));
+ [BILibUtils addMethodToClass:item.targetClass
+ selector:item.originalSel
+ imp:method_getImplementation(item.originalMethod)
+ typeEncoding:method_getTypeEncoding(item.originalMethod)
+ isClassMethod:item.isClassMethod];
+ [BILibUtils addMethodToClass:item.targetClass
+ selector:item.targetSel
+ imp:imp_implementationWithBlock(replaceBlock)
+ typeEncoding:method_getTypeEncoding(item.originalMethod)
+ isClassMethod:item.isClassMethod];
}
}
@@ -9,8 +9,8 @@
@interface BILibUtils : NSObject
+ (NSString*)saveNameForMethodName:(NSString*)methodName;
-+ (NSString*)preprocessNameForMethodName:(NSString*)methodName index:(int)index;
-+ (NSString*)postprocessNameForMethodName:(NSString*)methodName index:(int)index;
++ (NSString*)preprocessNameForClassName:(NSString*)className methodName:(NSString*)methodName index:(NSUInteger)index;
++ (NSString*)postprocessNameForClassName:(NSString*)className methodName:(NSString*)methodName index:(NSUInteger)index;
+ (NSString*)superNameForMethodName:(NSString*)methodName;
+ (Method)getMethodInClass:(Class)class selector:(SEL)selector;
@@ -15,14 +15,14 @@ + (NSString*)saveNameForMethodName:(NSString*)methodName
return [NSString stringWithFormat:@"__mi_save_%@", methodName];
}
-+ (NSString*)preprocessNameForMethodName:(NSString*)methodName index:(int)index
++ (NSString*)preprocessNameForClassName:(NSString*)className methodName:(NSString*)methodName index:(NSUInteger)index
{
- return [NSString stringWithFormat:@"__mi_pre_%d_%@", index, methodName];
+ return [NSString stringWithFormat:@"__mi_%@_pre_%lu_%@", className, (unsigned long)index, methodName];
}
-+ (NSString*)postprocessNameForMethodName:(NSString*)methodName index:(int)index
++ (NSString*)postprocessNameForClassName:(NSString*)className methodName:(NSString*)methodName index:(NSUInteger)index
{
- return [NSString stringWithFormat:@"__mi_post_%d_%@", index, methodName];
+ return [NSString stringWithFormat:@"__mi_%@_post_%lu_%@", className, (unsigned long)index, methodName];
}
+ (NSString*)superNameForMethodName:(NSString*)methodName
@@ -50,15 +50,13 @@ + (Method)getMethodInClass:(Class)class selector:(SEL)selector isClassMethod:(BO
+ (void)addMethodToClass:(Class)class selector:(SEL)selector imp:(IMP)imp typeEncoding:(const char*)typeEncoding isClassMethod:(BOOL)isClassMethod
{
- Method method = [BILibUtils getMethodInClass:class selector:selector];
- if (method) {
- method_setImplementation(method, imp);
- } else {
if (isClassMethod) {
- class = object_getClass(class);
+ class = object_getClass(class);
+ }
+ // The method is added first, to avoid setting the method implementation of a superclass.
+ if (!class_addMethod(class, selector, imp, typeEncoding)) {
+ method_setImplementation(class_getInstanceMethod(class, selector), imp);
}
- class_addMethod(class, selector, imp, typeEncoding);
- }
}
+ (NSArray*)classesWithRegex:(NSRegularExpression*)regex
@@ -60,18 +60,18 @@ - (void)testSuperClassMethod
++i;
}];
[BILib injectToClassWithNames:@[@"Parent"] methodNames:@[@"instanceMethod:"] preprocess:^{
- ++i;
+ i+=2;
}];
STAssertEquals(i, 0, @"i is invalid.");
[[Parent new] instanceMethod:@"hello!"];
- STAssertEquals(i, 1, @"i is invalid.");
+ STAssertEquals(i, 2, @"i is invalid.");
[[Child new] instanceMethod:@"hello!"];
- STAssertEquals(i, 3, @"i is invalid.");
+ STAssertEquals(i, 5, @"i is invalid.");
}
@end

0 comments on commit a1efd7c

Please sign in to comment.