Skip to content

Commit

Permalink
ast: Adding rule_head_refs capabilities feature flag (#6346)
Browse files Browse the repository at this point in the history
To signal/toggle support for general refs in rule heads.

Fixes: #6334
Signed-off-by: Johan Fylling <[email protected]>
  • Loading branch information
johanfylling authored Nov 15, 2023
1 parent 7933a40 commit 56f2350
Show file tree
Hide file tree
Showing 9 changed files with 5,145 additions and 11 deletions.
2 changes: 2 additions & 0 deletions ast/capabilities.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ var minVersionIndex = func() VersionIndex {
// OPA can work with policies that we're compiling -- if they don't know ref
// heads, they wouldn't be able to parse them.
const FeatureRefHeadStringPrefixes = "rule_head_ref_string_prefixes"
const FeatureRefHeads = "rule_head_refs"

// Capabilities defines a structure containing data that describes the capabilities
// or features supported by a particular version of OPA.
Expand Down Expand Up @@ -102,6 +103,7 @@ func CapabilitiesForThisVersion() *Capabilities {

f.Features = []string{
FeatureRefHeadStringPrefixes,
FeatureRefHeads,
}

return f
Expand Down
21 changes: 20 additions & 1 deletion ast/capabilities_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -181,14 +181,33 @@ func TestCapabilitiesMinimumCompatibleVersion(t *testing.T) {
version: "0.38.0",
},
{
note: "features",
note: "features (string prefix ref)",
module: `
package x
import future.keywords.if
p.a.b.c.d if { true }
`,
version: "0.46.0",
},
{
note: "features (general ref)",
module: `
package x
import future.keywords.if
p.a.b[c].d if { c := "foo" }
`,
version: "0.59.0",
},
{
note: "features (general ref + string prefix ref)",
module: `
package x
import future.keywords.if
p.a.b.c.d if { true }
p.a.b[c].d if { c := "foo" }
`,
version: "0.59.0",
},
}

for _, tc := range tests {
Expand Down
33 changes: 24 additions & 9 deletions ast/compile.go
Original file line number Diff line number Diff line change
Expand Up @@ -966,8 +966,13 @@ func (c *Compiler) buildRequiredCapabilities() {

for _, name := range c.sorted {
for _, rule := range c.Modules[name].Rules {
if len(rule.Head.Reference) >= 3 {
features[FeatureRefHeadStringPrefixes] = struct{}{}
refLen := len(rule.Head.Reference)
if refLen >= 3 {
if refLen > len(rule.Head.Reference.ConstantPrefix()) {
features[FeatureRefHeads] = struct{}{}
} else {
features[FeatureRefHeadStringPrefixes] = struct{}{}
}
}
}
}
Expand Down Expand Up @@ -1852,23 +1857,33 @@ func (c *Compiler) rewriteRuleHeadRefs() {
rule.Head.Reference = ref
}

cannotSpeakRefs := true
cannotSpeakStringPrefixRefs := true
cannotSpeakGeneralRefs := true
for _, f := range c.capabilities.Features {
if f == FeatureRefHeadStringPrefixes {
cannotSpeakRefs = false
break
switch f {
case FeatureRefHeadStringPrefixes:
cannotSpeakStringPrefixRefs = false
case FeatureRefHeads:
cannotSpeakGeneralRefs = false
}
}

if cannotSpeakRefs && rule.Head.Name == "" {
if cannotSpeakStringPrefixRefs && cannotSpeakGeneralRefs && rule.Head.Name == "" {
c.err(NewError(CompileErr, rule.Loc(), "rule heads with refs are not supported: %v", rule.Head.Reference))
return true
}

for i := 1; i < len(ref); i++ {
// Rewrite so that any non-scalar elements that in the last position of
// the rule are vars:
if cannotSpeakGeneralRefs && (rule.Head.RuleKind() == MultiValue || i != len(ref)-1) { // last
if _, ok := ref[i].Value.(String); !ok {
c.err(NewError(TypeErr, rule.Loc(), "rule heads with general refs (containing variables) are not supported: %v", rule.Head.Reference))
continue
}
}

// Rewrite so that any non-scalar elements in the rule's ref are vars:
// p.q.r[y.z] { ... } => p.q.r[__local0__] { __local0__ = y.z }
// p.q[a.b][c.d] { ... } => p.q[__local0__] { __local0__ = a.b; __local1__ = c.d }
// because that's what the RuleTree knows how to deal with.
if _, ok := ref[i].Value.(Var); !ok && !IsScalar(ref[i].Value) {
expr := f.Generate(ref[i])
Expand Down
176 changes: 176 additions & 0 deletions ast/compile_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -580,6 +580,22 @@ func TestCompilerCheckRuleHeadRefs(t *testing.T) {
),
expected: MustParseRule(`p.q.r[__local0__] { y := {"z": "a"}; __local0__ = y.z }`),
},
{
note: "rewrite: single-value with non-var ref term",
modules: modules(
`package x
p.q[y.z].r if y := {"z": "a"}`,
),
expected: MustParseRule(`p.q[__local0__].r { y := {"z": "a"}; __local0__ = y.z }`),
},
{
note: "rewrite: single-value with non-var ref term and key",
modules: modules(
`package x
p.q[a.b][c.d] if y := {"z": "a"}`,
),
expected: MustParseRule(`p.q[__local0__][__local1__] { y := {"z": "a"}; __local0__ = a.b; __local1__ = c.d }`),
},
}

for _, tc := range tests {
Expand Down Expand Up @@ -9750,6 +9766,166 @@ func runQueryCompilerTest(q, pkg string, imports []string, expected interface{})
}
}

func TestCompilerCapabilitiesFeatures(t *testing.T) {
cases := []struct {
note string
module string
features []string
expectedErr string
}{
{
note: "no features, no ref-head rules",
module: `package test
p := 42`,
},
{
note: "no features, ref-head rule",
module: `package test
p.q.r := 42`,
expectedErr: "rego_compile_error: rule heads with refs are not supported: p.q.r",
},
{
note: "no features, general-ref-head rule",
module: `package test
p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
expectedErr: "rego_compile_error: rule heads with refs are not supported: p[q].r[s]",
},
{
note: "string-prefix-ref-head feature, no ref-head rules",
features: []string{
FeatureRefHeadStringPrefixes,
},
module: `package test
p := 42`,
},
{
note: "string-prefix-ref-head feature, ref-head rule",
features: []string{
FeatureRefHeadStringPrefixes,
},
module: `package test
p.q.r := 42`,
},
{
note: "ref-head feature, ref-head rule",
features: []string{
FeatureRefHeads,
},
module: `package test
p.q.r := 42`,
},
{
note: "string-prefix-ref-head feature, general-ref-head rule",
features: []string{
FeatureRefHeadStringPrefixes,
},
module: `package test
p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
expectedErr: "rego_type_error: rule heads with general refs (containing variables) are not supported: p[q].r[s]",
},
{
note: "ref-head feature, general-ref-head rule",
features: []string{
FeatureRefHeads,
},
module: `package test
p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
},
{
note: "string-prefix-ref-head & ref-head features, general-ref-head rule",
features: []string{
FeatureRefHeadStringPrefixes,
FeatureRefHeads,
},
module: `package test
p[q].r[s] := 42 { q := "foo"; s := "bar" }`,
},
{
note: "string-prefix-ref-head & ref-head features, ref-head rule",
features: []string{
FeatureRefHeadStringPrefixes,
FeatureRefHeads,
},
module: `package test
p.q.r := 42`,
},
{
note: "no features, string-prefix-ref-head with contains kw",
features: []string{},
module: `package test
import future.keywords.contains
p.x contains 1`,
expectedErr: "rego_compile_error: rule heads with refs are not supported: p.x",
},
{
note: "string-prefix-ref-head feature, string-prefix-ref-head with contains kw",
features: []string{
FeatureRefHeadStringPrefixes,
},
module: `package test
import future.keywords.contains
p.x contains 1`,
},
{
note: "ref-head feature, string-prefix-ref-head with contains kw",
features: []string{
FeatureRefHeads,
},
module: `package test
import future.keywords.contains
p.x contains 1`,
},

{
note: "no features, general-ref-head with contains kw",
features: []string{},
module: `package test
import future.keywords
p[x] contains 1 if x = "foo"`,
expectedErr: "rego_compile_error: rule heads with refs are not supported: p[x]",
},
{
note: "string-prefix-ref-head feature, general-ref-head with contains kw",
features: []string{
FeatureRefHeadStringPrefixes,
},
module: `package test
import future.keywords
p[x] contains 1 if x = "foo"`,
expectedErr: "rego_type_error: rule heads with general refs (containing variables) are not supported: p[x]",
},
{
note: "ref-head feature, general-ref-head with contains kw",
features: []string{
FeatureRefHeads,
},
module: `package test
import future.keywords
p[x] contains 1 if x = "foo"`,
},
}

for _, tc := range cases {
t.Run(tc.note, func(t *testing.T) {
capabilities := CapabilitiesForThisVersion()
capabilities.Features = tc.features

compiler := NewCompiler().WithCapabilities(capabilities)
compiler.Compile(map[string]*Module{"test": MustParseModule(tc.module)})
if tc.expectedErr != "" {
if !compiler.Failed() {
t.Fatal("expected error but got success")
}
if !strings.Contains(compiler.Errors.Error(), tc.expectedErr) {
t.Fatalf("expected error:\n\n%s\n\nbut got:\n\n%v", tc.expectedErr, compiler.Errors)
}
} else if compiler.Failed() {
t.Fatalf("unexpected error(s): %v", compiler.Errors)
}
})
}
}

func TestCompilerCapabilitiesExtendedWithCustomBuiltins(t *testing.T) {

compiler := NewCompiler().WithCapabilities(&Capabilities{
Expand Down
7 changes: 7 additions & 0 deletions ast/version_index.json
Original file line number Diff line number Diff line change
Expand Up @@ -1373,6 +1373,13 @@
"Patch": 0,
"PreRelease": "",
"Metadata": ""
},
"rule_head_refs": {
"Major": 0,
"Minor": 59,
"Patch": 0,
"PreRelease": "",
"Metadata": ""
}
},
"keywords": {
Expand Down
Loading

0 comments on commit 56f2350

Please sign in to comment.