proc: simplify and generalize runtime.mallocgc workaround

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

@@ -857,8 +857,8 @@ 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
+	dwarfTreeCache  *simplelru.LRU
+	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 {
-		return image.runtimeMallocgcTree, nil
+	if image.workaroundCache[off] != 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,15 +573,19 @@ 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" {
-		// Debug info for function arguments on optimized functions is currently
-		// too incomplete to attempt injecting calls to arbitrary optimized
-		// functions.
-		// Prior to regabi we could do this because the ABI was simple enough to
-		// manually encode it in Delve.
-		// Runtime.mallocgc is an exception, we specifically patch it's DIE to be
-		// correct for call injection purposes.
-		return 0, nil, fmt.Errorf("can not call optimized function %s when regabi is in use", fn.Name)
+	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.
+			// Prior to regabi we could do this because the ABI was simple enough to
+			// manually encode it in Delve.
+			// Runtime.mallocgc is an exception, we specifically patch it's DIE to be
+			// 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
-
-func (e fakeEntry) Val(attr dwarf.Attr) interface{} {
-	if e[attr] == nil {
-		return nil
-	}
-
-	return e[attr].Val
-}
-
-func (e fakeEntry) AttrField(attr dwarf.Attr) *dwarf.Field {
-	return e[attr]
+// runtimeWhitelist is a list of functions in the runtime that we can call
+// (through call injection) even if they are optimized.
+var runtimeWhitelist = map[string]bool{
+	"runtime.mallocgc": true,
 }
 
-func regabiMallocgcWorkaround(bi *BinaryInfo) ([]*godwarf.Tree, error) {
-	ptrToRuntimeType := "*" + bi.runtimeTypeTypename()
-
-	var err1 error
-
-	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
+// 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
 
-	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,
-		}
-	}
+	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++
+			}
 
-	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
+			newlocfield := dwarf.Field{Attr: dwarf.AttrLocation, Val: []byte{byte(op.DW_OP_reg0) + byte(reg)}, Class: dwarf.ClassBlock}
 
-	default:
-		// do nothing
-		return nil, nil
+			locfield := childEntry.AttrField(dwarf.AttrLocation)
+			if locfield != nil {
+				*locfield = newlocfield
+			} else {
+				childEntry.Field = append(childEntry.Field, newlocfield)
+			}
+		}
 	}
 }

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},
 	}
 }