ast/parser: Detect function rule head + contains keyword.

Previously, we did not detect misuses of the `contains` keyword, which
is only valid as a syntactic sugar for partial set definition rules.

This commit adds some logic to detect when a function rule head has been
paired with the `contains` keyword, which should generally be an error
condition.

Fixes: open-policy-agent#5525

Signed-off-by: Philip Conrad <philipaconrad@gmail.com>

ast/check.go

@@ -213,7 +213,6 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) {
 	var tpe types.Type
 
 	if len(rule.Head.Args) > 0 {
-
 		// If args are not referred to in body, infer as any.
 		WalkVars(rule.Head.Args, func(v Var) bool {
 			if cpy.Get(v) == nil {

ast/compile_test.go

@@ -1734,6 +1734,27 @@ p {
 	}
 }
 
+func TestCompilerCheckSafetyFunctionAndContainsKeyword(t *testing.T) {
+	_, err := CompileModules(map[string]string{"test.rego": `package play
+
+			import future.keywords.contains
+
+			p(id) contains x {
+				x := id
+			}`})
+	if err == nil {
+		t.Fatal("expected error")
+	}
+
+	errs := err.(Errors)
+	if !strings.Contains(errs[0].Message, "the contains keyword can only be used with multi-value rule definitions (e.g., p contains <VALUE> { ... })") {
+		t.Fatal("wrong error message:", err)
+	}
+	if errs[0].Location.Row != 5 {
+		t.Fatal("expected error on line 5 but got:", errs[0].Location.Row)
+	}
+}
+
 func TestCompilerCheckTypes(t *testing.T) {
 	c := NewCompiler()
 	modules := getCompilerTestModules()

ast/parser.go

@@ -831,6 +831,10 @@ func (p *Parser) parseHead(defaultRule bool) (*Head, bool) {
 
 	switch p.s.tok {
 	case tokens.Contains: // NOTE: no Value for `contains` heads, we return here
+		// Catch error case of using 'contains' with a function definition rule head.
+		if head.Args != nil {
+			p.illegal("the contains keyword can only be used with multi-value rule definitions (e.g., %s contains <VALUE> { ... })", name)
+		}
 		p.scan()
 		head.Key = p.parseTermInfixCall()
 		if head.Key == nil {

ast/parser_test.go

@@ -1751,6 +1751,32 @@ func TestRuleContains(t *testing.T) {
 	}
 }
 
+func TestRuleContainsFail(t *testing.T) {
+	opts := ParserOptions{FutureKeywords: []string{"contains", "if", "every"}}
+
+	tests := []struct {
+		note     string
+		rule     string
+		expected string
+	}{
+		{
+			note:     "contains used with a 1+ argument function",
+			rule:     "p(a) contains x { x := a }",
+			expected: "the contains keyword can only be used with multi-value rule definitions (e.g., p contains <VALUE> { ... })",
+		},
+		{
+			note:     "contains used with a 0 argument function",
+			rule:     "p() contains x { x := 1 }",
+			expected: "the contains keyword can only be used with multi-value rule definitions (e.g., p contains <VALUE> { ... })",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			assertParseErrorContains(t, tc.note, tc.rule, tc.expected, opts)
+		})
+	}
+}
+
 func TestRuleIf(t *testing.T) {
 	opts := ParserOptions{FutureKeywords: []string{"contains", "if", "every"}}