From b3a79ff1944014ac59ecf580379a5072217f5d75 Mon Sep 17 00:00:00 2001 From: Johan Fylling Date: Thu, 26 Jan 2023 12:53:56 +0100 Subject: [PATCH] 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: #5588 Signed-off-by: Johan Fylling --- ast/annotations.go | 10 +++++++ ast/parser.go | 1 + bundle/bundle.go | 15 ++++++---- compile/compile.go | 75 +++++++++++++++++++++++++++++++++++++++++++--- 4 files changed, 92 insertions(+), 9 deletions(-) diff --git a/ast/annotations.go b/ast/annotations.go index b50635299b..1ea6930936 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. @@ -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 { 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..afa0fa70d8 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"` + 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 diff --git a/compile/compile.go b/compile/compile.go index 03d915adab..26c90e2d24 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 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, }) } @@ -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 } }