Skip to content

Commit 03646dd

Browse files
authored
topdown: Fix PE not namespacing vars in comprehensions nested inside every (#8816)
Partial-eval doesn't properly namespace in-scope vars inside comprehensions when they're nested inside an `every` statement. E.g. PE on `data.test.p = true` for the policy: ```rego package test p if { every x in input.x { {y | y := input.y; y < x} } } ``` will emit: ```rego every __local0__1, __local1__1 in input.x { {__local2__ | __local2__ = input.y; lt(__local2__, __local1__)} } ``` Notice how the comprehension makes a reference to `__local1__`, which has been namespaced to `__local1__1` in the outer scope, making the result query invalid. This fix checks for comprehension terms inside the `every`-body and amends them. Which gives us the updated result query: ```rego every __local0__1, __local1__1 in input.x { {__local2__1 | __local2__1 = input.y; lt(__local2__1, __local1__1)} } ``` where vars inside the comprehension are now namespaced. Note: this is a pretty narrow edge-case, so I wouldn't expect many real cases in the wild. --------- Signed-off-by: Johan Fylling <johan.dev@fylling.se>
1 parent bf2bb52 commit 03646dd

2 files changed

Lines changed: 173 additions & 9 deletions

File tree

v1/topdown/eval.go

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4204,36 +4204,64 @@ func isIterableValue(x ast.Value) bool {
42044204
}
42054205

42064206
func (e *evalEvery) save(iter unifyIterator) error {
4207-
return e.e.saveExpr(e.plug(e.expr), e.e.bindings, iter)
4207+
plugged, err := e.plug(e.expr)
4208+
if err != nil {
4209+
return err
4210+
}
4211+
return e.e.saveExpr(plugged, e.e.bindings, iter)
42084212
}
42094213

4210-
func (e *evalEvery) plug(expr *ast.Expr) *ast.Expr {
4214+
func (e *evalEvery) plug(expr *ast.Expr) (*ast.Expr, error) {
42114215
cpy := expr.Copy()
42124216
every := cpy.Terms.(*ast.Every)
4213-
e.plugBody(every.Body)
4217+
if err := e.plugBody(every.Body); err != nil {
4218+
return nil, err
4219+
}
42144220

42154221
every.Key = e.e.bindings.PlugNamespaced(every.Key, e.e.caller.bindings)
42164222
every.Value = e.e.bindings.PlugNamespaced(every.Value, e.e.caller.bindings)
42174223
every.Domain = e.e.bindings.PlugNamespaced(every.Domain, e.e.caller.bindings)
42184224
cpy.Terms = every
4219-
return cpy
4225+
return cpy, nil
42204226
}
42214227

4222-
func (e *evalEvery) plugBody(body ast.Body) {
4228+
func (e *evalEvery) plugBody(body ast.Body) error {
42234229
for i := range body {
42244230
switch t := body[i].Terms.(type) {
42254231
case *ast.Term:
4226-
body[i].Terms = e.e.bindings.PlugNamespaced(t, e.e.caller.bindings)
4232+
plugged, err := e.plugTerm(t)
4233+
if err != nil {
4234+
return err
4235+
}
4236+
body[i].Terms = plugged
42274237
case []*ast.Term:
42284238
for j := 1; j < len(t); j++ { // don't plug operator, t[0]
4229-
t[j] = e.e.bindings.PlugNamespaced(t[j], e.e.caller.bindings)
4239+
plugged, err := e.plugTerm(t[j])
4240+
if err != nil {
4241+
return err
4242+
}
4243+
t[j] = plugged
42304244
}
42314245
case *ast.Every:
4232-
body[i] = e.plug(body[i])
4246+
plugged, err := e.plug(body[i])
4247+
if err != nil {
4248+
return err
4249+
}
4250+
body[i] = plugged
42334251
case *ast.Not:
4234-
e.plugBody(t.Body)
4252+
if err := e.plugBody(t.Body); err != nil {
4253+
return err
4254+
}
42354255
}
42364256
}
4257+
return nil
4258+
}
4259+
4260+
func (e *evalEvery) plugTerm(t *ast.Term) (*ast.Term, error) {
4261+
if ast.IsComprehension(t.Value) {
4262+
return e.e.amendComprehension(t, e.e.bindings)
4263+
}
4264+
return e.e.bindings.PlugNamespaced(t, e.e.caller.bindings), nil
42374265
}
42384266

42394267
type evalNot struct {

v1/topdown/topdown_partial_test.go

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3051,6 +3051,142 @@ func TestTopDownPartialEval(t *testing.T) {
30513051
a2 = input; __local4__2 > __local2__2 }
30523052
}`},
30533053
},
3054+
{
3055+
note: "every: in-scope var in set comprehension",
3056+
query: "data.test.p = true",
3057+
modules: []string{`package test
3058+
p if {
3059+
every a in input.z {
3060+
a == 1
3061+
x := input.x
3062+
{y | y := input.y; y < x}
3063+
}
3064+
}`},
3065+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3066+
__local1__1 = 1
3067+
__local2__1 = input.x
3068+
{__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)}
3069+
}`},
3070+
},
3071+
{
3072+
note: "every: in-scope var in nested set comprehensions",
3073+
query: "data.test.p = true",
3074+
modules: []string{`package test
3075+
p if {
3076+
every x in input.x {
3077+
{y |
3078+
y := input.y
3079+
y < x
3080+
{z |
3081+
z := input.z
3082+
z < x
3083+
z > y
3084+
}
3085+
}
3086+
}
3087+
}`},
3088+
wantQueries: []string{`every __local0__1, __local1__1 in input.x {
3089+
{__local2__1 |
3090+
__local2__1 = input.y
3091+
lt(__local2__1, __local1__1)
3092+
{__local3__1 |
3093+
__local3__1 = input.z
3094+
lt(__local3__1, __local1__1)
3095+
gt(__local3__1, __local2__1)}
3096+
}
3097+
}`},
3098+
},
3099+
{
3100+
note: "every: in-scope var in set comprehension, assigned",
3101+
query: "data.test.p = true",
3102+
modules: []string{`package test
3103+
p if {
3104+
every a in input.z {
3105+
a == 1
3106+
x := input.x
3107+
z := {y | y := input.y; y < x}
3108+
z != {x}
3109+
}
3110+
}`},
3111+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3112+
__local1__1 = 1
3113+
__local2__1 = input.x
3114+
__local4__1 = {__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)}
3115+
neq(__local4__1, {__local2__1})
3116+
}`},
3117+
},
3118+
{
3119+
note: "every: in-scope var in array comprehension",
3120+
query: "data.test.p = true",
3121+
modules: []string{`package test
3122+
p if {
3123+
every a in input.z {
3124+
a == 1
3125+
x := input.x
3126+
[y | y := input.y; y < x]
3127+
}
3128+
}`},
3129+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3130+
__local1__1 = 1
3131+
__local2__1 = input.x
3132+
[__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)]
3133+
}`},
3134+
},
3135+
{
3136+
note: "every: in-scope var in array comprehension, assigned",
3137+
query: "data.test.p = true",
3138+
modules: []string{`package test
3139+
p if {
3140+
every a in input.z {
3141+
a == 1
3142+
x := input.x
3143+
z := [y | y := input.y; y < x]
3144+
z != [x]
3145+
}
3146+
}`},
3147+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3148+
__local1__1 = 1
3149+
__local2__1 = input.x
3150+
__local4__1 = [__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)]
3151+
neq(__local4__1, [__local2__1])
3152+
}`},
3153+
},
3154+
{
3155+
note: "every: in-scope var in object comprehension",
3156+
query: "data.test.p = true",
3157+
modules: []string{`package test
3158+
p if {
3159+
every a in input.z {
3160+
a == 1
3161+
x := input.x
3162+
{k: v | k := input.k; v := x}
3163+
}
3164+
}`},
3165+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3166+
__local1__1 = 1
3167+
__local2__1 = input.x
3168+
{__local3__1: __local4__1 | __local3__1 = input.k; __local4__1 = __local2__1}
3169+
}`},
3170+
},
3171+
{
3172+
note: "every: in-scope var in object comprehension, assigned",
3173+
query: "data.test.p = true",
3174+
modules: []string{`package test
3175+
p if {
3176+
every a in input.z {
3177+
a == 1
3178+
x := input.x
3179+
z := {k: v | k := input.k; v := x}
3180+
z != {"_": x}
3181+
}
3182+
}`},
3183+
wantQueries: []string{`every __local0__1, __local1__1 in input.z {
3184+
__local1__1 = 1
3185+
__local2__1 = input.x
3186+
__local5__1 = {__local3__1: __local4__1 | __local3__1 = input.k; __local4__1 = __local2__1}
3187+
neq(__local5__1, {"_": __local2__1})
3188+
}`},
3189+
},
30543190
{ // https://github.com/open-policy-agent/opa/issues/5367
30553191
note: "copypropagation: keep equations that are only found in comprehensions, inlined function call",
30563192
query: "data.test.p",

0 commit comments

Comments
 (0)