Golang Mock 原理分析

package main

//go:generate mockgen -source=main.go -destination=foo_mock.go -package=main Foo

type Foo interface {
	Say(string, []string) string
}

func main() {
}

// MockFoo is a mock of Foo interface.
type MockFoo struct {
   ctrl     *gomock.Controller
   recorder *MockFooMockRecorder
}

// MockFooMockRecorder is the mock recorder for MockFoo.
type MockFooMockRecorder struct {
   mock *MockFoo
}

// NewMockFoo creates a new mock instance.
func NewMockFoo(ctrl *gomock.Controller) *MockFoo {
   mock := &MockFoo{ctrl: ctrl}
   mock.recorder = &MockFooMockRecorder{mock}
   return mock
}

// Call represents an expected call to a mock.
type Call struct {
   t TestHelper // for triggering test failures on invalid call setup

   receiver   interface{}  // the receiver of the method call
   method     string       // the name of the method
   methodType reflect.Type // the type of the method
   args       []Matcher    // the args
   origin     string       // file and line number of call setup

   preReqs []*Call // prerequisite calls

   // Expectations
   minCalls, maxCalls int

   numCalls int // actual number made

   // actions are called when this Call is called. Each action gets the args and
   // can set the return values by returning a non-nil slice. Actions run in the
   // order they are created.
   actions []func([]interface{}) []interface{}
}

type Controller struct {
   // T should only be called within a generated mock. It is not intended to
   // be used in user code and may be changed in future versions. T is the
   // TestReporter passed in when creating the Controller via NewController.
   // If the TestReporter does not implement a TestHelper it will be wrapped
   // with a nopTestHelper.
   T             TestHelper
   mu            sync.Mutex
   expectedCalls *callSet
   finished      bool
}

// callSet represents a set of expected calls, indexed by receiver and method
// name.
type callSet struct {
	// Calls that are still expected.
	expected map[callSetKey][]*Call
	// Calls that have been exhausted.
	exhausted map[callSetKey][]*Call
}

// callSetKey is the key in the maps in callSet
type callSetKey struct {
	receiver interface{}
	fname    string
}

func TestMockFoo(t *testing.T) {
   ctrl := gomock.NewController(t)
   mockFoo := NewMockFoo(ctrl)
   mockFoo.EXPECT().
      Say("foo", []string{"a", "b"}). // 创建匹配 Say("foo", []string{"a", "b"}) 的 gomock.Call
      Do(func(arg string, arg2 []string) string { // 只执行，不会使用其返回值
         t.Logf("Do(%s, %v)", arg, arg2)
         return arg
      }).
      SetArg(1, []string{"c", "d"}). // 修改第 2 个参数
      DoAndReturn(func(arg string, arg2 []string) string { // 执行，并使用其返回值
         t.Logf("DoAndReturn(%s, %v)", arg, arg2)
         return arg
      }).
      Return("Return"). // 设置返回值为 "Return"
      AnyTimes()
   t.Log(mockFoo.Say("foo", []string{"a", "b"})) // 调用 Say，从 gomock.Controller.expectedCalls 中查找匹配的 gomock.Call，顺序调用 gomock.Call.actions
}

bar := Bar{field: "field"}
say := bar.Say // 必须把方法复制给一个变量
stub = gostub.StubFunc(&say, "bar")
println(bar.Say("foo"), say("foo")) // fieldfoo bar
stub.Reset()
println(bar.Say("foo"), say("foo")) // fieldfoo fieldfoo

// Stubs represents a set of stubbed variables that can be reset.
type Stubs struct {
   // stubs is a map from the variable pointer (being stubbed) to the original value.
   stubs   map[reflect.Value]reflect.Value
   origEnv map[string]envVal
}

// StubFunc replaces a function variable with a function that returns stubVal.
// funcVarToStub must be a pointer to a function variable. If the function
// returns multiple values, then multiple values should be passed to stubFunc.
// The values must match be assignable to the return values' types.
func (s *Stubs) StubFunc(funcVarToStub interface{}, stubVal ...interface{}) *Stubs {
   funcPtrType := reflect.TypeOf(funcVarToStub)
   if funcPtrType.Kind() != reflect.Ptr ||
      funcPtrType.Elem().Kind() != reflect.Func {
      panic("func variable to stub must be a pointer to a function")
   }
   funcType := funcPtrType.Elem()
   if funcType.NumOut() != len(stubVal) {
      panic(fmt.Sprintf("func type has %v return values, but only %v stub values provided",
         funcType.NumOut(), len(stubVal)))
   }

   return s.Stub(funcVarToStub, FuncReturning(funcPtrType.Elem(), stubVal...).Interface())
}

// Stub replaces the value stored at varToStub with stubVal.
// varToStub must be a pointer to the variable. stubVal should have a type
// that is assignable to the variable.
func (s *Stubs) Stub(varToStub interface{}, stubVal interface{}) *Stubs {
   v := reflect.ValueOf(varToStub)
   stub := reflect.ValueOf(stubVal)

   // varToStub 必须为变量的指针
   if v.Type().Kind() != reflect.Ptr {
      panic("variable to stub is expected to be a pointer")
   }

   if _, ok := s.stubs[v]; !ok {
      // 存储 varToStub 原始值
      s.stubs[v] = reflect.ValueOf(v.Elem().Interface())
   }

   // 设置为新的值
   // *varToStub = stubVal
   v.Elem().Set(stub)
   return s
}

type Bar struct {
   field string
}

func (b Bar) Say(arg string) string {
   return b.field + arg
}

func fn(i int) int {
   return i+1
}

func main() {
   foo := "foo"
   stub := gostub.Stub(&foo, "stubFoo")
   println(foo) // stubFoo
   stub.Reset()
   println(foo) // foo

   slice := []int{1, 2}
   stub = gostub.Stub(&slice, []int{3})
   fmt.Println(slice) // [3]
   stub.Reset()
   fmt.Println(slice) // [1 2]

   ff := fn
   stub = gostub.Stub(&ff, func(i int) int{return 0})
   println(ff(1)) // 0
   stub.Reset()
   println(ff(1)) // 2

   stub = gostub.StubFunc(&ff, 1)
   println(ff(1)) // 1
   stub.Reset()
   println(ff(1)) // 2

   bar := Bar{field: "field"}
   say := bar.Say
   stub = gostub.StubFunc(&say, "bar")
   println(bar.Say("foo"), say("foo")) // fieldfoo bar
   stub.Reset()
   println(bar.Say("foo"), say("foo")) //fieldfoo fieldfoo

   os.Setenv("GOSTUB_VAR", "value")
   stub = gostub.New()
   stub.SetEnv("GOSTUB_VAR", "stub_value")
   println(os.Getenv("GOSTUB_VAR")) // stub_value
   stub.Reset()
   println(os.Getenv("GOSTUB_VAR")) // value
}

func getDoubleFunc(funcType reflect.Type, outputs []OutputCell) reflect.Value {
    // 判断返回值个数是否正确
   if funcType.NumOut() != len(outputs[0].Values) {
      panic(fmt.Sprintf("func type has %v return values, but only %v values provided as double",
         funcType.NumOut(), len(outputs[0].Values)))
   }

    // 构造返回结果集
   slice := make([]Params, 0)
   for _, output := range outputs {
      t := 0
      if output.Times <= 1 { // 每个值至少被返回一次
         t = 1
      } else {
         t = output.Times
      }
      for j := 0; j < t; j++ { // 根据返回次数将值追加到结果集中
         slice = append(slice, output.Values)
      }
   }

   i := 0 // 调用次数统计
   len := len(slice)
   return reflect.MakeFunc(funcType, func(_ []reflect.Value) []reflect.Value {
      if i < len {
         i++
         return GetResultValues(funcType, slice[i-1]...)
      }
      panic("double seq is less than call seq")
   })
}

//go:noinline
func bar() string {
   return "bar"
}

func main() {
   fmt.Println(bar) // 0x10a2e20
   println(unsafe.Pointer(reflect.ValueOf(bar).Pointer())) // 0x10a2e20
}

func (this *Patches) ApplyCore(target, double reflect.Value) *Patches {
   this.check(target, double) // 类型检查
   if _, ok := this.originals[target]; ok {
      panic("patch has been existed")
   }

   this.valueHolders[double] = double // 因为 mock 函数通常是一个闭包，也就是个局部作用域的对象，为了防止 mock 函数被 GC 回收掉，需要增加引用
   // 替换 target 的机器指令，返回的 origin 是 target 会被覆盖的机器指令
   original := replace(*(*uintptr)(getPointer(target)), uintptr(getPointer(double)))
   // 保存 target 被覆盖的机器指令，用于恢复 target
   this.originals[target] = original
   return this
}

type funcValue struct {
	_ uintptr
	p unsafe.Pointer
}

func getPointer(v reflect.Value) unsafe.Pointer {
	return (*funcValue)(unsafe.Pointer(&v)).p
}

type Value struct {
   typ *rtype
   ptr unsafe.Pointer
   flag
}

package main

//go:noinline
func foo() string {
   return "foo"
}

func main() {
   funcVar := foo
   println(funcVar())
   funcVar2 := foo
   println(funcVar2())
}

"".main STEXT size=205 args=0x0 locals=0x28 funcid=0x0
		        ......
        0x0021 00033 (ff.go:10) MOVQ    "".foo·f(SB), AX ; funcVar()
        0x0028 00040 (ff.go:10) LEAQ    "".foo·f(SB), DX
        0x002f 00047 (ff.go:10) PCDATA  $1, $0
        0x002f 00047 (ff.go:10) CALL    AX
		        ......
        0x006f 00111 (ff.go:12) MOVQ    "".foo·f(SB), AX ; funcVar2()
        0x0076 00118 (ff.go:12) LEAQ    "".foo·f(SB), DX
        0x007d 00125 (ff.go:12) CALL    AX
		        ......
"".foo·f SRODATA dupok size=8
        0x0000 00 00 00 00 00 00 00 00                          ........
        rel 0+8 t=1 "".foo+0

// target 目标函数地址
// double mock 函数的指针
func replace(target, double uintptr) []byte {
	code := buildJmpDirective(double) // 生成跳转到 mock 函数的机器指令
	bytes := entryAddress(target, len(code))
	original := make([]byte, len(bytes))
	copy(original, bytes)
	modifyBinary(target, code)
	return original
}

func buildJmpDirective(double uintptr) []byte {
    d0 := byte(double)
    d1 := byte(double >> 8)
    d2 := byte(double >> 16)
    d3 := byte(double >> 24)
    d4 := byte(double >> 32)
    d5 := byte(double >> 40)
    d6 := byte(double >> 48)
    d7 := byte(double >> 56)

    // 返回跳转的机器指令
    return []byte{
        0x48, 0xBA, d0, d1, d2, d3, d4, d5, d6, d7, // MOV rdx, double 将 mock 函数的指针值放到 rdx 中
        0xFF, 0x22,     // JMP [rdx] 因为rdx 中存储的是 mock 函数的指针，所以需要使用[]，从内存中获得 mock 函数的地址，然后跳转
    }
}

func modifyBinary(target uintptr, bytes []byte) {
    function := entryAddress(target, len(bytes))
    
    // 默认情况下，代码段的内存页是不可写的，需要调用 Mprotect 修改 target 所在页的权限为可写
    page := entryAddress(pageStart(target), syscall.Getpagesize())
    err := syscall.Mprotect(page, syscall.PROT_READ|syscall.PROT_WRITE|syscall.PROT_EXEC)
    if err != nil {
        panic(err)
    }
    // 替换 target 函数开头的指令为跳转指令
    copy(function, bytes)
    // 恢复为读和执行权限
    err = syscall.Mprotect(page, syscall.PROT_READ|syscall.PROT_EXEC)
    if err != nil {
        panic(err)
    }
}

func init() {
   content, _ := nm.Parse(os.Args[0])

   lines := strings.Split(content, "\n")
   for _, line := range lines {
      line := strings.TrimSpace(line)
      arr := strings.Split(line, " ")
      if len(arr) < 3 {
         continue
      }

      funcSymbol, addr := arr[2], arr[0]
      addrUint, _ := strconv.ParseUint(addr, 16, 64) // addrUint 为函数地址
      symbolTable[funcSymbol] = uintptr(addrUint) // funcSymbol 函数符号，比如 supermonkey/pkg.Foo.say
   }
}

func main() {
   f := pkg.Foo{}
   patchGuard := sm.PatchByFullSymbolName("supermonkey/pkg.Foo.say", func() string {
      return "mock say"
   })
   println(f.Say())
   patchGuard.Unpatch()
}