2020-05-18

TrampolineHook - 解决栈污染问题支持变参 Hook

在之前的文章《基于桥的全量方法 Hook 方案（3）- TrampolineHook》的文末，我说如果对汇编熟悉的同学可能会发现我之前实现的一个错误 - 关于上下文污染

一提到到上下文污染，可能我们绝大多数人想到的都是寄存器污染，但是实际上还有一个不容我们忽视的上下文资源：栈，过去可能大家常见的 Hook 代码关注比较少，正好这次在借助 TrampolineHook 修复这个方面的问题，我们来一起探讨下。

先看一个例子

假设有这样一个类 TestObject和不定参函数 method，定义如下：

@interface TestObject : NSObject
@end

@implementation TestObject

- (void)method:(int *)value,...
{
     va_list list;
    va_start(list, value);

    while (value) {
        NSLog(@"orig value is %d", *value);
        value = va_arg(list, int *);
    }
    va_end(list);
}

@end

如果要使用 TrampolineHook 来拦截 method 的调用，也非常简单。如下所示：

THInterceptor *sharedInterceptor = [THInterceptor sharedInterceptorWithFunction:(IMP)wzq_check_variadic];

Method m = class_getInstanceMethod([TestObject class], @selector(method:));
IMP imp = method_getImplementation(m);

THInterceptorResult *result = [sharedInterceptor interceptFunction:(IMP)imp];
if (result.state == THInterceptStateSuccess) {
    method_setImplementation(m, (IMP)result.replacedAddress);
}

// 拦截函数
void wzq_check_variadic(id a, char * methodName, int *v, ...)
{
        NSLog(@"haha checked %@ %s", a, methodName);
}

当我们使用如下方式调用 -[TestObject method:] 的时候，你会发现一切正常，毫无问题。

TestObject *obj = [[TestObject alloc] init];
int a = 0;
int b = 1;
int c = 2;
int d = 3;
int e = 4;
int f = 5;
int g = 6;
int h = 7;
int i = 8;
[obj method:&a, &b, &c, &d, &e, &f, &g, &h, &i, nil];

但是如果你将拦截函数中添加打印参数的语句后，如下所示：

void wzq_check_variadic(id a, char * methodName, int *v, ...)
{
    NSLog(@"haha checked %@ %s", a, methodName);
    va_list args;
    va_start(args, v);
    while (v != NULL) { 
        NSLog(@"v is %d", *v);
        v = va_arg(args, int *); // crash 
    }
    va_end(args);
}

你会发现出现了必现的崩溃情形，而且是必定崩溃在第二次读取变参列表中的参数的时候。
为什么添加了读取参数的代码就导致运行崩溃了？有点意思。

了解变参的传递过程。

为了避免优化的干扰，如下汇编生成的优化选项为 -O0

为了看运行时的栈结构是如何生成的，我们通过汇编结合图的形式来一探究竟没 Hook 的时候的调用情况。

首先先看 Caller 函数，即 [obj method:&a, &b, &c, &d, &e, &f, &g, &h, &i, nil]; 这段代码所处的函数，汇编如下：

 // prologue
0x1009dc428 <+0>:   sub    sp, sp, #0xb0             ; =0xb0 
0x1009dc42c <+4>:   stp    x29, x30, [sp, #0xa0]
0x1009dc430 <+8>:   add    x29, sp, #0xa0            ; =0xa0 

// 构造 int 变量 0 - 8
0x1009dc498 <+112>: stur   wzr, [x29, #-0x2c]
0x1009dc49c <+116>: mov    w11, #0x1
0x1009dc4a0 <+120>: stur   w11, [x29, #-0x30]
0x1009dc4a4 <+124>: mov    w11, #0x2
0x1009dc4a8 <+128>: stur   w11, [x29, #-0x34]
0x1009dc4ac <+132>: mov    w11, #0x3
0x1009dc4b0 <+136>: stur   w11, [x29, #-0x38]
0x1009dc4b4 <+140>: mov    w11, #0x4
0x1009dc4b8 <+144>: stur   w11, [x29, #-0x3c]
0x1009dc4bc <+148>: mov    w11, #0x5
0x1009dc4c0 <+152>: stur   w11, [x29, #-0x40]
0x1009dc4c4 <+156>: mov    w11, #0x6
0x1009dc4c8 <+160>: stur   w11, [x29, #-0x44]
0x1009dc4cc <+164>: mov    w11, #0x7
0x1009dc4d0 <+168>: stur   w11, [x29, #-0x48]
0x1009dc4d4 <+172>: mov    w11, #0x8
0x1009dc4d8 <+176>: stur   w11, [x29, #-0x4c]
0x1009dc4dc <+180>: ldur   x9, [x29, #-0x28]
0x1009dc4e0 <+184>: ldr    x1, [x8]
0x1009dc4e4 <+188>: mov    x8, sp
0x1009dc4e8 <+192>: mov    x10, #0x0

// 把对应 int 变量的地址存入栈中 
0x1009dc4ec <+196>: str    x10, [x8, #0x40]
0x1009dc4f0 <+200>: sub    x10, x29, #0x4c           ; =0x4c 
0x1009dc4f4 <+204>: str    x10, [x8, #0x38]
0x1009dc4f8 <+208>: sub    x10, x29, #0x48           ; =0x48 
0x1009dc4fc <+212>: str    x10, [x8, #0x30]
0x1009dc500 <+216>: sub    x10, x29, #0x44           ; =0x44 
0x1009dc504 <+220>: str    x10, [x8, #0x28]
0x1009dc508 <+224>: sub    x10, x29, #0x40           ; =0x40 
0x1009dc50c <+228>: str    x10, [x8, #0x20]
0x1009dc510 <+232>: sub    x10, x29, #0x3c           ; =0x3c 
0x1009dc514 <+236>: str    x10, [x8, #0x18]
0x1009dc518 <+240>: sub    x10, x29, #0x38           ; =0x38 
0x1009dc51c <+244>: str    x10, [x8, #0x10]
0x1009dc520 <+248>: sub    x10, x29, #0x34           ; =0x34 
0x1009dc524 <+252>: str    x10, [x8, #0x8]
0x1009dc528 <+256>: sub    x10, x29, #0x30           ; =0x30 
0x1009dc52c <+260>: str    x10, [x8]

 // 其余参数 x0, x1, x2
0x1009dc530 <+264>: sub    x2, x29, #0x2c            ; =0x2c 
0x1009dc534 <+268>: mov    x0, x9

// 调用 method 函数
0x1009dc538 <+272>: bl     0x1009e8d9c               ; symbol stub for: objc_msgSend

上述这段函数，简要而言，就是干了四件事：

分配 176 byte 的栈内存
在栈上分配 a = 0, b = 1 等等 9 个变量
把 &b, &c 等 8个 int 变量的地址压栈。
x0 (obj), x1 (method), x2(&a)

特别注意，变参列表的第一个参数也是通过寄存器来传递。

调用 method 函数

如果不理解, 可以参考这张图：

而当进入 method: 函数时，汇编如下：

重点看两行蓝色汇编断点的地方，其实是在暗示一种循环，也从底层实现上对应上了我们不断循环获取变参列表的逻辑。

简要来说，就是从变参列表的第一参数（寄存器中的值代表地址），开始读取，循环遍历。这里的循环利用了栈空间在函数调用间的连续性，不断将偏移地址从原来 caller 函数的 sp 回溯，读取处于高地址的 caller 栈空间中的 int 变量地址。

看到这，我想大家也知道了为什么是必定崩溃在第二次读取变参的时候。

x0, x1 不用说，是寄存器参数。和变参不变参函数无关，这也能解释为什么只读取 id obj 和 SEL selector 不会崩溃。
x2，即变参函数列表的第一个参数，我这里把他称为变参的锚点参数，它也是通过寄存器传递，所以读取的时候没问题。
变参列表的后续参数都是分配在调用函数（caller）中，而 TrampolineHook 在调用 interceptor 之前利用了栈（操作 SP）来保存上下文，如下所示，因此破坏了栈资源上下文，导致循环从栈地址获取参数的时候崩溃
```
stp q0,  q1,   [sp, #-32]!
stp q2,  q3,   [sp, #-32]!
stp q4,  q5,   [sp, #-32]!
stp q6,  q7,   [sp, #-32]!
stp lr,  x10,  [sp, #-16]!
stp x0,  x1,   [sp, #-16]!
```
而调用原函数的时候，由于栈已经复原了，所以就不会出现崩溃了。

解决方案

了解了问题出现的原因，解决办法就很简单了，我们要让调用 inteceptor 时候的上下文和调用原函数一样。

还是构造一堆的动态 trampoline ，让原函数替换到 trampoline，同时保存原函数的 IMP。
依然保存原先需要的上下文，比如通用寄存器、浮点寄存器，但是不能使用栈了。
调用 interceptor。
恢复上下文，调用到原函数。

其实整个步骤和原先基本一样，唯一需要考虑的就是如何在一点也不用栈的资源的前提下保存寄存器上下文？

堆上。
堆上。
堆上。

简单而言，我们把上下文一股脑都保存到堆上就行。需要保存的上下文大致类似于一个结构体：

typedef struct _THPageVariadicContext {
    int64_t gR[10];              // general registers x0-x8 + x13
    int64_t vR[16];              // float   registers q0-q7
    int64_t linkRegister;        // lr
    int64_t originIMPRegister;   // origin
} THPageVariadicContext;

当然，这里的结构体只是形象化表示内存中的数据顺序和含义，真正使用汇编操作内存的时候，没有结构体。

保存上下文解决了我们不污染栈的诉求，但是同时也引出了一个新的问题，堆分配的地址我们保存在哪？跨函数调用后恢复上下文必须要让我们分配出的堆地址得到“持久化”存储啊。

保存到栈上？这肯定不可能，自己打自己脸嘛。
保存到寄存器上？如果是 caller-saved 寄存器，那不能保证跨函数调用完后，寄存器里面的内容还是我们原先设定的那样；而如果是 callee-saved 寄存器，确实可以解决跨函数调用后数据还原成我们保存的那样。但是同样的，我们自身也是其他 caller 函数的 callee，我们侵占了一个寄存器，怎么在返回到 caller 函数之前复原这个 callee-saved 寄存器呢？

上面这段话有点绕。

所以，我们在分配堆内存的时候，要多分配一个 8 byte 的空间，把侵占的 callee-saved register 的值保存到堆内存中，然后再继续存我们原先要保留的上下文。

关键代码简要概括如下：

第一步，在拦截到函数调用后，先进入我们的 pre 操作，这里是在堆上对应上下文空间大小的地方。需要注意的是，调用分配内存的函数是使用 malloc，我们并不知道 malloc 究竟会破坏哪些寄存器，因为也需要作一次额外的寄存器上下文保存，不过这个保存时短暂的，分配结束后就恢复。然后将这些上下文都保存到堆上。

attribute((naked))
void THPageVariadicContextPre(void)
{

// 先保存，避免调用 malloc 破坏寄存器
saveRegs();

// 分配堆上内存 extra 16 byte + sizeof(THPageVariadicContext)
__asm volatile ("mov x0, #0xF0");
__asm volatile ("bl _malloc");

// 返回的分配内存地址保存起来 callee-saved
__asm volatile ("str x19, [x0]");
__asm volatile ("mov x19, x0");

// 恢复堆栈，避免影响变参所处在的堆栈
restoreRegs();

// 用堆上空间保存数据
__asm volatile ("stp x0, x1,  [x19, #(16 + 0 * 16)]");
__asm volatile ("stp x2, x3,  [x19, #(16 + 1 * 16)]");
__asm volatile ("stp x4, x5,  [x19, #(16 + 2 * 16)]");
__asm volatile ("stp x6, x7,  [x19, #(16 + 3 * 16)]");
__asm volatile ("stp x8, x13, [x19, #(16 + 4 * 16)]");

__asm volatile ("stp q0, q1,  [x19, #(16 + 5 * 16 + 0 * 32)]");
__asm volatile ("stp q2, q3,  [x19, #(16 + 5 * 16 + 1 * 32)]");
__asm volatile ("stp q4, q5,  [x19, #(16 + 5 * 16 + 2 * 32)]");
__asm volatile ("stp q6, q7,  [x19, #(16 + 5 * 16 + 3 * 32)]");

__asm volatile ("stp lr, x10, [x19, #(16 + 5 * 16 + 4 * 32)]");

__asm volatile ("ret");

}

调用完拦截函数，我们需要销毁堆空间，由于我们之前使用的是 callee-saved 的寄存器，我们能确保寄存器的值还是调用之前的。所以我们放心的将其中的值取出来，然后销毁对应的占空间，然后恢复寄存器即可。

__attribute__((__naked__))
void THPageVariadicContextPost(void)
{
    // x19 肯定是正确的地址，使用x19恢复对应的数据
    __asm volatile ("ldp lr, x10, [x19, #(16 + 5 * 16 + 4 * 32)]");
    __asm volatile ("ldp q6, q7,  [x19, #(16 + 5 * 16 + 3 * 32)]");
    __asm volatile ("ldp q4, q5,  [x19, #(16 + 5 * 16 + 2 * 32)]");
    __asm volatile ("ldp q2, q3,  [x19, #(16 + 5 * 16 + 1 * 32)]");
    __asm volatile ("ldp q0, q1,  [x19, #(16 + 5 * 16 + 0 * 32)]");

    __asm volatile ("ldp x8, x13, [x19, #(16 + 4 * 16)]");
    __asm volatile ("ldp x6, x7,  [x19, #(16 + 3 * 16)]");
    __asm volatile ("ldp x4, x5,  [x19, #(16 + 2 * 16)]");
    __asm volatile ("ldp x2, x3,  [x19, #(16 + 1 * 16)]");
    __asm volatile ("ldp x0, x1,  [x19, #(16 + 0 * 16)]");

    // 保存一下，避免 free 的影响。
    saveRegs();

    // 恢复原先的 x19, 调用free
    __asm volatile ("mov x0, x19");
    __asm volatile ("ldr x19, [x19]");
    __asm volatile ("bl _free");

    // 恢复堆栈
    restoreRegs();

    __asm volatile ("mov lr, x13");
    __asm volatile ("br x10");
}

需要注意的是，我们这里用了__attribute__((__naked__))，这个作用是为了让我们的函数不会额外的生成函数 prologue/epilogue 中的压栈消栈操作。

至此，变参 Hook 就完成了，大家可以前往 Github 查看最新的 THVaradicInterceptor 来使用。

后记

有的朋友会问，为什么很多网上常见的 Hook 方案，都不要这么复杂的上下文保存流程？

其实道理很简单，保存什么上下文取决你的拦截或者 Hook 函数的目的以及使用方式。

举个非常常见的统计函数调用耗时的例子，在这个情形中，一般只用关注 x0，x1 两个参数来记录是什么类什么函数的调用。这种情况下，你的上下文保存可以极简，甚至只要保存 x0, x1 即可。

而 TrampolineHook 想要提供的拦截器，是一个通用的拦截器，我不能保证其内部的实现，因为我需要保留的上下文就必须很完整。

后续 TrampolineHook 除了完善对 x86_64 的支持外，还有两个比较大的技术目标，也会慢慢完善。如果有什么使用中遇到的问题或者 Bug 也欢迎提交代码。