diff --git a/opa/eval_prepared_test.go b/opa/eval_prepared_test.go new file mode 100644 index 0000000..a425695 --- /dev/null +++ b/opa/eval_prepared_test.go @@ -0,0 +1,141 @@ +package opa + +import ( + "context" + "sync" + "testing" + + "github.com/boostsecurityio/poutine/models" + "github.com/boostsecurityio/poutine/results" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// These tests pin the invariants that Eval's cached query plan must preserve: +// +// 1. The same query on the same Opa, evaluated with DIFFERENT inputs, must +// return per-input-correct results — input must stay per-eval and must never +// be baked into a cached plan. +// 2. Concurrent Eval calls on a shared Opa (as AnalyzeOrg does across repo +// goroutines) must be race-free and correct. +// 3. Config changes (WithConfig) and rule-set changes (Compile) must be +// reflected by subsequent Evals — a stale cached plan bound to the old +// compiler would be a regression. + +func TestEvalSameQueryDifferentInputs(t *testing.T) { + o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}}) + noOpaErrors(t, err) + ctx := context.TODO() + + // Reuse the exact same query string across calls, varying only the input. + const q = `utils.job_uses_self_hosted_runner(input)` + cases := map[string]bool{ + "self-hosted": true, + "ubuntu-latest": false, + "random-name": true, + "windows-2022": false, + } + // Run each case twice, interleaved, to catch any input carried across calls. + for pass := 0; pass < 2; pass++ { + for runner, want := range cases { + var got bool + err := o.Eval(ctx, q, map[string]interface{}{"runs_on": []string{runner}}, &got) + noOpaErrors(t, err) + assert.Equal(t, want, got, "runner=%s pass=%d", runner, pass) + } + } +} + +func TestEvalConcurrentSharedOpa(t *testing.T) { + o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}}) + noOpaErrors(t, err) + ctx := context.TODO() + + const q = `utils.job_uses_self_hosted_runner(input)` + type tc struct { + runner string + want bool + } + cases := []tc{ + {"self-hosted", true}, + {"ubuntu-latest", false}, + {"random-name", true}, + {"windows-2019", false}, + } + + var wg sync.WaitGroup + for i := 0; i < 64; i++ { + c := cases[i%len(cases)] + wg.Add(1) + go func(c tc) { + defer wg.Done() + var got bool + if err := o.Eval(ctx, q, map[string]interface{}{"runs_on": []string{c.runner}}, &got); err != nil { + t.Errorf("eval error: %v", err) + return + } + assert.Equal(t, c.want, got, "runner=%s", c.runner) + }(c) + } + wg.Wait() +} + +func TestEvalReflectsRecompile(t *testing.T) { + o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}}) + noOpaErrors(t, err) + ctx := context.TODO() + + const q = `data.rules.pr_runs_on_self_hosted.rule` + + var before *results.Rule + err = o.Eval(ctx, q, nil, &before) + noOpaErrors(t, err) + require.NotNil(t, before) + assert.Equal(t, []interface{}{}, before.Config["allowed_runners"].Value) + + // Change config (which rewrites the store) and recompile the modules. + err = o.WithConfig(ctx, &models.Config{ + RulesConfig: map[string]map[string]interface{}{ + "pr_runs_on_self_hosted": {"allowed_runners": []string{"self-hosted"}}, + }, + }) + require.NoError(t, err) + require.NoError(t, o.Compile(ctx, nil, nil)) + + // The same query on the same Opa must now reflect the new config, not a + // stale cached result/plan. + var after *results.Rule + err = o.Eval(ctx, q, nil, &after) + noOpaErrors(t, err) + require.NotNil(t, after) + assert.Equal(t, []interface{}{"self-hosted"}, after.Config["allowed_runners"].Value) +} + +// TestEvalRecompileInvalidatesCache specifically guards the cache-invalidation in +// Compile. Unlike TestEvalReflectsRecompile (whose config change is read from the +// store on every eval and so would pass even without invalidation), this changes +// the *rule set*: it primes the cache with a query against a rule, then recompiles +// with that rule skipped. A cached plan bound to the old compiler would still +// resolve the rule; a correctly invalidated cache re-prepares and the rule is +// gone. This test FAILS if the invalidation in Compile is removed. +func TestEvalRecompileInvalidatesCache(t *testing.T) { + o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}}) + noOpaErrors(t, err) + ctx := context.TODO() + + const q = `data.rules.pr_runs_on_self_hosted.rule` + + // Prime the cache: the rule resolves. + var before *results.Rule + require.NoError(t, o.Eval(ctx, q, nil, &before)) + require.NotNil(t, before) + + // Recompile with the rule skipped → new compiler, rule removed from the set. + require.NoError(t, o.Compile(ctx, []string{"pr_runs_on_self_hosted"}, nil)) + + // With a correctly invalidated cache, the query now resolves to nothing and + // Eval reports the empty result set. A stale plan would still succeed. + var after *results.Rule + err = o.Eval(ctx, q, nil, &after) + require.Error(t, err, "skipped rule must not resolve after recompile; stale cache would still find it") +} diff --git a/opa/opa.go b/opa/opa.go index d8f9ab0..28c5e3c 100644 --- a/opa/opa.go +++ b/opa/opa.go @@ -10,6 +10,7 @@ import ( "path/filepath" "slices" "strings" + "sync" "github.com/boostsecurityio/poutine/models" "github.com/open-policy-agent/opa/v1/ast" @@ -37,6 +38,15 @@ type Opa struct { Store storage.Store LoadPaths []string customEmbeddedRules []embeddedSource + + // preparedMu guards prepared. prepared caches one compiled query plan + // (rego.PreparedEvalQuery) per query string so Eval does not re-run + // PrepareForEval — the query compile + evaluation-plan build — on every call. + // The cache is invalidated in Compile, which is the only thing that changes + // the rule set the plans are bound to (config lives in the store and is read + // fresh on every eval, so it needs no invalidation). + preparedMu sync.RWMutex + prepared map[string]rego.PreparedEvalQuery } func NewOpa(ctx context.Context, config *models.Config) (*Opa, error) { @@ -232,20 +242,26 @@ func (o *Opa) Compile(ctx context.Context, skip []string, allowed []string) erro } o.Compiler = compiler + + // Cached plans are bound to the previous compiler (and thus the previous rule + // set); drop them so subsequent Evals re-prepare against the new compiler. + o.preparedMu.Lock() + o.prepared = nil + o.preparedMu.Unlock() + return nil } func (o *Opa) Eval(ctx context.Context, query string, input map[string]interface{}, result interface{}) error { - regoInstance := rego.New( - rego.Query(query), - rego.Compiler(o.Compiler), - rego.PrintHook(o), - rego.Input(input), - rego.Imports([]string{"data.poutine.utils"}), - rego.Store(o.Store), - ) + pq, err := o.preparedQuery(ctx, query) + if err != nil { + return err + } - rs, err := regoInstance.Eval(ctx) + // Input is supplied per call via EvalInput; the cached plan holds no input, + // and each Eval opens a fresh store transaction, so config changes written by + // WithConfig are observed here without invalidating the cache. + rs, err := pq.Eval(ctx, rego.EvalInput(input)) if err != nil { return err } @@ -263,6 +279,43 @@ func (o *Opa) Eval(ctx context.Context, query string, input map[string]interface return json.Unmarshal(data, result) } +// preparedQuery returns the cached compiled plan for query, preparing (and +// caching) it on first use. The compiled plan is bound to the current compiler +// and is safe for concurrent evaluation; Compile clears the cache when it swaps +// the compiler out. +func (o *Opa) preparedQuery(ctx context.Context, query string) (rego.PreparedEvalQuery, error) { + o.preparedMu.RLock() + pq, ok := o.prepared[query] + o.preparedMu.RUnlock() + if ok { + return pq, nil + } + + pq, err := rego.New( + rego.Query(query), + rego.Compiler(o.Compiler), + rego.PrintHook(o), + rego.Imports([]string{"data.poutine.utils"}), + rego.Store(o.Store), + ).PrepareForEval(ctx) + if err != nil { + return rego.PreparedEvalQuery{}, fmt.Errorf("failed to prepare query: %w", err) + } + + o.preparedMu.Lock() + defer o.preparedMu.Unlock() + // Another goroutine may have prepared the same query while we were preparing; + // prefer the already-cached one so all callers share a single plan. + if existing, ok := o.prepared[query]; ok { + return existing, nil + } + if o.prepared == nil { + o.prepared = make(map[string]rego.PreparedEvalQuery) + } + o.prepared[query] = pq + return pq, nil +} + func Capabilities() (*ast.Capabilities, error) { capabilites := &ast.Capabilities{} err := json.Unmarshal(capabilitiesJson, capabilites)