Adds Module.SectionSize and constrains text parsing rules (tetratelab…

…s#231)

This consolidates more code around SectionID, notably consistently
getting and using section ID length. This also enforces rules in text
parsing such as up to one memory, and imported functions can't be
defined after regular ones.

Signed-off-by: Adrian Cole <adrian@tetrate.io>
Signed-off-by: r8d8 <ckryvomaz@gmail.com>

wasm/binary/decoder.go

@@ -27,6 +27,8 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 
 	m := &wasm.Module{}
 	for {
+		// TODO: except custom sections, all others are required to be in order, but we aren't checking yet.
+		// See https://www.w3.org/TR/wasm-core-1/#modules%E2%91%A0%E2%93%AA
 		sectionID := make([]byte, 1)
 		if _, err := io.ReadFull(r, sectionID); err == io.EOF {
 			break
@@ -36,7 +38,7 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 
 		sectionSize, _, err := leb128.DecodeUint32(r)
 		if err != nil {
-			return nil, fmt.Errorf("get size of section for id=%d: %v", sectionID[0], err)
+			return nil, fmt.Errorf("get size of section %s: %v", wasm.SectionIDName(sectionID[0]), err)
 		}
 
 		sectionContentStart := r.Len()
@@ -104,11 +106,11 @@ func DecodeModule(binary []byte) (*wasm.Module, error) {
 		}
 
 		if err != nil {
-			return nil, fmt.Errorf("section ID %d: %v", sectionID[0], err)
+			return nil, fmt.Errorf("section %s: %v", wasm.SectionIDName(sectionID[0]), err)
 		}
 	}
 
-	functionCount, codeCount := len(m.FunctionSection), len(m.CodeSection)
+	functionCount, codeCount := m.SectionSize(wasm.SectionIDFunction), m.SectionSize(wasm.SectionIDCode)
 	if functionCount != codeCount {
 		return nil, fmt.Errorf("function and code section have inconsistent lengths: %d != %d", functionCount, codeCount)
 	}

wasm/binary/decoder_test.go

@@ -134,7 +134,7 @@ func TestDecodeModule_Errors(t *testing.T) {
 				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
 				0x04, 'm', 'e', 'm', 'e',
 				subsectionIDModuleName, 0x03, 0x01, 'y'),
-			expectedErr: "section ID 0: redundant custom section meme",
+			expectedErr: "section custom: redundant custom section meme",
 		},
 		{
 			name: "redundant name section",
@@ -145,7 +145,7 @@ func TestDecodeModule_Errors(t *testing.T) {
 				wasm.SectionIDCustom, 0x09, // 9 bytes in this section
 				0x04, 'n', 'a', 'm', 'e',
 				subsectionIDModuleName, 0x03, 0x01, 'x'),
-			expectedErr: "section ID 0: redundant custom section name",
+			expectedErr: "section custom: redundant custom section name",
 		},
 	}
 

wasm/binary/encoder.go

@@ -11,47 +11,50 @@ var sizePrefixedName = []byte{4, 'n', 'a', 'm', 'e'}
 // See https://www.w3.org/TR/wasm-core-1/#binary-format%E2%91%A0
 func EncodeModule(m *wasm.Module) (bytes []byte) {
 	bytes = append(magic, version...)
-	for name, data := range m.CustomSections {
-		bytes = append(bytes, encodeCustomSection(name, data)...)
-	}
-	if len(m.TypeSection) > 0 {
+	if m.SectionSize(wasm.SectionIDType) > 0 {
 		bytes = append(bytes, encodeTypeSection(m.TypeSection)...)
 	}
-	if len(m.ImportSection) > 0 {
+	if m.SectionSize(wasm.SectionIDImport) > 0 {
 		bytes = append(bytes, encodeImportSection(m.ImportSection)...)
 	}
-	if len(m.FunctionSection) > 0 {
+	if m.SectionSize(wasm.SectionIDFunction) > 0 {
 		bytes = append(bytes, encodeFunctionSection(m.FunctionSection)...)
 	}
-	if len(m.TableSection) > 0 {
+	if m.SectionSize(wasm.SectionIDTable) > 0 {
 		panic("TODO: TableSection")
 	}
-	if len(m.MemorySection) > 0 {
+	if m.SectionSize(wasm.SectionIDMemory) > 0 {
 		bytes = append(bytes, encodeMemorySection(m.MemorySection)...)
 	}
-	if len(m.GlobalSection) > 0 {
+	if m.SectionSize(wasm.SectionIDGlobal) > 0 {
 		panic("TODO: GlobalSection")
 	}
-	if len(m.ExportSection) > 0 {
+	if m.SectionSize(wasm.SectionIDExport) > 0 {
 		bytes = append(bytes, encodeExportSection(m.ExportSection)...)
 	}
-	if m.StartSection != nil {
+	if m.SectionSize(wasm.SectionIDStart) > 0 {
 		bytes = append(bytes, encodeStartSection(*m.StartSection)...)
 	}
-	if len(m.ElementSection) > 0 {
+	if m.SectionSize(wasm.SectionIDElement) > 0 {
 		panic("TODO: ElementSection")
 	}
-	if len(m.CodeSection) > 0 {
+	if m.SectionSize(wasm.SectionIDCode) > 0 {
 		bytes = append(bytes, encodeCodeSection(m.CodeSection)...)
 	}
-	if len(m.DataSection) > 0 {
+	if m.SectionSize(wasm.SectionIDData) > 0 {
 		panic("TODO: DataSection")
 	}
-	// >> The name section should appear only once in a module, and only after the data section.
-	// See https://www.w3.org/TR/wasm-core-1/#binary-namesec
-	if m.NameSection != nil {
-		nameSection := append(sizePrefixedName, encodeNameSectionData(m.NameSection)...)
-		bytes = append(bytes, encodeSection(wasm.SectionIDCustom, nameSection)...)
+	if m.SectionSize(wasm.SectionIDCustom) > 0 {
+		for name, data := range m.CustomSections {
+			bytes = append(bytes, encodeCustomSection(name, data)...)
+		}
+
+		// >> The name section should appear only once in a module, and only after the data section.
+		// See https://www.w3.org/TR/wasm-core-1/#binary-namesec
+		if m.NameSection != nil {
+			nameSection := append(sizePrefixedName, encodeNameSectionData(m.NameSection)...)
+			bytes = append(bytes, encodeSection(wasm.SectionIDCustom, nameSection)...)
+		}
 	}
 	return
 }

wasm/binary/section.go

@@ -165,7 +165,7 @@ func decodeExportSection(r *bytes.Reader) (map[string]*wasm.Export, error) {
 	return exportSection, nil
 }
 
-func decodeStartSection(r *bytes.Reader) (*uint32, error) {
+func decodeStartSection(r *bytes.Reader) (*wasm.Index, error) {
 	vs, _, err := leb128.DecodeUint32(r)
 	if err != nil {
 		return nil, fmt.Errorf("get size of vector: %w", err)
@@ -260,10 +260,10 @@ func encodeImportSection(imports []*wasm.Import) []byte {
 // WebAssembly 1.0 (MVP) Binary Format.
 //
 // See https://www.w3.org/TR/wasm-core-1/#function-section%E2%91%A0
-func encodeFunctionSection(functions []wasm.Index) []byte {
-	contents := leb128.EncodeUint32(uint32(len(functions)))
-	for _, typeIndex := range functions {
-		contents = append(contents, leb128.EncodeUint32(typeIndex)...)
+func encodeFunctionSection(typeIndices []wasm.Index) []byte {
+	contents := leb128.EncodeUint32(uint32(len(typeIndices)))
+	for _, index := range typeIndices {
+		contents = append(contents, leb128.EncodeUint32(index)...)
 	}
 	return encodeSection(wasm.SectionIDFunction, contents)
 }
@@ -307,6 +307,6 @@ func encodeExportSection(exports map[string]*wasm.Export) []byte {
 // encodeStartSection encodes a SectionIDStart for the given function index in WebAssembly 1.0 (MVP) Binary Format.
 //
 // See https://www.w3.org/TR/wasm-core-1/#start-section%E2%91%A0
-func encodeStartSection(funcidx uint32) []byte {
+func encodeStartSection(funcidx wasm.Index) []byte {
 	return encodeSection(wasm.SectionIDStart, leb128.EncodeUint32(funcidx))
 }

wasm/binary/section_test.go

@@ -111,6 +111,11 @@ func TestDecodeExportSection_Errors(t *testing.T) {
 	}
 }
 
+func TestEncodeFunctionSection(t *testing.T) {
+	require.Equal(t, []byte{wasm.SectionIDFunction, 0x2, 0x01, 0x05}, encodeFunctionSection([]wasm.Index{5}))
+}
+
+// TestEncodeStartSection uses the same index as TestEncodeFunctionSection to highlight the encoding is different.
 func TestEncodeStartSection(t *testing.T) {
 	require.Equal(t, []byte{wasm.SectionIDStart, 0x01, 0x05}, encodeStartSection(5))
 }

wasm/module.go

@@ -1,6 +1,8 @@
 package wasm
 
-// DecodeModule parses the configured source into a wasm.Module. This function returns when the source is exhausted or
+import "fmt"
+
+// DecodeModule parses the configured source into a Module. This function returns when the source is exhausted or
 // an error occurs. The result can be initialized for use via Store.Instantiate.
 //
 // Here's a description of the return values:
@@ -115,7 +117,8 @@ type Module struct {
 	// this module at TableSection[0].
 	//
 	// Note: Version 1.0 (MVP) of the WebAssembly spec allows at most one table definition per module, so the length of
-	// the TableSection can be zero or one.
+	// the TableSection can be zero or one, and can only be one if there is no ImportKindTable.
+	//
 	// Note: In the Binary Format, this is SectionIDTable.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#table-section%E2%91%A0
@@ -128,7 +131,8 @@ type Module struct {
 	// this module at TableSection[0].
 	//
 	// Note: Version 1.0 (MVP) of the WebAssembly spec allows at most one memory definition per module, so the length of
-	// the MemorySection can be zero or one.
+	// the MemorySection can be zero or one, and can only be one if there is no ImportKindMemory.
+	//
 	// Note: In the Binary Format, this is SectionIDMemory.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#memory-section%E2%91%A0
@@ -154,6 +158,7 @@ type Module struct {
 	//
 	// Note: The index here is not the position in the FunctionSection, rather in the function index namespace, which
 	// begins with imported functions.
+	//
 	// Note: In the Binary Format, this is SectionIDStart.
 	//
 	// See https://www.w3.org/TR/wasm-core-1/#start-section%E2%91%A0
@@ -226,7 +231,7 @@ const (
 	ValueTypeF64 ValueType = 0x7c
 )
 
-// ValuTypeName returns the type name of the given ValueType as a string.
+// ValueTypeName returns the type name of the given ValueType as a string.
 // These type names match the names used in the WebAssembly text format.
 // Note that ValueTypeName returns "unknown", if an undefined ValueType value is passed.
 func ValueTypeName(t ValueType) string {
@@ -461,3 +466,47 @@ func (m *Module) allDeclarations() (functions []Index, globals []*GlobalType, me
 	tables = append(tables, m.TableSection...)
 	return
 }
+
+// SectionSize returns the count of items for a given section ID
+//
+// For example, given...
+// * SectionIDType this returns the count of FunctionType
+// * SectionIDCustom this returns the count of unique section names
+// * SectionIDExport this returns the count of unique export names
+func (m *Module) SectionSize(sectionID SectionID) uint32 {
+	switch sectionID {
+	case SectionIDCustom:
+		count := uint32(len(m.CustomSections))
+		if m.NameSection != nil {
+			return count + 1
+		}
+		return count
+	case SectionIDType:
+		return uint32(len(m.TypeSection))
+	case SectionIDImport:
+		return uint32(len(m.ImportSection))
+	case SectionIDFunction:
+		return uint32(len(m.FunctionSection))
+	case SectionIDTable:
+		return uint32(len(m.TableSection))
+	case SectionIDMemory:
+		return uint32(len(m.MemorySection))
+	case SectionIDGlobal:
+		return uint32(len(m.GlobalSection))
+	case SectionIDExport:
+		return uint32(len(m.ExportSection))
+	case SectionIDStart:
+		if m.StartSection != nil {
+			return 1
+		}
+		return 0
+	case SectionIDElement:
+		return uint32(len(m.ElementSection))
+	case SectionIDCode:
+		return uint32(len(m.CodeSection))
+	case SectionIDData:
+		return uint32(len(m.DataSection))
+	default:
+		panic(fmt.Errorf("BUG: unknown section: %d", sectionID))
+	}
+}

wasm/module_test.go

@@ -203,3 +203,120 @@ func TestModule_allDeclarations(t *testing.T) {
 		})
 	}
 }
+
+func TestModule_SectionSize(t *testing.T) {
+	i32, f32 := ValueTypeI32, ValueTypeF32
+	zero := uint32(0)
+	empty := &ConstantExpression{Opcode: OpcodeI32Const, Data: []byte{0x00}}
+
+	tests := []struct {
+		name     string
+		input    *Module
+		expected map[string]uint32
+	}{
+		{
+			name:     "empty",
+			input:    &Module{},
+			expected: map[string]uint32{},
+		},
+		{
+			name:     "only name section",
+			input:    &Module{NameSection: &NameSection{ModuleName: "simple"}},
+			expected: map[string]uint32{"custom": 1},
+		},
+		{
+			name: "only custom section",
+			input: &Module{CustomSections: map[string][]byte{
+				"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
+			}},
+			expected: map[string]uint32{"custom": 1},
+		},
+		{
+			name: "name section and a custom section",
+			input: &Module{
+				NameSection: &NameSection{ModuleName: "simple"},
+				CustomSections: map[string][]byte{
+					"meme": {1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
+				},
+			},
+			expected: map[string]uint32{"custom": 2},
+		},
+		{
+			name: "type section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{},
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{i32, i32, i32, i32}, Results: []ValueType{i32}},
+				},
+			},
+			expected: map[string]uint32{"type": 3},
+		},
+		{
+			name: "type and import section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{f32, f32}, Results: []ValueType{f32}},
+				},
+				ImportSection: []*Import{
+					{
+						Module: "Math", Name: "Mul",
+						Kind:     ImportKindFunc,
+						DescFunc: 1,
+					}, {
+						Module: "Math", Name: "Add",
+						Kind:     ImportKindFunc,
+						DescFunc: 0,
+					},
+				},
+			},
+			expected: map[string]uint32{"import": 2, "type": 2},
+		},
+		{
+			name: "type function and start section",
+			input: &Module{
+				TypeSection:     []*FunctionType{{}},
+				FunctionSection: []Index{0},
+				CodeSection: []*Code{
+					{Body: []byte{OpcodeLocalGet, 0, OpcodeLocalGet, 1, OpcodeI32Add, OpcodeEnd}},
+				},
+				ExportSection: map[string]*Export{
+					"AddInt": {Name: "AddInt", Kind: ExportKindFunc, Index: Index(0)},
+				},
+				StartSection: &zero,
+			},
+			expected: map[string]uint32{"code": 1, "export": 1, "function": 1, "start": 1, "type": 1},
+		},
+		{
+			name: "memory and data",
+			input: &Module{
+				MemorySection: []*MemoryType{{Min: 1}},
+				DataSection:   []*DataSegment{{MemoryIndex: 0, OffsetExpression: empty}},
+			},
+			expected: map[string]uint32{"data": 1, "memory": 1},
+		},
+		{
+			name: "table and element",
+			input: &Module{
+				TableSection:   []*TableType{{ElemType: 0x70, Limit: &LimitsType{Min: 1}}},
+				ElementSection: []*ElementSegment{{TableIndex: 0, OffsetExpr: empty}},
+			},
+			expected: map[string]uint32{"element": 1, "table": 1},
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			actual := map[string]uint32{}
+			for i := SectionID(0); i <= SectionIDData; i++ {
+				if size := tc.input.SectionSize(i); size > 0 {
+					actual[SectionIDName(i)] = size
+				}
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}