From 8ad9024ebae1db3a07242052370b2fd193a09e3d Mon Sep 17 00:00:00 2001 From: Philip Conrad Date: Tue, 6 Sep 2022 16:21:16 -0400 Subject: [PATCH 1/3] 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. Signed-off-by: Philip Conrad --- topdown/eval.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 { From 9984920b62cc5ef52a0993f679936c41de967012 Mon Sep 17 00:00:00 2001 From: Philip Conrad Date: Tue, 6 Sep 2022 16:41:29 -0400 Subject: [PATCH 2/3] 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. Signed-off-by: Philip Conrad --- rego/rego.go | 2 -- 1 file changed, 2 deletions(-) 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 { From 27f2762cef14ceb92c8830da3207056617c3676b Mon Sep 17 00:00:00 2001 From: Philip Conrad Date: Tue, 6 Sep 2022 17:06:50 -0400 Subject: [PATCH 3/3] 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 --- rego/rego_test.go | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) 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) } }