Go语言实现猴子补丁

猴子补丁（Monkey Patch）可以在程序运行期间动态修改函数行为。也有人把这种技术叫「打桩」。这种技术在Python或者Ruby这样的动态语言中比较常见。但这并不意味着静态语言不能实现类似的效果。本文就跟大家分享一种上在Go语言中实现猴子补丁效果的黑科技。核心思想来自Bouke，我原本是想直接翻译的。但是原文一方面有部分内容需要对二进制文件做反汇编处理，而这部分内容其实不是必须的；另一方面，Bouke 采用了间接寻址跳转，也让整个方案变得很费解。如果按照信、达、雅的方式翻译，会带来很多不必要的理解成本。所以我结合自己的理解和实验，优化（简化）了 Bouke 的方案，最终形成本文的方案。下面开始分析实现思路。

假设我们有如下代码：

package main

func a() int { return 1}
func b() int { return 2}

func main() {
  patch(a, b)
  fmt.Println(a())
}

我们希望实现一个patch(from, to interface{})函数，调用 patch(a, b) 之后，后续所有调用a()函数的地方都会自动执行b()函数的代码逻辑。也就是说，我们希望上面的代码可以输出2而不是的1。

明确了目标，接下来我们来分析一下如何实现这种效果。

我们知道，编译器会把函数中的每一条语句翻译成机器码，存到可执行文件的代码段。我们在执行函数的时候会把对应的代码段加载到内存的某个位置。机器码说白了也是内存里的数据。我们能不能在执行a()函数之前把它的机器码替换成一段跳转指令，让CPU跳转到b()函数的机器码继续执行呢？能！

为了实现这个效果，我们要做三件事：

• 找到a()和b()机器码所在的内存地址
• 构造跳转指令
• 修改a()函数的机器码

函数的机器码在内存中的地址说白了就是函数指针。Bouke 花了大量的篇幅介绍如何找到这个内存地址，而实际上使用反射就可以轻松获取函数指针：

reflect.ValueOf(f).Pointer()

对于任意函数名f（或者匿名函数变量），我们可以通过reflect.ValueOf(f)构造一个反射对象，然后调用它的Pointer()方法拿到对应的函数指针。这个指针变量保存的就是函数的第一条机器码的内存地址（也就是指向函数）。

拿到了内存地址，接下来需要构造跳转指令。这里得用到 amd64 架构的jmp指令（不同架构的 CPU 指令不一样，本文只讲 amd64 架构）：

movabs rdx, 0x?? ; 将内存地址存到寄存器 rdx
jmp rdx ; 跳转到寄存器 rdx 中存储的内存地址，继续执行后面的指令
; 下面这条是 Bouke 的版本，不好理解，不建议使用
jmp DWORD PTR [edx]

Intel 的 x86/amd64 处理器都不支持绝对跳转，只能先把目标地址存到寄存器，再根据寄存器的内容完成跳转。我在这里使用 rdx 暂存目标地址。那问题就来了。Go语言在函数调用的过程中不会用到这个 rdx 寄存器吗？如果有用到，我们写入跳转地址不就把原来的值覆盖掉了吗？

这是一个非常重要的问题。Go语言在1.17版本之前，没有官方的ABI文档，也不确定它有没有用这个 rdx 寄存器。但到了1.17版本，Go语言支持通过寄存器传递参数和返回值，那这种直接写寄存器的方式大概率会影响Go语言的运行。我仔细看了一下相关文档，发现Go语言在 amd64 下主要使用RAX, RBX, RCX, RDI, RSI, R8, R9, R10, R11这九个寄存器传递参数。另外，Go语言在函数调用的时候还会用到以下寄存器：

这里Scratch的意思大致是说寄存器只能用来存放临时计算结果，它们的内容可能会被覆盖。如果程序在后面还要使用里面的值，需要自行保存其他地方。其中 RDX 会在调用匿名函数的时候保存闭包的上下文指针。我认为覆写这个寄存器可能会有问题。但我给 Bouke 发邮件请教，他说这个寄存器只会在有闭包调用的时候才会用到，不会有问题。而我自己也实验了几次，确实没有问题。不过为了保险起见，我还是建议大家选用 R12 和 R13 这两个寄存器。它们俩是专门用于存放中间临时结果的，应该可以随便覆盖。

让我们回到刚才的跳转指令：

movabs rdx, 0x?? ; 将内存地址存到寄存器 rdx
jmp rdx ; 跳转到寄存器 rdx 中存储的内存地址，继续执行后面的指令

我们需要把函数b()的内存地址转成字节序列。Intel 处理器用的是小端顺序，所以我们需要做如下转换：

func jmp(to uintptr) []byte {
  return []byte{
    0x48, 0xBA,     // movabs rdx
    byte(to),
    byte(to >> 8),
    byte(to >> 16),
    byte(to >> 24),
    byte(to >> 32),
    byte(to >> 40),
    byte(to >> 48),
    byte(to >> 56), // movabs rdx,to
    0xFF, 0xE2,     // jmp rdx
  }
}

最前面的0x48,0xBA表示moveabs rdx，后面的八个字节存放目标地址。因为是小端，所以低位字节在前。

有了跳转指令，我们就得想办法修改a()函数的机器码了。有同学可能会说这还不简单，直接使用copy将jmp返回的[]byte拷贝到a()函数指向的内存不就完了吗？其实不然。

前面通过反射拿到了函数指针，但它的类型是uintptr，这相当于C语言中的void *。因为没有具体的类型信息，Go编译器没法直接操作它所指向的内存。另一方面，copy方法的签名是func copy(dst, src []Type) int，要求dst和src类型相同。这里src是跳转指令，其类型为[]byte，所以dst的类型也必须是[]byte。

可dst是函数指针，它的类型是uintptr呀。怎么将它转成[]byte呢？这就需要用到unsafe这个包了：

func rawMemoryAccess(p uintptr, length int) []byte {
  return *(*[]byte)(unsafe.Pointer(&reflect.SliceHeader{
    Data: p,
    Len:  length,
    Cap:  length,
  }))
}

[]byte说白了就是一个三元组(Data, Len, Cap)，其中Data保存内存区域开始的地址，其他字段含义请参考这篇文章。通过reflect.SliceHeader，我们动态构建了一个[]byte切片的底层对象。但这个对象没办法直接转换成[]byte。为此，我们需要借助unsafe.Ponter。但是我们只能把指针强转成unsafe.Pointer，所以我们需要对reflect.SliceHeader取地址。这样转成unsafe.Pointer之后得到的实际上是一个指向[]byte的指针。unsafe.Pointer理论上是可以强转成任意类型的指针，编译器不会再检查对应的类型，如果底层类型不一致，后果自负。所以我们可以强转成*[]byte然后通过*提取到指向的内容，也就是[]byte对象。

整个过程比较绕，而且unsafe.Pointer也很不常用，所以过程比较麻烦。Go官方也觉得不好理解，所以在 v1.17 引入了unsafe.Slice方法，上面的代码可以简化成：

func rawMemoryAccess(p uintptr, length int) []byte {
  return *(*[]byte)(unsafe.Slice(p, length))
}

最终，我们希望patch的代码是这个样子的：

func patch(target, replacement interface{}) {
  from := reflect.ValueOf(target).Pointer()
  to := reflect.ValueOf(replacement).Pointer()
  
  data := jmp(to)
  f := rawMemoryAccess(from, len(data))
  copy(f, data)
}

大家可以运行一下试试。执行的时候一定要加上-gcflags '-N -l'告诉编译器不要优化代码，也不要内联函数。不然没法实现「打桩」。

go run -gcflags '-N -l' foo.go

我在自己的 mac 上执行会直接报错：

unexpected fault address 0x10a32c0
fatal error: fault
[signal SIGBUS: bus error code=0x2 addr=0x10a32c0 pc=0x1064fbe]

goroutine 1 [running]:
...
exit status 2

这又是为何？跟文件一样，程序的内存数据也是有权限的概念的，有可读、可写和可执行三种权限。代码段的数据因为是从二进制文件加载的，一般是只有可读和可执行两种权限，默认不能在程序的运行过程中修改。如果直接修改，操作系统就会报错。那有办法改权限吗？当然有。这就得用到syscall.Mprotect这个系统调用。

先别高兴地太早，这个syscal.Mprotect没法随意修改内存的权限，必须要以内存的页(page)为单位进行修改。也就是说，我们想要修改a()的代码区，就必须找到一共要修改哪些页。要想确定有哪些页，最重要的是确定第一页内存地址。为此，我们需要使用下面的函数：

func pageStart(ptr uintptr) uintptr {
  return ptr & ^(uintptr(syscall.Getpagesize() - 1))
}

其中ptr是要修改的函数指针（内存地址），syscall.Getpagesize()是当前系统的页的长度。这个算法的本质是ptr &^ (pageSize - 1)。那怎么理解这个算法呢？（我不是计算机科班出身，搞了半天才弄明白😭）

假设ptr落在了某一页，它的页起始地址是p0。定义d = ptr - p0，则一定有d < pageSize，否则ptr就会落入下一页。要想得到p0，我们需要确定d。我们不妨假设pageSize = 2^3 = 8，对应二进制为0b1000。因为d小于8，所以d的范围一定是0b0000-0b0111，d的最大值是0b0111。也就是说d只有低n位有值（n = 3），其他位肯定都是零，否则d的值就会超过0b0111。

为了计算p0 = ptr - d，我会需要从ptr中减掉d。而d的值只占据了ptr的低n位（2^n = pageSize），所以我们只需要把这低n位清零就能得到p0。而2^n - 1的低n位正好全为1，取反之后再跟ptr按位做逻辑与操作（也就是&^），正好把低n位清零，得到ptr - d也就是p0。

一旦确定了第一页的地址，后面的页就好说了：

func mprotectCrossPage(addr uintptr, length int, prot int) {
  pageSize := syscall.Getpagesize()
  for p := pageStart(addr); p < addr+uintptr(length); p += uintptr(pageSize) {
    page := rawMemoryAccess(p, pageSize)
    if err := syscall.Mprotect(page, prot); err != nil {
      panic(err)
    }
  }
}

我们可以通过下面的函数来修改某一片内存的权限：

mprotectCrossPage(location, len(data), syscall.PROT_READ|syscall.PROT_EXEC|syscall.PROT_WRITE)

最终我们写了这样一个复制函数：

func copyTo(p uintptr, data []byte) {
  f := rawMemoryAccess(p, len(data))

  mprotectCrossPage(p, len(data), syscall.PROT_READ|syscall.PROT_EXEC|syscall.PROT_WRITE)
  copy(f, data[:])
  mprotectCrossPage(p, len(data), syscall.PROT_READ|syscall.PROT_EXEC)
}

copyTo把传入的函数指针转成var f []byte，然后调用mprotectCrossPage修改内存权限，最后使用copy将跳转指令拷贝到目标函数代码区。为了安全起见，函数结束需要把内存的权限改回来。

到这里，所有障碍都扫清了✌️🍾🍻😇。完整代码在这里。运行代码，会看到程序输出2而不是1。

最后总结一下实现猴子补丁的三板斧🪓：获取函数地址➔构造跳转指令➔覆盖原函数代码段。简单粗暴！

那这个黑科技是完美方案吗？显然不是，这个方案有很多问题：

• 不可移植。不同架构的处理器需要构造不同的跳转指令。好在需要支持的架构不多，顶多是再支持一下 arm64，这方面的工作可以参见gomonkey项目。

• 只能做全局打桩。不同协程打桩会相互影响，甚至都不是并发安全的。比如有多个测试用例都调用time.Now()函数。如果一个用例修改了time.Now()的行为，而且是并发跑测试用例，那其他测试用例都会受到影响。因此，我们内部的测试用例都是线性运行的。我正在着手研究如何解决这个问题，有一定的进展，后面会再写一篇文章介绍Go语言实现猴子补丁【二】。

• 可能破坏Go语言的运行环境，产生各种诡异的问题。相关的讨论可以参考 Bouke 原文给出的链接。不过，如果是只用在单元测试场景的话，应该问题不大。

• 可能因为Go后续内部实现的变更而失效。这是我最担心的问题。如果我们写了成千上万的测试用例，到某个新版本都不能用，那是一种怎样的体验。不过从Go1.17引入新ABI的变更来看，这种问题出现的可能性不大。

以上就是本文的全部内容。希望能给大家带来启发。