rego: Avoid re-using transactions in compiler

Any time we do a compilation in the Rego object (or helpers) we need
to be careful to setup a conflict check that is valid for the current
context (both literal golang ctx and current storage transaction).

This isn't much of a concern if the Rego instance owns the compiler,
but if an external one was provided we need to be careful to update
the conflict check before compiling.

This was already doing the "right thing" when activating bundles, but
for partial evaluation results that were being updated on the compiler
it was not.

Fixes: #2197
Signed-off-by: Patrick East <east.patrick@gmail.com>

rego/rego.go

@@ -1494,7 +1494,7 @@ func (r *Rego) compileModules(ctx context.Context, txn storage.Transaction, m me
 			Ctx:          ctx,
 			Store:        r.store,
 			Txn:          txn,
-			Compiler:     r.compiler.WithPathConflictsCheck(storage.NonEmpty(ctx, r.store, txn)),
+			Compiler:     r.compilerForTxn(ctx, r.store, txn),
 			Metrics:      m,
 			Bundles:      r.bundles,
 			ExtraModules: r.parsedModules,
@@ -1717,7 +1717,7 @@ func (r *Rego) partialResult(ctx context.Context, pCfg *PrepareConfig) (PartialR
 	}
 
 	r.metrics.Timer(metrics.RegoModuleCompile).Start()
-	r.compiler.Compile(r.compiler.Modules)
+	r.compilerForTxn(ctx, r.store, r.txn).Compile(r.compiler.Modules)
 	r.metrics.Timer(metrics.RegoModuleCompile).Stop()
 
 	if r.compiler.Failed() {
@@ -1947,6 +1947,12 @@ func (r *Rego) getTxn(ctx context.Context) (storage.Transaction, transactionClos
 	return txn, closer, nil
 }
 
+func (r *Rego) compilerForTxn(ctx context.Context, store storage.Store, txn storage.Transaction) *ast.Compiler {
+	// Update the compiler to have a valid path conflict check
+	// for the current context and transaction.
+	return r.compiler.WithPathConflictsCheck(storage.NonEmpty(ctx, store, txn))
+}
+
 func isTermVar(v ast.Var) bool {
 	return strings.HasPrefix(string(v), ast.WildcardPrefix+"term")
 }

rego/rego_test.go

@@ -1035,6 +1035,38 @@ func TestPreparedPartialResultWithTracer(t *testing.T) {
 	}
 }
 
+func TestPartialResultSetsValidConflictChecker(t *testing.T) {
+	mod := `
+	package test
+	p {
+		true
+	}
+	`
+
+	c := ast.NewCompiler().WithPathConflictsCheck(func(_ []string) (bool, error) {
+		t.Fatal("Conflict check should not have been called")
+		return false, nil
+	})
+
+	r := New(
+		Query("data.test.p"),
+		Module("test.rego", mod),
+		PartialNamespace("test_ns1"),
+		Compiler(c),
+	)
+
+	ctx := context.Background()
+	pr, err := r.PartialResult(ctx)
+
+	if err != nil {
+		t.Fatalf("unexpected error from Rego.PartialResult(): %s", err.Error())
+	}
+
+	r2 := pr.Rego()
+
+	assertEval(t, r2, "[[true]]")
+}
+
 func TestMissingLocation(t *testing.T) {
 
 	// Create a query programmatically and evaluate it. The Location information

server/server_test.go

@@ -3527,6 +3527,53 @@ func TestServerReloadTrigger(t *testing.T) {
 	}
 }
 
+func TestServerClearsCompilerConflictCheck(t *testing.T) {
+	f := newFixture(t)
+	store := f.server.store
+	ctx := context.Background()
+
+	// Make a new transaction
+	params := storage.WriteParams
+	params.Context = storage.NewContext()
+	txn := storage.NewTransactionOrDie(ctx, store, params)
+
+	// Fresh compiler we will swap on the manager
+	c := ast.NewCompiler()
+
+	// Add the policy we want to use
+	c.Compile(map[string]*ast.Module{"test": ast.MustParseModule("package test\np=1")})
+	if len(c.Errors) > 0 {
+		t.Fatalf("Unexpected compile errors: %v", c.Errors)
+	}
+
+	// Add in a "bad" conflict check
+	c = c.WithPathConflictsCheck(func(_ []string) (bool, error) {
+		t.Fatal("Conflict check should not have been called")
+		return false, nil
+	})
+
+	// Set the compiler on the transaction context and commit to trigger listeners
+	plugins.SetCompilerOnContext(params.Context, c)
+
+	if err := store.UpsertPolicy(ctx, txn, "test", []byte("package test\np = 1")); err != nil {
+		panic(err)
+	}
+	if err := store.Commit(ctx, txn); err != nil {
+		panic(err)
+	}
+
+	// internal helpers should now give the new compiler back
+	if f.server.getCompiler() != c {
+		t.Fatalf("Expected to get the updated compiler")
+	}
+
+	// If we request for partial evaluation it will end up using the compiler set from the manager. Ensure it
+	// is using a correct conflict checker.
+	if err := f.v1(http.MethodGet, "/data/test?partial", "", 200, `{"result": {"p": 1}}`); err != nil {
+		t.Fatalf("Unexpected error from server: %v", err)
+	}
+}
+
 type queryBindingErrStore struct {
 	storage.WritesNotSupported
 	storage.PolicyNotSupported