From 2f264b12f1e004b1bbfb225ef7d0894cc83cfb2b Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Fri, 14 Oct 2022 10:38:04 +0200
Subject: [PATCH] ast: add 'features' to capabilities, use it to check for ref
 heads (#5208)

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 ast/capabilities.go | 15 ++++++++++++
 ast/compile.go      | 13 ++++++++++
 ast/compile_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++
 capabilities.json   |  3 +++
 4 files changed, 89 insertions(+)

diff --git a/ast/capabilities.go b/ast/capabilities.go
index 4869767a42..217e99c999 100644
--- a/ast/capabilities.go
+++ b/ast/capabilities.go
@@ -17,6 +17,12 @@ import (
 	"github.com/open-policy-agent/opa/util"
 )
 
+// In the compiler, we used this to check that we're OK working with ref heads.
+// If this isn't present, we'll fail. This is to ensure that older versions of
+// 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"
+
 // Capabilities defines a structure containing data that describes the capabilities
 // or features supported by a particular version of OPA.
 type Capabilities struct {
@@ -24,6 +30,11 @@ type Capabilities struct {
 	FutureKeywords  []string         `json:"future_keywords"`
 	WasmABIVersions []WasmABIVersion `json:"wasm_abi_versions"`
 
+	// Features is a bit of a mixed bag for checking that an older version of OPA
+	// is able to do what needs to be done.
+	// TODO(sr): find better words ^^
+	Features []string `json:"features"`
+
 	// allow_net is an array of hostnames or IP addresses, that an OPA instance is
 	// allowed to connect to.
 	// If omitted, ANY host can be connected to. If empty, NO host can be connected to.
@@ -60,6 +71,10 @@ func CapabilitiesForThisVersion() *Capabilities {
 	}
 	sort.Strings(f.FutureKeywords)
 
+	f.Features = []string{
+		FeatureRefHeadStringPrefixes,
+	}
+
 	return f
 }
 
diff --git a/ast/compile.go b/ast/compile.go
index 81b3d6110c..fcc54bfdd6 100644
--- a/ast/compile.go
+++ b/ast/compile.go
@@ -1635,6 +1635,19 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 				rule.Head.Reference = ref
 			}
 
+			cannotSpeakRefs := true
+			for _, f := range c.capabilities.Features {
+				if f == FeatureRefHeadStringPrefixes {
+					cannotSpeakRefs = false
+					break
+				}
+			}
+
+			if cannotSpeakRefs && 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++ {
 				// NOTE(sr): In the first iteration, non-string values in the refs are forbidden
 				// except for the last position, e.g.
diff --git a/ast/compile_test.go b/ast/compile_test.go
index 8b02e49609..f28345a5c2 100644
--- a/ast/compile_test.go
+++ b/ast/compile_test.go
@@ -3418,6 +3418,64 @@ x.y.w contains bar[i] if true
 	}
 }
 
+func TestCompilerRefHeadsNeedCapability(t *testing.T) {
+	popts := ParserOptions{AllFutureKeywords: true, unreleasedKeywords: true}
+	for _, tc := range []struct {
+		note string
+		mod  *Module
+		err  string
+	}{
+		{
+			note: "one-dot ref, single-value rule, short+compat",
+			mod: MustParseModule(`package t
+p[1] = 2`),
+		},
+		{
+			note: "one-dot ref, single-value rule, compat",
+			mod: MustParseModuleWithOpts(`package t
+p[3] = 4 if true`, popts),
+		},
+		{
+			note: "multi-value non-ref head",
+			mod: MustParseModuleWithOpts(`package t
+p contains 1 if true`, popts),
+		},
+		{ // NOTE(sr): this was previously forbidden because we need the `if` for disambiguation
+			note: "one-dot ref head",
+			mod: MustParseModuleWithOpts(`package t
+p[1] if true`, popts),
+			err: "rule heads with refs are not supported: p[1]",
+		},
+		{
+			note: "single-value ref rule",
+			mod: MustParseModuleWithOpts(`package t
+a.b.c[x] if x := input`, popts),
+			err: "rule heads with refs are not supported: a.b.c[x]",
+		},
+		{
+			note: "multi-value ref rule",
+			mod: MustParseModuleWithOpts(`package t
+a.b.c contains x if x := input`, popts),
+			err: "rule heads with refs are not supported: a.b.c",
+		},
+	} {
+		t.Run(tc.note, func(t *testing.T) {
+			caps, err := LoadCapabilitiesVersion("v0.44.0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			c := NewCompiler().WithCapabilities(caps)
+			c.Modules["test"] = tc.mod
+			compileStages(c, c.rewriteRefsInHead)
+			if tc.err != "" {
+				assertErrorWithMessage(t, c.Errors, tc.err)
+			} else {
+				assertNotFailed(t, c)
+			}
+		})
+	}
+}
+
 func TestCompilerRewriteRegoMetadataCalls(t *testing.T) {
 	tests := []struct {
 		note   string
diff --git a/capabilities.json b/capabilities.json
index 1d28f21b65..2c2fb50ce1 100644
--- a/capabilities.json
+++ b/capabilities.json
@@ -4302,5 +4302,8 @@
       "version": 1,
       "minor_version": 2
     }
+  ],
+  "features": [
+    "rule_head_ref_string_prefixes"
   ]
 }