测试 - monkey - 《golang生态》

Go 语言中的函数值是如何工作的
运行期替换一个函数
monkey">monkey
gomonkey
实际解决问题：指定ip去处理连接

很多人认为 monkey 补丁只能在动态语言，比如 Ruby 和 Python 中才存在。但是，这并不对。因为计算机只是很笨的机器，我们总能让它做我们想让它做的事儿！让我们看看 Go 中的函数是怎么工作的，并且，我们如何在运行时修改它们。本文会用到大量的 Intel 汇编(为方便理解还是用AT&T)，所以，我假设你可以读汇编代码，或者在读本文时正拿着参考手册.

如果你对 monkey 补丁是怎么工作的不感兴趣，你只是想使用它的话，你可以在这里找到对应的库文件
让我们看看一下代码产生的汇编码：

package main
func a() int { return 1 }
func main() {
  print(a())
}

go build -gcflags=-l 来编译，以避免内联。在本文中我假设你的电脑架构是 64 位，并且你使用的是一个基于unix 的操作系统比如 Mac OSX 或者某个 Linux 系统。
当代码编译后，我们用 Hopper 来查看，可以看到如上代码会产生如下汇编代码：

函数调用过程：在0x105e1c0 调用了a()函数，然后在0x105e1c0 将调用a函数的返回值给了print

Go 语言中的函数值是如何工作的

package main
import (
  "fmt"
  "unsafe"
)
func a() int { return 1 }
func main() {
  f := a
  fmt.Printf("0x%x\n", *(*uintptr)(unsafe.Pointer(&f)))
  fmt.Printf("0x%x\n", **(**uintptr)(unsafe.Pointer(&f)))
}

我在第11行将 a 赋值给 f，这意味者，执行 f() 就会调用 a。然后我使用 Go 中的 unsafe 包来直接读出 f 中存储的值。如果你是有 C 语言的开发背景，你可以会觉得 f 就是一个简单的函数指针，并且这段代码第一个输出为 0x10d07d8。当我在我的机器上运行时，我得到的是 0x10a46a0, 这与代码13行的输出对应：

这说明f 并不是一个函数指针，而是一个指向函数指针的指针

让我们看看调用一个函数值是怎么工作的。我把上面的代码修改一下，在给 f 赋值后直接调用它。

package main
func a() int { return 1 }
func main() {
    f := a
    f()
}

首先将 f 加载到%rax，然后call %rax，注意：f是指向函数指针的指针，因此 %rax 才能代表函数的地址（0x105e180）

运行期替换一个函数

我们希望做到的是，让下面的代码输出 2：

package main
func a() int { return 1 }
func b() int { return 2 }
func main() {
    replace(a, b)
    print(a())
}

现在我们该怎么实现这种替换？我们需要修改函数 a 跳到 b 的代码，而不是执行它自己的函数体。本质上，我们需要这么替换，把 b 的函数值加载到 rdx 然后跳转到 rdx 所指向的地址。

mov rdx, main.b.f ; 48 C7 C2 ?? ?? ?? ??
jmp [rdx] ; FF 22

我将上述代码编译后产生的对应的机器码列出来了（用在线编译器，比如这个，你可以随意尝试编译）。很明显，我们需要写一个能产生这样机器码的函数，它应该看起来像这样：

func assembleJump(f func() int) []byte {
  funcVal := *(*uintptr)(unsafe.Pointer(&f))
  return []byte{
    0x48, 0xC7, 0xC2,
    byte(funcval >> 0),
    byte(funcval >> 8),
    byte(funcval >> 16),
    byte(funcval >> 24), // MOV rdx, funcVal
    0xFF, 0x22,          // JMP [rdx]
  }
}

现在万事俱备，我们已经准备好将 a 的函数体替换为从 a 跳转到 b了！下述代码尝试直接将机器码拷贝到函数体中。

package main
import (
    "syscall"
    "unsafe"
)
func a() int { return 1 }
func b() int { return 2 }
func rawMemoryAccess(b uintptr) []byte {
    return (*(*[0xFF]byte)(unsafe.Pointer(b)))[:]
}
func assembleJump(f func() int) []byte {
    funcVal := *(*uintptr)(unsafe.Pointer(&f))
    return []byte{
        0x48, 0xC7, 0xC2,
        byte(funcVal >> 0),
        byte(funcVal >> 8),
        byte(funcVal >> 16),
        byte(funcVal >> 24), // MOV rdx, funcVal
        0xFF, 0x22,          // JMP [rdx]
    }
}
func replace(orig, replacement func() int) {
    bytes := assembleJump(replacement)
    functionLocation := **(**uintptr)(unsafe.Pointer(&orig))
    window := rawMemoryAccess(functionLocation)
    copy(window, bytes)
}
func main() {
    replace(a, b)
    print(a())
}

然而，运行上述代码并没有达到我们的目的，实际上，它会产生一个段错误。这是因为默认情况下，已经加载的二进制代码是不可写的。我们可以用 mprotect 系统调用来取消这个保护，并且这个最终版本的代码就像我们期望的一样，把函数 a 替换成了 b，然后 ‘2’ 被打印出来。

package main
import (
    "syscall"
    "unsafe"
)
func a() int { return 1 }
func b() int { return 2 }
func getPage(p uintptr) []byte {
    return (*(*[0xFFFFFF]byte)(unsafe.Pointer(p & ^uintptr(syscall.Getpagesize()-1))))[:syscall.Getpagesize()]
}
func rawMemoryAccess(b uintptr) []byte {
    return (*(*[0xFF]byte)(unsafe.Pointer(b)))[:]
}
func assembleJump(f func() int) []byte {
    funcVal := *(*uintptr)(unsafe.Pointer(&f))
    return []byte{
        0x48, 0xC7, 0xC2,
        byte(funcVal >> 0),
        byte(funcVal >> 8),
        byte(funcVal >> 16),
        byte(funcVal >> 24), // MOV rdx, funcVal
        0xFF, 0x22,          // JMP rdx
    }
}
func replace(orig, replacement func() int) {
    bytes := assembleJump(replacement)
    functionLocation := **(**uintptr)(unsafe.Pointer(&orig))
    window := rawMemoryAccess(functionLocation)
    page := getPage(functionLocation)
    syscall.Mprotect(page, syscall.PROT_READ|syscall.PROT_WRITE|syscall.PROT_EXEC)
    copy(window, bytes)
}
func main() {
    replace(a, b)
    print(a())
}

monkey

注意：编译时，一定要使用 -gcflags=”all=-N -l”

monkey.Patch(,

package main
import (
    "fmt"
    "os"
    "strings"
    "bou.ke/monkey"
)
func main() {
    monkey.Patch(fmt.Println, func(a ...interface{}) (n int, err error) {
        s := make([]interface{}, len(a))
        for i, v := range a {
            s[i] = strings.Replace(fmt.Sprint(v), "hell", "*bleep*", -1)
        }
        return fmt.Fprintln(os.Stdout, s...)
    })
    fmt.Println("what the hell?") // what the *bleep*?
}

可使用 monkey.Unpatch() 取消补丁

monkey.PatchInstanceMethod(, , )

package main
import (
    "fmt"
    "net"
    "net/http"
    "reflect"
    "bou.ke/monkey"
)
func main() {
    var d *net.Dialer // Has to be a pointer to because `Dial` has a pointer receiver
    monkey.PatchInstanceMethod(reflect.TypeOf(d), "Dial", func(_ *net.Dialer, _, _ string) (net.Conn, error) {
        return nil, fmt.Errorf("no dialing allowed")
    })
    _, err := http.Get("http://google.com")
    fmt.Println(err) // Get http://google.com: no dialing allowed
}

注意，目前只对一个实例打补丁是不可能的，PatchInstanceMethod将为所有实例打补丁
不要尝试monkey.Patch(instance.Method, replacement)，它不会工作
可使用 monkey.UnpatchInstanceMethod(, ) 取消补丁

可使用 monkey.UnpatchAll 取消所有补丁

如果你想在替换函数中调用原始函数，你需要使用monkey.PatchGuard。patchguard可以让你轻松地删除和恢复补丁，这样你就可以调用原来的功能。

package main
import (
    "fmt"
    "net/http"
    "reflect"
    "strings"
    "bou.ke/monkey"
)
func main() {
    var guard *monkey.PatchGuard
    guard = monkey.PatchInstanceMethod(reflect.TypeOf(http.DefaultClient), "Get", func(c *http.Client, url string) (*http.Response, error) {
        guard.Unpatch()
        defer guard.Restore()
        if !strings.HasPrefix(url, "https://") {
            return nil, fmt.Errorf("only https requests allowed")
        }
        return c.Get(url)
    })
    _, err := http.Get("http://google.com")
    fmt.Println(err) // only https requests allowed
    resp, err := http.Get("https://google.com")
    fmt.Println(resp.Status, err) // 200 OK <nil>
}

gomonkey

monkey 归档了，转至：https://github.com/agiledragon/gomonkey

实际解决问题：指定ip去处理连接

一台机器配置了多个IP并不罕见，当你在这台机器连接其它的TCP服务器时，本地到底使用的是哪一个IP地址呢？如serverfault有人提出的问题，在默认的情况下，Linux会依照子网的分类，选择和服务器在相同的子网的本地地址，但是如果同一个子网中配置了多个IP地址，那么Linux会选择此子网的”主”IP地址作为本地Ip地址连接服务器。

在我这个项目中，会有很多的网络连接，比如连接mysql，连接clickhouse,连接第三方的HTTP API服务，连接Kafka、连接 Redis等。不幸的事，当使用第三方库比如go-sql-driver/mysql、go-redis时，Linux所选择的本地IP地址并不是我期望的本地IP地址，导致权限验证失败无法连接。

本质上，无论是go-sql-driver/mysql或者go-redis,都是基于net.Dial或者net.DialContext建立的TCP连接。go-sql-driver/mysql 提供了 RegisterDialContext用于定制化Dial,go-redis提供了Dialer字段用来定制，你如果想指定本地的IP地址，可以通过定制的net.Dialer来实现：

localAddrDialier := &net.Dialer{
    LocalAddr: localAddr,
}

这是一个不错的、传统的方法，唯一不爽的是，每种类型我都需要进行地址，访问mysql、访问redis、访问kafka、访问第三方库、访问服务器……, 有没有一劳永逸的方法呢？
有!
bouk/monkey是一个相当相当”骇客”的技术，当过运行时动态将方法的实现替换成JMP 新的函数, 来实现在运行时替换方法。经常我们会在单元测试的时候用来”Mock”一些方法，非常的有效，这一次，我要尝试使用它替换所有的net.Dialer.Dial或者net.Dialer.DialContext方法，来实现强制指定本地地址。

// 指定要使用的本地地址. // by https://colobu.com
   localAddr := &net.TCPAddr{
    IP:   net.ParseIP(localIP),
    Port: 0,
}
var d *net.Dialer
   // 替换Dialer.DialContext方法
dialContextGuard = monkey.PatchInstanceMethod(reflect.TypeOf(d), "DialContext", func(d *net.Dialer, ctx context.Context, network, address string) (net.Conn, error) {
       // 临时恢复
    dialContextGuard.Unpatch()
    defer dialContextGuard.Restore()
    if network == "tcp" || network == "tcp4" || network == "tcp6" {
        localAddrDialier := &net.Dialer{
            LocalAddr: localAddr,
        }
           // 使用指定本地地址的dialer
        return localAddrDialier.DialContext(ctx, network, address)
    }
       // 其它情况，比如UDP、UnixDomain等，使用标准库的方法
    return d.DialContext(ctx, network, address)
})
   // 替换Dail方法
dialGuard = monkey.PatchInstanceMethod(reflect.TypeOf(d), "Dial", func(d *net.Dialer, network, address string) (net.Conn, error) {
       // 临时恢复
    dialGuard.Unpatch()
    defer dialGuard.Restore()
    if network == "tcp" || network == "tcp4" || network == "tcp6" {
        localAddrDialier := &net.Dialer{
            LocalAddr: localAddr,
        }
           // 使用指定本地地址的dialer
        return localAddrDialier.Dial(network, address)
    }
       // 其它情况，比如UDP、UnixDomain等，使用标准库的方法
    return d.Dial(network, address)
})

替换了这两个方法后，之后即使新建立net.Dailer对象，也是使用替换后的方法执行。
这里并没有考虑并发定位情况，如果你的程序有并发的调用Dial或者DialContext,你需要加锁。

这样，我们就一劳永逸的解决了指定本地IP地址创建TCP连接的问题，无需改动标准库的代码，无需逐个定制Dial方法。

同样的，你可以更改标准库的net.DefaultResolver, 这是标准库用来进行域名解析的实现，支持Go自己的解析实现和CGO方式查询。本身它是一个struct,而不是一个接口，所以虽然它是一个单例的对象，但是通常情况下你也没有多少定制化的可能。比如在调用LookupIP方法时,你想使用自己的一个协议返回IP列表，而不是查询本地文件或者DNS服务器，你基本是没有办法的。但是通过bouk/monkey，你可以更改LookupIP方法，这样你就可以定制了。

所以，bouk/monkey不仅仅可以用来在单元测试中mock对象和方法，还可以在应用运行中替换一些常规没有办法更改的函数。