diff --git a/rego/rego.go b/rego/rego.go index 89b7bd159a..9d02e15c70 100644 --- a/rego/rego.go +++ b/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 { diff --git a/rego/rego_test.go b/rego/rego_test.go index 06214f5bde..3688b8cdce 100644 --- a/rego/rego_test.go +++ b/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) } } diff --git a/topdown/eval.go b/topdown/eval.go index dae6a35209..0ee1706124 100644 --- a/topdown/eval.go +++ b/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 {