使用 Linux signal handler 劫持内存访问

背景与需求

最近实验室有师兄做了一个高性能的用户态 page cache（主要面向 out-of-core 计算场景，相关论文已经被 OSDI’22 接收），工作原理大概如下：

劫持内存申请相关函数（如 mmap / malloc），在程序申请内存时提供有特殊 tag 的非法（如 kernel space）指针；
使用 LLVM Pass 对 IR 上所有访存指令进行 instrumentation，在前面插入指针的转换。即对于所有类似 %data = load type* %ptr 的指令，将其修改为 %raw_ptr = get_raw_ptr(%ptr, LOAD); %data = load type* %raw_ptr，对于 store 也类似；
get_raw_ptr 函数对于不带 tag 的指针不做任何处理，对于带 tag 的指针，使用一定的规则转换为有效地之后返回。

通过上述的处理，可以不修改负载程序的源码，只需重新编译，即可将访存劫持到 cache 系统中，从而（相比使用系统机制）获得显著的性能提升。

但这一方法并不能解决所有问题，因为负载程序中还可能将 cache 系统生成的 tagged 指针传递给外部库（最典型的例子包括使用 printf 等 libc 函数）。为了使得外部库能够正确进行访存，有两种方法：

源码级（使用上述的 LLVM Pass）编译所有外部依赖，使它们主动对指针进行转换。
不修改外部依赖的二进制文件，而尝试动态地解决问题，即在外部库访问非法指针出错时尝试修复。

第一种方案显然不太现实：多数程序依赖于 C/C++ 标准库，这一编译工作量相当大；此外程序的依赖是无法穷举的，很多库可能也无法获得源码。第二种方式显然会带来较大的性能损失，但考虑到在这一工作面向的场景下，程序调用外部库的频率并不高，也多半不是性能瓶颈，因此应当是可以接受的。

Signal handler 与 `ucontext`

在指令访问非法地址时，CPU 会产生 page fault 异常，而 Linux 的处理程序会将此异常转换为信号（singal）发送给程序。通常来说，可能收到 SIGSEGV (Segmentation Fault) 或者 SIGBUS (Bus Error) 信号。而 POSIX 允许程序使用 sigaction(2) 对除了 SIGKILL 之外的信号安装 signal handler。在收到注册的信号时，正常控制流会被打断，进入 handler 中。

Handler 的原型为：void (*sa_sigaction)(int, siginfo_t *, void *)，当注册 handler 时提供了 SA_SIGINFO 标志，则第三个指针事实上是 ucontext_t *，否则是 void *。ucontext_t 即类似于内核的 exception frame，保存了程序的所有状态，尤其是各个寄存器的值。在 X86_64 平台上，ucontext_t 的结构定义如下（摘录自 /usr/include/x86_64-linux-gnu/sys/ucontext.h）：

/* Context to describe whole processor state.  */
typedef struct
  {
    gregset_t __ctx(gregs);
    /* Due to Linux's history we have to use a pointer here.  The SysV/i386
       ABI requires a struct with the values.  */
    fpregset_t __ctx(fpregs);
    unsigned long int __ctx(oldmask);
    unsigned long int __ctx(cr2);
  } mcontext_t;

/* Userlevel context.  */
typedef struct ucontext_t
  {
    unsigned long int __ctx(uc_flags);
    struct ucontext_t *uc_link;
    stack_t uc_stack;
    mcontext_t uc_mcontext;
    sigset_t uc_sigmask;
    struct _libc_fpstate __fpregs_mem;
    unsigned long int __ssp[4];
  } ucontext_t;

其中 gregs、fpregs 为所有的通用寄存器、浮点寄存器，还有一些 libc 等软件的内部状态。通过这一机制，Linux 允许用户态程序保存、恢复和切换自己的运行状态。

具体实践

使用 context 机制后，我们可以如下设计：

在程序启动时，注册 SIGSEGV handler，以及 SIGTRAP handler。
在 SIGSEGV handler 中，获取当前的 context C，用于劫持程序的访存行为：
1. 解码当前 RIP 指向的指令 I，获得 CPU 的访存地址 A；如果 A 没有 cache 系统的 tag，则终止程序（意味着真的发生了非法访问）。
2. 计算 A' = get_raw_ptr(A)
3. 构造一片可执行的代码空间 S，将产生异常的指令 I 复制到该空间中，并在后面跟随 0xCC 指令，以便在执行完后触发 SIGTRAP 信号。
4. 复制一份程序的 context C'，但修改访存地址对应的寄存器的值为 A'，RIP 为 S 的起始地址
5. 使用 C' 恢复程序运行状态。
在 SIGTRAP handler 中，获得当前的 context D，用于恢复程序状态：
1. 将 D 中的上述第 4 步修改的寄存器还原为 A（来自 C），并将 RIP 指向 I 的下一条指令。
2. 使用 D' 恢复程序运行状态。

其中的关键点是使用两个不同的 handler。这是由于必须保证执行被劫持的访存指令 I 时，除了地址对应的寄存器外，其他状态均与原来一致；而在 I 使用新地址执行后，必须 使用 D（注意不是 C，因为 I 可能会修改很多状态，尤其是各类标志位）继续执行，并将修改过的地址寄存器恢复，而不可将被修改的地址保留在寄存器中。这是由 cache 系统的设计决定的：每个被转换的指针有效大小仅为一个 page（通常是 4K），如果不恢复地址，很可能在此后产生越界访问：考虑数组遍历的情况，通常编译器会使用 base + offset * mult 的访存模式，而复用 base 寄存器。这一越界访问可能无法被立刻捕获，但会造成严重的危害（如数据损坏）。

这一设计中还包含了现场生成代码和执行，这是由于出错的代码区很可能不可写。并且在执行完当条指令后，基于上述的考虑需要立刻恢复寄存器的值，因此需要插入代码劫持控制流，从而需要一些额外的代码空间。这里简单粗暴地使用 INT3 指令触发 SIGTRAP，当然也可以做得更优雅一些。

上面的叙述中还省略了一些细节。如 SIGSEGV handler 在返回前需要设置一个标志位，标志主动劫持行为的开始，后由 SIGTRAP handler 在返回前清除。如果进入 SIGSEGV handler 时此标志是 true，说明这一指令在劫持后的地址访存时再次发生了 SIGSEGV（即访存行为原本就是非法的），此时需要终止程序，否则会陷入死循环。

CISC 之痛

思路虽然比较清晰，但在实现时还是遇到了一些问题。特别地，“获得访存地址”和“指向下一条指令”都需要解码当前指令，以获得操作数和指令长度，在 X86 这样的 CISC 架构上并不简单。我在 GitHub 上找到了一个比较轻量的 decoder/encoder 框架 Fadec，用起来还算简明。对指令的处理大概如下（代码来自 Guanyu，有修改）：

int instr_length = fd_decode(instr_pointer, MAX_INSTR_LENGTH, 64, 0, &instr);
context->uc_mcontext.gregs[REG_RIP] = (uint64_t)trap_context.exec_buffer;

for (size_t i = 0; i < MAX_OP_SIZE; i++) {
    FdOpType op_type = FD_OP_TYPE(&instr, i);
    if (op_type == FD_OT_NONE) break;
    if (op_type == FD_OT_MEM && FD_OP_BASE(&instr, i) != FD_REG_NONE) {
        uint8_t *ptr = nullptr;

        // calculate base
        uint64_t base = context->uc_mcontext.gregs[FADEC_REG_TO_GNU_REG[FD_OP_BASE(&instr, i)]];
        int64_t offset = 0;
        if (FD_ADDRSIZE(&instr) == 2) base &= 0xffff;
        else if (FD_ADDRSIZE(&instr) == 4) base &= 0xffffffff;

        // calculate offset
        bool has_idx = FD_OP_INDEX(&instr, i) != FD_REG_NONE;
        if (has_idx) {
            offset = context->uc_mcontext.gregs[FADEC_REG_TO_GNU_REG[FD_OP_INDEX(&instr, i)]];
            if (FD_ADDRSIZE(&instr) == 2)
                offset &= 0xffff;
            else if (FD_ADDRSIZE(&instr) == 4)
                offset &= 0xffffffff;
            offset *= (1 << FD_OP_SCALE(&instr, i));
        }

        // calculate ptr = base + offset + disp
        int64_t disp = FD_OP_DISP(&instr, i);
        if (disp) offset += disp;
        ptr = (uint8_t *) base + offset;

        // convert ptr & hijack address
        if (is_cache_space_ptr(ptr)) {
            const void *raw_ptr = nullptr;
            raw_ptr = cache_get_raw_ptr_store(ptr);
            raw_ptr = (const uint8_t*)raw_ptr - offset;
            context->uc_mcontext.gregs[FADEC_REG_TO_GNU_REG[FD_OP_BASE(&instr, i)]] = (unsigned long long)(raw_ptr);
            trap_context.modified[FADEC_REG_TO_GNU_REG[FD_OP_BASE(&instr, i)]] = true;
        }
    }
}

由于框架并不能简单区分指令类型（读/写），因此默认将所有指令都视为写总是安全的。此外 Fadec 的寄存器列表和 ucontext_t 中的不一致，所以使用了 FADEC_REG_TO_GNU_REG 来进行映射。

测试与总结

在使用此方法劫持 RocksDB 的所有第三方库后，可以通过其自带的一系列 ctest 测试，并且动态劫持部分的性能下降是可接受的（一个数量级左右）。

有其他同学提出了可以使用 Capstone，它的功能更强大，也是跨平台的，但对于本文的需求可能有些杀鸡用牛刀了。此外 Intel Pin 也可以完成类似任务，但它也非常重，带来的性能开销很大。

还有同学指出或许可以使用 userfaultfd(2) 机制实现上述的功能，或许可以获得更高的性能。但根据我的经验，userfaultfd 限制较多，在这一场景下也不一定合适（比如并不一定能处理带 tag 的地址）。并且根据文献的 Table 1 的结论，使用 userfaultfd 并不比 signal handler 有本质的性能提升。

而我对这一工作的评价是：