From d07c40f5d621002d2631305539756ed95e4f97bf Mon Sep 17 00:00:00 2001 From: Edoardo Vacchi Date: Thu, 19 Jan 2023 13:39:43 +0100 Subject: [PATCH] CompiledModule: export Custom Sections (#1048) Signed-off-by: Edoardo Vacchi Co-authored-by: Takeshi Yoneda --- api/wasm.go | 8 +++++++ config.go | 47 ++++++++++++++++++++++++++++++++++++++ config_example_test.go | 43 ++++++++++++++++++++++++++++++++++ config_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++ runtime.go | 4 +++- 5 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 config_example_test.go diff --git a/api/wasm.go b/api/wasm.go index 1141a025..e9d33d6b 100644 --- a/api/wasm.go +++ b/api/wasm.go @@ -562,6 +562,14 @@ type Memory interface { WriteString(offset uint32, v string) bool } +// CustomSection contains the name and raw data of a custom section. +type CustomSection interface { + // Name is the name of the custom section + Name() string + // Data is the raw data of the custom section + Data() []byte +} + // EncodeExternref encodes the input as a ValueTypeExternref. // // See DecodeExternref diff --git a/config.go b/config.go index 65ae437a..bf0143fa 100644 --- a/config.go +++ b/config.go @@ -117,6 +117,16 @@ type RuntimeConfig interface { // foo := wazero.NewRuntimeWithConfig(context.Background(), config) // bar := wazero.NewRuntimeWithConfig(context.Background(), config) WithCompilationCache(CompilationCache) RuntimeConfig + + // WithCustomSections toggles parsing of "custom sections". Defaults to false. + // + // When enabled, it is possible to retrieve custom sections from a CompiledModule: + // + // config := wazero.NewRuntimeConfig().WithCustomSections(true) + // r := wazero.NewRuntimeWithConfig(ctx, config) + // c, err := r.CompileModule(ctx, wasm) + // customSections := c.CustomSections() + WithCustomSections(bool) RuntimeConfig } // NewRuntimeConfig returns a RuntimeConfig using the compiler if it is supported in this environment, @@ -135,6 +145,7 @@ type runtimeConfig struct { dwarfDisabled bool // negative as defaults to enabled newEngine newEngine cache CompilationCache + storeCustomSections bool } // engineLessConfig helps avoid copy/pasting the wrong defaults. @@ -227,6 +238,13 @@ func (c *runtimeConfig) WithDebugInfoEnabled(dwarfEnabled bool) RuntimeConfig { return ret } +// WithCustomSections implements RuntimeConfig.WithCustomSections +func (c *runtimeConfig) WithCustomSections(storeCustomSections bool) RuntimeConfig { + ret := c.clone() + ret.storeCustomSections = storeCustomSections + return ret +} + // CompiledModule is a WebAssembly module ready to be instantiated (Runtime.InstantiateModule) as an api.Module. // // In WebAssembly terminology, this is a decoded, validated, and possibly also compiled module. wazero avoids using @@ -265,6 +283,10 @@ type CompiledModule interface { // memory. ExportedMemories() map[string]api.MemoryDefinition + // CustomSections returns all the custom sections + // (api.CustomSection) in this module keyed on the section name. + CustomSections() []api.CustomSection + // Close releases all the allocated resources for this CompiledModule. // // Note: It is safe to call Close while having outstanding calls from an @@ -318,6 +340,31 @@ func (c *compiledModule) ExportedMemories() map[string]api.MemoryDefinition { return c.module.ExportedMemories() } +// CustomSections implements CompiledModule.CustomSections +func (c *compiledModule) CustomSections() []api.CustomSection { + ret := make([]api.CustomSection, len(c.module.CustomSections)) + for i, d := range c.module.CustomSections { + ret[i] = &customSection{data: d.Data, name: d.Name} + } + return ret +} + +// customSection implements wasm.CustomSection +type customSection struct { + name string + data []byte +} + +// Name implements wasm.CustomSection.Name +func (c *customSection) Name() string { + return c.name +} + +// Data implements wasm.CustomSection.Data +func (c *customSection) Data() []byte { + return c.data +} + // ModuleConfig configures resources needed by functions that have low-level interactions with the host operating // system. Using this, resources such as STDIN can be isolated, so that the same module can be safely instantiated // multiple times. diff --git a/config_example_test.go b/config_example_test.go new file mode 100644 index 00000000..c93cef89 --- /dev/null +++ b/config_example_test.go @@ -0,0 +1,43 @@ +package wazero_test + +import ( + "context" + _ "embed" + "log" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +// This is a basic example of retrieving custom sections using RuntimeConfig.WithCustomSections. +func Example_runtimeConfig_WithCustomSections() { + ctx := context.Background() + config := wazero.NewRuntimeConfig().WithCustomSections(true) + + r := wazero.NewRuntimeWithConfig(ctx, config) + defer r.Close(ctx) + + m, err := r.CompileModule(ctx, addWasm) + if err != nil { + log.Panicln(err) + } + + if m.CustomSections() == nil { + log.Panicln("Custom sections should not be nil") + } + + mustContain(m.CustomSections(), "producers") + mustContain(m.CustomSections(), "target_features") + + // Output: + // +} + +func mustContain(ss []api.CustomSection, name string) { + for _, s := range ss { + if s.Name() == name { + return + } + } + log.Panicf("Could not find section named %s\n", name) +} diff --git a/config_test.go b/config_test.go index c41e6654..b5fa17df 100644 --- a/config_test.go +++ b/config_test.go @@ -60,6 +60,15 @@ func TestRuntimeConfig(t *testing.T) { dwarfDisabled: true, // dwarf is a more technical name and ok here. }, }, + { + name: "WithCustomSections", + with: func(c RuntimeConfig) RuntimeConfig { + return c.WithCustomSections(true) + }, + expected: &runtimeConfig{ + storeCustomSections: true, + }, + }, } for _, tt := range tests { @@ -597,6 +606,49 @@ func Test_compiledModule_Name(t *testing.T) { } } +func Test_compiledModule_CustomSections(t *testing.T) { + tests := []struct { + name string + input *compiledModule + expected []string + }{ + { + name: "no custom section", + input: &compiledModule{module: &wasm.Module{}}, + expected: []string{}, + }, + { + name: "name", + input: &compiledModule{module: &wasm.Module{ + CustomSections: []*wasm.CustomSection{ + {Name: "custom1"}, + {Name: "custom2"}, + {Name: "customDup"}, + {Name: "customDup"}, + }, + }}, + expected: []string{ + "custom1", + "custom2", + "customDup", + "customDup", + }, + }, + } + + for _, tt := range tests { + tc := tt + + t.Run(tc.name, func(t *testing.T) { + customSections := tc.input.CustomSections() + require.Equal(t, len(tc.expected), len(customSections)) + for i := 0; i < len(tc.expected); i++ { + require.Equal(t, tc.expected[i], customSections[i].Name()) + } + }) + } +} + func Test_compiledModule_Close(t *testing.T) { for _, ctx := range []context.Context{nil, testCtx} { // Ensure it doesn't crash on nil! e := &mockEngine{name: "1", cachedModules: map[*wasm.Module]struct{}{}} diff --git a/runtime.go b/runtime.go index 1a6eafb1..d545de3a 100644 --- a/runtime.go +++ b/runtime.go @@ -129,6 +129,7 @@ func NewRuntimeWithConfig(ctx context.Context, rConfig RuntimeConfig) Runtime { memoryLimitPages: config.memoryLimitPages, memoryCapacityFromMax: config.memoryCapacityFromMax, dwarfDisabled: config.dwarfDisabled, + storeCustomSections: config.storeCustomSections, } } @@ -140,6 +141,7 @@ type runtime struct { memoryLimitPages uint32 memoryCapacityFromMax bool dwarfDisabled bool + storeCustomSections bool } // Module implements Runtime.Module. @@ -158,7 +160,7 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod } internal, err := binaryformat.DecodeModule(binary, r.enabledFeatures, - r.memoryLimitPages, r.memoryCapacityFromMax, !r.dwarfDisabled, false) + r.memoryLimitPages, r.memoryCapacityFromMax, !r.dwarfDisabled, r.storeCustomSections) if err != nil { return nil, err } else if err = internal.Validate(r.enabledFeatures); err != nil {