From 2e122c9e69e89474893d06e61be5bbe9de2c194f Mon Sep 17 00:00:00 2001 From: Stephan Renatus Date: Wed, 13 Jan 2021 09:30:39 +0100 Subject: [PATCH] wasm: read/write custom sections (incl. dwarf), deal with Name section specifically (#2900) * wasm: read/write custom sections (incl. dwarf) I've extended the roundtrip OPA test to have some assertions on custom sections that work for both debug and non-debug wasm builds. Without debugging enabled, we had still been stripping the `names` and `producers` sections. This changes therefore increases the size of wasm modules. For a simple policy, package test default allow = false allow { input.foo == "bar" } the size generated is 544K (master) vs 620K (this commit). With debugging enabled (and the opa binary rebuilt), the size becomes 6.1M. * wasm/encoding: treat 'name' custom section separately This roundtrips module, functions, and locals, as per https://webassembly.github.io/spec/core/appendix/custom.html#name-section There is currently no use of module and locals, afaict. * compiler/wasm: record function names in 'name' custom section Quality-of-life improvement when dealing with our generated WASM code. Before: local.get 2 local.get 3 call 1172 local.set 6 After: local.get 2 local.get 3 call $g0.data.foo.p local.set 6 Signed-off-by: Stephan Renatus --- internal/compiler/wasm/wasm.go | 18 +++- internal/wasm/constant/constant.go | 10 ++ internal/wasm/encoding/encoding_test.go | 31 +++++- internal/wasm/encoding/reader.go | 125 +++++++++++++++++++++++- internal/wasm/encoding/writer.go | 115 ++++++++++++++++++++++ internal/wasm/module/module.go | 28 ++++++ 6 files changed, 318 insertions(+), 9 deletions(-) diff --git a/internal/compiler/wasm/wasm.go b/internal/compiler/wasm/wasm.go index f5d6e81e30..f5e1d3ea90 100644 --- a/internal/compiler/wasm/wasm.go +++ b/internal/compiler/wasm/wasm.go @@ -1148,18 +1148,32 @@ func (c *Compiler) emitFunctionDecl(name string, tpe module.FunctionType, export typeIndex := c.emitFunctionType(tpe) c.module.Function.TypeIndices = append(c.module.Function.TypeIndices, typeIndex) c.module.Code.Segments = append(c.module.Code.Segments, module.RawCodeSegment{}) - c.funcs[name] = uint32((len(c.module.Function.TypeIndices) - 1) + c.functionImportCount()) + idx := uint32((len(c.module.Function.TypeIndices) - 1) + c.functionImportCount()) + c.funcs[name] = idx if export { c.module.Export.Exports = append(c.module.Export.Exports, module.Export{ Name: name, Descriptor: module.ExportDescriptor{ Type: module.FunctionExportType, - Index: c.funcs[name], + Index: idx, }, }) } + // add functions 'name' entry + var found bool + for _, m := range c.module.Names.Functions { + if m.Index == idx { + found = true + } + } + if !found { + c.module.Names.Functions = append(c.module.Names.Functions, module.NameMap{ + Index: idx, + Name: name, + }) + } } func (c *Compiler) emitFunctionType(tpe module.FunctionType) uint32 { diff --git a/internal/wasm/constant/constant.go b/internal/wasm/constant/constant.go index 84e4d4746c..878979fb6e 100644 --- a/internal/wasm/constant/constant.go +++ b/internal/wasm/constant/constant.go @@ -65,3 +65,13 @@ const ( Const byte = iota Mutable ) + +// NameSectionCustomID is the ID of the "Name" section Custom Section +const NameSectionCustomID = "name" + +// Subtypes of the 'name' custom section +const ( + NameSectionModuleType byte = iota + NameSectionFunctionsType + NameSectionLocalsType +) diff --git a/internal/wasm/encoding/encoding_test.go b/internal/wasm/encoding/encoding_test.go index 487bce275d..60d4944005 100644 --- a/internal/wasm/encoding/encoding_test.go +++ b/internal/wasm/encoding/encoding_test.go @@ -12,6 +12,7 @@ import ( "testing" "github.com/open-policy-agent/opa/internal/compiler/wasm/opa" + "github.com/open-policy-agent/opa/internal/wasm/module" ) func TestRoundtrip(t *testing.T) { @@ -67,16 +68,33 @@ func TestRoundtripOPA(t *testing.T) { t.Fatal(err) } - module, err := ReadModule(bytes.NewBuffer(bs)) + module1, err := ReadModule(bytes.NewBuffer(bs)) if err != nil { t.Fatal(err) } + // When using a WASM module with or without debug, the custom sections differ. + // Both variants have 'producers'. + customSections := map[string]int{} + for _, s := range module1.Customs { + customSections[s.Name]++ + } + if expected, actual := 1, customSections["producers"]; expected != actual { + t.Errorf("expected %d 'producers' custom sections, found %d", expected, actual) + } + if len(module1.Names.Functions) == 0 { + t.Errorf("expected non-zero function names in 'name' custom sections") + } + + // Note(sr): We don't have this set by any other means, so manually set it, and + // check the write->read roundtrip at least. + module1.Names.Module = "foo" + module1.Names.Locals = []module.LocalNameMap{{FuncIndex: 1172, NameMap: module.NameMap{Index: 0, Name: "data"}}} // TODO(tsandall): when all instructions are handled by reader, add logic to // check code section contents. var buf2 bytes.Buffer - if err := WriteModule(&buf2, module); err != nil { + if err := WriteModule(&buf2, module1); err != nil { t.Fatal(err) } @@ -85,9 +103,12 @@ func TestRoundtripOPA(t *testing.T) { t.Fatal(err) } - // TODO(tsandall): how to make this more debuggable - if !reflect.DeepEqual(module, module2) { - t.Fatal("modules are not equal") + if expected, actual := len(module1.Names.Functions), len(module2.Names.Functions); expected != actual { + t.Errorf("expected %d function names in 'name' custom sections, found %d", expected, actual) } + // TODO(tsandall): how to make this more debuggable + if !reflect.DeepEqual(module1, module2) { + t.Fatal("modules are not equal") + } } diff --git a/internal/wasm/encoding/reader.go b/internal/wasm/encoding/reader.go index 9ecb8a43f7..6ae6ca3cf1 100644 --- a/internal/wasm/encoding/reader.go +++ b/internal/wasm/encoding/reader.go @@ -9,6 +9,7 @@ import ( "encoding/binary" "fmt" "io" + "io/ioutil" "github.com/pkg/errors" @@ -142,8 +143,22 @@ func readSections(r io.Reader, m *module.Module) error { bufr := bytes.NewReader(buf) switch id { - case constant.CustomSectionID, constant.StartSectionID, constant.MemorySectionID: + case constant.StartSectionID, constant.MemorySectionID: continue + case constant.CustomSectionID: + var name string + if err := readByteVectorString(bufr, &name); err != nil { + return errors.Wrap(err, "read custom section type") + } + if name == "name" { + if err := readCustomNameSections(bufr, &m.Names); err != nil { + return errors.Wrap(err, "custom 'name' section") + } + } else { + if err := readCustomSection(bufr, name, &m.Customs); err != nil { + return errors.Wrap(err, "custom section") + } + } case constant.TypeSectionID: if err := readTypeSection(bufr, &m.Type); err != nil { return errors.Wrap(err, "type section") @@ -186,6 +201,112 @@ func readSections(r io.Reader, m *module.Module) error { } } +func readCustomSection(r io.Reader, name string, s *[]module.CustomSection) error { + buf, err := ioutil.ReadAll(r) + if err != nil { + return err + } + + *s = append(*s, module.CustomSection{ + Name: name, + Data: buf, + }) + return nil +} + +func readCustomNameSections(r io.Reader, s *module.NameSection) error { + for { + id, err := readByte(r) + if err != nil { + if err == io.EOF { + break + } + return err + } + n, err := leb128.ReadVarUint32(r) + if err != nil { + return err + } + buf := make([]byte, n) + if _, err := io.ReadFull(r, buf); err != nil { + return err + } + bufr := bytes.NewReader(buf) + switch id { + case constant.NameSectionModuleType: + err = readNameSectionModule(bufr, s) + case constant.NameSectionFunctionsType: + err = readNameSectionFunctions(bufr, s) + case constant.NameSectionLocalsType: + err = readNameSectionLocals(bufr, s) + } + if err != nil { + return err + } + } + return nil +} + +func readNameSectionModule(r io.Reader, s *module.NameSection) error { + return readByteVectorString(r, &s.Module) +} + +func readNameSectionFunctions(r io.Reader, s *module.NameSection) error { + nm, err := readNameMap(r) + if err != nil { + return err + } + s.Functions = nm + return nil +} + +func readNameMap(r io.Reader) ([]module.NameMap, error) { + n, err := leb128.ReadVarUint32(r) + if err != nil { + return nil, err + } + nm := make([]module.NameMap, n) + for i := uint32(0); i < n; i++ { + var name string + id, err := leb128.ReadVarUint32(r) + if err != nil { + return nil, err + } + + if err := readByteVectorString(r, &name); err != nil { + return nil, err + } + nm[i] = module.NameMap{Index: id, Name: name} + } + return nm, nil +} + +func readNameSectionLocals(r io.Reader, s *module.NameSection) error { + n, err := leb128.ReadVarUint32(r) // length of vec(indirectnameassoc) + if err != nil { + return err + } + for i := uint32(0); i < n; i++ { + id, err := leb128.ReadVarUint32(r) // func index + if err != nil { + return err + } + nm, err := readNameMap(r) + if err != nil { + return err + } + for _, m := range nm { + s.Locals = append(s.Locals, module.LocalNameMap{ + FuncIndex: id, + NameMap: module.NameMap{ + Index: m.Index, + Name: m.Name, + }}) + } + } + return nil +} + func readTypeSection(r io.Reader, s *module.TypeSection) error { n, err := leb128.ReadVarUint32(r) @@ -804,6 +925,6 @@ func readBlockValueType(r io.Reader, v *types.ValueType) error { func readByte(r io.Reader) (byte, error) { buf := make([]byte, 1) - _, err := r.Read(buf) + _, err := io.ReadFull(r, buf) return buf[0], err } diff --git a/internal/wasm/encoding/writer.go b/internal/wasm/encoding/writer.go index c95045b05d..413e3d9d68 100644 --- a/internal/wasm/encoding/writer.go +++ b/internal/wasm/encoding/writer.go @@ -70,6 +70,16 @@ func WriteModule(w io.Writer, module *module.Module) error { return err } + if err := writeNameSection(w, module.Names); err != nil { + return err + } + + for _, custom := range module.Customs { + if err := writeCustomSection(w, custom); err != nil { + return err + } + } + return nil } @@ -370,6 +380,111 @@ func writeDataSection(w io.Writer, s module.DataSection) error { return writeRawSection(w, &buf) } +func writeNameSection(w io.Writer, s module.NameSection) error { + if s.Module == "" && len(s.Functions) == 0 && len(s.Locals) == 0 { + return nil + } + + if err := writeByte(w, constant.CustomSectionID); err != nil { + return err + } + + var buf bytes.Buffer + if err := writeByteVector(&buf, []byte(constant.NameSectionCustomID)); err != nil { + return err + } + + // "module" subsection + if s.Module != "" { + if err := writeByte(&buf, constant.NameSectionModuleType); err != nil { + return err + } + var mbuf bytes.Buffer + if err := writeByteVector(&mbuf, []byte(s.Module)); err != nil { + return err + } + if err := writeRawSection(&buf, &mbuf); err != nil { + return err + } + } + + // "functions" subsection + if len(s.Functions) != 0 { + if err := writeByte(&buf, constant.NameSectionFunctionsType); err != nil { + return err + } + + var fbuf bytes.Buffer + if err := writeNameMap(&fbuf, s.Functions); err != nil { + return err + } + if err := writeRawSection(&buf, &fbuf); err != nil { + return err + } + } + + // "locals" subsection + if len(s.Locals) != 0 { + if err := writeByte(&buf, constant.NameSectionLocalsType); err != nil { + return err + } + funs := map[uint32][]module.NameMap{} + for _, e := range s.Locals { + funs[e.FuncIndex] = append(funs[e.FuncIndex], module.NameMap{Index: e.Index, Name: e.Name}) + } + var lbuf bytes.Buffer + if err := leb128.WriteVarUint32(&lbuf, uint32(len(funs))); err != nil { + return err + } + for fidx, e := range funs { + if err := leb128.WriteVarUint32(&lbuf, fidx); err != nil { + return err + } + if err := writeNameMap(&lbuf, e); err != nil { + return err + } + } + if err := writeRawSection(&buf, &lbuf); err != nil { + return err + } + } + + return writeRawSection(w, &buf) +} + +func writeNameMap(buf io.Writer, nm []module.NameMap) error { + if err := leb128.WriteVarUint32(buf, uint32(len(nm))); err != nil { + return err + } + for _, m := range nm { + if err := leb128.WriteVarUint32(buf, m.Index); err != nil { + return err + } + if err := writeByteVector(buf, []byte(m.Name)); err != nil { + return err + } + } + return nil +} + +func writeCustomSection(w io.Writer, s module.CustomSection) error { + + if err := writeByte(w, constant.CustomSectionID); err != nil { + return err + } + + var buf bytes.Buffer + if err := writeByteVector(&buf, []byte(s.Name)); err != nil { + return err + } + + if _, err := io.Copy(&buf, bytes.NewReader(s.Data)); err != nil { + return err + } + + return writeRawSection(w, &buf) +} + func writeFunctionType(w io.Writer, fsig module.FunctionType) error { if err := writeByte(w, constant.FunctionTypeID); err != nil { diff --git a/internal/wasm/module/module.go b/internal/wasm/module/module.go index b55094c8c0..6d1edcd03d 100644 --- a/internal/wasm/module/module.go +++ b/internal/wasm/module/module.go @@ -25,6 +25,8 @@ type ( Export ExportSection Code RawCodeSection Data DataSection + Customs []CustomSection + Names NameSection } // TypeSection represents a WASM type section. @@ -73,6 +75,32 @@ type ( Segments []DataSegment } + // CustomSection represents a WASM custom section. + CustomSection struct { + Name string + Data []byte + } + + // NameSection represents the WASM custom section "name". + NameSection struct { + Module string + Functions []NameMap + Locals []LocalNameMap + } + + // NameMap maps function or local arg indices to their names. + NameMap struct { + Index uint32 + Name string + } + + // LocalNameMap maps function indices, and argument indices for the args + // of the indexed function to their names. + LocalNameMap struct { + FuncIndex uint32 + NameMap + } + // FunctionType represents a WASM function type definition. FunctionType struct { Params []types.ValueType