proc: generalize escapeCheck and allocString

pkg/proc/eval.go

@@ -1045,6 +1045,10 @@ func (stack *evalStack) executeOp() {
 		}
 		stack.push(v)
 
+	case *evalop.PushLen:
+		v := stack.peek()
+		stack.push(newConstant(constant.MakeInt64(v.Len), scope.Mem))
+
 	case *evalop.Select:
 		scope.evalStructSelector(op, stack)
 
@@ -1112,8 +1116,11 @@ func (stack *evalStack) executeOp() {
 		stack.fncallPeek().undoInjection = nil
 		stack.callInjectionContinue = true
 
-	case *evalop.CallInjectionAllocString:
-		stack.callInjectionContinue = scope.allocString(op.Phase, stack, curthread)
+	case *evalop.CallInjectionStartSpecial:
+		stack.callInjectionContinue = scope.callInjectionStartSpecial(stack, op, curthread)
+
+	case *evalop.ConvertAllocToString:
+		scope.convertAllocToString(stack)
 
 	case *evalop.SetValue:
 		lhv := stack.pop()
@@ -1155,6 +1162,14 @@ func (scope *EvalScope) evalJump(op *evalop.Jump, stack *evalStack) {
 		v = true
 	case evalop.JumpIfFalse:
 		v = false
+	case evalop.JumpIfAllocStringChecksFail:
+		stringChecksFailed := x.Kind != reflect.String || x.Addr != 0 || x.Flags&VariableConstant == 0 || x.Len <= 0
+		nilCallCtx := scope.callCtx == nil // do not complain here, setValue will if no other errors happen
+		if stringChecksFailed || nilCallCtx {
+			stack.opidx = op.Target - 1
+			return
+		}
+		return
 	}
 
 	if x.Kind != reflect.Bool {

pkg/proc/evalop/evalcompile.go

@@ -112,9 +112,38 @@ func CompileSet(lookup evalLookup, lhexpr, rhexpr string) ([]Op, error) {
 }
 
 func (ctx *compileCtx) compileAllocLiteralString() {
-	ctx.pushOp(&CallInjectionAllocString{Phase: 0})
-	ctx.pushOp(&CallInjectionAllocString{Phase: 1})
-	ctx.pushOp(&CallInjectionAllocString{Phase: 2})
+	jmp := &Jump{When: JumpIfAllocStringChecksFail}
+	ctx.pushOp(jmp)
+
+	ctx.compileSpecialCall("runtime.mallocgc", []ast.Expr{
+		&ast.BasicLit{Kind: token.INT, Value: "0"},
+		&ast.Ident{Name: "nil"},
+		&ast.Ident{Name: "false"},
+	}, []Op{
+		&PushLen{},
+		&PushNil{},
+		&PushConst{constant.MakeBool(false)},
+	})
+
+	ctx.pushOp(&ConvertAllocToString{})
+	jmp.Target = len(ctx.ops)
+}
+
+func (ctx *compileCtx) compileSpecialCall(fnname string, argAst []ast.Expr, args []Op) {
+	id := ctx.curCall
+	ctx.curCall++
+	ctx.pushOp(&CallInjectionStartSpecial{
+		id:     id,
+		FnName: fnname,
+		ArgAst: argAst})
+	ctx.pushOp(&CallInjectionSetTarget{id: id})
+
+	for i := range args {
+		ctx.pushOp(args[i])
+		ctx.pushOp(&CallInjectionCopyArg{id: id, ArgNum: i})
+	}
+
+	ctx.pushOp(&CallInjectionComplete{id: id})
 }
 
 func (ctx *compileCtx) pushOp(op Op) {

pkg/proc/evalop/ops.go

@@ -66,6 +66,13 @@ type PushPackageVar struct {
 
 func (*PushPackageVar) depthCheck() (npop, npush int) { return 0, 1 }
 
+// PushLen pushes the length of the variable at the top of the stack into
+// the stack.
+type PushLen struct {
+}
+
+func (*PushLen) depthCheck() (npop, npush int) { return 1, 2 }
+
 // Select replaces the topmost stack variable v with v.Name.
 type Select struct {
 	Name string
@@ -160,6 +167,7 @@ type JumpCond uint8
 const (
 	JumpIfFalse JumpCond = iota
 	JumpIfTrue
+	JumpIfAllocStringChecksFail
 )
 
 // Binary pops two variables from the stack, applies the specified binary
@@ -229,20 +237,24 @@ type CallInjectionComplete struct {
 
 func (*CallInjectionComplete) depthCheck() (npop, npush int) { return 0, 1 }
 
-// CallInjectionAllocString uses the call injection protocol to allocate the
-// value of a string literal somewhere on the target's memory so that it can
-// be assigned to a variable (or passed to a function).
-// There are three phases to CallInjectionAllocString, distinguished by the
-// Phase field. They must always appear in sequence in the program:
-//
-//	CallInjectionAllocString{Phase: 0}
-//	CallInjectionAllocString{Phase: 1}
-//	CallInjectionAllocString{Phase: 2}
-type CallInjectionAllocString struct {
-	Phase int
+// CallInjectionStartSpecial starts call injection for a function with a
+// name and arguments known at compile time.
+type CallInjectionStartSpecial struct {
+	id     int
+	FnName string
+	ArgAst []ast.Expr
+}
+
+func (*CallInjectionStartSpecial) depthCheck() (npop, npush int) { return 0, 1 }
+
+// ConvertAllocToString pops two variables from the stack, a constant string
+// and the return value of runtime.mallocgc (mallocv), copies the contents
+// of the string at the address in mallocv and pushes on the stack a new
+// string value that uses the backing storage of mallocv.
+type ConvertAllocToString struct {
 }
 
-func (op *CallInjectionAllocString) depthCheck() (npop, npush int) { return 1, 1 }
+func (*ConvertAllocToString) depthCheck() (npop, npush int) { return 2, 1 }
 
 // SetValue pops to variables from the stack, lhv and rhv, and sets lhv to
 // rhv.

pkg/proc/fncall.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"go/ast"
 	"go/constant"
-	"go/token"
 	"reflect"
 	"sort"
 	"strconv"
@@ -532,13 +531,14 @@ type funcCallArg struct {
 func funcCallCopyOneArg(scope *EvalScope, fncall *functionCallState, actualArg *Variable, formalArg *funcCallArg, thread Thread) error {
 	if scope.callCtx.checkEscape {
 		//TODO(aarzilli): only apply the escapeCheck to leaking parameters.
-		if err := escapeCheck(actualArg, formalArg.name, scope.g.stack); err != nil {
-			return fmt.Errorf("cannot use %s as argument %s in function %s: %v", actualArg.Name, formalArg.name, fncall.fn.Name, err)
-		}
-		for _, stack := range scope.callCtx.stacks {
-			if err := escapeCheck(actualArg, formalArg.name, stack); err != nil {
-				return fmt.Errorf("cannot use %s as argument %s in function %s: %v", actualArg.Name, formalArg.name, fncall.fn.Name, err)
+		err := allPointers(actualArg, formalArg.name, func(addr uint64, name string) error {
+			if !pointerEscapes(addr, scope.g.stack, scope.callCtx.stacks) {
+				return fmt.Errorf("cannot use %s as argument %s in function %s: stack object passed to escaping pointer: %s", actualArg.Name, formalArg.name, fncall.fn.Name, name)
 			}
+			return nil
+		})
+		if err != nil {
+			return err
 		}
 	}
 
@@ -691,7 +691,8 @@ func alignAddr(addr, align int64) int64 {
 	return (addr + align - 1) &^ (align - 1)
 }
 
-func escapeCheck(v *Variable, name string, stack stack) error {
+// allPointers calls f on every pointer contained in v
+func allPointers(v *Variable, name string, f func(addr uint64, name string) error) error {
 	if v.Unreadable != nil {
 		return fmt.Errorf("escape check for %s failed, variable unreadable: %v", name, v.Unreadable)
 	}
@@ -704,43 +705,48 @@ func escapeCheck(v *Variable, name string, stack stack) error {
 		} else {
 			w = v.maybeDereference()
 		}
-		return escapeCheckPointer(w.Addr, name, stack)
+		return f(w.Addr, name)
 	case reflect.Chan, reflect.String, reflect.Slice:
-		return escapeCheckPointer(v.Base, name, stack)
+		return f(v.Base, name)
 	case reflect.Map:
 		sv := v.clone()
 		sv.RealType = resolveTypedef(&(v.RealType.(*godwarf.MapType).TypedefType))
 		sv = sv.maybeDereference()
-		return escapeCheckPointer(sv.Addr, name, stack)
+		return f(sv.Addr, name)
 	case reflect.Struct:
 		t := v.RealType.(*godwarf.StructType)
 		for _, field := range t.Field {
 			fv, _ := v.toField(field)
-			if err := escapeCheck(fv, fmt.Sprintf("%s.%s", name, field.Name), stack); err != nil {
+			if err := allPointers(fv, fmt.Sprintf("%s.%s", name, field.Name), f); err != nil {
 				return err
 			}
 		}
 	case reflect.Array:
 		for i := int64(0); i < v.Len; i++ {
 			sv, _ := v.sliceAccess(int(i))
-			if err := escapeCheck(sv, fmt.Sprintf("%s[%d]", name, i), stack); err != nil {
+			if err := allPointers(sv, fmt.Sprintf("%s[%d]", name, i), f); err != nil {
 				return err
 			}
 		}
 	case reflect.Func:
-		if err := escapeCheckPointer(v.funcvalAddr(), name, stack); err != nil {
+		if err := f(v.funcvalAddr(), name); err != nil {
 			return err
 		}
 	}
 
 	return nil
 }
 
-func escapeCheckPointer(addr uint64, name string, stack stack) error {
+func pointerEscapes(addr uint64, stack stack, stacks []stack) bool {
 	if addr >= stack.lo && addr < stack.hi {
-		return fmt.Errorf("stack object passed to escaping pointer: %s", name)
+		return false
 	}
-	return nil
+	for _, stack := range stacks {
+		if addr >= stack.lo && addr < stack.hi {
+			return false
+		}
+	}
+	return true
 }
 
 const (
@@ -986,92 +992,45 @@ func fakeFunctionEntryScope(scope *EvalScope, fn *Function, cfa int64, sp uint64
 	return nil
 }
 
-func (scope *EvalScope) allocString(phase int, stack *evalStack, curthread Thread) bool {
-	switch phase {
-	case 0:
-		x := stack.peek()
-		if !(x.Kind == reflect.String && x.Addr == 0 && (x.Flags&VariableConstant) != 0 && x.Len > 0) {
-			stack.opidx += 2 // skip the next two allocString phases, we don't need to do an allocation
-			return false
-		}
-		if scope.callCtx == nil {
-			// do not complain here, setValue will if no other errors happen
-			stack.opidx += 2
-			return false
-		}
-		mallocv, err := scope.findGlobal("runtime", "mallocgc")
-		if mallocv == nil {
-			stack.err = err
-			return false
-		}
-		stack.push(mallocv)
-		scope.evalCallInjectionStart(&evalop.CallInjectionStart{HasFunc: true, Node: &ast.CallExpr{
-			Fun: &ast.SelectorExpr{
-				X:   &ast.Ident{Name: "runtime"},
-				Sel: &ast.Ident{Name: "mallocgc"},
-			},
-			Args: []ast.Expr{
-				&ast.BasicLit{Kind: token.INT, Value: "0"},
-				&ast.Ident{Name: "nil"},
-				&ast.Ident{Name: "false"},
-			},
-		}}, stack)
-		if stack.err == nil {
-			stack.pop() // return value of evalop.CallInjectionStart
-		}
-		return true
-
-	case 1:
-		fncall := stack.fncallPeek()
-		savedLoadCfg := scope.callCtx.retLoadCfg
-		scope.callCtx.retLoadCfg = loadFullValue
-		defer func() {
-			scope.callCtx.retLoadCfg = savedLoadCfg
-		}()
-
-		scope.evalCallInjectionSetTarget(nil, stack, curthread)
-
-		strvar := stack.peek()
-
-		stack.err = funcCallCopyOneArg(scope, fncall, newConstant(constant.MakeInt64(strvar.Len), scope.Mem), &fncall.formalArgs[0], curthread)
-		if stack.err != nil {
-			return false
-		}
-		stack.err = funcCallCopyOneArg(scope, fncall, nilVariable, &fncall.formalArgs[1], curthread)
-		if stack.err != nil {
-			return false
-		}
-		stack.err = funcCallCopyOneArg(scope, fncall, newConstant(constant.MakeBool(false), scope.Mem), &fncall.formalArgs[2], curthread)
-		if stack.err != nil {
-			return false
-		}
+func (scope *EvalScope) callInjectionStartSpecial(stack *evalStack, op *evalop.CallInjectionStartSpecial, curthread Thread) bool {
+	fnv, err := scope.findGlobalInternal(op.FnName)
+	if fnv == nil {
+		stack.err = err
+		return false
+	}
+	stack.push(fnv)
+	scope.evalCallInjectionStart(&evalop.CallInjectionStart{HasFunc: true, Node: &ast.CallExpr{
+		Fun:  &ast.Ident{Name: op.FnName},
+		Args: op.ArgAst,
+	}}, stack)
+	if stack.err == nil {
+		stack.pop() // return value of evalop.CallInjectionStart
 		return true
+	}
+	return false
+}
 
-	case 2:
-		mallocv := stack.pop()
-		v := stack.pop()
-		if mallocv.Unreadable != nil {
-			stack.err = mallocv.Unreadable
-			return false
-		}
-
-		if mallocv.DwarfType.String() != "*void" {
-			stack.err = fmt.Errorf("unexpected return type for mallocgc call: %v", mallocv.DwarfType.String())
-			return false
-		}
+func (scope *EvalScope) convertAllocToString(stack *evalStack) {
+	mallocv := stack.pop()
+	v := stack.pop()
+	if mallocv.Unreadable != nil {
+		stack.err = mallocv.Unreadable
+		return
+	}
 
-		if len(mallocv.Children) != 1 {
-			stack.err = errors.New("internal error, could not interpret return value of mallocgc call")
-			return false
-		}
+	if mallocv.DwarfType.String() != "*void" {
+		stack.err = fmt.Errorf("unexpected return type for mallocgc call: %v", mallocv.DwarfType.String())
+		return
+	}
 
-		v.Base = mallocv.Children[0].Addr
-		_, stack.err = scope.Mem.WriteMemory(v.Base, []byte(constant.StringVal(v.Value)))
-		stack.push(v)
-		return false
+	if len(mallocv.Children) != 1 {
+		stack.err = errors.New("internal error, could not interpret return value of mallocgc call")
+		return
 	}
 
-	panic("unreachable")
+	v.Base = mallocv.Children[0].Addr
+	_, stack.err = scope.Mem.WriteMemory(v.Base, []byte(constant.StringVal(v.Value)))
+	stack.push(v)
 }
 
 func isCallInjectionStop(t *Target, thread Thread, loc *Location) bool {