From 4e5c36d05ced824f1529d196b9c0c3ac76db5eee Mon Sep 17 00:00:00 2001 From: Ashutosh Narkar Date: Thu, 30 May 2024 05:54:48 -0700 Subject: [PATCH] Include annotations in rule AST (#6771) Signed-off-by: Ashutosh Narkar Co-authored-by: Johan Fylling --- ast/annotations.go | 28 +++ ast/compile.go | 2 + ast/marshal_test.go | 16 ++ ast/parser_ext.go | 2 + ast/parser_test.go | 477 ++++++++++++++++++++++++++++++++++++++++++++ ast/policy.go | 30 ++- 6 files changed, 548 insertions(+), 7 deletions(-) diff --git a/ast/annotations.go b/ast/annotations.go index 1868685542..f7a5c78f76 100644 --- a/ast/annotations.go +++ b/ast/annotations.go @@ -509,6 +509,34 @@ func (a *Annotations) toObject() (*Object, *Error) { return &obj, nil } +func attachRuleAnnotations(mod *Module) { + // make a copy of the annotations + cpy := make([]*Annotations, len(mod.Annotations)) + for i, a := range mod.Annotations { + cpy[i] = a.Copy(a.node) + } + + for _, rule := range mod.Rules { + var j int + var found bool + for i, a := range cpy { + if rule.Ref().Equal(a.GetTargetPath()) { + if a.Scope == annotationScopeDocument { + rule.Annotations = append(rule.Annotations, a) + } else if a.Scope == annotationScopeRule && rule.Loc().Row > a.Location.Row { + j = i + found = true + rule.Annotations = append(rule.Annotations, a) + } + } + } + + if found && j < len(cpy) { + cpy = append(cpy[:j], cpy[j+1:]...) + } + } +} + func attachAnnotationsNodes(mod *Module) Errors { var errs Errors diff --git a/ast/compile.go b/ast/compile.go index 422ba468de..6afa648d0f 100644 --- a/ast/compile.go +++ b/ast/compile.go @@ -2196,6 +2196,8 @@ func (c *Compiler) parseMetadataBlocks() { for _, err := range errs { c.err(err) } + + attachRuleAnnotations(mod) } } } diff --git a/ast/marshal_test.go b/ast/marshal_test.go index 4d71ebcd1c..2b5952745c 100644 --- a/ast/marshal_test.go +++ b/ast/marshal_test.go @@ -444,6 +444,22 @@ func TestRule_MarshalJSON(t *testing.T) { }(), ExpectedJSON: `{"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]},"location":{"file":"example.rego","row":6,"col":2}}`, }, + "annotations included": { + Rule: func() *Rule { + r := rule.Copy() + r.Annotations = []*Annotations{{ + Scope: "rule", + Title: "My rule", + Entrypoint: true, + Organizations: []string{"org1"}, + Description: "My desc", + Custom: map[string]interface{}{ + "foo": "bar", + }}} + return r + }(), + ExpectedJSON: `{"annotations":[{"custom":{"foo":"bar"},"description":"My desc","entrypoint":true,"organizations":["org1"],"scope":"rule","title":"My rule"}],"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]}}`, + }, } for name, data := range testCases { diff --git a/ast/parser_ext.go b/ast/parser_ext.go index 19af82f5b2..afaa1d890c 100644 --- a/ast/parser_ext.go +++ b/ast/parser_ext.go @@ -713,6 +713,8 @@ func parseModule(filename string, stmts []Statement, comments []*Comment, regoCo return nil, errs } + attachRuleAnnotations(mod) + return mod, nil } diff --git a/ast/parser_test.go b/ast/parser_test.go index 8895ca3ebd..3e9743fcc5 100644 --- a/ast/parser_test.go +++ b/ast/parser_test.go @@ -5297,6 +5297,483 @@ p { input = "str" }`, } } +func TestAnnotationsAttachedToRule(t *testing.T) { + + tests := []struct { + note string + module string + expAnnotations map[int][]*Annotations + }{ + { + note: "single metadata block for rule (implied rule scope)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# title: p +# description: p +p := 1`, + expAnnotations: map[int][]*Annotations{9: {{ + Description: "p", + Scope: "rule", + Title: "p", + }}}, + }, + { + note: "single metadata block for rule (explicit rule scope)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# title: p +# description: p +# scope: rule +p := 1`, + expAnnotations: map[int][]*Annotations{10: {{ + Description: "p", + Scope: "rule", + Title: "p", + }}}, + }, + { + note: "multiple metadata blocks for single rule", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# title: One + +# METADATA +# title: Two + +# METADATA +# title: Three + +# METADATA +# title: Four +p := 1`, + expAnnotations: map[int][]*Annotations{17: { + { + Scope: "rule", + Title: "One", + }, + { + Scope: "rule", + Title: "Two", + }, { + Scope: "rule", + Title: "Three", + }, + { + Scope: "rule", + Title: "Four", + }, + }}, + }, + { + note: "document scope", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# scope: document +# title: doc +# description: doc + +p := 1`, + expAnnotations: map[int][]*Annotations{11: {{ + Description: "doc", + Scope: "document", + Title: "doc", + }}}, + }, + { + note: "document and rule scope (single rule)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# scope: document +# title: doc +# description: doc + +# METADATA +# title: p +# description: p +p := 1`, + expAnnotations: map[int][]*Annotations{14: { + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + { + Description: "p", + Scope: "rule", + Title: "p", + }, + }}, + }, + { + note: "document and rule scope (multiple rules)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# scope: document +# title: doc +# description: doc + +# METADATA +# title: p +# description: p +p := 1 + +# METADATA +# title: q +# description: q +q := 1`, + expAnnotations: map[int][]*Annotations{ + 14: { + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + { + Description: "p", + Scope: "rule", + Title: "p", + }, + }, + 19: { + { + Description: "q", + Scope: "rule", + Title: "q", + }, + }, + }, + }, + { + note: "document and rule scope (unordered annotations, multiple rules)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# scope: document +# title: p-rules + +# METADATA +# title: p-1 +# description: p-1 +p[1] + +# METADATA +# title: p-2 +# description: p-2 +p[2] + +# METADATA +# title: q +# description: q +q := 1`, + expAnnotations: map[int][]*Annotations{ + 13: { + { + Scope: "document", + Title: "p-rules", + }, + { + Description: "p-1", + Scope: "rule", + Title: "p-1", + }, + }, + 18: { + { + Scope: "document", + Title: "p-rules", + }, + { + Description: "p-2", + Scope: "rule", + Title: "p-2", + }, + }, + 23: { + { + Description: "q", + Scope: "rule", + Title: "q", + }, + }, + }, + }, + { + note: "document and rule scope (unordered annotations, multiple unordered rules)", + module: `# METADATA +# title: pkg +# description: pkg +package test + +# METADATA +# scope: document +# title: p-rules + +# METADATA +# title: p-1 +# description: p-1 +p[1] + +# METADATA +# title: q +# description: q +q := 1 + +# METADATA +# title: p-2 +# description: p-2 +p[2] +`, + expAnnotations: map[int][]*Annotations{ + 13: { + { + Scope: "document", + Title: "p-rules", + }, + { + Description: "p-1", + Scope: "rule", + Title: "p-1", + }, + }, + 18: { + { + Description: "q", + Scope: "rule", + Title: "q", + }, + }, + 23: { + { + Scope: "document", + Title: "p-rules", + }, + { + Description: "p-2", + Scope: "rule", + Title: "p-2", + }, + }, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.note, func(t *testing.T) { + + pm, err := ParseModuleWithOpts("test.rego", tc.module, ParserOptions{ProcessAnnotation: true}) + if err != nil { + t.Fatal(err) + } + + for _, rule := range pm.Rules { + annotations, ok := tc.expAnnotations[rule.Location.Row] + if !ok { + t.Fatalf("No annotations for rule on row %v", rule.Location.Row) + } + + if annotationsCompare(annotations, rule.Annotations) != 0 { + t.Fatalf("expected rule on row %d to have annotations:\n\n%v\n\nbut got:\n\n%v", + rule.Location.Row, annotations, rule.Annotations) + } + } + }) + } +} + +func TestAnnotationsAttachedToRuleMixScope(t *testing.T) { + + module := `# METADATA +# title: pkg +# description: pkg +package test + +import rego.v1 + +# METADATA +# scope: document +# title: doc +# description: doc + +# METADATA +# title: p1 +# description: p1 +p contains x if { + input.x == 1 + x := "hello" +} + +# METADATA +# title: p2 +# description: p2 +p contains x if { + input.x == 2 + x := "world" +} + +# METADATA +# title: q +# description: q +q := 1` + + pm, err := ParseModuleWithOpts("test.rego", module, ParserOptions{ProcessAnnotation: true}) + if err != nil { + t.Fatal(err) + } + + a1 := []*Annotations{ + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + { + Description: "p1", + Scope: "rule", + Title: "p1", + }, + } + + a2 := []*Annotations{ + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + { + Description: "p2", + Scope: "rule", + Title: "p2", + }, + } + + a3 := []*Annotations{ + { + Description: "q", + Scope: "rule", + Title: "q", + }, + } + + expAnnotations := [][]*Annotations{a1, a2, a3} + + for i, rule := range pm.Rules { + if annotationsCompare(expAnnotations[i], rule.Annotations) != 0 { + t.Fatalf("expected %v but got %v", expAnnotations[i], rule.Annotations) + } + } +} + +func TestAnnotationsAttachedToRuleDocScopeBeforeRule(t *testing.T) { + + module := `# METADATA +# title: pkg +# description: pkg +package test + +import rego.v1 + +# METADATA +# title: p1 +# description: p1 + +# METADATA +# scope: document +# title: doc +# description: doc + +p contains x if { + input.x == 1 + x := "hello" +} + +# METADATA +# title: p2 +# description: p2 +p contains x if { + input.x == 2 + x := "world" +} + +# METADATA +# title: q +# description: q +q := 1` + + pm, err := ParseModuleWithOpts("test.rego", module, ParserOptions{ProcessAnnotation: true}) + if err != nil { + t.Fatal(err) + } + + a1 := []*Annotations{ + { + Description: "p1", + Scope: "rule", + Title: "p1", + }, + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + } + + a2 := []*Annotations{ + { + Description: "doc", + Scope: "document", + Title: "doc", + }, + { + Description: "p2", + Scope: "rule", + Title: "p2", + }, + } + + a3 := []*Annotations{ + { + Description: "q", + Scope: "rule", + Title: "q", + }, + } + + expAnnotations := [][]*Annotations{a1, a2, a3} + + for i, rule := range pm.Rules { + if annotationsCompare(expAnnotations[i], rule.Annotations) != 0 { + t.Fatalf("expected %v but got %v", expAnnotations[i], rule.Annotations) + } + } +} + func TestAnnotationsAugmentedError(t *testing.T) { tests := []struct { note string diff --git a/ast/policy.go b/ast/policy.go index 051eccc1e6..505b7abd7f 100644 --- a/ast/policy.go +++ b/ast/policy.go @@ -185,11 +185,12 @@ type ( // Rule represents a rule as defined in the language. Rules define the // content of documents that represent policy decisions. Rule struct { - Default bool `json:"default,omitempty"` - Head *Head `json:"head"` - Body Body `json:"body"` - Else *Rule `json:"else,omitempty"` - Location *Location `json:"location,omitempty"` + Default bool `json:"default,omitempty"` + Head *Head `json:"head"` + Body Body `json:"body"` + Else *Rule `json:"else,omitempty"` + Location *Location `json:"location,omitempty"` + Annotations []*Annotations `json:"annotations,omitempty"` // Module is a pointer to the module containing this rule. If the rule // was NOT created while parsing/constructing a module, this should be @@ -309,8 +310,8 @@ func (mod *Module) Copy() *Module { nodes[mod.Package] = cpy.Package cpy.Annotations = make([]*Annotations, len(mod.Annotations)) - for i := range mod.Annotations { - cpy.Annotations[i] = mod.Annotations[i].Copy(nodes[mod.Annotations[i].node]) + for i, a := range mod.Annotations { + cpy.Annotations[i] = a.Copy(nodes[a.node]) } cpy.Comments = make([]*Comment, len(mod.Comments)) @@ -663,6 +664,11 @@ func (rule *Rule) Compare(other *Rule) int { if cmp := rule.Body.Compare(other.Body); cmp != 0 { return cmp } + + if cmp := annotationsCompare(rule.Annotations, other.Annotations); cmp != 0 { + return cmp + } + return rule.Else.Compare(other.Else) } @@ -671,6 +677,12 @@ func (rule *Rule) Copy() *Rule { cpy := *rule cpy.Head = rule.Head.Copy() cpy.Body = rule.Body.Copy() + + cpy.Annotations = make([]*Annotations, len(rule.Annotations)) + for i, a := range rule.Annotations { + cpy.Annotations[i] = a.Copy(&cpy) + } + if cpy.Else != nil { cpy.Else = rule.Else.Copy() } @@ -763,6 +775,10 @@ func (rule *Rule) MarshalJSON() ([]byte, error) { } } + if len(rule.Annotations) != 0 { + data["annotations"] = rule.Annotations + } + return json.Marshal(data) }