From 492aeb3cbb3bbc5cfb6054aed2fcea3aca8d4338 Mon Sep 17 00:00:00 2001
From: Johan Fylling <johan.dev@fylling.se>
Date: Wed, 13 Apr 2022 14:07:10 +0200
Subject: [PATCH] ast: Respecting capabilities built-ins for rego.metadata.*
 functions (#4581)

Signed-off-by: Johan Fylling <johan.dev@fylling.se>
---
 ast/compile.go   | 19 +++++++------
 cmd/eval_test.go | 74 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/ast/compile.go b/ast/compile.go
index 4f23549fa4..f04dcbb076 100644
--- a/ast/compile.go
+++ b/ast/compile.go
@@ -1771,6 +1771,9 @@ func (c *Compiler) parseMetadataBlocks() {
 func (c *Compiler) rewriteRegoMetadataCalls() {
 	eqFactory := newEqualityFactory(c.localvargen)
 
+	_, chainFuncAllowed := c.builtins[RegoMetadataChain.Name]
+	_, ruleFuncAllowed := c.builtins[RegoMetadataRule.Name]
+
 	for _, name := range c.sorted {
 		mod := c.Modules[name]
 
@@ -1779,9 +1782,9 @@ func (c *Compiler) rewriteRegoMetadataCalls() {
 			var ruleCalled bool
 
 			WalkExprs(rule, func(expr *Expr) bool {
-				if isRegoMetadataChainCall(expr) {
+				if chainFuncAllowed && isRegoMetadataChainCall(expr) {
 					chainCalled = true
-				} else if isRegoMetadataRuleCall(expr) {
+				} else if ruleFuncAllowed && isRegoMetadataRuleCall(expr) {
 					ruleCalled = true
 				}
 				return chainCalled && ruleCalled
@@ -1835,7 +1838,7 @@ func (c *Compiler) rewriteRegoMetadataCalls() {
 				rule.Body = body
 
 				vis := func(b Body) bool {
-					for _, err := range rewriteRegoMetadataCalls(metadataChainVar, metadataRuleVar, b, &c.RewrittenVars) {
+					for _, err := range rewriteRegoMetadataCalls(&metadataChainVar, &metadataRuleVar, b, &c.RewrittenVars) {
 						c.err(err)
 					}
 					return false
@@ -1864,7 +1867,7 @@ func getPrimaryRuleAnnotations(as *AnnotationSet, rule *Rule) *Annotations {
 	return annots[0]
 }
 
-func rewriteRegoMetadataCalls(metadataChainVar Var, metadataRuleVar Var, body Body, rewrittenVars *map[Var]Var) Errors {
+func rewriteRegoMetadataCalls(metadataChainVar *Var, metadataRuleVar *Var, body Body, rewrittenVars *map[Var]Var) Errors {
 	var errs Errors
 
 	WalkClosures(body, func(x interface{}) bool {
@@ -1885,10 +1888,10 @@ func rewriteRegoMetadataCalls(metadataChainVar Var, metadataRuleVar Var, body Bo
 		expr := body[i]
 		var metadataVar Var
 
-		if isRegoMetadataChainCall(expr) {
-			metadataVar = metadataChainVar
-		} else if isRegoMetadataRuleCall(expr) {
-			metadataVar = metadataRuleVar
+		if metadataChainVar != nil && isRegoMetadataChainCall(expr) {
+			metadataVar = *metadataChainVar
+		} else if metadataRuleVar != nil && isRegoMetadataRuleCall(expr) {
+			metadataVar = *metadataRuleVar
 		} else {
 			continue
 		}
diff --git a/cmd/eval_test.go b/cmd/eval_test.go
index 099c1c0c65..40730c8e11 100755
--- a/cmd/eval_test.go
+++ b/cmd/eval_test.go
@@ -500,6 +500,80 @@ func TestEvalWithSchemaFileWithRemoteRef(t *testing.T) {
 	})
 }
 
+func TestBuiltinsCapabilities(t *testing.T) {
+	tests := []struct {
+		note            string
+		policy          string
+		query           string
+		ruleName        string
+		expectedCode    string
+		expectedMessage string
+	}{
+		{
+			note:            "rego.metadata.chain() not allowed",
+			policy:          "package p\n r := rego.metadata.chain()",
+			query:           "data.p",
+			ruleName:        "rego.metadata.chain",
+			expectedCode:    "rego_type_error",
+			expectedMessage: "undefined function rego.metadata.chain",
+		},
+		{
+			note:            "rego.metadata.rule() not allowed",
+			policy:          "package p\n r := rego.metadata.rule()",
+			query:           "data.p",
+			ruleName:        "rego.metadata.rule",
+			expectedCode:    "rego_type_error",
+			expectedMessage: "undefined function rego.metadata.rule",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+
+			files := map[string]string{
+				"p.rego": tc.policy,
+			}
+
+			test.WithTempFS(files, func(path string) {
+				params := newEvalCommandParams()
+				params.capabilities.C = ast.CapabilitiesForThisVersion()
+				params.capabilities.C.Builtins = removeBuiltin(params.capabilities.C.Builtins, tc.ruleName)
+
+				_ = params.dataPaths.Set(filepath.Join(path, "p.rego"))
+
+				var buf bytes.Buffer
+				_, err := eval([]string{tc.query}, params, &buf)
+				if err == nil {
+					t.Fatal("expected error, got nil")
+				}
+				var output presentation.Output
+				if err := util.NewJSONDecoder(&buf).Decode(&output); err != nil {
+					t.Fatal(err)
+				}
+				if exp, act := 1, len(output.Errors); exp != act {
+					t.Fatalf("expected %d errors, got %d", exp, act)
+				}
+				if code := output.Errors[0].Code; code != tc.expectedCode {
+					t.Errorf("expected code '%v', got '%v'", tc.expectedCode, code)
+				}
+				if msg := output.Errors[0].Message; msg != tc.expectedMessage {
+					t.Errorf("expected message '%v', got '%v'", tc.expectedMessage, msg)
+				}
+			})
+		})
+	}
+}
+
+func removeBuiltin(builtins []*ast.Builtin, name string) []*ast.Builtin {
+	var cpy []*ast.Builtin
+	for _, builtin := range builtins {
+		if builtin.Name != name {
+			cpy = append(cpy, builtin)
+		}
+	}
+	return cpy
+}
+
 func TestEvalReturnsRegoError(t *testing.T) {
 	buf := new(bytes.Buffer)
 	_, err := eval([]string{`{k: v | k = ["a", "a"][_]; v = [0,1][_]}`}, newEvalCommandParams(), buf)