修正monkey补丁在windows下失败的问题

2022-02-22

起因

公司内正在做的一个项目最近正在写接口集成测试的脚本，但是由于项目内部分代码使用了全局单例，这部分很难进行mock测试，虽然通过gomonkey可以进行一定程度的patch，但是由于gomonkey当前是非协程安全的，也就是没办法做并行的测试，然而我们的项目用例很多，如果进行串行测试，整个CI流程走下来要接近10分钟，这会非常影响开发的效率，故我开始寻求解决方案。

转机

最终在网上发现了有人做了一个支持多协程的monkey补丁go-kiss/monkey，看了作者的实现文章后，其基本原理就是通过写入一段汇编实现的switch代码段，通过当前协程ID的比较，进行分支跳转到不同的patch上，虽然当前的版本比较粗糙，安全性据说也没有完善处理，但是作为测试我觉得是够用了，所以打算在项目里先试试。

问题

测试这个monkey库的时候发现了一个问题，在AMD64 + Linux 的环境下是能够正常跑通独立patch的逻辑，但是AMD64 + Windows下却失效了，然而我司不少小伙伴的开发机还是使用Windows平台，如果在Windows下无法正常使用，就需要用wsl等方式来跑自测，这样会影响到原先的开发流程。既然作者把基本原理都提供了，故我打算尝试解决这个Bug。

排查

排查Hook流程

写了一段测试代码用于排查，执行后停在input处，拿x86debug附加进去，找到协程判断的汇编段

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

//go:noinline
func Add(a, b int) int {
    return a + b
}

func Add1(a, b int) int {
    return 10
}


func main() {
    monkey.Patch(Add, Add1)
    var input string
    fmt.Scanln(&input)
    fmt.Printf("g1 add = %d\n", Add(99, 99))
}

发现其判断协程ID的逻辑没有通过，然而测试代码只有一个协程，明显和我们的预期相悖，按理想逻辑，cmp应该成功，不走jne跳转，实际上却判断失败，跳过了patch，手工将跳转nop掉之后，能够正常的执行patch代码，所以初步定位是协程ID获取的地方有错误，导致程序中获取的协程ID和汇编代码中获取的协程ID不一致。

// 错误的代码
func getg() []byte {
    return []byte{
        // mov r12,QWORD PTR gs:0x28
        0x65, 0x4C, 0x8B, 0x24, 0x25, 0x28, 0x00, 0x00, 0x00,
    }
}

比较Linux和Windows下获取协程ID的差异

按照原文作者的方法，写了一段代码用于原始获取汇编指令

import "github.com/huandu/go-tls/g"

func main() {
    _ = g.G()   // 这里其实没有必要按原文那样修改源码，直接这样也能找到的
}

在Linux和Windows下分别编译

go build -ldflags=-w -gcflags '-N -l' main.go

dump

// linux
go tool objdump main > main.txt
// windows
go tool objdump main.exe > main.txt

找到getg的汇编代码比较
Linux

TEXT github.com/huandu/go-tls/g.getg.abi0(SB) /root/go/pkg/mod/github.com/huandu/go-tls@v1.0.1/g/getg_amd64.s
  getg_amd64.s:10    0x455420        64488b0425f8ffffff    MOVQ FS:0xfffffff8, AX    
  getg_amd64.s:11    0x455429        4889442408        MOVQ AX, 0x8(SP)    
  getg_amd64.s:12    0x45542e        c3            RET

Windows

TEXT github.com/huandu/go-tls/g.getg.abi0(SB) C:/Users/lzq/go/pkg/mod/github.com/huandu/go-tls@v1.0.1/g/getg_amd64.s
  getg_amd64.s:9    0x45ab80        65488b0c2528000000    MOVQ GS:0x28, CX    
  getg_amd64.s:10    0x45ab89        488b8100000000        MOVQ 0(CX), AX        
  getg_amd64.s:11    0x45ab90        4889442408        MOVQ AX, 0x8(SP)    
  getg_amd64.s:12    0x45ab95        c3            RET

可以发现其实非原文所说的单纯只是偏移和寄存器不一致，Windows下还多了一行

  getg_amd64.s:10    0x45ab89        488b8100000000        MOVQ 0(CX), AX

也就是说在Windows下其实是2级指针，从GS:0x28获取出来的是一个二级指针，还需要一次move才能获取到g的指针

修复和测试

修复

知道问题了，于是尝试修改一波，不过我对汇编不是很熟，大致想法是先将二级指针移到r13，在二次取值回r12，这样只需要修改getg的逻辑，其他的代码可以保持原有的逻辑，写入的协程ID在后面的代码中会覆盖掉r12，照我的理解应该是没有问题。

func getg() []byte {
    return []byte{
        // mov r13,QWORD PTR gs:0x28
        0x65, 0x4C, 0x8B, 0x2C, 0x25, 0x28, 0x00, 0x00, 0x00,
        // mov r12,QWORD PTR [r13]
        0x4D, 0x8B, 0x65, 0x00,
    }
}

修复后的代码我也提了个pr给原仓库

测试

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

//go:noinline
func Add(a, b int) int {
    return a + b
}

func Add1(a, b int) int {
    return 10
}

func Add2(a, b int) int {
    return 100
}

func main() {

    go func() {
        monkey.Patch(Add, Add1)
        fmt.Printf("g1 add = %d\n", Add(99, 99))
    }()

    go func() {
        monkey.Patch(Add, Add2)
        fmt.Printf("g2 add = %d\n", Add(99, 99))
    }()
    
    var input string
    fmt.Scanln(&input)
}

输出

C:\Users\lzq\AppData\Local\Temp\GoLand\___go_build_monkeyDemo.exe
g1 add = 10
g2 add = 100

测试符合预期，收工。

总结

目前monkey patch这块还是比较hack的行为，由于和底层汇编实现强关联，需要单独去适配不同的平台、不同的架构还是挺蛋疼的，不知道是否有更好的方式来实现，比如使用go的中间层汇编在执行的时候转换成字节码再patch进去，不过目前好像没有找到实现方式，待定吧~