Merge pull request #35 from bytedance/fix/generics

fix: got wrong original function addr when generics params or returns…

internal/monkey/inst/disasm_amd64.go

@@ -17,13 +17,57 @@
 package inst
 
 import (
+	"reflect"
 	"unsafe"
 
 	"github.com/bytedance/mockey/internal/monkey/common"
 	"github.com/bytedance/mockey/internal/tool"
 	"golang.org/x/arch/x86/x86asm"
 )
 
+//go:linkname duffcopy runtime.duffcopy
+func duffcopy()
+
+//go:linkname duffzero runtime.duffzero
+func duffzero()
+
+var duffcopyStart, duffcopyEnd, duffzeroStart, duffzeroEnd uintptr
+
+func init() {
+	duffcopyStart, duffcopyEnd = calcFnAddrRange("duffcopy", duffcopy)
+	duffzeroStart, duffzeroEnd = calcFnAddrRange("duffzero", duffzero)
+}
+
+func calcFnAddrRange(name string, fn func()) (uintptr, uintptr) {
+	v := reflect.ValueOf(fn)
+	var start, end uintptr
+	start = v.Pointer()
+	maxScan := 2000
+	code := common.BytesOf(v.Pointer(), 2000)
+	pos := 0
+
+	for pos < maxScan {
+		inst, err := x86asm.Decode(code[pos:], 64)
+		tool.Assert(err == nil, err)
+
+		args := []interface{}{name, inst.Op}
+		for i := range inst.Args {
+			args = append(args, inst.Args[i])
+		}
+		tool.DebugPrintf("init: <%v>\t%v\t%v\t%v\t%v\t%v\t%v\n", args...)
+
+		if inst.Op == x86asm.RET {
+			end = start + uintptr(pos)
+			tool.DebugPrintf("init: %v(%v,%v)\n", name, start, end)
+			return start, end
+		}
+
+		pos += int(inst.Len)
+	}
+	tool.Assert(false, "%v ret not found", name)
+	return 0, 0
+}
+
 func Disassemble(code []byte, required int, checkLen bool) int {
 	var pos int
 	var err error
@@ -45,26 +89,53 @@ func GetGenericJumpAddr(addr uintptr, maxScan uint64) uintptr {
 	var err error
 	var inst x86asm.Inst
 
+	allAddrs := []uintptr{}
 	for pos < maxScan {
 		inst, err = x86asm.Decode(code[pos:], 64)
 		tool.Assert(err == nil, err)
-		// if inst.Op == arm64asm.BL {
+
 		args := []interface{}{inst.Op}
 		for i := range inst.Args {
 			args = append(args, inst.Args[i])
 		}
 		tool.DebugPrintf("%v\t%v\t%v\t%v\t%v\t%v\n", args...)
 
+		if inst.Op == x86asm.RET {
+			break
+		}
+
 		if inst.Op == x86asm.CALL {
 			rel := int32(inst.Args[0].(x86asm.Rel))
-			tool.DebugPrintf("found: CALL, raw is: %x, rel: %v\n", inst.String(), rel)
-			return calcAddr(uintptr(unsafe.Pointer(&code[0]))+uintptr(pos+uint64(inst.Len)), rel)
+			fnAddr := calcAddr(uintptr(unsafe.Pointer(&code[0]))+uintptr(pos+uint64(inst.Len)), rel)
+
+			/*
+				When function argument size is too big, golang will use duff-copy
+				to get better performance. Thus we will see more call-instruction
+				in asm code.
+				For example, assume struct type like this:
+					```
+					type Large15 struct _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _ string}
+					```
+				when we use `Large15` as generic function's argument or return types,
+				the wrapper `function[Large15]` will call `duffcopy` and `duffzero`
+				before passing arguments and after receiving returns.
+
+				Notice that `duff` functions are very special, they are always called
+				in the middle of function body(not at beginning). So we should check
+				the `call` instruction's target address with a range.
+			*/
+
+			isDuffCopy := (fnAddr >= duffcopyStart && fnAddr <= duffcopyEnd)
+			isDuffZero := (fnAddr >= duffzeroStart && fnAddr <= duffzeroEnd)
+			tool.DebugPrintf("found: CALL, raw is: %x, rel: %v  isDuffCopy: %v, isDuffZero: %v fnAddr: %v\n", inst.String(), rel, isDuffCopy, isDuffZero, fnAddr)
+			if !isDuffCopy && !isDuffZero {
+				allAddrs = append(allAddrs, fnAddr)
+			}
 		}
-		tool.Assert(inst.Op != x86asm.RET, "!!!FOUND RET!!!")
 		pos += uint64(inst.Len)
 	}
-	tool.Assert(false, "CALL op not found")
-	return 0
+	tool.Assert(len(allAddrs) == 1, "invalid callAddr: %v", allAddrs)
+	return allAddrs[0]
 }
 
 func calcAddr(from uintptr, rel int32) uintptr {

internal/monkey/inst/disasm_arm64.go

@@ -17,13 +17,56 @@
 package inst
 
 import (
+	"reflect"
 	"unsafe"
 
 	"github.com/bytedance/mockey/internal/monkey/common"
 	"github.com/bytedance/mockey/internal/tool"
 	"golang.org/x/arch/arm64/arm64asm"
 )
 
+//go:linkname duffcopy runtime.duffcopy
+func duffcopy()
+
+//go:linkname duffzero runtime.duffzero
+func duffzero()
+
+var duffcopyStart, duffcopyEnd, duffzeroStart, duffzeroEnd uintptr
+
+func init() {
+	duffcopyStart, duffcopyEnd = calcFnAddrRange("duffcopy", duffcopy)
+	duffzeroStart, duffzeroEnd = calcFnAddrRange("duffzero", duffzero)
+}
+
+func calcFnAddrRange(name string, fn func()) (uintptr, uintptr) {
+	v := reflect.ValueOf(fn)
+	var start, end uintptr
+	start = v.Pointer()
+	maxScan := 2000
+	code := common.BytesOf(start, 2000)
+	pos := 0
+	for pos < maxScan {
+		inst, err := arm64asm.Decode(code[pos:])
+		tool.Assert(err == nil, err)
+
+		args := []interface{}{name, inst.Op}
+		for i := range inst.Args {
+			args = append(args, inst.Args[i])
+		}
+		tool.DebugPrintf("init: <%v>\t%v\t%v\t%v\t%v\t%v\t%v\n", args...)
+
+		if inst.Op == arm64asm.RET {
+			end = start + uintptr(pos)
+			tool.DebugPrintf("init: duffcopy(%v,%v)\n", name, start, end)
+			return start, end
+		}
+
+		pos += int(unsafe.Sizeof(inst.Enc))
+	}
+	tool.Assert(false, "%v end not found", name)
+	return 0, 0
+}
+
 func Disassemble(code []byte, required int, checkLen bool) int {
 	tool.Assert(len(code) > required, "function is too short to patch")
 	return required
@@ -35,25 +78,51 @@ func GetGenericJumpAddr(addr uintptr, maxScan uint64) uintptr {
 	var err error
 	var inst arm64asm.Inst
 
+	allAddrs := []uintptr{}
 	for pos < maxScan {
 		inst, err = arm64asm.Decode(code[pos:])
 		tool.Assert(err == nil, err)
-		// if inst.Op == arm64asm.BL {
 		args := []interface{}{inst.Op}
 		for i := range inst.Args {
 			args = append(args, inst.Args[i])
 		}
 		tool.DebugPrintf("%v\t%v\t%v\t%v\t%v\t%v\n", args...)
 
+		if inst.Op == arm64asm.RET {
+			break
+		}
+
 		if inst.Op == arm64asm.BL {
-			tool.DebugPrintf("found: BL, raw is: %x\n", inst.Enc)
-			return calcAddr(uintptr(unsafe.Pointer(&code[0]))+uintptr(pos), inst.Enc)
+			fnAddr := calcAddr(uintptr(unsafe.Pointer(&code[0]))+uintptr(pos), inst.Enc)
+
+			/*
+				When function argument size is too big, golang will use duff-copy
+				to get better performance. Thus we will see more call-instruction
+				in asm code.
+				For example, assume struct type like this:
+					```
+					type Large15 struct _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _ string}
+					```
+				when we use `Large15` as generic function's argument or return types,
+				the wrapper `function[Large15]` will call `duffcopy` and `duffzero`
+				before passing arguments and after receiving returns.
+
+				Notice that `duff` functions are very special, they are always called
+				in the middle of function body(not at beginning). So we should check
+				the `call` instruction's target address with a range.
+			*/
+
+			isDuffCopy := (fnAddr >= duffcopyStart && fnAddr <= duffcopyEnd)
+			isDuffZero := (fnAddr >= duffzeroStart && fnAddr <= duffzeroEnd)
+			tool.DebugPrintf("found: BL, raw is: %x,  isDuffCopy: %v, isDuffZero: %v fnAddr: %v\n", inst.String(), isDuffCopy, isDuffZero, fnAddr)
+			if !isDuffCopy && !isDuffZero {
+				allAddrs = append(allAddrs, fnAddr)
+			}
 		}
 		pos += uint64(unsafe.Sizeof(inst.Enc))
-		tool.Assert(inst.Op != arm64asm.RET, "!!!FOUND RET!!!")
 	}
-	tool.Assert(false, "BL op not found")
-	return 0
+	tool.Assert(len(allAddrs) == 1, "invalid callAddr: %v", allAddrs)
+	return allAddrs[0]
 }
 
 func calcAddr(from uintptr, bl uint32) uintptr {

internal/monkey/patch.go

@@ -50,7 +50,7 @@ func PatchValue(target, hook, proxy reflect.Value, unsafe, generic bool) *Patch
 	targetAddr := target.Pointer()
 	if generic {
 		// we assume that generic call/bl op is located in first 200 bytes of codes from targetAddr
-		targetAddr = inst.GetGenericJumpAddr(targetAddr, 200)
+		targetAddr = inst.GetGenericJumpAddr(targetAddr, 10000)
 	}
 	// The first few bytes of the target function code
 	const bufSize = 64

internal/tool/call.go

@@ -34,13 +34,8 @@ func ReflectCall(f reflect.Value, args []reflect.Value) []reflect.Value {
 		for i := 0; i < len(args)-1; i++ {
 			newArgs = append(newArgs, args[i])
 		}
-		// FIXME: There are a few cases that lastArgs.cap == 0 but lastArg.len == MAX_INT (arm64).
-		// 		Currently we have no idea how it happens.
-		//      In that case we assume that i should less than lastArg.cap
-		if lastArg.Len() < lastArg.Cap() {
-			DebugPrintf("ReflectCall: broken variadic params: len is %d but cap is %d", lastArg.Len(), lastArg.Cap())
-		}
-		for i := 0; i < lastArg.Len() && i < lastArg.Cap(); i++ {
+
+		for i := 0; i < lastArg.Len(); i++ {
 			newArgs = append(newArgs, lastArg.Index(i))
 		}
 		return f.Call(newArgs)

internal/tool/caller.go

@@ -36,7 +36,14 @@ func (c CallerInfo) String() string {
 //	func innerBar() CallerInfo { /*do some thing*/ return Caller() }
 //
 // The return value of innerBar should represent the line in a/b/foo.go where a/b/foo.Foo calls a/c/bar.Bar
-func OuterCaller() CallerInfo {
+func OuterCaller() (info CallerInfo) {
+	defer func() {
+		if err := recover(); err != nil {
+			DebugPrintf("OuterCaller: get stack failed, err: %v", err)
+			info = CallerInfo(runtime.Frame{File: "Nan"})
+		}
+	}()
+
 	caller, _, _, _ := runtime.Caller(1)
 	oriPkg, _ := getPackageAndFunction(caller)
 
@@ -65,7 +72,11 @@ func getPackageAndFunction(pc uintptr) (string, string) {
 	packageName := ""
 	funcName := parts[pl-1]
 
-	if parts[pl-2][0] == '(' {
+	// if mock run in an anonymous function of a global variable,
+	// the stack will looks like a.b.c.glob..func1(), so the
+	// second last part of the caller stack would not be guaranteed
+	// always to be non-empty.
+	if len(parts[pl-2]) > 0 && parts[pl-2][0] == '(' {
 		funcName = parts[pl-2] + "." + funcName
 		packageName = strings.Join(parts[0:pl-2], ".")
 	} else {