Fix safety check for nested function calls

Refactor how output vars are computed for call expressions. Previously
the number of rule args were not used to determine which args were
considered outputs. Instead, it was assumed the last arg in the call was
an output. This meant that if an arg in the output position was omitted,
an input arg would be incorrectly marked safe.

With these changes, the compiler looks up the number of args (arity) of
the rule when checking whether an arg is an output.

ast/compile.go

@@ -532,7 +532,7 @@ func (c *Compiler) checkSafetyRuleBodies() {
 }
 
 func (c *Compiler) checkBodySafety(safe VarSet, m *Module, b Body, l *Location) Body {
-	reordered, unsafe := reorderBodyForSafety(safe, b)
+	reordered, unsafe := reorderBodyForSafety(getRuleArgArity(c), safe, b)
 	if len(unsafe) != 0 {
 		for v := range unsafe.Vars() {
 			if !c.generatedVars[m].Contains(v) {
@@ -981,7 +981,7 @@ func (qc *queryCompiler) rewriteLocalAssignments(_ *QueryContext, body Body) (Bo
 func (qc *queryCompiler) checkSafety(_ *QueryContext, body Body) (Body, error) {
 
 	safe := ReservedVars.Copy()
-	reordered, unsafe := reorderBodyForSafety(safe, body)
+	reordered, unsafe := reorderBodyForSafety(getRuleArgArity(qc.compiler), safe, body)
 
 	if len(unsafe) != 0 {
 		var err Errors
@@ -1401,6 +1401,19 @@ func (vs unsafeVars) Vars() VarSet {
 	return r
 }
 
+// getArity defines a function to return the arity of the rule referred to by a ref.
+type getArity func(Ref) int
+
+func getRuleArgArity(c *Compiler) getArity {
+	return func(ref Ref) int {
+		rules := c.GetRulesExact(ref)
+		if len(rules) == 0 {
+			return 0
+		}
+		return len(rules[0].Head.Args)
+	}
+}
+
 // reorderBodyForSafety returns a copy of the body ordered such that
 // left to right evaluation of the body will not encounter unbound variables
 // in input positions or negated expressions.
@@ -1411,9 +1424,9 @@ func (vs unsafeVars) Vars() VarSet {
 //
 // If the body cannot be reordered to ensure safety, the second return value
 // contains a mapping of expressions to unsafe variables in those expressions.
-func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) {
+func reorderBodyForSafety(arity getArity, globals VarSet, body Body) (Body, unsafeVars) {
 
-	body, unsafe := reorderBodyForClosures(globals, body)
+	body, unsafe := reorderBodyForClosures(arity, globals, body)
 	if len(unsafe) != 0 {
 		return nil, unsafe
 	}
@@ -1439,7 +1452,7 @@ func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) {
 				continue
 			}
 
-			safe.Update(e.OutputVars(safe))
+			safe.Update(outputVarsForExpr(e, arity, safe))
 
 			for v := range unsafe[e] {
 				if safe.Contains(v) {
@@ -1467,6 +1480,7 @@ func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) {
 			g.Update(reordered[i-1].Vars(safetyCheckVarVisitorParams))
 		}
 		vis := &bodySafetyVisitor{
+			arity:   arity,
 			current: e,
 			globals: g,
 			unsafe:  unsafe,
@@ -1482,6 +1496,7 @@ func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) {
 }
 
 type bodySafetyVisitor struct {
+	arity   getArity
 	current *Expr
 	globals VarSet
 	unsafe  unsafeVars
@@ -1516,7 +1531,7 @@ func (vis *bodySafetyVisitor) checkComprehensionSafety(tv VarSet, body Body) Bod
 	}
 
 	// Check body for safety, reordering as necessary.
-	r, u := reorderBodyForSafety(vis.globals, body)
+	r, u := reorderBodyForSafety(vis.arity, vis.globals, body)
 	if len(u) == 0 {
 		return r
 	}
@@ -1542,7 +1557,7 @@ func (vis *bodySafetyVisitor) checkSetComprehensionSafety(sc *SetComprehension)
 // reorderBodyForClosures returns a copy of the body ordered such that
 // expressions (such as array comprehensions) that close over variables are ordered
 // after other expressions that contain the same variable in an output position.
-func reorderBodyForClosures(globals VarSet, body Body) (Body, unsafeVars) {
+func reorderBodyForClosures(arity getArity, globals VarSet, body Body) (Body, unsafeVars) {
 
 	reordered := Body{}
 	unsafe := unsafeVars{}
@@ -1568,7 +1583,7 @@ func reorderBodyForClosures(globals VarSet, body Body) (Body, unsafeVars) {
 			// contained in the output position of an expression in the reordered
 			// body. These vars are considered unsafe.
 			cv := vs.Intersect(body.Vars(safetyCheckVarVisitorParams)).Diff(globals)
-			uv := cv.Diff(reordered.OutputVars(globals))
+			uv := cv.Diff(outputVarsForBody(reordered, arity, globals))
 
 			if len(uv) == 0 {
 				reordered = append(reordered, e)
@@ -1586,6 +1601,159 @@ func reorderBodyForClosures(globals VarSet, body Body) (Body, unsafeVars) {
 	return reordered, unsafe
 }
 
+func outputVarsForBody(body Body, arity getArity, safe VarSet) VarSet {
+	o := safe.Copy()
+	for _, e := range body {
+		o.Update(outputVarsForExpr(e, arity, o))
+	}
+	return o.Diff(safe)
+}
+
+func outputVarsForExpr(expr *Expr, arity getArity, safe VarSet) VarSet {
+
+	// Negated expressions must be safe.
+	if expr.Negated {
+		return VarSet{}
+	}
+
+	// With modifier inputs must be safe.
+	for _, with := range expr.With {
+		unsafe := false
+		WalkVars(with, func(v Var) bool {
+			if !safe.Contains(v) {
+				unsafe = true
+				return true
+			}
+			return false
+		})
+		if unsafe {
+			return VarSet{}
+		}
+	}
+
+	if !expr.IsCall() {
+		return outputVarsForExprRefs(expr, safe)
+	}
+
+	terms := expr.Terms.([]*Term)
+	name := terms[0].String()
+
+	if b := BuiltinMap[name]; b != nil {
+		if b.Name == Equality.Name {
+			return outputVarsForExprEq(expr, safe)
+		}
+		return outputVarsForExprBuiltin(expr, b, safe)
+	}
+
+	return outputVarsForExprCall(expr, arity, safe, terms)
+}
+
+func outputVarsForExprBuiltin(expr *Expr, b *Builtin, safe VarSet) VarSet {
+
+	output := outputVarsForExprRefs(expr, safe)
+	terms := expr.Terms.([]*Term)
+
+	// Check that all input terms are safe.
+	for i, t := range terms[1:] {
+		if b.IsTargetPos(i) {
+			continue
+		}
+		vis := NewVarVisitor().WithParams(VarVisitorParams{
+			SkipClosures:   true,
+			SkipSets:       true,
+			SkipObjectKeys: true,
+			SkipRefHead:    true,
+		})
+		Walk(vis, t)
+		unsafe := vis.Vars().Diff(output).Diff(safe)
+		if len(unsafe) > 0 {
+			return VarSet{}
+		}
+	}
+
+	// Add vars in target positions to result.
+	for i, t := range terms[1:] {
+		if b.IsTargetPos(i) {
+			vis := NewVarVisitor().WithParams(VarVisitorParams{
+				SkipRefHead:    true,
+				SkipSets:       true,
+				SkipObjectKeys: true,
+				SkipClosures:   true,
+			})
+			Walk(vis, t)
+			output.Update(vis.vars)
+		}
+	}
+
+	return output
+}
+
+func outputVarsForExprEq(expr *Expr, safe VarSet) VarSet {
+	ts := expr.Terms.([]*Term)
+	output := outputVarsForExprRefs(expr, safe)
+	output.Update(safe)
+	output.Update(Unify(output, ts[1], ts[2]))
+	return output.Diff(safe)
+}
+
+func outputVarsForExprCall(expr *Expr, arity getArity, safe VarSet, terms []*Term) VarSet {
+
+	output := outputVarsForExprRefs(expr, safe)
+
+	ref, ok := terms[0].Value.(Ref)
+	if !ok {
+		return VarSet{}
+	}
+
+	numArgs := arity(ref)
+	if numArgs == 0 {
+		return VarSet{}
+	}
+
+	numInputTerms := numArgs + 1
+
+	if numInputTerms >= len(terms) {
+		return VarSet{}
+	}
+
+	vis := NewVarVisitor().WithParams(VarVisitorParams{
+		SkipClosures:   true,
+		SkipSets:       true,
+		SkipObjectKeys: true,
+		SkipRefHead:    true,
+	})
+
+	Walk(vis, Args(terms[:numInputTerms]))
+	unsafe := vis.Vars().Diff(output).Diff(safe)
+
+	if len(unsafe) > 0 {
+		return VarSet{}
+	}
+
+	vis = NewVarVisitor().WithParams(VarVisitorParams{
+		SkipRefHead:    true,
+		SkipSets:       true,
+		SkipObjectKeys: true,
+		SkipClosures:   true,
+	})
+
+	Walk(vis, Args(terms[numInputTerms:]))
+	output.Update(vis.vars)
+	return output
+}
+
+func outputVarsForExprRefs(expr *Expr, safe VarSet) VarSet {
+	output := VarSet{}
+	WalkRefs(expr, func(r Ref) bool {
+		if safe.Contains(r[0].Value.(Var)) {
+			output.Update(r.OutputVars())
+			return false
+		}
+		return true
+	})
+	return output
+}
+
 type equalityFactory struct {
 	gen *localVarGenerator
 }

ast/compile_test.go

@@ -237,9 +237,9 @@ func TestCompilerFunctions(t *testing.T) {
 				import data.x
 				import data.x.f as g
 
-				p { g(1, _) }
-				p { x.f(1, _) }
-				p { data.x.f(1, _) }
+				p { g(1, a) }
+				p { x.f(1, b) }
+				p { data.x.f(1, c) }
 				`,
 			},
 		},
@@ -499,6 +499,8 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 		{"builtin-input", `p { count([1, 2, x], x) }`, `{x,}`},
 		{"builtin-input-name", `p { count(eq, 1) }`, `{eq,}`},
 		{"builtin-multiple", `p { x > 0; x <= 3; x != 2 }`, `{x,}`},
+		{"unordered-object-keys", `p { x = "a"; [{x: y, z: a}] = [{"a": 1, "b": 2}]}`, `{a,y,z}`},
+		{"unordered-sets", `p { x = "a"; [{x, y}] = [{1, 2}]}`, `{y,}`},
 		{"array-compr", `p { _ = [x | x = data.a[_]; y > 1] }`, `{y,}`},
 		{"array-compr-nested", `p { _ = [x | x = a[_]; a = [y | y = data.a[_]; z > 1]] }`, `{z,}`},
 		{"array-compr-closure", `p { _ = [v | v = [x | x = data.a[_]]; x > 1] }`, `{x,}`},
@@ -514,10 +516,11 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 		{"with-value", `p { data.a.b.d.t with input as x }`, `{x,}`},
 		{"with-value-2", `p { x = data.a.b.d.t with input as x }`, `{x,}`},
 		{"else-kw", "p { false } else { count(x, 1) }", `{x,}`},
-		{"userfunc", "foo(x) = [y, z] { split(x, y, z) }", `{y,z}`},
+		{"function", "foo(x) = [y, z] { split(x, y, z) }", `{y,z}`},
 		{"call-vars", "p { f[i].g[j](1) }", `{i, j}`},
 		{"call-vars-input", "p { f(x, x) } f(x) = x { true }", `{x,}`},
 		{"call-no-output", "p { f(x) } f(x) = x { true }", `{x,}`},
+		{"call-too-few", "p { f(1,x) } f(x,y) { true }", "{x,}"},
 	}
 
 	makeErrMsg := func(varName string) string {

ast/policy.go

@@ -702,16 +702,6 @@ func (body Body) Loc() *Location {
 	return body[0].Location
 }
 
-// OutputVars returns a VarSet containing the variables that would be bound by evaluating
-// the body.
-func (body Body) OutputVars(safe VarSet) VarSet {
-	o := safe.Copy()
-	for _, e := range body {
-		o.Update(e.OutputVars(o))
-	}
-	return o.Diff(safe)
-}
-
 func (body Body) String() string {
 	var buf []string
 	for _, v := range body {
@@ -937,35 +927,6 @@ func (expr *Expr) IsGround() bool {
 	return true
 }
 
-// OutputVars returns a VarSet containing variables that would be bound by evaluating
-// this expression.
-func (expr *Expr) OutputVars(safe VarSet) VarSet {
-	if !expr.Negated {
-
-		// Currently the with modifier does not produce any outputs. Any
-		// variables in the value must be safe before the expression can be
-		// evaluated.
-		if !expr.withModifierSafe(safe) {
-			return VarSet{}
-		}
-
-		switch terms := expr.Terms.(type) {
-		case *Term:
-			return expr.outputVarsRefs(safe)
-		case []*Term:
-			name := terms[0].String()
-			if b := BuiltinMap[name]; b != nil {
-				if b.Name == Equality.Name {
-					return expr.outputVarsEquality(safe)
-				}
-				return expr.outputVarsBuiltins(b, safe)
-			}
-			return expr.outputVarsFunc(safe, terms)
-		}
-	}
-	return VarSet{}
-}
-
 // SetOperator sets the expr's operator and returns the expr itself. If expr is
 // not a call expr, this function will panic.
 func (expr *Expr) SetOperator(term *Term) *Expr {
@@ -1019,120 +980,6 @@ func (expr *Expr) Vars(params VarVisitorParams) VarSet {
 	return vis.Vars()
 }
 
-func (expr *Expr) outputVarsBuiltins(b *Builtin, safe VarSet) VarSet {
-
-	o := expr.outputVarsRefs(safe)
-	terms := expr.Terms.([]*Term)
-
-	// Check that all input terms are ground or safe.
-	for i, t := range terms[1:] {
-		if b.IsTargetPos(i) {
-			continue
-		}
-		if t.Value.IsGround() {
-			continue
-		}
-		vis := NewVarVisitor().WithParams(VarVisitorParams{
-			SkipClosures:   true,
-			SkipObjectKeys: true,
-			SkipRefHead:    true,
-		})
-		Walk(vis, t)
-		unsafe := vis.Vars().Diff(o).Diff(safe)
-		if len(unsafe) > 0 {
-			return VarSet{}
-		}
-	}
-
-	// Add vars in target positions to result.
-	for i, t := range terms[1:] {
-		if b.IsTargetPos(i) {
-			vis := NewVarVisitor().WithParams(VarVisitorParams{
-				SkipRefHead:    true,
-				SkipSets:       true,
-				SkipObjectKeys: true,
-				SkipClosures:   true,
-			})
-			Walk(vis, t)
-			o.Update(vis.vars)
-		}
-	}
-
-	return o
-}
-
-func (expr *Expr) outputVarsEquality(safe VarSet) VarSet {
-	ts := expr.Terms.([]*Term)
-	o := expr.outputVarsRefs(safe)
-	o.Update(safe)
-	o.Update(Unify(o, ts[1], ts[2]))
-	return o.Diff(safe)
-}
-
-func (expr *Expr) outputVarsFunc(safe VarSet, terms []*Term) VarSet {
-
-	// Functions called with 0 or 1 args cannot produce output vars.
-	if len(expr.Operands()) < 2 {
-		return VarSet{}
-	}
-
-	o := expr.outputVarsRefs(safe)
-
-	// Find unsafe input vars.
-	args := Args(terms[:len(terms)-1])
-	vis := NewVarVisitor().WithParams(VarVisitorParams{
-		SkipClosures:   true,
-		SkipObjectKeys: true,
-		SkipRefHead:    true,
-	})
-	Walk(vis, args)
-	unsafe := vis.Vars().Diff(o).Diff(safe)
-	if len(unsafe) > 0 {
-		return VarSet{}
-	}
-
-	// Find safe output vars.
-	vis = NewVarVisitor().WithParams(VarVisitorParams{
-		SkipRefHead:    true,
-		SkipSets:       true,
-		SkipObjectKeys: true,
-		SkipClosures:   true,
-	})
-	Walk(vis, terms[len(terms)-1])
-	o.Update(vis.vars)
-
-	return o
-}
-
-func (expr *Expr) outputVarsRefs(safe VarSet) VarSet {
-	o := VarSet{}
-	WalkRefs(expr, func(r Ref) bool {
-		if safe.Contains(r[0].Value.(Var)) {
-			o.Update(r.OutputVars())
-			return false
-		}
-		return true
-	})
-	return o
-}
-
-func (expr *Expr) withModifierSafe(safe VarSet) bool {
-	for _, with := range expr.With {
-		unsafe := false
-		WalkVars(with, func(v Var) bool {
-			if !safe.Contains(v) {
-				unsafe = true
-				return true
-			}
-			return false
-		})
-		if unsafe {
-			return false
-		}
-	}
-	return true
-}
-
 // NewBuiltinExpr creates a new Expr object with the supplied terms.
 // The builtin operator must be the first term.
 func NewBuiltinExpr(terms ...*Term) *Expr {

ast/policy_test.go

@@ -10,7 +10,6 @@ import (
 	"reflect"
 	"testing"
 
-	"github.com/open-policy-agent/opa/types"
 	"github.com/open-policy-agent/opa/util"
 )
 
@@ -217,71 +216,6 @@ func TestBodyIsGround(t *testing.T) {
 	}
 }
 
-func TestExprOutputVars(t *testing.T) {
-
-	RegisterBuiltin(&Builtin{
-		Name: "test_out_array",
-		Decl: types.NewFunction(
-			nil, types.NewArray(nil, types.N),
-		),
-	})
-
-	RegisterBuiltin(&Builtin{
-		Name: "test_out_set",
-		Decl: types.NewFunction(
-			nil, types.NewArray(nil, types.N),
-		),
-	})
-
-	RegisterBuiltin(&Builtin{
-		Name: "foo",
-		Decl: types.NewFunction(
-			types.Args(types.A),
-			types.A,
-		),
-	})
-
-	tests := []struct {
-		note     string
-		expr     string
-		safe     string
-		expected string
-	}{
-		{"ref 1", "a[i].b[j]", "[a]", "[i, j]"},
-		{"ref 2", "[1,2,a[i]]", "[a]", "[i]"},
-		{"simple unify", `{"a": [{x: y}, b[z]]} = c[i]`, "[b, c]", "[y, z, i]"},
-		{"built-in", "count([], x)", "[]", "[x]"},
-		{"built-in-array", "test_out_array([x])", "[]", "[x]"},
-		{"built-in-set", "test_out_set({x})", "[]", "[]"},
-		{"with", "data.foo[x] with input as bar", "[bar]", "[x]"},
-		{"with unsafe", "data.foo[x] with input as x", "[]", "[]"},
-		{"userfunc", "foo(x, y)", "[x]", "[y]"},
-	}
-
-	for i, tc := range tests {
-
-		expr := MustParseBody(tc.expr)[0]
-		safe := ReservedVars.Copy()
-
-		for _, x := range MustParseTerm(tc.safe).Value.(Array) {
-			safe.Add(x.Value.(Var))
-		}
-
-		result := expr.OutputVars(safe)
-
-		expected := VarSet{}
-		for _, x := range MustParseTerm(tc.expected).Value.(Array) {
-			expected.Add(x.Value.(Var))
-		}
-
-		missing := expected.Diff(result)
-		extra := result.Diff(expected)
-		if len(missing) != 0 || len(extra) != 0 {
-			t.Errorf("%s (%d): Missing output vars: %v, extra output vars: %v", tc.note, i, missing, extra)
-		}
-	}
-}
-
 func TestExprString(t *testing.T) {
 	expr1 := &Expr{
 		Terms: RefTerm(VarTerm("q"), StringTerm("r"), VarTerm("x")),

repl/repl_test.go

@@ -412,7 +412,7 @@ func TestUnset(t *testing.T) {
 	repl.OneShot(ctx, "p(x) = y { y = x }")
 	repl.OneShot(ctx, "unset p")
 
-	err = repl.OneShot(ctx, "data.repl.p(5, y)")
+	err = repl.OneShot(ctx, "data.repl.p(1, 2)")
 	if err == nil || err.Error() != `1 error occurred: 1:1: rego_type_error: undefined function data.repl.p` {
 		t.Fatalf("Expected eval error (undefined built-in) but got err: '%v'", err)
 	}
@@ -422,7 +422,7 @@ func TestUnset(t *testing.T) {
 	repl.OneShot(ctx, "p(2, x) = y { y = x+1 }")
 	repl.OneShot(ctx, "unset p")
 
-	err = repl.OneShot(ctx, "data.repl.p(1, 2, y)")
+	err = repl.OneShot(ctx, "data.repl.p(1, 2, 3)")
 	if err == nil || err.Error() != `1 error occurred: 1:1: rego_type_error: undefined function data.repl.p` {
 		t.Fatalf("Expected eval error (undefined built-in) but got err: '%v'", err)
 	}