proc: simplify and generalize runtime.mallocgc workaround (#3571)

Instead of having a different version for each architecture have a
single version that uses an architecture specific list of registers.
Also generalize it so that, if we want, we can extend the workaround to
other runtime functions we might want to call (for example the channel
send/receive functions).

pkg/proc/amd64_arch.go

@@ -43,6 +43,7 @@ func AMD64Arch(goos string) *Arch {
 		RegnumToString:                   regnum.AMD64ToName,
 		debugCallMinStackSize:            256,
 		maxRegArgBytes:                   9*8 + 15*8,
+		argumentRegs:                     []int{regnum.AMD64_Rax, regnum.AMD64_Rbx, regnum.AMD64_Rcx},
 	}
 }
 

pkg/proc/arch.go

@@ -55,6 +55,8 @@ type Arch struct {
 	// maxRegArgBytes is extra padding for ABI1 call injections, equivalent to
 	// the maximum space occupied by register arguments.
 	maxRegArgBytes int
+	// argumentRegs are function call injection registers for runtimeOptimizedWorkaround
+	argumentRegs []int
 
 	// asmRegisters maps assembly register numbers to dwarf registers.
 	asmRegisters map[int]asmRegister

pkg/proc/arm64_arch.go

@@ -53,6 +53,7 @@ func ARM64Arch(goos string) *Arch {
 		RegnumToString:                   regnum.ARM64ToName,
 		debugCallMinStackSize:            288,
 		maxRegArgBytes:                   16*8 + 16*8, // 16 int argument registers plus 16 float argument registers
+		argumentRegs:                     []int{regnum.ARM64_X0, regnum.ARM64_X0 + 1, regnum.ARM64_X0 + 2},
 	}
 }
 

pkg/proc/bininfo.go

@@ -858,7 +858,7 @@ type Image struct {
 	compileUnits []*compileUnit // compileUnits is sorted by increasing DWARF offset
 
 	dwarfTreeCache  *simplelru.LRU
-	runtimeMallocgcTree *godwarf.Tree // patched version of runtime.mallocgc's DIE
+	workaroundCache map[dwarf.Offset]*godwarf.Tree
 
 	// runtimeTypeToDIE maps between the offset of a runtime._type in
 	// runtime.moduledata.types and the offset of the DIE in debug_info. This
@@ -1012,8 +1012,8 @@ func (image *Image) LoadError() error {
 }
 
 func (image *Image) getDwarfTree(off dwarf.Offset) (*godwarf.Tree, error) {
-	if image.runtimeMallocgcTree != nil && off == image.runtimeMallocgcTree.Offset {
+	if image.workaroundCache[off] != nil {
-		return image.runtimeMallocgcTree, nil
+		return image.workaroundCache[off], nil
 	}
 	if r, ok := image.dwarfTreeCache.Get(off); ok {
 		return r.(*godwarf.Tree), nil
@@ -2257,23 +2257,6 @@ func (bi *BinaryInfo) loadDebugInfoMaps(image *Image, debugInfoBytes, debugLineB
 	sort.Strings(bi.Sources)
 	bi.Sources = uniq(bi.Sources)
 
-	if bi.regabi {
-		// prepare patch for runtime.mallocgc's DIE
-		fn := bi.lookupOneFunc("runtime.mallocgc")
-		if fn != nil && fn.cu.image == image {
-			tree, err := image.getDwarfTree(fn.offset)
-			if err == nil {
-				children, err := regabiMallocgcWorkaround(bi)
-				if err != nil {
-					bi.logger.Errorf("could not patch runtime.mallocgc: %v", err)
-				} else {
-					tree.Children = children
-					image.runtimeMallocgcTree = tree
-				}
-			}
-		}
-	}
-
 	if cont != nil {
 		cont()
 	}

pkg/proc/fncall.go

@@ -573,7 +573,10 @@ func funcCallArgs(fn *Function, bi *BinaryInfo, includeRet bool) (argFrameSize i
 		return 0, nil, fmt.Errorf("DWARF read error: %v", err)
 	}
 
-	if bi.regabi && fn.cu.optimized && fn.Name != "runtime.mallocgc" {
+	if bi.regabi && fn.cu.optimized {
+		if runtimeWhitelist[fn.Name] {
+			runtimeOptimizedWorkaround(bi, fn.cu.image, dwarfTree)
+		} else {
 			// Debug info for function arguments on optimized functions is currently
 			// too incomplete to attempt injecting calls to arbitrary optimized
 			// functions.
@@ -583,6 +586,7 @@ func funcCallArgs(fn *Function, bi *BinaryInfo, includeRet bool) (argFrameSize i
 			// correct for call injection purposes.
 			return 0, nil, fmt.Errorf("can not call optimized function %s when regabi is in use", fn.Name)
 		}
+	}
 
 	varEntries := reader.Variables(dwarfTree, fn.Entry, int(^uint(0)>>1), reader.VariablesSkipInlinedSubroutines)
 
@@ -1214,82 +1218,53 @@ func debugCallProtocolReg(archName string, version int) (uint64, bool) {
 	}
 }
 
-type fakeEntry map[dwarf.Attr]*dwarf.Field
+// runtimeWhitelist is a list of functions in the runtime that we can call
-
+// (through call injection) even if they are optimized.
-func (e fakeEntry) Val(attr dwarf.Attr) interface{} {
+var runtimeWhitelist = map[string]bool{
-	if e[attr] == nil {
+	"runtime.mallocgc": true,
-		return nil
 }
 
-	return e[attr].Val
+// runtimeOptimizedWorkaround modifies the input DIE so that arguments and
+// return variables have the appropriate registers for call injection.
+// This function can not be called on arbitrary DIEs, it is only valid for
+// the functions specified in runtimeWhitelist.
+// In particular this will fail if any of the arguments of the function
+// passed in input does not fit in an integer CPU register.
+func runtimeOptimizedWorkaround(bi *BinaryInfo, image *Image, in *godwarf.Tree) {
+	if image.workaroundCache == nil {
+		image.workaroundCache = make(map[dwarf.Offset]*godwarf.Tree)
+	}
+	if image.workaroundCache[in.Offset] == in {
+		return
+	}
+	image.workaroundCache[in.Offset] = in
+
+	curArg, curRet := 0, 0
+	for _, child := range in.Children {
+		if child.Tag == dwarf.TagFormalParameter {
+			childEntry, ok := child.Entry.(*dwarf.Entry)
+			if !ok {
+				panic("internal error: bad DIE for runtimeOptimizedWorkaround")
+			}
+			isret, _ := child.Entry.Val(dwarf.AttrVarParam).(bool)
+
+			var reg int
+			if isret {
+				reg = bi.Arch.argumentRegs[curRet]
+				curRet++
+			} else {
+				reg = bi.Arch.argumentRegs[curArg]
+				curArg++
 			}
 
-func (e fakeEntry) AttrField(attr dwarf.Attr) *dwarf.Field {
+			newlocfield := dwarf.Field{Attr: dwarf.AttrLocation, Val: []byte{byte(op.DW_OP_reg0) + byte(reg)}, Class: dwarf.ClassBlock}
-	return e[attr]
-}
 
-func regabiMallocgcWorkaround(bi *BinaryInfo) ([]*godwarf.Tree, error) {
+			locfield := childEntry.AttrField(dwarf.AttrLocation)
-	ptrToRuntimeType := "*" + bi.runtimeTypeTypename()
+			if locfield != nil {
-
+				*locfield = newlocfield
-	var err1 error
+			} else {
-
+				childEntry.Field = append(childEntry.Field, newlocfield)
-	t := func(name string) godwarf.Type {
-		if err1 != nil {
-			return nil
-		}
-		typ, err := bi.findType(name)
-		if err != nil {
-			err1 = err
-			return nil
-		}
-		return typ
-	}
-
-	m := func(name string, typ godwarf.Type, reg int, isret bool) *godwarf.Tree {
-		if err1 != nil {
-			return nil
-		}
-		var e fakeEntry = map[dwarf.Attr]*dwarf.Field{
-			dwarf.AttrName:     &dwarf.Field{Attr: dwarf.AttrName, Val: name, Class: dwarf.ClassString},
-			dwarf.AttrType:     &dwarf.Field{Attr: dwarf.AttrType, Val: typ.Common().Offset, Class: dwarf.ClassReference},
-			dwarf.AttrLocation: &dwarf.Field{Attr: dwarf.AttrLocation, Val: []byte{byte(op.DW_OP_reg0) + byte(reg)}, Class: dwarf.ClassBlock},
-			dwarf.AttrVarParam: &dwarf.Field{Attr: dwarf.AttrVarParam, Val: isret, Class: dwarf.ClassFlag},
-		}
-
-		return &godwarf.Tree{
-			Entry: e,
-			Tag:   dwarf.TagFormalParameter,
 			}
 		}
-
-	switch bi.Arch.Name {
-	case "amd64":
-		r := []*godwarf.Tree{
-			m("size", t("uintptr"), regnum.AMD64_Rax, false),
-			m("typ", t(ptrToRuntimeType), regnum.AMD64_Rbx, false),
-			m("needzero", t("bool"), regnum.AMD64_Rcx, false),
-			m("~r1", t("unsafe.Pointer"), regnum.AMD64_Rax, true),
-		}
-		return r, err1
-	case "arm64":
-		r := []*godwarf.Tree{
-			m("size", t("uintptr"), regnum.ARM64_X0, false),
-			m("typ", t(ptrToRuntimeType), regnum.ARM64_X0+1, false),
-			m("needzero", t("bool"), regnum.ARM64_X0+2, false),
-			m("~r1", t("unsafe.Pointer"), regnum.ARM64_X0, true),
-		}
-		return r, err1
-	case "ppc64le":
-		r := []*godwarf.Tree{
-			m("size", t("uintptr"), regnum.PPC64LE_R0+3, false),
-			m("typ", t(ptrToRuntimeType), regnum.PPC64LE_R0+4, false),
-			m("needzero", t("bool"), regnum.PPC64LE_R0+5, false),
-			m("~r1", t("unsafe.Pointer"), regnum.PPC64LE_R0+3, true),
-		}
-		return r, err1
-
-	default:
-		// do nothing
-		return nil, nil
 	}
 }

pkg/proc/ppc64le_arch.go

@@ -41,6 +41,7 @@ func PPC64LEArch(goos string) *Arch {
 		RegisterNameToDwarf:              nameToDwarfFunc(regnum.PPC64LENameToDwarf),
 		debugCallMinStackSize:            320,
 		maxRegArgBytes:                   13*8 + 13*8,
+		argumentRegs:                     []int{regnum.PPC64LE_R0 + 3, regnum.PPC64LE_R0 + 4, regnum.PPC64LE_R0 + 5},
 	}
 }