ast: Adding rule_head_general_refs capabilities feature flag

ast/capabilities.go

@@ -22,6 +22,7 @@ import (
 // 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 FeatureGeneralRefHeads = "rule_head_general_refs"
 
 // Capabilities defines a structure containing data that describes the capabilities
 // or features supported by a particular version of OPA.
@@ -73,6 +74,7 @@ func CapabilitiesForThisVersion() *Capabilities {
 
 	f.Features = []string{
 		FeatureRefHeadStringPrefixes,
+		FeatureGeneralRefHeads,
 	}
 
 	return f

ast/compile.go

@@ -1753,10 +1753,13 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 			}
 
 			cannotSpeakRefs := true
+			cannotSpeakGeneralRefs := true
 			for _, f := range c.capabilities.Features {
-				if f == FeatureRefHeadStringPrefixes {
+				switch f {
+				case FeatureRefHeadStringPrefixes:
 					cannotSpeakRefs = false
-					break
+				case FeatureGeneralRefHeads:
+					cannotSpeakGeneralRefs = false
 				}
 			}
 
@@ -1766,9 +1769,16 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 			}
 
 			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 && i != len(ref)-1 { // last
+					if _, ok := ref[i].Value.(String); !ok {
+						c.err(NewError(TypeErr, rule.Loc(), "rule heads with general refs (variables outside of last term) 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 {
@@ -9544,6 +9560,106 @@ 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: "ref-head feature, no ref-head rules",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				p := 42`,
+		},
+		{
+			note: "ref-head feature, ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+			},
+			module: `package test
+				p.q.r := 42`,
+		},
+		{
+			note: "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 (variables outside of last term) are not supported: p[q].r[s]",
+		},
+		{
+			note: "general-ref-head feature, general-ref-head rule",
+			features: []string{
+				FeatureGeneralRefHeads,
+			},
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+			// Both FeatureRefHeadStringPrefixes and FeatureGeneralRefHeads flags are required for general-ref heads
+			expectedErr: "rego_compile_error: rule heads with refs are not supported: p[q].r[s]",
+		},
+		{
+			note: "ref-head & general-ref-head features, general-ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+				FeatureGeneralRefHeads,
+			},
+			module: `package test
+				p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
+		},
+		{
+			note: "ref-head & general-ref-head features, ref-head rule",
+			features: []string{
+				FeatureRefHeadStringPrefixes,
+				FeatureGeneralRefHeads,
+			},
+			module: `package test
+				p.q.r := 42`,
+		},
+	}
+
+	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{

capabilities.json

@@ -4705,6 +4705,7 @@
     }
   ],
   "features": [
-    "rule_head_ref_string_prefixes"
+    "rule_head_ref_string_prefixes",
+    "rule_head_general_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_general_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. Only the last element of the ref (the key) is allowed to be a variable.
+* `rule_head_general_refs`: Enables the use of [variables at arbitrary locations](../policy-language/#variables-in-rule-head-references) in a rule's ref (i.e. a general ref). Also requires `rule_head_ref_string_prefixes` to be enabled.
+
 ### Future keywords
 
 The availability of future keywords in an OPA version can also be controlled using the capabilities file: