ast: don't leak generated locals in ref type errors (#8902)

Fixes #8897 When a reference has a composite subject (e.g. [1, 2][i]) or a dynamic index term (e.g. [1, 2][input.x]), the compiler hoists that part of the ref into a generated local. Previously these locals leaked verbatim into type errors, e.g.: undefined ref: __localq0__[i][j] The type checker's var rewriter only knew how to map user variables back (via RewrittenVars), so anonymous generated locals were rendered as-is. Record a mapping from each generated ref-subject/dynamic-operand local back to the original term value (in localVarGenerator.subjects) and use it when rendering refs in type errors, so the original expression is shown instead: undefined ref: [1, 2][i][j] --------- Signed-off-by: Sebastian Spaink <sebastianspaink@gmail.com>
2026-08-12 19:32:48 -06:00 · 2026-08-10 11:30:10 -05:00
parent 413903e8cc
commit 990061876a
4 changed files with 191 additions and 31 deletions
@@ -468,7 +468,7 @@ func NewCompiler() *Compiler {
 		{StageCheckSafetyRuleHeads, "compile_stage_check_safety_rule_heads", c.checkSafetyRuleHeads},
 		{StageCheckSafetyRuleBodies, "compile_stage_check_safety_rule_bodies", c.checkSafetyRuleBodies},
 		{StageRewriteEquals, "compile_stage_rewrite_equals", c.rewriteEquals},
-		{StageRewriteDynamicTerms, "compile_stage_rewrite_dynamic_terms", c.rewriteDynamicTerms},
+		{StageRewriteDynamicTerms, "compile_stage_rewrite_dynamic_terms", c.rewriteDynamicTerms},                       // stages before CheckTypes must not rewrite hoisted terms, see recordSubjectNoCopy
 		{StageRewriteTestRulesForTracing, "compile_stage_rewrite_test_rules_for_tracing", c.rewriteTestRuleEqualities}, // must run after RewriteDynamicTerms
 		{StageCheckRecursion, "compile_stage_check_recursion", c.checkRecursion},
 		{StageCheckTypes, "compile_stage_check_types", c.checkTypes}, // must be run after CheckRecursion
@@ -1906,7 +1906,7 @@ func (c *Compiler) checkTypes() {
 		WithInputType(c.inputType).
 		WithBuiltins(c.builtins).
 		WithRequiredCapabilities(c.Required).
-		WithVarRewriter(rewriteVarsInRef(c.RewrittenVars)).
+		WithVarRewriter(rewriteRefErrVars(c.localvargen.subjects, c.RewrittenVars)).
 		WithAllowUndefinedFunctionCalls(c.allowUndefinedFuncCalls)
 	var as *AnnotationSet
 	if c.useTypeCheckAnnotations {
@@ -3625,6 +3625,7 @@ type queryCompiler struct {
 	qctx                  *QueryContext
 	typeEnv               *TypeEnv
 	rewritten             map[Var]Var
+	refSubjects           map[Var]Value
 	after                 map[string][]QueryCompilerStageDefinition
 	unsafeBuiltins        map[string]struct{}
 	comprehensionIndices  map[*Term]*ComprehensionIndex
@@ -3724,7 +3725,7 @@ func (qc *queryCompiler) Compile(query Body) (Body, error) {
 		{StageRewriteWithValues, "query_compile_stage_rewrite_with_values", qc.rewriteWithModifiers},
 		{StageCheckUndefinedFuncs, "query_compile_stage_check_undefined_funcs", qc.checkUndefinedFuncs},
 		{StageCheckSafety, "query_compile_stage_check_safety", qc.checkSafety},
-		{StageRewriteDynamicTerms, "query_compile_stage_rewrite_dynamic_terms", qc.rewriteDynamicTerms},
+		{StageRewriteDynamicTerms, "query_compile_stage_rewrite_dynamic_terms", qc.rewriteDynamicTerms}, // see recordSubjectNoCopy
 		{StageCheckTypes, "query_compile_stage_check_types", qc.checkTypes},
 		{StageCheckUnsafeBuiltins, "query_compile_stage_check_unsafe_builtins", qc.checkUnsafeBuiltins},
 		{StageCheckDeprecatedBuiltins, "query_compile_stage_check_deprecated_builtins", qc.checkDeprecatedBuiltins},
@@ -3813,15 +3814,19 @@ func (*queryCompiler) rewriteComprehensionTerms(_ *QueryContext, body Body) (Bod
 	return node.(Body), nil
 }

-func (*queryCompiler) rewriteDynamicTerms(_ *QueryContext, body Body) (Body, error) {
+func (qc *queryCompiler) rewriteDynamicTerms(_ *QueryContext, body Body) (Body, error) {
 	gen := newLocalVarGenerator("q", body)
 	f := newEqualityFactory(gen)
-	return rewriteDynamics(f, body), nil
+	body = rewriteDynamics(f, body)
+	qc.refSubjects = mergeRefSubjects(qc.refSubjects, gen.subjects)
+	return body, nil
 }

-func (*queryCompiler) rewriteExprTerms(_ *QueryContext, body Body) (Body, error) {
+func (qc *queryCompiler) rewriteExprTerms(_ *QueryContext, body Body) (Body, error) {
 	gen := newLocalVarGenerator("q", body)
-	return rewriteExprTermsInBody(gen, body), nil
+	body = rewriteExprTermsInBody(gen, body)
+	qc.refSubjects = gen.subjects
+	return body, nil
 }

 func (qc *queryCompiler) rewriteLocalVars(_ *QueryContext, body Body) (Body, error) {
@@ -3889,7 +3894,7 @@ func (qc *queryCompiler) checkTypes(_ *QueryContext, body Body) (Body, error) {
 	checker := newTypeChecker().
 		WithSchemaSet(qc.compiler.schemaSet).
 		WithInputType(qc.compiler.inputType).
-		WithVarRewriter(rewriteVarsInRef(qc.rewritten, qc.compiler.RewrittenVars))
+		WithVarRewriter(rewriteRefErrVars(qc.refSubjects, qc.rewritten, qc.compiler.RewrittenVars))
 	qc.typeEnv, errs = checker.CheckBody(qc.compiler.TypeEnv, body)
 	if len(errs) > 0 {
 		return nil, errs
@@ -5504,6 +5509,38 @@ type localVarGenerator struct {
 	exclude VarSet
 	suffix  string
 	next    int
+
+	// subjects maps a generated local back to the original term it replaced,
+	// so type errors can render the original expression (e.g. [1, 2][i]
+	// instead of __local0__[i]). Populated lazily.
+	subjects map[Var]Value
+}
+
+// recordSubject records that local stands in for value. The value is copied, as
+// stages running between the caller and CheckTypes may rewrite it in place: a
+// composite subject recorded in RewriteExprTerms, say [x, input.y][i], has its
+// dynamic elements hoisted by the later RewriteDynamicTerms stage, which would
+// otherwise turn the recorded value into [__local5__, __local6__].
+func (l *localVarGenerator) recordSubject(local Var, value *Term) {
+	l.putSubject(local, CopyValue(value.Value))
+}
+
+// recordSubjectNoCopy records that local stands in for value, aliasing value
+// rather than copying it. Only callers in the RewriteDynamicTerms stage may use
+// this: only RewriteTestRulesForTracing and CheckRecursion run between that
+// stage and CheckTypes, and neither rewrites hoisted terms, so nothing can
+// mutate value before the mapping is read. Copying here instead would allocate
+// on every hoisted ref of every compile, for a map only read when a type error
+// is rendered.
+func (l *localVarGenerator) recordSubjectNoCopy(local Var, value *Term) {
+	l.putSubject(local, value.Value)
+}
+
+func (l *localVarGenerator) putSubject(local Var, value Value) {
+	if l.subjects == nil {
+		l.subjects = map[Var]Value{}
+	}
+	l.subjects[local] = value
 }

 func newLocalVarGeneratorForModuleSet(sorted []string, modules map[string]*Module) *localVarGenerator {
@@ -6127,6 +6164,7 @@ func rewriteDynamicsOne(original *Expr, f *equalityFactory, term *Term, result B
 		generated.With = original.With
 		result.Append(generated)
 		connectGeneratedExprs(original, generated)
+		f.gen.recordSubjectNoCopy(generated.Operand(0).Value.(Var), term)
 		return result, result[len(result)-1].Operand(0)
 	case *Array:
 		for i := range v.Len() {
@@ -6156,18 +6194,21 @@ func rewriteDynamicsOne(original *Expr, f *equalityFactory, term *Term, result B
 		v.Body, extra = rewriteDynamicsComprehensionBody(original, f, v.Body, term)
 		result.Append(extra)
 		connectGeneratedExprs(original, extra)
+		f.gen.recordSubjectNoCopy(extra.Operand(0).Value.(Var), term)
 		return result, result[len(result)-1].Operand(0)
 	case *SetComprehension:
 		var extra *Expr
 		v.Body, extra = rewriteDynamicsComprehensionBody(original, f, v.Body, term)
 		result.Append(extra)
 		connectGeneratedExprs(original, extra)
+		f.gen.recordSubjectNoCopy(extra.Operand(0).Value.(Var), term)
 		return result, result[len(result)-1].Operand(0)
 	case *ObjectComprehension:
 		var extra *Expr
 		v.Body, extra = rewriteDynamicsComprehensionBody(original, f, v.Body, term)
 		result.Append(extra)
 		connectGeneratedExprs(original, extra)
+		f.gen.recordSubjectNoCopy(extra.Operand(0).Value.(Var), term)
 		return result, result[len(result)-1].Operand(0)
 	}
 	return result, term
@@ -6399,6 +6440,7 @@ func expandExprRef(gen *localVarGenerator, v []*Term) (support []*Expr) {
 		assignToLocal := f.Generate(subject)
 		support = append(support, assignToLocal)
 		v[0] = assignToLocal.Operand(0)
+		gen.recordSubject(v[0].Value.(Var), subject)
 	}
 	return
 }
@@ -7401,6 +7443,39 @@ func rewriteVarsInRef(vars ...map[Var]Var) varRewriter {
 	}
 }

+// mergeRefSubjects merges src into dst, allocating dst if needed.
+func mergeRefSubjects(dst, src map[Var]Value) map[Var]Value {
+	if len(src) == 0 {
+		return dst
+	}
+	if dst == nil {
+		dst = make(map[Var]Value, len(src))
+	}
+	maps.Copy(dst, src)
+	return dst
+}
+
+// rewriteRefErrVars returns a varRewriter for rendering refs in type errors.
+// Beyond the var-to-var mappings of rewriteVarsInRef, it substitutes generated
+// locals recorded in localVarGenerator.subjects with the original term (so
+// errors show [1, 2][i] rather than __local0__[i]). It operates on a copy.
+func rewriteRefErrVars(subjects map[Var]Value, vars ...map[Var]Var) varRewriter {
+	return func(node Ref) Ref {
+		i, _ := TransformVars(node.Copy(), func(v Var) (Value, error) {
+			if val, ok := subjects[v]; ok {
+				return CopyValue(val), nil
+			}
+			for _, m := range vars {
+				if u, ok := m[v]; ok {
+					return u, nil
+				}
+			}
+			return v, nil
+		})
+		return i.(Ref)
+	}
+}
+
 type ruleRefSet struct {
 	s []ruleRef
 }
@@ -1,10 +1,65 @@
 package ast

 import (
+	"fmt"
 	"strconv"
 	"testing"
 )

+// Cost of recording generated locals for ref type errors, at 100 modules:
+// 9452905 ns/op   6487160 B/op   162824 allocs/op // not recorded
+// 9944902 ns/op   6751259 B/op   168040 allocs/op // every subject copied
+// 9578020 ns/op   6580108 B/op   163241 allocs/op // copied only where a later stage can rewrite it
+func BenchmarkCompileModules(b *testing.B) {
+	// The choice of module set is somewhat arbitrary. These rules are
+	// representative of the ones that exercise the term-rewriting stages:
+	// composite ref subjects, dynamic ref operands and comprehensions all get
+	// hoisted into generated locals, so per-hoist work in those stages shows up
+	// here. BenchmarkRewriteDynamics covers one of them in isolation, but reuses
+	// the same bodies across iterations, so they are already rewritten after the
+	// first pass.
+	sizes := []int{1, 10, 100}
+
+	for _, size := range sizes {
+		b.Run(strconv.Itoa(size), func(b *testing.B) {
+			base := make(map[string]*Module, size)
+			for i := range size {
+				base[fmt.Sprintf("mod%d.rego", i)] = MustParseModule(fmt.Sprintf(`package bench.p%d
+
+allow if {
+	some x
+	input.users[x].roles[_] == "admin"
+	data.perms[input.tenant][x]
+	count([y | y := input.items[_]; y.n > 0]) > 2
+	input.a.b.c[input.i].d == data.z.w[input.j]
+}
+
+deny contains msg if {
+	msg := input.msgs[input.i].text
+	[1, 2][input.k]
+}
+`, i))
+			}
+
+			for b.Loop() {
+				// Compile rewrites modules in place, so every iteration needs
+				// its own copies. Copying is not what we're measuring.
+				b.StopTimer()
+				modules := make(map[string]*Module, len(base))
+				for name, module := range base {
+					modules[name] = module.Copy()
+				}
+				b.StartTimer()
+
+				c := NewCompiler()
+				if c.Compile(modules); c.Failed() {
+					b.Fatal(c.Errors)
+				}
+			}
+		})
+	}
+}
+
 func BenchmarkRewriteDynamics(b *testing.B) {
 	// The choice of query to use is somewhat arbitrary. This query is
 	// representative of the ones that result from partial evaluation on IAM
@@ -11522,6 +11522,27 @@ func TestQueryCompiler(t *testing.T) {
 			q:        "data.deadbeef(x)",
 			expected: errors.New("rego_type_error: undefined function data.deadbeef"),
 		},
+		{
+			// Regression test for https://github.com/open-policy-agent/opa/issues/8897:
+			// the generated local for a composite ref subject must not leak.
+			note:     "composite ref subject not leaked in type error",
+			q:        `[1, 2][i][j]`,
+			expected: errors.New("1 error occurred: 1:1: rego_type_error: undefined ref: [1, 2][i][j]"),
+		},
+		{
+			// A dynamic index term hoisted into its own local must not leak either.
+			note:     "dynamic index local not leaked in type error",
+			q:        `[1, 2][input.x][j]`,
+			expected: errors.New("1 error occurred: 1:1: rego_type_error: undefined ref: [1, 2][input.x][j]"),
+		},
+		{
+			// A dynamic index term is rewritten before the local standing in for
+			// it is recorded, so a nested one must be rendered through both
+			// mappings rather than leaking the inner local.
+			note:     "nested dynamic index locals not leaked in type error",
+			q:        `[1, 2][input.a[input.b]][j]`,
+			expected: errors.New("1 error occurred: 1:1: rego_type_error: undefined ref: [1, 2][input.a[input.b]][j]"),
+		},
 		{
 			note:     "imports resolved without package",
 			q:        "abc",
@@ -349,33 +349,42 @@ func (term *Term) Copy() *Term {
 	}

 	cpy := *term
-
-	switch v := term.Value.(type) {
-	case Null, Boolean, Number, String, Var:
-		cpy.Value = v
-	case Ref:
-		cpy.Value = v.Copy()
-	case *Array:
-		cpy.Value = v.Copy()
-	case Set:
-		cpy.Value = v.Copy()
-	case *object:
-		cpy.Value = v.Copy()
-	case *ArrayComprehension:
-		cpy.Value = v.Copy()
-	case *ObjectComprehension:
-		cpy.Value = v.Copy()
-	case *SetComprehension:
-		cpy.Value = v.Copy()
-	case *TemplateString:
-		cpy.Value = v.Copy()
-	case Call:
-		cpy.Value = v.Copy()
-	}
+	cpy.Value = CopyValue(term.Value)

 	return &cpy
 }

+// CopyValue returns a deep copy of v. The Value interface doesn't require a
+// Copy method, so this dispatches on the known value types. Values of any other
+// type are returned as-is.
+func CopyValue(v Value) Value {
+	switch v := v.(type) {
+	case Null, Boolean, Number, String, Var:
+		// Scalars are immutable, no copy needed.
+		return v
+	case Ref:
+		return v.Copy()
+	case *Array:
+		return v.Copy()
+	case Set:
+		return v.Copy()
+	case *object:
+		return v.Copy()
+	case *ArrayComprehension:
+		return v.Copy()
+	case *ObjectComprehension:
+		return v.Copy()
+	case *SetComprehension:
+		return v.Copy()
+	case *TemplateString:
+		return v.Copy()
+	case Call:
+		return v.Copy()
+	}
+
+	return v
+}
+
 // Equal returns true if this term equals the other term. Equality is
 // defined for each kind of term, and does not compare the Location.
 func (term *Term) Equal(other *Term) bool {