CompiledModule: export Custom Sections (#1048)

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>
Co-authored-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

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

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.

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)
+}

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{}{}}

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 {