topdown: Fix key construction for NDBCache; rego: avoid NDBCache init (…

…#5097)

* topdown/eval: Fix key construction for NDBCache.

This commit fixes the construction process for keys used to look up
entries in the non-deterministic builtins cache. The original cache
lookup process could use refs that were not fully-grounded, and this
meant that calling a builtin twice with differing parameters could
appear identical to the cache if they involved refs as parameters.

We now use fully-grounded terms for the cache keys during insertion/lookup,
meaning that non-deterministic builtin calls with different parameters
are now guaranteed to result in unique cache entries, even if hidden
behind a ref.

* rego/rego: Fix unneeded NDBCache initializations.

The NDBCache is intended as an opt-in feature, and the initializations
present in the `rego` module meant that it was forced on virtually every
eval. This commit removes those initializations, so that the only way an
NDBCache will ever make it to the evaluator is if it's provided from
outside the evaluator.

* rego/rego_test: Add ND builtin iteration test.

This commit adds a test to ensure that the NDBCache correctly
handles iterative calls of a builtin with different input parameters.

Signed-off-by: Philip Conrad <philipaconrad@gmail.com>

rego/rego.go

@@ -320,7 +320,6 @@ func (pq preparedQuery) newEvalContext(ctx context.Context, options []EvalOption
 		compiledQuery:    compiledQuery{},
 		indexing:         true,
 		earlyExit:        true,
-		ndBuiltinCache:   builtins.NDBCache{},
 		resolvers:        pq.r.resolvers,
 		printHook:        pq.r.printHook,
 		capabilities:     pq.r.capabilities,
@@ -1112,7 +1111,6 @@ func New(options ...func(r *Rego)) *Rego {
 		builtinDecls:    map[string]*ast.Builtin{},
 		builtinFuncs:    map[string]*topdown.Builtin{},
 		bundles:         map[string]*bundle.Bundle{},
-		ndBuiltinCache:  builtins.NDBCache{},
 	}
 
 	for _, option := range options {

rego/rego_test.go

@@ -2105,7 +2105,47 @@ p {
 	}
 	_, ok := ndBC["http.send"]
 	if !ok {
-		t.Errorf("expected http.send cache entry")
+		t.Fatalf("expected http.send cache entry")
+	}
+}
+
+// Catches issues around iteration with ND builtins.
+func TestNDBCacheWithRuleBodyAndIteration(t *testing.T) {
+	ctx := context.Background()
+	ndBC := builtins.NDBCache{}
+	query := "data.foo.results = x"
+	_, err := New(
+		Query(query),
+		NDBuiltinCache(ndBC),
+		Module("test.rego", `package foo
+
+import future.keywords
+
+urls := [
+	"https://httpbin.org/headers",
+	"https://httpbin.org/ip",
+	"https://httpbin.org/user-agent"
+]
+
+results[response] {
+	some url in urls
+	response := http.send({
+		"method": "GET",
+		"url": url
+	})
+}`),
+	).Eval(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Ensure that the cache exists, and has exactly 3 entries.
+	entries, ok := ndBC["http.send"]
+	if !ok {
+		t.Fatalf("expected http.send cache entry")
+	}
+	if entries.Len() != 3 {
+		t.Fatalf("expected 3 http.send cache entries, received:\n%v", ndBC)
 	}
 }
 

topdown/eval.go

@@ -1707,7 +1707,7 @@ func (e evalBuiltin) eval(iter unifyIterator) error {
 		e.e.instr.stopTimer(evalOpBuiltinCall)
 
 		// Unify against the NDBCache result if present.
-		if v, ok := e.bctx.NDBuiltinCache.Get(e.bi.Name, ast.NewArray(e.terms[:endIndex]...)); ok {
+		if v, ok := e.bctx.NDBuiltinCache.Get(e.bi.Name, ast.NewArray(operands[:endIndex]...)); ok {
 			switch {
 			case e.bi.Decl.Result() == nil:
 				err = iter()
@@ -1753,7 +1753,7 @@ func (e evalBuiltin) eval(iter unifyIterator) error {
 		// call was not cached earlier.
 		if e.canUseNDBCache(e.bi) {
 			// Populate the NDBCache from the output term.
-			e.bctx.NDBuiltinCache.Put(e.bi.Name, ast.NewArray(e.terms[:endIndex]...), output.Value)
+			e.bctx.NDBuiltinCache.Put(e.bi.Name, ast.NewArray(operands[:endIndex]...), output.Value)
 		}
 
 		if err != nil {