ast: Adding rule_head_refs capabilities feature flag (#6346)

To signal/toggle support for general refs in rule heads.

Fixes: #6334
Signed-off-by: Johan Fylling <johan.dev@fylling.se>

ast/capabilities.go

@@ -51,6 +51,7 @@ var minVersionIndex = func() VersionIndex {
 // OPA can work with policies that we're compiling -- if they don't know ref
 // heads, they wouldn't be able to parse them.
 const FeatureRefHeadStringPrefixes = "rule_head_ref_string_prefixes"
+const FeatureRefHeads = "rule_head_refs"
 
 // Capabilities defines a structure containing data that describes the capabilities
 // or features supported by a particular version of OPA.
@@ -102,6 +103,7 @@ func CapabilitiesForThisVersion() *Capabilities {
 
 	f.Features = []string{
 		FeatureRefHeadStringPrefixes,
+		FeatureRefHeads,
 	}
 
 	return f

ast/capabilities_test.go

@@ -181,7 +181,7 @@ func TestCapabilitiesMinimumCompatibleVersion(t *testing.T) {
 			version: "0.38.0",
 		},
 		{
-			note: "features",
+			note: "features (string prefix ref)",
 			module: `
 				package x
 				import future.keywords.if
@@ -189,6 +189,25 @@ func TestCapabilitiesMinimumCompatibleVersion(t *testing.T) {
 			`,
 			version: "0.46.0",
 		},
+		{
+			note: "features (general ref)",
+			module: `
+				package x
+				import future.keywords.if
+				p.a.b[c].d if { c := "foo" }
+			`,
+			version: "0.59.0",
+		},
+		{
+			note: "features (general ref + string prefix ref)",
+			module: `
+				package x
+				import future.keywords.if
+				p.a.b.c.d if { true }
+				p.a.b[c].d if { c := "foo" }
+			`,
+			version: "0.59.0",
+		},
 	}
 
 	for _, tc := range tests {

ast/compile.go

@@ -966,8 +966,13 @@ func (c *Compiler) buildRequiredCapabilities() {
 
 	for _, name := range c.sorted {
 		for _, rule := range c.Modules[name].Rules {
-			if len(rule.Head.Reference) >= 3 {
-				features[FeatureRefHeadStringPrefixes] = struct{}{}
+			refLen := len(rule.Head.Reference)
+			if refLen >= 3 {
+				if refLen > len(rule.Head.Reference.ConstantPrefix()) {
+					features[FeatureRefHeads] = struct{}{}
+				} else {
+					features[FeatureRefHeadStringPrefixes] = struct{}{}
+				}
 			}
 		}
 	}
@@ -1852,23 +1857,33 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 				rule.Head.Reference = ref
 			}
 
-			cannotSpeakRefs := true
+			cannotSpeakStringPrefixRefs := true
+			cannotSpeakGeneralRefs := true
 			for _, f := range c.capabilities.Features {
-				if f == FeatureRefHeadStringPrefixes {
-					cannotSpeakRefs = false
-					break
+				switch f {
+				case FeatureRefHeadStringPrefixes:
+					cannotSpeakStringPrefixRefs = false
+				case FeatureRefHeads:
+					cannotSpeakGeneralRefs = false
 				}
 			}
 
-			if cannotSpeakRefs && rule.Head.Name == "" {
+			if cannotSpeakStringPrefixRefs && cannotSpeakGeneralRefs && rule.Head.Name == "" {
 				c.err(NewError(CompileErr, rule.Loc(), "rule heads with refs are not supported: %v", rule.Head.Reference))
 				return true
 			}
 
 			for i := 1; i < len(ref); i++ {
-				// Rewrite so that any non-scalar elements that in the last position of
-				// the rule are vars:
+				if cannotSpeakGeneralRefs && (rule.Head.RuleKind() == MultiValue || i != len(ref)-1) { // last
+					if _, ok := ref[i].Value.(String); !ok {
+						c.err(NewError(TypeErr, rule.Loc(), "rule heads with general refs (containing variables) are not supported: %v", rule.Head.Reference))
+						continue
+					}
+				}
+
+				// Rewrite so that any non-scalar elements in the rule's ref are vars:
 				//     p.q.r[y.z] { ... }  =>  p.q.r[__local0__] { __local0__ = y.z }
+				//     p.q[a.b][c.d] { ... }  =>  p.q[__local0__] { __local0__ = a.b; __local1__ = c.d }
 				// because that's what the RuleTree knows how to deal with.
 				if _, ok := ref[i].Value.(Var); !ok && !IsScalar(ref[i].Value) {
 					expr := f.Generate(ref[i])

ast/compile_test.go

@@ -580,6 +580,22 @@ func TestCompilerCheckRuleHeadRefs(t *testing.T) {
 			),
 			expected: MustParseRule(`p.q.r[__local0__]  { y := {"z": "a"}; __local0__ = y.z }`),
 		},
+		{
+			note: "rewrite: single-value with non-var ref term",
+			modules: modules(
+				`package x
+				p.q[y.z].r if y := {"z": "a"}`,
+			),
+			expected: MustParseRule(`p.q[__local0__].r  { y := {"z": "a"}; __local0__ = y.z }`),
+		},
+		{
+			note: "rewrite: single-value with non-var ref term and key",
+			modules: modules(
+				`package x
+				p.q[a.b][c.d] if y := {"z": "a"}`,
+			),
+			expected: MustParseRule(`p.q[__local0__][__local1__]  { y := {"z": "a"}; __local0__ = a.b; __local1__ = c.d }`),
+		},
 	}
 
 	for _, tc := range tests {
@@ -9750,6 +9766,166 @@ func runQueryCompilerTest(q, pkg string, imports []string, expected interface{})
 	}
 }
 
+func TestCompilerCapabilitiesFeatures(t *testing.T) {
+	cases := []struct {
+		note        string
+		module      string
+		features    []string
+		expectedErr string
+	}{
+		{
+			note: "no features, no ref-head rules",
+			module: `package test
+				p := 42`,
+		},
+		{
+			note: "no features, ref-head rule",
+			module: `package test
+				p.q.r := 42`,
+			expectedErr: "rego_compile_error: rule heads with refs are not supported: p.q.r",
+		},
+		{
+			note: "no features, general-ref-head rule",
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+			expectedErr: "rego_compile_error: rule heads with refs are not supported: p[q].r[s]",
+		},
+		{
+			note: "string-prefix-ref-head feature, no ref-head rules",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				p := 42`,
+		},
+		{
+			note: "string-prefix-ref-head feature, ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				p.q.r := 42`,
+		},
+		{
+			note: "ref-head feature, ref-head rule",
+			features: []string{
+				FeatureRefHeads,
+			},
+			module: `package test
+				p.q.r := 42`,
+		},
+		{
+			note: "string-prefix-ref-head feature, general-ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+			expectedErr: "rego_type_error: rule heads with general refs (containing variables) are not supported: p[q].r[s]",
+		},
+		{
+			note: "ref-head feature, general-ref-head rule",
+			features: []string{
+				FeatureRefHeads,
+			},
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+		},
+		{
+			note: "string-prefix-ref-head & ref-head features, general-ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+				FeatureRefHeads,
+			},
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+		},
+		{
+			note: "string-prefix-ref-head & ref-head features, ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+				FeatureRefHeads,
+			},
+			module: `package test
+				p.q.r := 42`,
+		},
+		{
+			note:     "no features, string-prefix-ref-head with contains kw",
+			features: []string{},
+			module: `package test
+				import future.keywords.contains
+				p.x contains 1`,
+			expectedErr: "rego_compile_error: rule heads with refs are not supported: p.x",
+		},
+		{
+			note: "string-prefix-ref-head feature, string-prefix-ref-head with contains kw",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				import future.keywords.contains
+				p.x contains 1`,
+		},
+		{
+			note: "ref-head feature, string-prefix-ref-head with contains kw",
+			features: []string{
+				FeatureRefHeads,
+			},
+			module: `package test
+				import future.keywords.contains
+				p.x contains 1`,
+		},
+
+		{
+			note:     "no features, general-ref-head with contains kw",
+			features: []string{},
+			module: `package test
+				import future.keywords
+				p[x] contains 1 if x = "foo"`,
+			expectedErr: "rego_compile_error: rule heads with refs are not supported: p[x]",
+		},
+		{
+			note: "string-prefix-ref-head feature, general-ref-head with contains kw",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				import future.keywords
+				p[x] contains 1 if x = "foo"`,
+			expectedErr: "rego_type_error: rule heads with general refs (containing variables) are not supported: p[x]",
+		},
+		{
+			note: "ref-head feature, general-ref-head with contains kw",
+			features: []string{
+				FeatureRefHeads,
+			},
+			module: `package test
+				import future.keywords
+				p[x] contains 1 if x = "foo"`,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			capabilities := CapabilitiesForThisVersion()
+			capabilities.Features = tc.features
+
+			compiler := NewCompiler().WithCapabilities(capabilities)
+			compiler.Compile(map[string]*Module{"test": MustParseModule(tc.module)})
+			if tc.expectedErr != "" {
+				if !compiler.Failed() {
+					t.Fatal("expected error but got success")
+				}
+				if !strings.Contains(compiler.Errors.Error(), tc.expectedErr) {
+					t.Fatalf("expected error:\n\n%s\n\nbut got:\n\n%v", tc.expectedErr, compiler.Errors)
+				}
+			} else if compiler.Failed() {
+				t.Fatalf("unexpected error(s): %v", compiler.Errors)
+			}
+		})
+	}
+}
+
 func TestCompilerCapabilitiesExtendedWithCustomBuiltins(t *testing.T) {
 
 	compiler := NewCompiler().WithCapabilities(&Capabilities{

ast/version_index.json

@@ -1373,6 +1373,13 @@
       "Patch": 0,
       "PreRelease": "",
       "Metadata": ""
+    },
+    "rule_head_refs": {
+      "Major": 0,
+      "Minor": 59,
+      "Patch": 0,
+      "PreRelease": "",
+      "Metadata": ""
     }
   },
   "keywords": {

capabilities.json

@@ -4705,6 +4705,7 @@
     }
   ],
   "features": [
-    "rule_head_ref_string_prefixes"
+    "rule_head_ref_string_prefixes",
+    "rule_head_refs"
   ]
 }

docs/content/deployments.md

@@ -449,6 +449,24 @@ Note that the metaschemas [http://json-schema.org/draft-04/schema](http://json-s
 
 Similarly, the `allow_net` capability restricts what hosts the `http.send` built-in function may send requests to, and what hosts the `net.lookup_ip_addr` built-in function may resolve IP addresses for.
 
+### Features
+
+Some features of OPA can be toggled on and off through the `features` list:
+
+```json
+{
+    "features": [ 
+      "rule_head_ref_string_prefixes", 
+      "rule_head_refs" 
+    ]
+}
+```
+
+Features present in the list are enabled, while features not present are disabled. The following features are available:
+
+* `rule_head_ref_string_prefixes`: Enables the use of a [reference in place of name](../policy-language/#rule-heads-containing-references) in the head of rules. This is a subset of `rule_head_refs`, and only covers references where all terms are primitive types, or where only the last element of the ref (the key in the generated object or set) is allowed to be a variable.
+* `rule_head_refs`: Enables general support for [references in rule heads](../policy-language/#rule-heads-containing-references), including [variables at arbitrary locations](../policy-language/#variables-in-rule-head-references). This feature also covers the functionality of `rule_head_ref_string_prefixes`. 
+
 ### Future keywords
 
 The availability of future keywords in an OPA version can also be controlled using the capabilities file: