Include annotations in rule AST (open-policy-agent#6771)

Signed-off-by: Ashutosh Narkar <[email protected]>
Co-authored-by: Johan Fylling <[email protected]>

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
 

ast/compile.go

@@ -2196,6 +2196,8 @@ func (c *Compiler) parseMetadataBlocks() {
 				for _, err := range errs {
 					c.err(err)
 				}
+
+				attachRuleAnnotations(mod)
 			}
 		}
 	}

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 {

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
 }