bundle: Retain metadata annotations for Wasm entrypoints during inspection (#5603)

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

Fixes: #5588
Signed-off-by: Johan Fylling <johan.dev@fylling.se>

compile/compile.go

@@ -642,13 +642,21 @@ 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 annotations []*ast.Annotations
+		if !c.isPackage(e) {
+			annotations = findAnnotationsForTerm(e, flattenedAnnotations)
+		}
+
 		c.bundle.Manifest.WasmResolvers = append(c.bundle.Manifest.WasmResolvers, bundle.WasmResolver{
-			Module:     "/" + strings.TrimLeft(modulePath, "/"),
-			Entrypoint: entrypointPath,
+			Module:      "/" + strings.TrimLeft(modulePath, "/"),
+			Entrypoint:  entrypointPath,
+			Annotations: annotations,
 		})
 	}
 
@@ -656,6 +664,33 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 	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
 			}
 		}

compile/compile_test.go

@@ -627,6 +627,89 @@ func TestCompilerWasmTargetMultipleEntrypoints(t *testing.T) {
 	})
 }
 
+func TestCompilerWasmTargetAnnotations(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+# METADATA
+# title: My test package
+package test
+
+# METADATA
+# title: My P rule
+# entrypoint: true
+p = true`,
+		"policy.rego": `
+package policy
+
+# METADATA
+# title: All my Q rules
+# scope: document
+
+# METADATA
+# title: My Q rule
+q = true`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().WithPaths(root).WithTarget("wasm").
+			WithEntrypoints("test", "policy/q").
+			WithRegoAnnotationEntrypoints(true)
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(compiler.bundle.WasmModules) != 1 {
+			t.Fatalf("expected 1 Wasm modules, got: %d", len(compiler.bundle.WasmModules))
+		}
+
+		expWasmResolvers := []bundle.WasmResolver{
+			{
+				Entrypoint: "test",
+				Module:     "/policy.wasm",
+			},
+			{
+				Entrypoint: "policy/q",
+				Module:     "/policy.wasm",
+				Annotations: []*ast.Annotations{
+					{
+						Title: "All my Q rules",
+						Scope: "document",
+					},
+					{
+						Title: "My Q rule",
+						Scope: "rule",
+					},
+				},
+			},
+			{
+				Entrypoint: "test/p",
+				Module:     "/policy.wasm",
+				Annotations: []*ast.Annotations{
+					{
+						Title:      "My P rule",
+						Scope:      "rule",
+						Entrypoint: true,
+					},
+				},
+			},
+		}
+
+		if len(expWasmResolvers) != len(compiler.bundle.Manifest.WasmResolvers) {
+			t.Fatalf("\nExpected WasmResolvers:\n  %+v\nGot:\n  %+v\n", expWasmResolvers, compiler.bundle.Manifest.WasmResolvers)
+		}
+
+		for i, expWasmResolver := range expWasmResolvers {
+			if !expWasmResolver.Equal(&compiler.bundle.Manifest.WasmResolvers[i]) {
+				t.Fatalf("WasmResolver at index %v mismatch\n\nExpected WasmResolvers:\n  %+v\nGot:\n  %+v\n",
+					i, expWasmResolvers, compiler.bundle.Manifest.WasmResolvers)
+			}
+		}
+	})
+}
+
 func TestCompilerWasmTargetEntrypointDependents(t *testing.T) {
 	files := map[string]string{
 		"test.rego": `package test