任意方法的Swizzle的应用之一AOP

作者: 吸血鬼de晚餐 | 来源:发表于2018-11-12 19:14 被阅读11次

在上篇博客曾聊过对任意方法Swizzle有多种应用，其中之一就是对多个方法的开始或者结束添加统一的切面调用。很有名的Aspect库利用类似于KVO实现原理来实现，在运行时给需要Swizzle的类动态添加子类，同时该对象isa指针指向创建的子类，然后hook子类的forwardInvocation关联到函数__ASPECTS_ARE_BEING_CALLED__。之后再将要hook的method的实现替换为_objc_msgForward，这样对该method调用会就转发到__ASPECTS_ARE_BEING_CALLED__，然后利用NSInvocation转发该调用即可。而有了对任意方法Swizzle的方案后就可以另辟蹊径来解决这个问题。

上篇博客给了个简单的demo实现，不过功能较弱，这次花了一天多时间，写了一个较为成熟的实现，总共三百行左右，其中一半是汇编。

先上代码，先吓走一票人再说。

#import <objc/runtime.h>

typedef NS_OPTIONS(NSInteger, ZWInvocationOption) {
    ZWInvocationOptionReplace = 1,
    ZWInvocationOptionBefore = 1 << 1,
    ZWInvocationOptionAfter = 1 << 2,
    ZWInvocationOptionOnly = 1 << 3,
};

#if defined(__arm64__)

static NSMutableDictionary  *_ZWBeforeIMP;
static NSMutableDictionary  *_ZWOriginIMP;
static NSMutableDictionary  *_ZWAfterIMP;
static NSLock  *_ZWLock;

OS_ALWAYS_INLINE void ZWStoreParams(void) {
    asm volatile("str    d7, [x11, #0x88]\n\
                 str    d6, [x11, #0x80]\n\
                 str    d5, [x11, #0x78]\n\
                 str    d4, [x11, #0x70]\n\
                 str    d3, [x11, #0x68]\n\
                 str    d2, [x11, #0x60]\n\
                 str    d1, [x11, #0x58]\n\
                 str    d0, [x11, #0x50]\n\
                 str    x8, [x11, #0x40]\n\
                 str    x7, [x11, #0x38]\n\
                 str    x6, [x11, #0x30]\n\
                 str    x5, [x11, #0x28]\n\
                 str    x4, [x11, #0x20]\n\
                 str    x3, [x11, #0x18]\n\
                 str    x2, [x11, #0x10]\n\
                 str    x1, [x11, #0x8]\n\
                 str    x0, [x11]\n\
                 ");
}
OS_ALWAYS_INLINE void ZWLoadParams(void) {
    asm volatile("ldr    d7, [x11, #0x88]\n\
                 ldr    d6, [x11, #0x80]\n\
                 ldr    d5, [x11, #0x78]\n\
                 ldr    d4, [x11, #0x70]\n\
                 ldr    d3, [x11, #0x68]\n\
                 ldr    d2, [x11, #0x60]\n\
                 ldr    d1, [x11, #0x58]\n\
                 ldr    d0, [x11, #0x50]\n\
                 ldr    x8, [x11, #0x40]\n\
                 ldr    x7, [x11, #0x38]\n\
                 ldr    x6, [x11, #0x30]\n\
                 ldr    x5, [x11, #0x28]\n\
                 ldr    x4, [x11, #0x20]\n\
                 ldr    x3, [x11, #0x18]\n\
                 ldr    x2, [x11, #0x10]\n\
                 ldr    x1, [x11, #0x8]\n\
                 ldr    x0, [x11]\n\
                 ");
}

OS_ALWAYS_INLINE void ZWGlobalOCSwizzle(void) {
    asm volatile("stp    x29, x30, [sp, #-0x10]!");
    
    asm volatile("mov    x29, sp\n\
                 sub    sp, sp, #0xb0");
    
    asm volatile("mov    x11, sp");
    asm volatile("bl    _ZWStoreParams");

    asm volatile("mov    x0, sp");
    asm volatile("bl    _ZWBeforeInvocation");
    
    asm volatile("mov    x0, sp");
    asm volatile("bl    _ZWInvocation");
    
    asm volatile("str    x0, [sp, #0xa0]");
    asm volatile("str    d0, [sp, #0xa8]");

    asm volatile("mov    x0, sp");
    asm volatile("bl    _ZWAfterInvocation");
    
    asm volatile("ldr    x0, [sp, #0xa0]");
    asm volatile("ldr    d0, [sp, #0xa8]");
    
    asm volatile("mov    sp, x29");
    asm volatile("ldp    x29, x30, [sp], #0x10");
}

OS_ALWAYS_INLINE NSString *ZWGetMetaSelName(SEL sel) {
    return [@"__META_" stringByAppendingString:NSStringFromSelector(sel)];
}

OS_ALWAYS_INLINE id ZWGetInvocation(NSDictionary *dict, id obj, SEL sel) {
    Class class = object_getClass(obj);
    NSString *className = NSStringFromClass(class);
    NSString *selName = class_isMetaClass(class) ? ZWGetMetaSelName(sel) : NSStringFromSelector(sel);
    [_ZWLock lock];
    id Invocation = dict[className][selName];
    [_ZWLock unlock];
    return Invocation;
}

OS_ALWAYS_INLINE NSUInteger ZWGetInvocationCount(NSDictionary *dict, id obj, SEL sel) {
    id ret = ZWGetInvocation(dict, obj, sel);
    if ([ret isKindOfClass:[NSArray class]]) {
        return [ret count];
    } else if ([ret isKindOfClass:NSClassFromString(@"NSBlock")]) {
        return 1;
    }
    return 0;
}

IMP ZWGetOriginImp(id obj, SEL sel) {
    id Invocation = ZWGetInvocation(_ZWOriginIMP, obj, sel);
    if ([Invocation isKindOfClass:[NSValue class]]) {
        return [Invocation pointerValue];
    }
    return NULL;
}

IMP ZWGetCurrentImp(id obj, SEL sel) {
    id Invocation = ZWGetInvocation(_ZWOriginIMP, obj, sel);
    if ([Invocation isKindOfClass:NSClassFromString(@"NSBlock")]) {
        if (!Invocation) return NULL;
        uint64_t *p = (__bridge void *)(Invocation);
        return (IMP)*(p + 2);
    }
    return NULL;
}

IMP ZWGetAopImp(NSDictionary *Invocation, id obj, SEL sel, NSUInteger index) {
    if (!obj || !sel) return NULL;
    id block = ZWGetInvocation(Invocation, obj, sel);
    if ([block isKindOfClass:[NSArray class]]) {
        block = block[index];
    }
    if (!block) return NULL;
    uint64_t *p = (__bridge void *)(block);
    return (IMP)*(p + 2);
}


void ZWAopInvocation(void **sp, NSDictionary *Invocation, ZWInvocationOption option) {
    id obj = (__bridge id)(*sp);
    SEL sel = *(sp + 1);
    if (!obj || !sel) return;
    NSInteger count = ZWGetInvocationCount(Invocation, obj, sel);
    __autoreleasing NSArray *arr = @[obj, [NSValue valueWithPointer:sel], @(option)];
    for (int i = 0; i < count; ++i) {
        ZWGetAopImp(Invocation, obj, sel, i);
        asm volatile("cbz    x0, LBB_20181107");
        asm volatile("mov    x17, x0");
        asm volatile("ldr    x11, %0": "=m"(sp));
        asm volatile("ldr    x12, %0": "=m"(arr));
        asm volatile("bl    _ZWLoadParams");
        asm volatile("mov    x1, x12");
        asm volatile("blr    x17");
        asm volatile("LBB_20181107:");
    }
}

void ZWAfterInvocation(void **sp) {
    ZWAopInvocation(sp, _ZWAfterIMP, ZWInvocationOptionAfter);
}
void ZWInvocation(void **sp) {
    __autoreleasing id obj;
    SEL sel;
    void *obj_p = &obj;
    void *sel_p = &sel;
    
    asm volatile("ldr    x11, %0": "=m"(sp));
    asm volatile("ldr    x10, %0": "=m"(obj_p));
    asm volatile("ldr    x0, [x11]");
    asm volatile("str    x0, [x10]");
    asm volatile("ldr    x10, %0": "=m"(sel_p));
    asm volatile("ldr    x0, [x11, #0x8]");
    asm volatile("str    x0, [x10]");
    
    asm volatile("ldr    x11, %0": "=m"(sp));
    asm volatile("ldr    x0, [x11]");
    asm volatile("ldr    x1, [x11, #0x8]");
    asm volatile("bl    _ZWGetOriginImp");
    asm volatile("cbnz    x0, LZW_20181105");
    
    __autoreleasing NSArray *arr = @[obj, [NSValue valueWithPointer:sel], @(ZWInvocationOptionReplace)];

    asm volatile("ldr    x11, %0": "=m"(sp));
    asm volatile("ldr    x0, [x11]");
    asm volatile("ldr    x1, [x11, #0x8]");
    asm volatile("bl    _ZWGetCurrentImp");
    asm volatile("cbz    x0, LZW_20181106");
    asm volatile("mov    x17, x0");
    asm volatile("ldr    x12, %0": "=m"(arr));
    asm volatile("ldr    x11, %0": "=m"(sp));
    asm volatile("bl    _ZWLoadParams");
    asm volatile("mov    x1, x12");
    asm volatile("blr    x17");
    asm volatile("b LZW_20181106");

    asm volatile("LZW_20181105:");
    asm volatile("mov    x17, x0");
    asm volatile("ldr    x11, %0": "=m"(sp));
    asm volatile("bl    _ZWLoadParams");
    asm volatile("blr    x17");
    asm volatile("LZW_20181106:");
}
void ZWBeforeInvocation(void **sp) {
    ZWAopInvocation(sp, _ZWBeforeIMP, ZWInvocationOptionBefore);
}

OS_ALWAYS_INLINE Method ZWGetMethod(Class cls, SEL sel) {
    unsigned int count = 0;
    Method retMethod = NULL;
    Method *list = class_copyMethodList(cls, &count);
    for (int i = 0; i < count; ++i) {
        Method m = list[i];
        SEL s = method_getName(m);
        if (sel_isEqual(s, sel)) {
            retMethod = m;
        }
    }
    free(list);
    return retMethod;
}

OS_ALWAYS_INLINE void ZWAddInvocation(NSDictionary *dict, NSString *className, NSString *selName, id block,  ZWInvocationOption options) {
    NSArray *tmp = dict[className][selName];
    if (options & ZWInvocationOptionOnly) {
        dict[className][selName] = block;
    } else {
        if ([tmp isKindOfClass:[NSArray class]]) {
            dict[className][selName] = [tmp arrayByAddingObject:block];
        } else if (!tmp) {
            dict[className][selName] = @[block];
        }
    }
}

void ZWAddAop(id obj, SEL sel, ZWInvocationOption options, id block) {
    if (options == ZWInvocationOptionOnly || !obj || !sel || !block) return;
    
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        _ZWOriginIMP = [NSMutableDictionary dictionary];
        _ZWBeforeIMP = [NSMutableDictionary dictionary];
        _ZWAfterIMP = [NSMutableDictionary dictionary];
        _ZWLock = [[NSLock alloc] init];
    });
    
    Class class = object_getClass(obj);
    Method method = NULL;
    NSString *className = NSStringFromClass(class);
    ZWGetMethod(class, sel, &method);//class_getInstanceMethod(class, sel)会获取父类的方法
    NSString *selName = class_isMetaClass(class) ? ZWGetMetaSelName(sel) : NSStringFromSelector(sel);
    
    [_ZWLock lock];
    if (!_ZWOriginIMP[className]) {
        _ZWOriginIMP[className] = [NSMutableDictionary dictionary];
        _ZWBeforeIMP[className] = [NSMutableDictionary dictionary];
        _ZWAfterIMP[className] = [NSMutableDictionary dictionary];
    }

    if (options & ZWInvocationOptionReplace) {
        _ZWOriginIMP[className][selName] = block;
    } else {
        IMP originImp = method_getImplementation(method);
        if (originImp != ZWGlobalOCSwizzle) {
            _ZWOriginIMP[className][selName] = [NSValue valueWithPointer:originImp];
        }
    }
    
    if (options & ZWInvocationOptionBefore) {
        ZWAddInvocation(_ZWBeforeIMP, className, selName, block, options);
    }
    if (options & ZWInvocationOptionAfter) {
        ZWAddInvocation(_ZWAfterIMP, className, selName, block, options);
    }
    method_setImplementation(method, ZWGlobalOCSwizzle);
    [_ZWLock unlock];
}
#else
void ZWAddAop(id obj, SEL sel, ZWInvocationOption options, id block) {
}
#endif

这里只写了arm64下的实现，arm32机器已经比较少，x86_64汇编比较麻烦，就不实现了。

入口函数ZWAddAop，前俩参数就不多说了，第三个参数是一个block，这个block在定义的时候，一定要注意参数类型要对应，Method的前两个参数签名为固定的“@:”，而block第一个参数签名是固定的“@?”，就是block本身，所以这里本来是需要一个占位的参数SEL和原始方法对应，后续的参数一一对应即可，如果不需要使用某些参数可以省略声明。但需要注意的是这里的占位参数被我改变成了NSArray，其会将当前调用的对象和方法，以及调用位置封装成数组传递，方便使用者。（简易实现，所以直接用了公共容器）。另外需要说明的是这里我并没有校验block和原始方法签名是否对应，需要的可以自行添加。

- (int )aMethod3:(NSString *)str;
对应
^(NSArray *info, NSString *str) {}；
如果不使用str， ^(NSArray *info){}；不使用info，也可以不声明

第四个参数ZWInvocationOption，表示需要在哪些位置插入该block调用。其中ZWInvocationOptionOnly不能单独使用，其他情况都可以单独使用或者复合使用。Only和Before，After复合使用时表示当前添加的block是当前位置唯一的调用，其他调用会被替换掉，否则在同一个方法的同一个位置可以添加多个block调用。Only和Replace复合没有意义，因为我觉得没有必要替换多次，因为从功能上讲Before就可以完成这项需求，如果有需求的话可以自行修改。

ZWAddAop函数实现比较容易看懂，需要说明一下的是，这里我操作NSMutableDictionary时我加了锁，如果有大量方法需要替换时，可以采用串行调用而不是加锁的方式，这样速度会提高。如果是元类方法会在存储的时候添加上__META_头和实例方法区分开。另外class_getInstanceMethod会获取父类的实现，所以这里class_copyMethodList来搜索，因为存在内存拷贝，所以效率不太好。

注意：优化点，将所有NSDictionary的selector key由NSString变成NSNumber，该地址是一个常量，而class对象是可以直接作为key的。AOP切面调用极为频繁，这就会带来极大的性能提升，优化后甚至可以提高一个数量级。

以上是添加AOP调用的过程，接下来说明一下调用过程。

所有添加AOP的方法，IMP会修改为ZWGlobalOCSwizzle方法，其声明没有参数没有返回值，也不涉及任何OC和C的代码，所以编译器会给一个空的实现，不会添加多余代码，可以方便的添加汇编代码（这种叫C桩(stub)函数）。其逻辑还是很容易看懂的。

我大概说一下，这里我抽离了参数入栈的的汇编到ZWStoreParams，这个函数会被多个地方调用，所以不能直接使用sp，我这里使用了x11，调用的时候需要设置x11的值（复杂的情况下使用寄存器，需要先暂存其值到栈上，方便之后恢复供后续使用，只不过比较麻烦，会涉及到初始栈空间的分配，这里其他地方不会使用x11存数据就简单处理了），调用ZWStoreParams，将参数存在栈上，同时将栈指针放入x0，传递给ZWBeforeInvocation使用。

接下来调用ZWBeforeInvocation，其接收了栈指针sp，上面有所有参数，其内部调用ZWAopInvocation。

在ZWAopInvocation中获枚举当前方法当前位置所有添加的AOP block。下面代码会将一些信息封装到数组中，后面写x1也就是block调用的第二个参数（info）。

__autoreleasing NSArray *arr = @[obj, [NSValue valueWithPointer:sel], @(option)];

for循环所有调用，并依次调用ZWGetAopImp获取所有block入口地址，ZWGetAopImp在字典中_ZWBeforeIMP或者_ZWAfterIMP获取存入的Block，并计算出入口地址返回。

使用cbz指令，如果x0==0就跳转标签LBB_20181107，就是末尾。否则将入口地址写入x17，接下来加载sp指针到x11，加载arr到x12，调用ZWLoadParams加载参数到寄存器，将arr写入x1，最后跳转到x17执行block方法。

note:标签LZW_20181107似乎有一些要求，不能随便乱写，似乎要L打头（L表示局部标签），后面就随便了，确保唯一即可。

回到ZWGlobalOCSwizzle函数，调用完ZWBeforeInvocation后，调用ZWInvocation，也就是原方法或是其替换方法。该方法也接收sp参数，然后我们声明obj，sel，及其指针obj_p，sel_p，接下来将sp和obj，sel的地址加载到寄存器，然后将ps，ps+8，也就是栈上的第一第二个参数写入obj，sel中。这样就把汇编的数据转移到了C语言的变量里面，这也是内联汇编和OC/C要数据交互方法。

接下来将sp，sp+8写入x0，x1，然后调用ZWGetOriginImp，其会尝试获取_ZWOriginIMP字典中的原始IMP并返回，汇编cbnz x0, LZW_20181105如果原始IMP不为NULL，则跳转标签LZW_20181105。

再看标签LZW_20181105，其实现也比较简单，将原始IMP写入x17，同时加载sp到x11，调用ZWLoadParams恢复参数到寄存器上，然后跳转到x17执行原始IMP后就结束了（变量声明都是autoreleasing的，所以不会插入release也就不会破坏x0，d0）。

如果原始IMP为NULL，则表明其被替换过。则将obj，sel，调用位置ZWInvocationOptionReplace封装成数组arr待用。加载sp，sp+8到x0，x1，然后调用ZWGetCurrentImp，获取替换后的实现。其是一个block，计算出IMP并返回。cbz x0, LZW_20181106，如果返回值为NULL，则直接跳转到标签LZW_20181106结束调用，否则继续执行，将IMP写入x17，加载arr，sp到x12，x11，调用ZWLoadParams，恢复所有的参数，这里和调用原始IMP不一样的时候将第二个参数x1替换为数组arr。然后blr x17执行IMP，之后跳转到标签LZW_20181106，这里不能直接使用汇编的"ret"返回。因为该句之后会调用___stack_chk_guard来检查栈有没有被破坏，不过幸运的是该调用没有破坏x0，d0寄存器。

最后再次回到ZWGlobalOCSwizzle函数，调用完ZWInvocation后，将可能存储返回值的寄存器x0，d0写入sp的+0xa0,+0xa8的位置，然后加载sp为参数调用ZWAfterInvocation，其逻辑和ZWBeforeInvocation基本一致就不再多说。

然后从sp的+0xa0,+0xa8的位置恢复数据到x0，d0，然后恢复sp，x29，x30寄存器，最后返回。

怎么使用呢？这里给一些例子

- (void)viewDidLoad {
    [super viewDidLoad];

    ZWAddAop(self, @selector(aMethod1:), ZWInvocationOptionBefore | ZWInvocationOptionOnly, ^(NSArray *info){
        NSLog(@"before1 : ");
    });
    ZWAddAop(self, @selector(aMethod1:), ZWInvocationOptionBefore | ZWInvocationOptionReplace  | ZWInvocationOptionAfter , ^(NSArray *info, int a){
        NSLog(@"before1 | replace1 | after1 : %d", a);
    });

    ZWAddAop(self, @selector(aMethod2:), ZWInvocationOptionAfter, ^(NSArray *info, NSString *str){
        NSLog(@"after2: %@", str);
    });

    ZWAddAop(self, @selector(aMethod2:), ZWInvocationOptionReplace | ZWInvocationOptionOnly, ^(NSArray *info, NSString *str){
        NSLog(@"replace2 | after2: %@ \n%@",info, str);
    });

    ZWAddAop(self, @selector(aMethod3::), ZWInvocationOptionReplace, ^int (NSArray *info, NSString *str){
        NSLog(@"replace3: %@", str);
        return 11034;
    });

    ZWAddAop(self, @selector(aMethod3::), ZWInvocationOptionAfter, ^(NSArray *info, NSString *str){
        NSLog(@"after31: %@ \n %@", info, str);
    });

    ZWAddAop(self, @selector(aMethod3::), ZWInvocationOptionAfter, ^(NSArray *info, NSString *str){
        NSLog(@"after32: %@", str);
    });
    
    ZWAddAop([self class], @selector(aMethod4:), ZWInvocationOptionReplace, ^(id info , int a){
        NSLog(@"META replace4:");
    });

    [self aMethod1:8848];
    [self aMethod2:@"test str"];
    int r = [self aMethod3:@"咋不上天呢" :@[@1,@2]];
    NSLog(@"return int: %d",r);
    [ViewController aMethod4:12358];
}

- (NSRange)aMethod1:(int)a {
    NSLog(@"method1: %d",a);
    return (NSRange){0,1};
}
- (void)aMethod2:(NSString *)str {
    NSLog(@"method2: %@", str);
}
- (int )aMethod3:(NSString *)str :(NSArray *)array{
    NSLog(@"method3: %@", str);
    return 11;
}
+ (void)aMethod4:(int )obj {
    NSLog(@"method4: %d", obj);
}

本来只是想写一个demo供大家参考的，但发现很多东西还是要过手了才有说服力，于是前前后后（主要是想简单实现基本功能即可，后来发现很多功能特性还是得添加，于是乎就不断涂涂改改...）花了两天时间写了个较完善些的实现。如果对稳定性要求较高，可以多测试一下，汇编和OC/C混编确实容易出bug，写的时候稍不注意就出现莫名其妙的问题，需要不断的修改和改进，另外我测试的case也不太足。

最后说一下这套方案的优缺点吧。

优点：1、简洁，虽然可能汇编细节不太懂，但大致的实现思路还是很容易看懂。
2、效率较高，大量使用汇编，减少函数调用，当然也有一个龟速代表class_copyMethodList拖慢了效率，当然也仅仅只是指针拷贝，而且大部分类的方法也就几个十几个。主要是没有其他类似的方法，要替换class_copyMethodList就得自己访问Class的MethodList，这比较麻烦，需要大量的代码才能实现。

缺点，改进点：
1、追求简洁，抛弃了一些细节问题。比如传入的block和原始方法的签名没有校验是否匹配，再比如：block的传回的第一个参数使用NSArray来存储obj，sel和调用位置，其实最好可以使用model对象，同时也可以传回更多的信息。
2、整型参数大于6个或者浮点参数大于8个则无法处理。因为多出来的参数在栈上，目前我还没有想出除了拷贝栈以外的简洁解决方案，而拷贝栈比较麻烦（我大概思考了一下，比较麻烦的部分是确定栈上存了多少数据，这个得根据签名和参数传递的规则去计算）。如果需要处理这类函数就需要自己另外处理了。

Github源码地址

任意方法的Swizzle的应用之一AOP

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读