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

2,033

前言

开源地址:MessageMock

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

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

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

核心原理

借助 fishhook 把指向objc_msgSend的指针指向自定义函数实现 Hook,在原函数调用前后分别插入两个切口,Hook 操作可以参考戴明老师的处理代码

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

修改和检查返回值

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

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

...
bl origin_msgSend
缓存 x0-x8 q0-q8bl 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_countreference

那什么时候能释放?合适时机就是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

后语

考虑到这份代码的落地场景并不是很多,至少还需要支持 x86 机器和大于指针类型的数据处理, 才可能替代 OCMock,考虑时间成本,笔者目前就做了一个雏形,供大家一乐。另外,源码中的 C++ / Assembly 不是专业的、性能和设计也不是最优的,望各大佬指点一二不胜感激 😁。