From f93d0f8fea011af9929c7d5d485224f04eeb0e77 Mon Sep 17 00:00:00 2001 From: Johan Fylling Date: Fri, 3 Feb 2023 12:15:21 +0100 Subject: [PATCH] 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 --- ast/annotations.go | 65 ++++++- ast/parser.go | 1 + bundle/bundle.go | 46 ++++- cmd/build_test.go | 300 ++++++++++++++++++++++++++++- cmd/inspect.go | 13 +- cmd/inspect_test.go | 203 +++++++++++++++++++ compile/compile.go | 79 +++++++- compile/compile_test.go | 83 ++++++++ internal/bundle/inspect/inspect.go | 19 +- 9 files changed, 782 insertions(+), 27 deletions(-) diff --git a/ast/annotations.go b/ast/annotations.go index b50635299b..7c40919a42 100644 --- a/ast/annotations.go +++ b/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. @@ -74,6 +75,10 @@ type ( Annotations *Annotations `json:"annotations,omitempty"` node Node // The node the annotations are applied to } + + AnnotationsRefSet []*AnnotationsRef + + FlatAnnotationsRefSet AnnotationsRefSet ) func (a *Annotations) String() string { @@ -91,6 +96,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 { @@ -162,8 +176,13 @@ func (a *Annotations) GetTargetPath() Ref { } func NewAnnotationsRef(a *Annotations) *AnnotationsRef { + var loc *Location + if a.node != nil { + loc = a.node.Loc() + } + return &AnnotationsRef{ - Location: a.node.Loc(), + Location: loc, Path: a.GetTargetPath(), Annotations: a, node: a.node, @@ -668,7 +687,7 @@ func (as *AnnotationSet) GetPackageScope(pkg *Package) *Annotations { // Flatten returns a flattened list view of this AnnotationSet. // The returned slice is sorted, first by the annotations' target path, then by their target location -func (as *AnnotationSet) Flatten() []*AnnotationsRef { +func (as *AnnotationSet) Flatten() FlatAnnotationsRefSet { // This preallocation often won't be optimal, but it's superior to starting with a nil slice. refs := make([]*AnnotationsRef, 0, len(as.byPath.Children)+len(as.byRule)+len(as.byPackage)) @@ -686,13 +705,7 @@ func (as *AnnotationSet) Flatten() []*AnnotationsRef { // Sort by path, then annotation location, for stable output sort.SliceStable(refs, func(i, j int) bool { - if refs[i].Path.Compare(refs[j].Path) < 0 { - return true - } - if refs[i].Annotations.Location.Compare(refs[j].Annotations.Location) < 0 { - return true - } - return false + return refs[i].Compare(refs[j]) < 0 }) return refs @@ -705,7 +718,7 @@ func (as *AnnotationSet) Flatten() []*AnnotationsRef { // 2. The 'package' scope entry, if any // 3. Entries for the 'subpackages' scope, if any; ordered from the closest package path to the fartest. E.g.: 'do.re.mi', 'do.re', 'do' // The returned slice is guaranteed to always contain at least one entry, corresponding to the given rule. -func (as *AnnotationSet) Chain(rule *Rule) []*AnnotationsRef { +func (as *AnnotationSet) Chain(rule *Rule) AnnotationsRefSet { var refs []*AnnotationsRef ruleAnnots := as.GetRuleScope(rule) @@ -751,6 +764,26 @@ func (as *AnnotationSet) Chain(rule *Rule) []*AnnotationsRef { return refs } +func (ars FlatAnnotationsRefSet) Insert(ar *AnnotationsRef) FlatAnnotationsRefSet { + result := make(FlatAnnotationsRefSet, 0, len(ars)+1) + + // insertion sort, first by path, then location + for i, current := range ars { + if ar.Compare(current) < 0 { + result = append(result, ar) + result = append(result, ars[i:]...) + break + } + result = append(result, current) + } + + if len(result) < len(ars)+1 { + result = append(result, ar) + } + + return result +} + func newAnnotationTree() *annotationTreeNode { return &annotationTreeNode{ Value: nil, @@ -814,3 +847,15 @@ func (t *annotationTreeNode) flatten(refs []*AnnotationsRef) []*AnnotationsRef { } return refs } + +func (ar *AnnotationsRef) Compare(other *AnnotationsRef) int { + if c := ar.Path.Compare(other.Path); c != 0 { + return c + } + + if c := ar.Annotations.Location.Compare(other.Annotations.Location); c != 0 { + return c + } + + return ar.Annotations.Compare(other.Annotations) +} diff --git a/ast/parser.go b/ast/parser.go index c4a53b8cb8..1d700825de 100644 --- a/ast/parser.go +++ b/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 diff --git a/bundle/bundle.go b/bundle/bundle.go index b5ff532fbf..af05f7c1d4 100644 --- a/bundle/bundle.go +++ b/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"` + Annotations []*ast.Annotations `json:"annotations,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 !m.WasmResolvers[i].Equal(&other.WasmResolvers[i]) { return false } } @@ -218,6 +223,37 @@ func (m Manifest) equalWasmResolversAndRoots(other Manifest) bool { return m.rootSet().Equal(other.rootSet()) } +func (wr *WasmResolver) Equal(other *WasmResolver) bool { + if wr == nil && other == nil { + return true + } + + if wr == nil || other == nil { + return false + } + + if wr.Module != other.Module { + return false + } + + if wr.Entrypoint != other.Entrypoint { + return false + } + + annotLen := len(wr.Annotations) + if annotLen != len(other.Annotations) { + return false + } + + for i := 0; i < annotLen; i++ { + if wr.Annotations[i].Compare(other.Annotations[i]) != 0 { + return false + } + } + + return true +} + type stringSet map[string]struct{} func (ss stringSet) Equal(other stringSet) bool { @@ -849,7 +885,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 +962,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 diff --git a/cmd/build_test.go b/cmd/build_test.go index 394b2a0b21..b2c062d497 100644 --- a/cmd/build_test.go +++ b/cmd/build_test.go @@ -9,12 +9,13 @@ import ( "os" "path" "path/filepath" + "reflect" "strings" "testing" "github.com/open-policy-agent/opa/ast" "github.com/open-policy-agent/opa/loader" - + "github.com/open-policy-agent/opa/util" "github.com/open-policy-agent/opa/util/test" ) @@ -585,7 +586,7 @@ package test p[1] f(x) { p[x] } - `, +`, }, err: fmt.Errorf("plan compilation requires at least one entrypoint"), }, @@ -619,6 +620,301 @@ f(x) { p[x] } } } +func TestBuildWasmWithAnnotations(t *testing.T) { + tests := []struct { + note string + files map[string]string + entrypoints []string + manifest string + }{ + { + note: "last rule is annotated entrypoint", + files: map[string]string{ + "test.rego": ` +package test + +# METADATA +# title: P1 +p1 := 1 + +# METADATA +# title: P2 +# entrypoint: true +p2 := 2 +`, + }, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test/p2", + "module":"/policy.wasm", + "annotations":[{ + "scope":"rule", + "title":"P2", + "entrypoint":true + }] + }] +} +`, + }, + { + note: "last rule is (not annotated) entrypoint", + files: map[string]string{ + "test.rego": ` +package test + +# METADATA +# title: P1 +p1 := 1 + +# METADATA +# title: P2 +p2 := 2 +`, + }, + entrypoints: []string{"test/p2"}, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test/p2", + "module":"/policy.wasm", + "annotations":[{ + "scope":"rule", + "title":"P2" + }] + }] +} +`, + }, + { + note: "rules in multiple files are entrypoints", + files: map[string]string{ + "test1.rego": ` +package test + +# METADATA +# title: P1 +p1 := 1 + +# METADATA +# title: P2 +# entrypoint: true +p2 := 2 +`, + "test2.rego": ` +package test + +# METADATA +# title: P3 +p3 := 3 + +# METADATA +# title: P4 +p4 := 4 +`, + "test3.rego": ` +package test.foo + +# METADATA +# title: BAR +# entrypoint: true +bar := "baz" +`, + }, + entrypoints: []string{"test/p3"}, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test/p3", + "annotations":[{"scope":"rule","title":"P3"}], + "module":"/policy.wasm" + },{ + "entrypoint":"test/foo/bar", + "module":"/policy.wasm", + "annotations":[{ + "scope":"rule", + "title":"BAR", + "entrypoint":true + }] + },{ + "entrypoint":"test/p2", + "module":"/policy.wasm", + "annotations":[{ + "scope":"rule", + "title":"P2", + "entrypoint":true + }] + }] +} +`, + }, + { + note: "rule with multiple metadata blocks", + files: map[string]string{ + "test.rego": ` +package test + +# METADATA +# title: P doc +# scope: document + +# METADATA +# title: P +# entrypoint: true +p := 1 +`, + }, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test/p", + "module":"/policy.wasm", + "annotations":[{ + "scope":"document", + "title":"P doc" + },{ + "scope":"rule", + "title":"P", + "entrypoint":true + }] + }] +} +`, + }, + + // Package annotations are not injected into manifest, as package definition is always retained in Rego source. + { + note: "package is annotated entrypoint", + files: map[string]string{ + "test.rego": ` +# METADATA +# title: PKG +# entrypoint: true +package test + +# METADATA +# title: P1 +p1 := 1 + +# METADATA +# title: P2 +p2 := 2 +`, + }, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test", + "module":"/policy.wasm" + }] +} +`, + }, + { + note: "package is (not annotated) entrypoint", + files: map[string]string{ + "test.rego": ` +package test + +# METADATA +# title: P1 +p1 := 1 + +# METADATA +# title: P2 +p2 := 2 +`, + }, + entrypoints: []string{"test"}, + manifest: ` +{ + "revision":"", + "roots":[""], + "wasm":[{ + "entrypoint":"test", + "module":"/policy.wasm" + }] +} +`, + }, + } + + for _, tc := range tests { + t.Run(tc.note, func(t *testing.T) { + test.WithTempFS(tc.files, func(root string) { + params := newBuildParams() + if err := params.target.Set("wasm"); err != nil { + t.Fatal(err) + } + params.pruneUnused = true + params.outputFile = path.Join(root, "bundle.tar.gz") + params.entrypoints.v = tc.entrypoints + + // Build should fail if entrypoint is not discovered from annotations. + err := dobuild(params, []string{root}) + if err != nil { + t.Fatal(err) + } + + _, err = loader.NewFileLoader().AsBundle(params.outputFile) + if err != nil { + t.Fatal(err) + } + + // Check that manifest has expected content + f, err := os.Open(params.outputFile) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + gr, err := gzip.NewReader(f) + if err != nil { + t.Fatal(err) + } + + tr := tar.NewReader(gr) + + found := false + for { + f, err := tr.Next() + if err == io.EOF { + break + } else if err != nil { + t.Fatal(err) + } + if f.Name == "/.manifest" { + found = true + data, err := io.ReadAll(tr) + if err != nil { + t.Fatal(err) + } + manifest := util.MustUnmarshalJSON(data) + if !reflect.DeepEqual(manifest, util.MustUnmarshalJSON([]byte(tc.manifest))) { + t.Fatalf("expected manifest\n\n%v\n\nbut got\n\n%v", tc.manifest, string(util.MustMarshalJSON(manifest))) + } + break + } + } + + if !found { + t.Fatal("no manifest found in bundle") + } + }) + }) + } +} + func TestBuildBundleModeIgnoreFlag(t *testing.T) { files := map[string]string{ diff --git a/cmd/inspect.go b/cmd/inspect.go index 031a0d0335..27a42f4dd0 100644 --- a/cmd/inspect.go +++ b/cmd/inspect.go @@ -216,9 +216,16 @@ func populateAnnotations(out io.Writer, refs []*ast.AnnotationsRef) error { if r := ref.GetRule(); r != nil { fmt.Fprintln(out, "Rule: ", r.Head.Name) } - fmt.Fprintln(out, "Location:", ref.Location.String()) - if a := ref.Annotations; a != nil && a.Entrypoint { - fmt.Fprintln(out, "Entrypoint:", a.Entrypoint) + if loc := ref.Location; loc != nil { + fmt.Fprintln(out, "Location:", loc.String()) + } + if a := ref.Annotations; a != nil { + if len(a.Scope) > 0 { + fmt.Fprintln(out, "Scope:", a.Scope) + } + if a.Entrypoint { + fmt.Fprintln(out, "Entrypoint:", a.Entrypoint) + } } fmt.Fprintln(out) diff --git a/cmd/inspect_test.go b/cmd/inspect_test.go index eefe206f37..8ac578df27 100644 --- a/cmd/inspect_test.go +++ b/cmd/inspect_test.go @@ -288,6 +288,7 @@ pkg-descr Package: test Location: %[1]s/x.rego:16 +Scope: package Organizations: pkg-org @@ -313,6 +314,7 @@ doc-descr Package: test Rule: p Location: %[1]s/x.rego:50 +Scope: document Organizations: doc-org @@ -338,6 +340,7 @@ rule-title Package: test Rule: p Location: %[1]s/x.rego:50 +Scope: rule Organizations: rule-org @@ -361,3 +364,203 @@ Custom: }) } + +func TestDoInspectTarballPrettyWithAnnotations(t *testing.T) { + + files := [][2]string{ + {"x.rego", `# METADATA +# title: pkg-title +# description: pkg-descr +# organizations: +# - pkg-org +# related_resources: +# - https://pkg +# - ref: https://pkg +# description: rr-pkg-note +# authors: +# - pkg-author +# schemas: +# - input: {"type": "boolean"} +# custom: +# pkg: pkg-custom +package test + +# METADATA +# scope: document +# title: doc-title +# description: doc-descr +# organizations: +# - doc-org +# related_resources: +# - https://doc +# - ref: https://doc +# description: rr-doc-note +# authors: +# - doc-author +# schemas: +# - input: {"type": "integer"} +# custom: +# doc: doc-custom + +# METADATA +# title: rule-title +# description: rule-title +# organizations: +# - rule-org +# related_resources: +# - https://rule +# - ref: https://rule +# description: rr-rule-note +# authors: +# - rule-author +# schemas: +# - input: {"type": "string"} +# custom: +# rule: rule-custom +p = 1`}, + {".manifest", ` +{ + "revision": "", + "roots": [ + "" + ], + "wasm": [ + { + "entrypoint": "test/a", + "module": "/policy.wasm" + }, + { + "entrypoint": "test/b", + "module": "/policy.wasm", + "annotations": [ + { + "scope": "rule", + "title": "WASM RULE B", + "entrypoint": true + } + ] + } + ] +}`}, + {"policy.wasm", ""}, + } + + buf := archive.MustWriteTarGz(files) + + test.WithTempFS(nil, func(rootDir string) { + bundleFile := filepath.Join(rootDir, "bundle.tar.gz") + + bf, err := os.Create(bundleFile) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + _, err = bf.Write(buf.Bytes()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + ps := newInspectCommandParams() + ps.listAnnotations = true + var out bytes.Buffer + + err = doInspect(ps, bundleFile, &out) + if err != nil { + t.Fatalf("Unexpected error %v", err) + } + + bs := out.Bytes() + idx := bytes.Index(bs, []byte(`ANNOTATIONS`)) // skip NAMESPACE box + output := strings.TrimSpace(string(bs[idx:])) + expected := strings.TrimSpace(` +ANNOTATIONS: +pkg-title +========= + +pkg-descr + +Package: test +Location: /x.rego:16 +Scope: package + +Organizations: + pkg-org + +Authors: + pkg-author + +Schemas: + input: {"type":"boolean"} + +Related Resources: + https://pkg + https://pkg rr-pkg-note + +Custom: + pkg: "pkg-custom" + +WASM RULE B +=========== + +Location: /policy.wasm:0 +Scope: rule +Entrypoint: true + +doc-title +========= + +doc-descr + +Package: test +Rule: p +Location: /x.rego:50 +Scope: document + +Organizations: + doc-org + +Authors: + doc-author + +Schemas: + input: {"type":"integer"} + +Related Resources: + https://doc + https://doc rr-doc-note + +Custom: + doc: "doc-custom" + +rule-title +========== + +rule-title + +Package: test +Rule: p +Location: /x.rego:50 +Scope: rule + +Organizations: + rule-org + +Authors: + rule-author + +Schemas: + input: {"type":"string"} + +Related Resources: + https://rule + https://rule rr-rule-note + +Custom: + rule: "rule-custom"`) + + if output != expected { + t.Fatalf("Unexpected output. Expected:\n\n%q\n\nGot:\n\n%q", expected, output) + } + + }) +} diff --git a/compile/compile.go b/compile/compile.go index 03d915adab..e294463af8 100644 --- a/compile/compile.go +++ b/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 } } diff --git a/compile/compile_test.go b/compile/compile_test.go index ad71ffb1a1..5783069417 100644 --- a/compile/compile_test.go +++ b/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 diff --git a/internal/bundle/inspect/inspect.go b/internal/bundle/inspect/inspect.go index 2636915a1a..cae88ad17b 100644 --- a/internal/bundle/inspect/inspect.go +++ b/internal/bundle/inspect/inspect.go @@ -52,7 +52,24 @@ func File(path string, includeAnnotations bool) (*Info, error) { if len(errs) > 0 { return nil, errs } - bi.Annotations = as.Flatten() + flattened := as.Flatten() + + for _, wr := range bi.Manifest.WasmResolvers { + if as := wr.Annotations; len(as) > 0 { + path, err := ast.PtrRef(ast.DefaultRootDocument, wr.Entrypoint) + if err != nil { + return nil, fmt.Errorf("failed to parse Wasm entrypoint in manifest: %s", err) + } + for _, a := range as { + ar := ast.NewAnnotationsRef(a) + ar.Path = path + ar.Location = ast.NewLocation(nil, wr.Module, 0, 0) + flattened = flattened.Insert(ar) + } + } + } + + bi.Annotations = flattened } err = bi.getBundleDataWasmAndSignatures(path)