MessageMock : 优雅的模拟 Objective-C 方法

发表于 4年以前 | 总阅读数：628 次

前言

开源地址：MessageMock https://github.com/indulgeIn/MessageMock

我们在调试代码或编写单元测试时，为了触发特定场景，往往需要通过一系列前置操作，或者直接修改源代码数据。实际上更期望有一种不需侵入源码且更快捷的方式，知名的 OCMock 正是为了解决这些问题，不过它有不支持多线程、接口怪异、重复调用、类型处理复杂等问题，笔者看了源码过后决定换一种思路，基于objc_msgSend来进行方法的“模拟”和“校验”。

MessageMock通过任意[target selector]调用命中目标方法：

修改目标方法返回值、参数
验证目标方法返回值、参数
跳过目标方法调用
获取目标方法命中次数

核心原理

借助 fishhook 把指向objc_msgSend的指针指向自定义函数实现 Hook，在原函数调用前后分别插入两个切口，类似的代码业界已经很多了，可以参考戴明老师的处理代码 https://github.com/ming1016/GCDFetchFeed，

拿到切面过后，就可以拦截到所有的 Objective-C 方法调用，具备了做任何“坏事”的条件。但值得注意的是，MessageMock 代码必经路径不能包含任何的 Objective-C 方法调用，不然会死循环，所以源码大部分是使用 C++ / Assembly 实现的。

修改和检查返回值

目前只考虑小于等于指针类型的返回值，那浮点型就在d0，其它就在x0。

修改返回值就在origin_msgSend调用后修改掉对应寄存器的值：

...
bl origin_msgSend
缓存 x0-x8 q0-q8 值
bl after_msgSend    //把需要改的值存 x2
...
str x2, [sp, #160]  //修改栈上 x0 对应位置的数据
...
str x2, [sp]        //修改栈上 d0 对应位置的数据
...
恢复 x0-x8 q0-q8 值  //若前面修改了栈上数据，这里就完成了修改返回值操作

返回值的检查回调，只需要在after_msgSend函数里调用一下外部传入的函数指针。

修改和检查参数

目前只考虑小于等于指针类型的参数，大致测试了一下方法调用仅使用寄存器的情况：

通用寄存器参数最多 6 个（x2 - x7）
浮点寄存器参数最多 8 个（d0 - d7 编译器限制不能连续超过 6 个）

而参数到寄存器的分配比较简单，就是 x2 - x7 / d0 - d7 挨着用，用完为止。

修改方法入参就在origin_msgSend调用之前处理：

缓存 x0-x8 q0-q8 值
bl before_msgSend   //把需要修改的数据写入 x2 - x7 / d0 - d7
...
str d0, [sp, #0]
str d1, [sp, #16]
... 直到 d7          //修改栈上 d0 - d7 对应位置的数据
...
str x2, [sp, #176]
str x3, [sp, #184]
... 直到 x7          //修改栈上 x2 - x7 对应位置的数据
...
恢复 x0-x8 q0-q8 值  //若前面修改了栈上数据，这里就完成了修改入参操作
bl origin_msgSend
...

参数的检查回调，只需要在before_msgSend函数里面挨着调用一下外部传入的函数指针。

跳过原方法调用

处理很简单：

...
b 1f
...
bl origin_msgSend
...
1:
...

只是需要注意跳过原方法调用后修改入参就无效了。

析构带来的问题

代码里面用了一些内嵌汇编，由于作用域结束时会触发析构函数，可能会影响目标函数末尾的汇编代码，导致寄存器状态变化从而引发 Crash，多使用{...}限制作用域就能解决这个问题。

数据安全

底层设计上使用的一个 C++ 类来进行各种处理的配置：

class MethodMatcher {
public:
    ...
    /// 被引用的次数（用于上层代码不期望该内存释放）
    long reference = 0;
    /// 正在使用的数量
    long using_count = 0;
    /// 目标 id 地址
    uintptr_t target = 0;
    /// 目标 SEL 地址
    uintptr_t selector = 0;
    ...
};

用了一个映射表存储所有的配置数据：

typedef std::unordered_map<uintptr_t, std::unordered_map<uintptr_t, MethodMatcher *>> MethodMatcherMap;
static MethodMatcherMap *matcher_map = NULL;

问题的根源

首先MethodMatcher *指针的访问安全使用一个互斥锁就行了，关键是 MessageMock 有两个重要的能力是修改返回值和入参，当这些自定义返回值是 Objective-C 对象时，代码里面直接通过汇编指令操作，编译器不能在合适的地方插入retain，那这些 Objective-C 对象就可能提前释放（比如当前作用域结束）。

当自定义的方法返回值和入参是 Objective-C 对象时，这里称之为游离对象便于理解。

游离对象的生命周期

对于游离对象，目前是通过__bridge_retained将目标对象引用计数加一。由于这些对象都是依附于MethodMatcher *存在，所以这些引用计数被加一的 Objective-C 对象不释放，那MethodMatcher *也不能释放。

一旦游离对象被某个方法使用，最好的方式是持续到origin_msgSend方法调用结束再release。

临界区的考虑

可能第一想法在把before_msgSend/origin_msgSend/after_msgSend整个代码作为临界区保护起来，这样做肯定是不合适的。对于这种问题，可以利用读写安全的“标记”来最小化临界区。

这里的标记就是using_count和reference。

那什么时候能释放？合适时机就是origin_msgSend调用完成后。所以在origin_msgSend调用之前如果用到了某个MethodMatcher *的游离对象，其using_count属性就++，在origin_msgSend调用之后如果用到了某个MethodMatcher *的游离对象，其using_count属性就--。

上层使用的考虑

而考虑到上层接口是在 Objective-C 环境中运行，若一个作用域还未结束，这个MethodMatcher *就被释放了就会 Crash，所以上层接口层面是这样设计的：

@implementation MessageMocker {
    MethodMatcher *_matcher;
    ...
}
init {
    _matcher = new MethodMatcher();
    ++_matcher->reference;
    ...
}
dealloc {
    --_matcher->reference;
    // 尝试移除 matcher
}

如此一来，从并发数据安全角度考虑，释放一个 matcher 需要满足：reference == 0 && using_count == 0。

接口设计

使用链式语法并不是笔者的初衷，主要是基于一些特殊的考虑。

我们通常所涉及的泛型实际上是id类型，难以通过常规的手段实现真正的泛型，那比如修改返回值的接口就得很多个：

- (void)mockReturnObject:(id)value;
- (void)mockReturnInt:(int)value;
- (void)mockReturnFloat:(float)value;
...

考虑到接口和实现的简洁，还是希望能做一个真正的泛型接口，最好是能支持编译器的索引，能想到的有两点：C 多参和宏。

使用 C 多参实现：

- (void)mockReturn:(const char *)type, ...;

其实这样用户还是需要传一个前置参数 type，非常别扭，更期望的是类似于这样：

- (void)mockReturn:...;

但编译器不支持，所以考虑利用宏来处理，而宏的调用方式都是类似于macro(arg)，可以使用宏来简化参数：

#define mockReturn(arg) mockReturn:@encode(typeof(arg)), arg, nil

但编译器是不会索引出这个宏的，所以又改进一下：

@property (nonatomic, copy, readonly) MessageMocker *(^mockReturn)(const char *type, ...);
#define mockReturn(arg) mockReturn(@encode(typeof(arg)), arg, nil)

如此过后，这个宏就能索引出来了（可以在代码里面试一下），达到了简化参数的目的。

而其它的接口也顺势都做成链式调用了，使用起来也是比较优雅的，放一个简单的例子：

// 跳过 NSObject 的 new 方法调用，直接返回 nil
MessageMocker.build(NSObject.self, @selector(new)).mockReturn(nil).start();
// 当使用时始终返回 nil
id res = [NSObject new];
//res == nil