bundle: Retain metadata annotations for Wasm entrypoints

* Pruning METADATA blocks associated with Wasm compiled entrypoints from Rego source
* Adding metadata annotations to wasm entrypoint declarations in bundle .manifest file

Fixes: open-policy-agent#5588

Signed-off-by: Johan Fylling <[email protected]>

ast/annotations.go

@@ -37,6 +37,7 @@ type (
 		Schemas          []*SchemaAnnotation          `json:"schemas,omitempty"`
 		Custom           map[string]interface{}       `json:"custom,omitempty"`
 		node             Node
+		comments         []*Comment
 	}
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
@@ -91,6 +92,15 @@ func (a *Annotations) SetLoc(l *Location) {
 	a.Location = l
 }
 
+// EndLoc returns the location of this annotation's last comment line.
+func (a *Annotations) EndLoc() *Location {
+	count := len(a.comments)
+	if count == 0 {
+		return a.Location
+	}
+	return a.comments[count-1].Location
+}
+
 // Compare returns an integer indicating if a is less than, equal to, or greater
 // than other.
 func (a *Annotations) Compare(other *Annotations) int {

ast/parser.go

@@ -2188,6 +2188,7 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 	}
 
 	var result Annotations
+	result.comments = b.comments
 	result.Scope = raw.Scope
 	result.Entrypoint = raw.Entrypoint
 	result.Title = raw.Title

bundle/bundle.go

@@ -127,8 +127,9 @@ type Manifest struct {
 
 // WasmResolver maps a wasm module to an entrypoint ref.
 type WasmResolver struct {
-	Entrypoint string `json:"entrypoint,omitempty"`
-	Module     string `json:"module,omitempty"`
+	Entrypoint string             `json:"entrypoint,omitempty"`
+	Module     string             `json:"module,omitempty"`
+	Metadata   []*ast.Annotations `json:"metadata,omitempty"`
 }
 
 // Init initializes the manifest. If you instantiate a manifest
@@ -166,6 +167,10 @@ func (m Manifest) Equal(other Manifest) bool {
 	return m.equalWasmResolversAndRoots(other)
 }
 
+func (m Manifest) Empty() bool {
+	return m.Equal(Manifest{})
+}
+
 // Copy returns a deep copy of the manifest.
 func (m Manifest) Copy() Manifest {
 	m.Init()
@@ -210,7 +215,7 @@ func (m Manifest) equalWasmResolversAndRoots(other Manifest) bool {
 	}
 
 	for i := 0; i < len(m.WasmResolvers); i++ {
-		if m.WasmResolvers[i] != other.WasmResolvers[i] {
+		if reflect.DeepEqual(m.WasmResolvers[i], other.WasmResolvers[i]) {
 			return false
 		}
 	}
@@ -849,7 +854,7 @@ func (w *Writer) writePlan(tw *tar.Writer, bundle Bundle) error {
 
 func writeManifest(tw *tar.Writer, bundle Bundle) error {
 
-	if bundle.Manifest.Equal(Manifest{}) {
+	if bundle.Manifest.Empty() {
 		return nil
 	}
 
@@ -926,7 +931,7 @@ func hashBundleFiles(hash SignatureHasher, b *Bundle) ([]FileInfo, error) {
 	// parse the manifest into a JSON structure;
 	// then recursively order the fields of all objects alphabetically and then apply
 	// the hash function to result to compute the hash.
-	if !b.Manifest.Equal(Manifest{}) {
+	if !b.Manifest.Empty() {
 		mbs, err := json.Marshal(b.Manifest)
 		if err != nil {
 			return files, err

compile/compile.go

@@ -642,20 +642,55 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 		Raw:  buf.Bytes(),
 	}}
 
+	flattenedAnnotations := c.compiler.GetAnnotationSet().Flatten()
+
 	// Each entrypoint needs an entry in the manifest
-	for i := range c.entrypointrefs {
+	for i, e := range c.entrypointrefs {
 		entrypointPath := c.entrypoints[i]
 
+		var metadata []*ast.Annotations
+		if !c.isPackage(e) {
+			metadata = findAnnotationsForTerm(e, flattenedAnnotations)
+		}
+
 		c.bundle.Manifest.WasmResolvers = append(c.bundle.Manifest.WasmResolvers, bundle.WasmResolver{
 			Module:     "/" + strings.TrimLeft(modulePath, "/"),
 			Entrypoint: entrypointPath,
+			Metadata:   metadata,
 		})
 	}
 
 	// Remove the entrypoints from remaining source rego files
 	return pruneBundleEntrypoints(c.bundle, c.entrypointrefs)
 }
 
+func (c *Compiler) isPackage(term *ast.Term) bool {
+	for _, m := range c.compiler.Modules {
+		if m.Package.Path.Equal(term.Value) {
+			return true
+		}
+	}
+	return false
+}
+
+// findAnnotationsForTerm returns a slice of all annotations directly associated with the given term.
+func findAnnotationsForTerm(term *ast.Term, annotationRefs []*ast.AnnotationsRef) []*ast.Annotations {
+	r, ok := term.Value.(ast.Ref)
+	if !ok {
+		return nil
+	}
+
+	var result []*ast.Annotations
+
+	for _, ar := range annotationRefs {
+		if r.Equal(ar.Path) {
+			result = append(result, ar.Annotations)
+		}
+	}
+
+	return result
+}
+
 // pruneBundleEntrypoints will modify modules in the provided bundle to remove
 // rules matching the entrypoints along with injecting import statements to
 // preserve their ability to compile.
@@ -691,11 +726,43 @@ func pruneBundleEntrypoints(b *bundle.Bundle, entrypointrefs []*ast.Term) error
 				}
 			}
 
-			// If any rules were dropped update the module accordingly
-			if len(rules) != len(mf.Parsed.Rules) {
+			// Drop any Annotations for rules matching the entrypoint path
+			var annotations []*ast.Annotations
+			var prunedAnnotations []*ast.Annotations
+			for _, annotation := range mf.Parsed.Annotations {
+				p := annotation.GetTargetPath()
+				// We prune annotations of dropped rules, but not packages, as the Rego file is always retained
+				if p.Equal(entrypoint.Value) && !mf.Parsed.Package.Path.Equal(entrypoint.Value) {
+					prunedAnnotations = append(prunedAnnotations, annotation)
+				} else {
+					annotations = append(annotations, annotation)
+				}
+			}
+
+			// Drop comments associated with pruned annotations
+			var comments []*ast.Comment
+			for _, comment := range mf.Parsed.Comments {
+				pruned := false
+				for _, annotation := range prunedAnnotations {
+					if comment.Location.Row >= annotation.Location.Row &&
+						comment.Location.Row <= annotation.EndLoc().Row {
+						pruned = true
+						break
+					}
+				}
+
+				if !pruned {
+					comments = append(comments, comment)
+				}
+			}
+
+			// If any rules or annotations were dropped update the module accordingly
+			if len(rules) != len(mf.Parsed.Rules) || len(comments) != len(mf.Parsed.Comments) {
 				mf.Parsed.Rules = rules
+				mf.Parsed.Annotations = annotations
+				mf.Parsed.Comments = comments
 				// Remove the original raw source, we're editing the AST
-				// directly so it wont be in sync anymore.
+				// directly, so it won't be in sync anymore.
 				mf.Raw = nil
 			}
 		}