ast: forbid dynamic recursion with hidden document

ast/compile.go

@@ -121,6 +121,14 @@ type CompilerStageDefinition struct {
 	Stage      CompilerStage
 }
 
+// RulesOptions defines the options for retrieving rules by Ref from the
+// compiler.
+type RulesOptions struct {
+	// IncludeHiddenModules determines if the result contains hidden modules,
+	// currently only the "system" namespace, i.e. "data.system.*".
+	IncludeHiddenModules bool
+}
+
 // QueryContext contains contextual information for running an ad-hoc query.
 //
 // Ad-hoc queries can be run in the context of a package and imports may be
@@ -554,6 +562,14 @@ func (c *Compiler) GetRules(ref Ref) (rules []*Rule) {
 }
 
 // GetRulesDynamic returns a slice of rules that could be referred to by a ref.
+//
+// Deprecated: use GetRulesDynamicWithOpts
+func (c *Compiler) GetRulesDynamic(ref Ref) []*Rule {
+	return c.GetRulesDynamicWithOpts(ref, RulesOptions{})
+}
+
+// GetRulesDynamicWithOpts returns a slice of rules that could be referred to by
+// a ref.
 // When parts of the ref are statically known, we use that information to narrow
 // down which rules the ref could refer to, but in the most general case this
 // will be an over-approximation.
@@ -572,10 +588,25 @@ func (c *Compiler) GetRules(ref Ref) (rules []*Rule) {
 //
 // The following calls yield the rules on the right.
 //
-//  GetRulesDynamic("data.a[x].c[y]") => [rule1, rule2]
-//  GetRulesDynamic("data.a[x].c.r2") => [rule2]
-//  GetRulesDynamic("data.a.b[x][y]") => [rule1]
-func (c *Compiler) GetRulesDynamic(ref Ref) (rules []*Rule) {
+//  GetRulesDynamicWithOpts("data.a[x].c[y]", opts) => [rule1, rule2]
+//  GetRulesDynamicWithOpts("data.a[x].c.r2", opts) => [rule2]
+//  GetRulesDynamicWithOpts("data.a.b[x][y]", opts) => [rule1]
+//
+// Using the RulesOptions parameter, the inclusion of hidden modules can be
+// controlled:
+//
+// With
+//
+//  package system.main
+//
+//  r3 = 3 # rule3
+//
+// We'd get this result:
+//
+//  GetRulesDynamicWithOpts("data[x]", RulesOptions{IncludeHiddenModules: true}) => [rule1, rule2, rule3]
+//
+// Without the options, it would be excluded.
+func (c *Compiler) GetRulesDynamicWithOpts(ref Ref, opts RulesOptions) []*Rule {
 	node := c.RuleTree
 
 	set := map[*Rule]struct{}{}
@@ -586,6 +617,9 @@ func (c *Compiler) GetRulesDynamic(ref Ref) (rules []*Rule) {
 			// under this "prefix".
 			node.DepthFirst(func(descendant *TreeNode) bool {
 				insertRules(set, descendant.Values)
+				if opts.IncludeHiddenModules {
+					return false
+				}
 				return descendant.Hide
 			})
 		} else if i == 0 || IsConstant(ref[i].Value) {
@@ -606,7 +640,7 @@ func (c *Compiler) GetRulesDynamic(ref Ref) (rules []*Rule) {
 			// This part of the ref is a dynamic term.  We can't know what it refers
 			// to and will just need to try all of the children.
 			for _, child := range node.Children {
-				if child.Hide {
+				if child.Hide && !opts.IncludeHiddenModules {
 					continue
 				}
 				insertRules(set, child.Values)
@@ -616,6 +650,7 @@ func (c *Compiler) GetRulesDynamic(ref Ref) (rules []*Rule) {
 	}
 
 	walk(node, 0)
+	rules := make([]*Rule, 0, len(set))
 	for rule := range set {
 		rules = append(rules, rule)
 	}
@@ -1764,7 +1799,10 @@ func (c *Compiler) setRuleTree() {
 }
 
 func (c *Compiler) setGraph() {
-	c.Graph = NewGraph(c.Modules, c.GetRulesDynamic)
+	list := func(r Ref) []*Rule {
+		return c.GetRulesDynamicWithOpts(r, RulesOptions{IncludeHiddenModules: true})
+	}
+	c.Graph = NewGraph(c.Modules, list)
 }
 
 type queryCompiler struct {

ast/compile_test.go

@@ -3462,24 +3462,34 @@ func TestCompilerCheckDynamicRecursion(t *testing.T) {
 	// This test tries to circumvent the recursion check by using dynamic
 	// references.  For more background info, see
 	// <https://github.com/open-policy-agent/opa/issues/1565>.
-	c := NewCompiler()
-	c.Modules = map[string]*Module{
-		"recursion": MustParseModule(`package recursion
 
+	for note, mod := range map[string]*Module{
+		"recursion": MustParseModule(`
+package recursion
 pkg = "recursion"
-
 foo[x] {
-  data[pkg]["foo"][x]
-}`),
-	}
-
-	compileStages(c, c.checkRecursion)
+	data[pkg]["foo"][x]
+}
+`),
+		"system.main": MustParseModule(`
+package system.main
+foo {
+  data[input]
+}
+`),
+	} {
+		t.Run(note, func(t *testing.T) {
+			c := NewCompiler()
+			c.Modules = map[string]*Module{note: mod}
+			compileStages(c, c.checkRecursion)
 
-	result := compilerErrsToStringSlice(c.Errors)
-	expected := "rego_recursion_error: rule foo is recursive: foo -> foo"
+			result := compilerErrsToStringSlice(c.Errors)
+			expected := "rego_recursion_error: rule foo is recursive: foo -> foo"
 
-	if len(result) != 1 || result[0] != expected {
-		t.Errorf("Expected %v but got: %v", expected, result)
+			if len(result) != 1 || result[0] != expected {
+				t.Errorf("Expected %v but got: %v", expected, result)
+			}
+		})
 	}
 }
 
@@ -3753,32 +3763,40 @@ r1 = 1`,
 r2 = 2`,
 		"mod3": `package a.b
 r3 = 3`,
+		"hidden": `package system.hidden
+r4 = 4`,
 	})
 
 	compileStages(compiler, nil)
 
 	rule1 := compiler.Modules["mod1"].Rules[0]
 	rule2 := compiler.Modules["mod2"].Rules[0]
 	rule3 := compiler.Modules["mod3"].Rules[0]
+	rule4 := compiler.Modules["hidden"].Rules[0]
 
 	tests := []struct {
-		input    string
-		expected []*Rule
+		input         string
+		expected      []*Rule
+		excludeHidden bool
 	}{
-		{"data.a.b.c.d.r1", []*Rule{rule1}},
-		{"data.a.b[x]", []*Rule{rule1, rule2, rule3}},
-		{"data.a.b[x].d", []*Rule{rule1, rule3}},
-		{"data.a.b.c", []*Rule{rule1, rule2}},
-		{"data.a.b.d", nil},
-		{"data[x]", []*Rule{rule1, rule2, rule3}},
-		{"data[data.complex_computation].b[y]", []*Rule{rule1, rule2, rule3}},
-		{"data[x][y].c.e", []*Rule{rule2}},
-		{"data[x][y].r3", []*Rule{rule3}},
+		{input: "data.a.b.c.d.r1", expected: []*Rule{rule1}},
+		{input: "data.a.b[x]", expected: []*Rule{rule1, rule2, rule3}},
+		{input: "data.a.b[x].d", expected: []*Rule{rule1, rule3}},
+		{input: "data.a.b.c", expected: []*Rule{rule1, rule2}},
+		{input: "data.a.b.d"},
+		{input: "data[x]", expected: []*Rule{rule1, rule2, rule3, rule4}},
+		{input: "data[data.complex_computation].b[y]", expected: []*Rule{rule1, rule2, rule3}},
+		{input: "data[x][y].c.e", expected: []*Rule{rule2}},
+		{input: "data[x][y].r3", expected: []*Rule{rule3}},
+		{input: "data[x][y]", expected: []*Rule{rule1, rule2, rule3}, excludeHidden: true}, // old behaviour of GetRulesDynamic
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.input, func(t *testing.T) {
-			result := compiler.GetRulesDynamic(MustParseRef(tc.input))
+			result := compiler.GetRulesDynamicWithOpts(
+				MustParseRef(tc.input),
+				RulesOptions{IncludeHiddenModules: !tc.excludeHidden},
+			)
 
 			if len(result) != len(tc.expected) {
 				t.Fatalf("Expected %v but got: %v", tc.expected, result)

rego/rego.go

@@ -1469,22 +1469,22 @@ func (r *Rego) PrepareForEval(ctx context.Context, opts ...PrepareOption) (Prepa
 
 		queries := []ast.Body{r.compiledQueries[evalQueryType].query}
 
-		// nolint: staticcheck // SA4006 false positive
-		cr, err := r.compileWasm(modules, queries, evalQueryType)
+		e, err := opa.LookupEngine(targetWasm)
 		if err != nil {
-			_ = txnClose(ctx, err) // Ignore error
 			return PreparedEvalQuery{}, err
 		}
 
 		// nolint: staticcheck // SA4006 false positive
-		data, err := r.store.Read(ctx, r.txn, storage.Path{})
+		cr, err := r.compileWasm(modules, queries, evalQueryType)
 		if err != nil {
 			_ = txnClose(ctx, err) // Ignore error
 			return PreparedEvalQuery{}, err
 		}
 
-		e, err := opa.LookupEngine(targetWasm)
+		// nolint: staticcheck // SA4006 false positive
+		data, err := r.store.Read(ctx, r.txn, storage.Path{})
 		if err != nil {
+			_ = txnClose(ctx, err) // Ignore error
 			return PreparedEvalQuery{}, err
 		}
 

topdown/save.go

@@ -347,7 +347,7 @@ func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, s
 				} else if ic.Disabled(v.ConstantPrefix(), icIgnoreInternal) {
 					found = true
 				} else {
-					for _, rule := range c.GetRulesDynamic(v) {
+					for _, rule := range c.GetRulesDynamicWithOpts(v, ast.RulesOptions{IncludeHiddenModules: false}) {
 						if saveRequired(c, ic, icIgnoreInternal, ss, b, rule, true) {
 							found = true
 							break