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

2021-09-12

这是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 设备,只能寄希望有志同道合的朋友贡献代码了。

以上就是今天的全部内容。后面准备把部门内部的测试用例都迁移到新库上。希望能给大家一些启发。

「taoshu」微信公众号