Go语言实现猴子补丁【三】

涛叔 2021-09-12 ⏳4.8分钟(1.9千字)

这是Go语言实现猴子补丁的最后一篇。在上一篇中，我分享了一种协程隔离的打桩思路。但在文章结尾部分也出了该方案的几个问题。其中最核心的问题是不支持闭包引用，而这个功能在写单元测试又极为常用，所以不得不花时间研究解决方案。期间我还给 Bouke 发邮件看能否接手他的项目继续开发，他也没回复，估计是不愿意。所以，我把自己的工作整理了一下，最终形成go-kiss这个项目。好了，今天就给大家分享一下修复闭包引用问题的过程。

本文假定读者已经理解 monkey 的工作原理。如果你是第一次阅读本文，请先阅读Go语言实现猴子补丁和Go语言实现猴子补丁【二】。

我在第一篇文章中说 Bouke 采用了间接跳转的方案，代码如下：

movabs rdx, 0x?? ; 将内存地址存到寄存器 rdx
jmp DWORD PTR [rdx]

而且 Bouke 的方案使用了 rdx 寄存器。Go语言在运行闭包函数的时候会使用这个寄存器存储闭包上下文指针（Closure context pointer）。我当时也纳闷用这个寄存器到底会不会有问题🤨我也给 Bouke 写信，他回复说没有问题，但没细讲原因。

因为有这样一个疑问，我自己把跳转方式改成了直接跳转。代码如下：

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

然后所有的函数指针都可以使用使用下面的方式读取：

v := reflection.ValueOf(f)
p := v.Pointer()

不需要处理一级指针和二级指针了，我以为很美好。但这一改动却导致闭包调用出错。示例代码如下：

package main

import (
  "fmt"
  "github.com/go-kiss/monkey"
)

func foo() bool { retrun false }

func main() {
  t := true
  monkey.Patch(foo, func() bool { return t })
  fmt.Println(foo())
}

输出结果居然为false，有问题。当时，我虽然改用了直接跳转方案，但寄存器还是用了 rdx。所以发现问题后，我第一反应是可能是因为覆盖 rdx 寄存器触发了问题。于是，我将寄存器换成了 r13，结果非但没好，而是直接报空指针错误！陷入僵局💣。

我在想到底是哪里出了问题。我试着自己脑补一下。在前面的代码中，foo 是普通函数，补丁函数是闭包函数。闭包函数在调用的时候需要准备好上下文，但我们调的是 foo 函数，不是闭包。是不是因为 foo 不是闭包函数，导致补丁函数执行的时候闭包上下文没有初始化，进而导致空指针错误呢？

如果真是这个错误，那我能不能在执行补丁函数之前手工初始化闭包上下文指针呢？可问题是要从哪里找这个指针呢？谷歌了半天也没有找到直接资料。我甚至想看看 Go 编译器是怎么编译这段代码的。当然，一打开代码就被劝退了🤣

最后，我不得不静下心来研究一下 Go 函数调用原理。读了好几遍The Go low-level calling convention on x86-64，发现没有提到闭包函数调用。最后只能读Go 1.1 Function Calls，这是Russ Cox在 2013 发布的 Go 函数重构方案。里面最主要的内容是说 Go 中的函数是二级指针（其实 Bouke 早就指出这一点了）。函数指向一个结构体，结构体的第一个字段指向真正的代码段，后面的字段存放其他相关信息。到底有哪些信息呢？其实我也没弄明白。但我猜闭包上下文信息可能就在这个结构体中。

为了验证我的猜想，我构造如下代码：

package main

func call(f func()) { f() }

func test() {
  x := 1
  call(func() { x += 2 })
}

func main() { test() }

这里 call 接受并执行一个函数。而 test 中给 call 传入了一个闭包函数。这个闭包上下文指针一定是在 test 中确定的，而且一定可以通过函数指针获取。因为 call 在执行 f 的时候已经不知道它是不是一个闭包函数了。然后，我们查看汇编代码：

go build -ldflags=-w -gcflags '-N -l' main.go
go tool objdump main|less

我们找到main.call的汇编代码如下：

MOVQ AX, 0x10(SP)
MOVQ 0(AX), CX
MOVQ AX, DX
CALL CX
MOVQ 0(SP), BP
ADDQ $0x8, SP

Go语言使用 DX 寄存器保存闭包上下文。这是一个虚拟寄存器，具体到 amd64 平台就是 rdx 寄存器。这里的 AX 寄存器保存 main.call 参数f，也就是指向函数对象的指针。第三行的 MOVQ AX, DX 直接把这个指针保存到 DX 寄存器了！这一定是在为后面的调用恢复装包上下文！换句话说，DX 寄存器存放的闭包上下文指针应该就是函数对象的指针。

为了验证我的猜想，我又仔细看了 Bouke 的原始代码和测试用例，发现他的方案居然能正常工作！那他为什么能工作呢？他的方案是二次跳转方案，说白了就是把函数对象的地址存放到 rdx，然后根据 rdx 内容（函数对象第一个字段存放的地址）二次跳转到真正的代码区执行。因为他用了 rdx 寄存器，也就顺便完成了闭包上下文的准备工作！一切都对上了。

所以，我微调了在第二篇文章中讲的jmpTable函数：

func jmpToGoFn(to uintptr) []byte {
  return []byte{
    0x48, 0xBA,
    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, 0x22,     // jmp QWORD PTR [rdx]
  }
}

func jmpTable(g, to uintptr) []byte {
  b := []byte{
    // movq r13, g
    0x49, 0xBD,
    byte(g),
    byte(g >> 8),
    byte(g >> 16),
    byte(g >> 24),
    byte(g >> 32),
    byte(g >> 40),
    byte(g >> 48),
    byte(g >> 56),
    // cmp r12, r13
    0x4D, 0x39, 0xEC,
    // jne $+(2+12)
    0x75, 0x0c,
  }
  b = append(b, jmpToGoFn(to)...)
  return b
}

这样就修复了闭包补丁函数运行报错的问题🎉🎊🍾👏

除了闭包问题外，我还试着解决并发执行可能相互影响的问题。核心思路是在首次 patch 之后，后续的 patch 只更新八个字节的跳转地址。

func (p *patch) Apply() {
  p.patch = p.Marshal()

  v := reflect.ValueOf(p.patch)
  allowExec(v.Pointer(), len(p.patch))

  if p.patched {
    // 后续只更新八个字节的跳转地址
    data := littleEndian(v.Pointer())
    copyToLocation(p.from+2, data)
  } else {
    jumpData := jmpToFunctionValue(v.Pointer())
    copyToLocation(p.from, jumpData)
    p.patched = true
  }
}

我看了 copy 的源码（对应 runtime/slice.go 的 slicecopy 函数，进一步对应 runtime/memmove_amd64.s 的 memmove 函数）。如果是指针对象（八个字节），会直接转换成 movq 指令。理论上是原子操作，但这需要目标地址能被 8 整除（鸣谢 @雷鸣指教）。现在的方案好像不能保证跳转地址所在的地址一定能被 8 整除，所以可能出问题。但是我加了一个测试用例，起了两个协程循环 mock 同一个函数，并没有报错。所以这个问题还需要继续观察和研究。

另外，创建 go-kiss/monkey 项目后，发现 GitHub 开始自动运行单元测试。我的方案在 linux 居然不能运行🤣通过 objdump 工具，最终发现 linux 下获取当前协程的寄存器需要用 fs，而且 windows 下的寄存器用 gs 但偏移量不同。最终修好 linux 下的问题。~~本来也想顺手支持 windows 平台，代码改好了，但单元测试过不了。奈何没有 windows 设备，只能寄希望有志同道合的朋友贡献代码了。~~现在已经支持 windows 平台，具体问题请参考我的第二篇文章。