binary: adds custom section decoder (#877)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

internal/integration_test/bench/codec_test.go

@@ -89,7 +89,7 @@ func newExample() *wasm.Module {
 
 func TestExampleUpToDate(t *testing.T) {
 	t.Run("binary.DecodeModule", func(t *testing.T) {
-		m, err := binary.DecodeModule(exampleWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false)
+		m, err := binary.DecodeModule(exampleWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false, false)
 		require.NoError(t, err)
 		require.Equal(t, example, m)
 	})
@@ -116,7 +116,7 @@ func BenchmarkCodec(b *testing.B) {
 	b.Run("binary.DecodeModule", func(b *testing.B) {
 		b.ReportAllocs()
 		for i := 0; i < b.N; i++ {
-			if _, err := binary.DecodeModule(exampleWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false); err != nil {
+			if _, err := binary.DecodeModule(exampleWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false, false); err != nil {
 				b.Fatal(err)
 			}
 		}

internal/integration_test/spectest/spectest.go

@@ -325,7 +325,7 @@ var spectestWasm []byte
 // See https://github.com/WebAssembly/spec/blob/wg-1.0/test/core/imports.wast
 // See https://github.com/WebAssembly/spec/blob/wg-1.0/interpreter/script/js.ml#L13-L25
 func addSpectestModule(t *testing.T, ctx context.Context, s *wasm.Store, ns *wasm.Namespace, enabledFeatures api.CoreFeatures) {
-	mod, err := binaryformat.DecodeModule(spectestWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false)
+	mod, err := binaryformat.DecodeModule(spectestWasm, api.CoreFeaturesV2, wasm.MemoryLimitPages, false, false)
 	require.NoError(t, err)
 
 	// (global (export "global_i32") i32 (i32.const 666))
@@ -420,7 +420,7 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, newEngine func(
 					case "module":
 						buf, err := testDataFS.ReadFile(testdataPath(c.Filename))
 						require.NoError(t, err, msg)
-						mod, err := binaryformat.DecodeModule(buf, enabledFeatures, wasm.MemoryLimitPages, false)
+						mod, err := binaryformat.DecodeModule(buf, enabledFeatures, wasm.MemoryLimitPages, false, false)
 						require.NoError(t, err, msg)
 						require.NoError(t, mod.Validate(enabledFeatures))
 						mod.AssignModuleID(buf)
@@ -561,7 +561,7 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, newEngine func(
 							//
 							// In practice, such a module instance can be used for invoking functions without any issue. In addition, we have to
 							// retain functions after the expected "instantiation" failure, so in wazero we choose to not raise error in that case.
-							mod, err := binaryformat.DecodeModule(buf, s.EnabledFeatures, wasm.MemoryLimitPages, false)
+							mod, err := binaryformat.DecodeModule(buf, s.EnabledFeatures, wasm.MemoryLimitPages, false, false)
 							require.NoError(t, err, msg)
 
 							err = mod.Validate(s.EnabledFeatures)
@@ -590,7 +590,7 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, newEngine func(
 }
 
 func requireInstantiationError(t *testing.T, ctx context.Context, s *wasm.Store, ns *wasm.Namespace, buf []byte, msg string) {
-	mod, err := binaryformat.DecodeModule(buf, s.EnabledFeatures, wasm.MemoryLimitPages, false)
+	mod, err := binaryformat.DecodeModule(buf, s.EnabledFeatures, wasm.MemoryLimitPages, false, false)
 	if err != nil {
 		return
 	}

internal/wasm/binary/custom.go

@@ -0,0 +1,22 @@
+package binary
+
+import (
+	"bytes"
+
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+// decodeCustomSection deserializes the data **not** associated with the "name" key in SectionIDCustom.
+//
+// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#custom-section%E2%91%A0
+func decodeCustomSection(r *bytes.Reader, name string, limit uint64) (result *wasm.CustomSection, err error) {
+	buf := make([]byte, limit)
+	_, err = r.Read(buf)
+
+	result = &wasm.CustomSection{
+		Name: name,
+		Data: buf,
+	}
+
+	return
+}

internal/wasm/binary/decoder.go

@@ -17,7 +17,8 @@ func DecodeModule(
 	binary []byte,
 	enabledFeatures api.CoreFeatures,
 	memoryLimitPages uint32,
-	memoryCapacityFromMax bool,
+	memoryCapacityFromMax,
+	storeCustomSections bool,
 ) (*wasm.Module, error) {
 	r := bytes.NewReader(binary)
 
@@ -66,17 +67,22 @@ func DecodeModule(
 				break
 			}
 
-			// Now, either decode the NameSection or skip an unsupported one
+			// Now, either decode the NameSection or CustomSection
 			limit := sectionSize - nameSize
 			if name == "name" {
 				m.NameSection, err = decodeNameSection(r, uint64(limit))
+			} else if storeCustomSections {
+				custom, err := decodeCustomSection(r, name, uint64(limit))
+				if err != nil {
+					return nil, fmt.Errorf("failed to read custom section name[%s]: %w", name, err)
+				}
+				m.CustomSections = append(m.CustomSections, custom)
 			} else {
 				// Note: Not Seek because it doesn't err when given an offset past EOF. Rather, it leads to undefined state.
 				if _, err = io.CopyN(io.Discard, r, int64(limit)); err != nil {
 					return nil, fmt.Errorf("failed to skip name[%s]: %w", name, err)
 				}
 			}
-
 		case wasm.SectionIDType:
 			m.TypeSection, err = decodeTypeSection(enabledFeatures, r)
 		case wasm.SectionIDImport:

internal/wasm/binary/decoder_test.go

@@ -81,7 +81,7 @@ func TestDecodeModule(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			m, e := DecodeModule(EncodeModule(tc.input), api.CoreFeaturesV1, wasm.MemoryLimitPages, false)
+			m, e := DecodeModule(EncodeModule(tc.input), api.CoreFeaturesV1, wasm.MemoryLimitPages, false, false)
 			require.NoError(t, e)
 			require.Equal(t, tc.input, m)
 		})
@@ -92,11 +92,28 @@ func TestDecodeModule(t *testing.T) {
 			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
 			0x04, 'm', 'e', 'm', 'e',
 			1, 2, 3, 4, 5, 6, 7, 8, 9, 0)
-		m, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false)
+		m, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false, false)
 		require.NoError(t, e)
 		require.Equal(t, &wasm.Module{}, m)
 	})
 
+	t.Run("reads custom sections", func(t *testing.T) {
+		input := append(append(Magic, version...),
+			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
+			0x04, 'm', 'e', 'm', 'e',
+			1, 2, 3, 4, 5, 6, 7, 8, 9, 0)
+		m, e := DecodeModule(input, api.CoreFeaturesV2, wasm.MemoryLimitPages, false, true)
+		require.NoError(t, e)
+		require.Equal(t, &wasm.Module{
+			CustomSections: []*wasm.CustomSection{
+				{
+					Name: "meme",
+					Data: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
+				},
+			},
+		}, m)
+	})
+
 	t.Run("skips custom section, but not name", func(t *testing.T) {
 		input := append(append(Magic, version...),
 			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
@@ -107,14 +124,38 @@ func TestDecodeModule(t *testing.T) {
 			subsectionIDModuleName, 0x07, // 7 bytes in this subsection
 			0x06, // the Module name simple is 6 bytes long
 			's', 'i', 'm', 'p', 'l', 'e')
-		m, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false)
+		m, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false, false)
 		require.NoError(t, e)
 		require.Equal(t, &wasm.Module{NameSection: &wasm.NameSection{ModuleName: "simple"}}, m)
 	})
+
+	t.Run("read custom sections and name separately", func(t *testing.T) {
+		input := append(append(Magic, version...),
+			wasm.SectionIDCustom, 0xf, // 15 bytes in this section
+			0x04, 'm', 'e', 'm', 'e',
+			1, 2, 3, 4, 5, 6, 7, 8, 9, 0,
+			wasm.SectionIDCustom, 0x0e, // 14 bytes in this section
+			0x04, 'n', 'a', 'm', 'e',
+			subsectionIDModuleName, 0x07, // 7 bytes in this subsection
+			0x06, // the Module name simple is 6 bytes long
+			's', 'i', 'm', 'p', 'l', 'e')
+		m, e := DecodeModule(input, api.CoreFeaturesV2, wasm.MemoryLimitPages, false, true)
+		require.NoError(t, e)
+		require.Equal(t, &wasm.Module{
+			NameSection: &wasm.NameSection{ModuleName: "simple"},
+			CustomSections: []*wasm.CustomSection{
+				{
+					Name: "meme",
+					Data: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
+				},
+			},
+		}, m)
+	})
+
 	t.Run("data count section disabled", func(t *testing.T) {
 		input := append(append(Magic, version...),
 			wasm.SectionIDDataCount, 1, 0)
-		_, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false)
+		_, e := DecodeModule(input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false, false)
 		require.EqualError(t, e, `data count section not supported as feature "bulk-memory-operations" is disabled`)
 	})
 }
@@ -164,7 +205,7 @@ func TestDecodeModule_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			_, e := DecodeModule(tc.input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false)
+			_, e := DecodeModule(tc.input, api.CoreFeaturesV1, wasm.MemoryLimitPages, false, false)
 			require.EqualError(t, e, tc.expectedErr)
 		})
 	}

internal/wasm/counts.go

@@ -41,16 +41,17 @@ func (m *Module) importCount(et ExternType) (res uint32) {
 //
 // For example...
 // * SectionIDType returns the count of FunctionType
-// * SectionIDCustom returns one if the NameSection is present
+// * SectionIDCustom returns the count of CustomSections plus one if NameSection is present
 // * SectionIDHostFunction returns the count of HostFunctionSection
 // * SectionIDExport returns the count of unique export names
 func (m *Module) SectionElementCount(sectionID SectionID) uint32 { // element as in vector elements!
 	switch sectionID {
 	case SectionIDCustom:
+		numCustomSections := uint32(len(m.CustomSections))
 		if m.NameSection != nil {
-			return 1
+			numCustomSections++
 		}
-		return 0
+		return numCustomSections
 	case SectionIDType:
 		return uint32(len(m.TypeSection))
 	case SectionIDImport:

internal/wasm/module.go

@@ -154,6 +154,11 @@ type Module struct {
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#custom-section%E2%91%A0
 	NameSection *NameSection
 
+	// CustomSections are set when the SectionIDCustom other than "name" were successfully decoded from the binary format.
+	//
+	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#custom-section%E2%91%A0
+	CustomSections []*CustomSection
+
 	// validatedActiveElementSegments are built on Validate when
 	// SectionIDElement is non-empty and all inputs are valid.
 	//
@@ -877,6 +882,12 @@ type NameSection struct {
 	LocalNames IndirectNameMap
 }
 
+// CustomSection contains the name and raw data of a custom section.
+type CustomSection struct {
+	Name string
+	Data []byte
+}
+
 // NameMap associates an index with any associated names.
 //
 // Note: Often the index namespace bridges multiple sections. For example, the function index namespace starts with any

runtime.go

@@ -172,7 +172,7 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 		return nil, errors.New("invalid binary")
 	}
 
-	internal, err := binaryformat.DecodeModule(binary, r.enabledFeatures, r.memoryLimitPages, r.memoryCapacityFromMax)
+	internal, err := binaryformat.DecodeModule(binary, r.enabledFeatures, r.memoryLimitPages, r.memoryCapacityFromMax, false)
 	if err != nil {
 		return nil, err
 	} else if err = internal.Validate(r.enabledFeatures); err != nil {