Repository: tetratelabs/wazero
Pull request #137 (2022-01-17 01:53:46)

Finishes binary custom name section codec
This completes the codec for the binary name section, leaving some TODOs
for a rainy day.


codefromthecrypt: Thanks for the look!

Patch:
diff --git a/wasm/names.go b/wasm/names.go
index fbc6a39e..99cd1601 100644
--- a/wasm/names.go
+++ b/wasm/names.go
@@ -14,8 +14,10 @@ import (
 // See https://github.com/tetratelabs/wazero/issues/134 about adding this to Module
 type CustomNameSection struct {
 	// ModuleName is the symbolic identifier for a module. Ex. math
+	// The corresponding subsection is subsectionIDModuleName.
 	ModuleName string
 	// FunctionNames is an association of a function index to its symbolic identifier. Ex. add
+	// The corresponding subsection is subsectionIDFunctionNames.
 	//
 	// * the key (idx) is in the function namespace, where module defined functions are preceded by imported ones.
 	//
@@ -27,6 +29,7 @@ type CustomNameSection struct {
 	FunctionNames map[uint32]string
 
 	// LocalNames is an association of a function index to any locals which have a symbolic identifier. Ex. add x
+	// The corresponding subsection is subsectionIDLocalNames.
 	//
 	// * the key (funcIndex) is in the function namespace, where module defined functions are preceded by imported ones.
 	// * the second key (idx) is in the local namespace, where parameters precede any function locals.
@@ -39,24 +42,30 @@ type CustomNameSection struct {
 	LocalNames map[uint32]map[uint32]string
 }
 
+const (
+	// subsectionIDModuleName contains only the module name.
+	subsectionIDModuleName = uint8(0)
+	// subsectionIDFunctionNames is a map of indices to function names, in ascending order by function index
+	subsectionIDFunctionNames = uint8(1)
+	// subsectionIDLocalNames contain a map of function indices to a map of local indices to their names, in ascending
+	// order by function and local index
+	subsectionIDLocalNames = uint8(2)
+)
+
 // EncodeData serializes the data associated with the "name" key in SectionIDCustom according to the standard:
 //
-// * ModuleName encode as subsection 0
-// * FunctionNames encode as subsection 1 in ascending order by function index
-// * LocalNames encode as subsection 2 in ascending order by function and local index
-//
 // Note: The result can be nil because this does not encode empty subsections
 //
 // See https://www.w3.org/TR/wasm-core-1/#binary-namesec
 func (n *CustomNameSection) EncodeData() (data []byte) {
 	if n.ModuleName != "" {
-		data = append(data, encodeNameSubsection(uint8(0), encodeSizePrefixed([]byte(n.ModuleName)))...)
+		data = append(data, encodeNameSubsection(subsectionIDModuleName, encodeSizePrefixed([]byte(n.ModuleName)))...)
 	}
 	if fd := n.encodeFunctionNameData(); len(fd) > 0 {
-		data = append(data, encodeNameSubsection(uint8(1), fd)...)
+		data = append(data, encodeNameSubsection(subsectionIDFunctionNames, fd)...)
 	}
 	if ld := n.encodeLocalNameData(); len(ld) > 0 {
-		data = append(data, encodeNameSubsection(uint8(2), ld)...)
+		data = append(data, encodeNameSubsection(subsectionIDLocalNames, ld)...)
 	}
 	return
 }
@@ -141,54 +150,133 @@ func encodeSizePrefixed(data []byte) []byte {
 // * LocalNames decode from subsection 2
 //
 // See https://www.w3.org/TR/wasm-core-1/#binary-namesec
-func DecodeCustomNameSection(data []byte) (*CustomNameSection, error) {
+func DecodeCustomNameSection(data []byte) (result *CustomNameSection, err error) {
+	// TODO: add leb128 functions that work on []byte and offset. While using a reader allows us to reuse reader-based
+	// leb128 functions, it is less efficient, causes untestable code and in some cases more complex vs plain []byte.
 	r := bytes.NewReader(data)
+	result = &CustomNameSection{}
+
+	// subsectionID is decoded if known, and skipped if not
+	var subsectionID uint8
+	// subsectionSize is the length to skip when the subsectionID is unknown
+	var subsectionSize uint32
 	for {
-		id, err := r.ReadByte()
-		if err != nil {
-			return nil, fmt.Errorf("failed to read subsection ID: %w", err)
+		if subsectionID, err = r.ReadByte(); err != nil {
+			if err == io.EOF {
+				return result, nil
+			}
+			// TODO: untestable as this can't fail for a reason beside EOF reading a byte from a buffer
+			return nil, fmt.Errorf("failed to read a subsection ID: %w", err)
 		}
 
-		size, _, err := leb128.DecodeUint32(r)
-		if err != nil {
-			return nil, fmt.Errorf("failed to read the size of subsection %d: %w", id, err)
+		// TODO: unused except when skipping. This means we can pass on a corrupt length of a known subsection
+		if subsectionSize, _, err = leb128.DecodeUint32(r); err != nil {
+			return nil, fmt.Errorf("failed to read the size of subsection[%d]: %w", subsectionID, err)
 		}
 
-		if id == 1 { // the function name subsection.
-			break
-		} else { // TODO: ModuleName and LocalNames!
-			// Skip other subsections.
-			_, err := r.Seek(int64(size), io.SeekCurrent)
-			if err != nil {
-				return nil, fmt.Errorf("failed to skip subsection %d: %w", id, err)
+		switch subsectionID {
+		case subsectionIDModuleName:
+			if result.ModuleName, err = decodeName(r, "module"); err != nil {
+				return nil, err
+			}
+		case subsectionIDFunctionNames:
+			if result.FunctionNames, err = decodeFunctionNames(r); err != nil {
+				return nil, err
+			}
+		case subsectionIDLocalNames:
+			if result.LocalNames, err = decodeLocalNames(r); err != nil {
+				return nil, err
+			}
+		default: // Skip other subsections.
+			// 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(subsectionSize)); err != nil {
+				return nil, fmt.Errorf("failed to skip subsection[%d]: %w", subsectionID, err)
 			}
 		}
 	}
+}
 
-	nameVectorSize, _, err := leb128.DecodeUint32(r)
+func decodeFunctionNames(r *bytes.Reader) (map[uint32]string, error) {
+	functionCount, err := decodeFunctionCount(r, subsectionIDFunctionNames)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read the size of name vector: %w", err)
+		return nil, err
 	}
 
-	funcNames := make(map[uint32]string, nameVectorSize)
-	for i := uint32(0); i < nameVectorSize; i++ {
-		functionIndex, _, err := leb128.DecodeUint32(r)
+	result := make(map[uint32]string, functionCount)
+	for i := uint32(0); i < functionCount; i++ {
+		functionIndex, err := decodeFunctionIndex(r, subsectionIDFunctionNames)
 		if err != nil {
-			return nil, fmt.Errorf("failed to read function index: %w", err)
+			return nil, err
 		}
 
-		functionNameSize, _, err := leb128.DecodeUint32(r)
+		if result[functionIndex], err = decodeName(r, "function[%d]", functionIndex); err != nil {
+			return nil, err
+		}
+	}
+	return result, nil
+}
+
+func decodeLocalNames(r *bytes.Reader) (map[uint32]map[uint32]string, error) {
+	functionCount, err := decodeFunctionCount(r, subsectionIDLocalNames)
+	if err != nil {
+		return nil, err
+	}
+
+	result := make(map[uint32]map[uint32]string, functionCount)
+	for i := uint32(0); i < functionCount; i++ {
+		functionIndex, err := decodeFunctionIndex(r, subsectionIDLocalNames)
 		if err != nil {
-			return nil, fmt.Errorf("failed to read function name size: %w", err)
+			return nil, err
 		}
 
-		namebuf := make([]byte, functionNameSize)
-		_, err = io.ReadFull(r, namebuf)
+		localCount, _, err := leb128.DecodeUint32(r)
 		if err != nil {
-			return nil, fmt.Errorf("failed to read function name: %w", err)
+			return nil, fmt.Errorf("failed to read the local count for function[%d]: %w", functionIndex, err)
 		}
-		funcNames[functionIndex] = string(namebuf)
+
+		locals := make(map[uint32]string, localCount)
+		for i := uint32(0); i < localCount; i++ {
+			localIndex, _, err := leb128.DecodeUint32(r)
+			if err != nil {
+				return nil, fmt.Errorf("failed to read a local index of function[%d]: %w", functionIndex, err)
+			}
+			locals[localIndex], err = decodeName(r, "function[%d] local[%d]", functionIndex, localIndex)
+			if err != nil {
+				return nil, err
+			}
+		}
+		result[functionIndex] = locals
 	}
+	return result, nil
+}
 
-	return &CustomNameSection{FunctionNames: funcNames}, nil
+func decodeFunctionIndex(r *bytes.Reader, subsectionID uint8) (uint32, error) {
+	functionIndex, _, err := leb128.DecodeUint32(r)
+	if err != nil {
+		return 0, fmt.Errorf("failed to read a function index in subsection[%d]: %w", subsectionID, err)
+	}
+	return functionIndex, nil
+}
+
+func decodeFunctionCount(r *bytes.Reader, subsectionID uint8) (uint32, error) {
+	functionCount, _, err := leb128.DecodeUint32(r)
+	if err != nil {
+		return 0, fmt.Errorf("failed to read the function count of subsection[%d]: %w", subsectionID, err)
+	}
+	return functionCount, nil
+}
+
+// decodeName decodes a size prefixed string from the reader.
+// contextFormat and contextArgs apply an error format when present
+func decodeName(r *bytes.Reader, contextFormat string, contextArgs ...interface{}) (string, error) {
+	size, _, err := leb128.DecodeUint32(r)
+	if err != nil {
+		return "", fmt.Errorf("failed to read %s name size: %w", fmt.Sprintf(contextFormat, contextArgs...), err)
+	}
+
+	namebuf := make([]byte, size)
+	if _, err = io.ReadFull(r, namebuf); err != nil {
+		return "", fmt.Errorf("failed to read %s name: %w", fmt.Sprintf(contextFormat, contextArgs...), err)
+	}
+	return string(namebuf), nil
 }
diff --git a/wasm/section.go b/wasm/section.go
index c20a56ff..eba7a1ef 100644
--- a/wasm/section.go
+++ b/wasm/section.go
@@ -11,6 +11,7 @@ import (
 type SectionID = byte
 
 const (
+	// SectionIDCustom includes the standard defined CustomNameSection and possibly others not defined in the standard.
 	SectionIDCustom   SectionID = 0
 	SectionIDType     SectionID = 1
 	SectionIDImport   SectionID = 2
@@ -30,9 +31,6 @@ const (
 // See https://www.w3.org/TR/wasm-core-1/#modules%E2%91%A0%E2%93%AA
 func (m *Module) encodeSections(buffer []byte) (bytes []byte) {
 	bytes = buffer
-	for name, data := range m.CustomSections {
-		bytes = append(bytes, encodeCustomSection(name, data)...)
-	}
 	if len(m.TypeSection) > 0 {
 		panic("TODO: TypeSection")
 	}
@@ -66,6 +64,13 @@ func (m *Module) encodeSections(buffer []byte) (bytes []byte) {
 	if len(m.DataSection) > 0 {
 		panic("TODO: DataSection")
 	}
+
+	// We encode custom sections after data as that ensures the correct order in the only section where order matters:
+	// >> 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
+	for name, data := range m.CustomSections {
+		bytes = append(bytes, encodeCustomSection(name, data)...)
+	}
 	return
 }
 


Test patch:
diff --git a/wasm/names_test.go b/wasm/names_test.go
index d13a2c47..0f140ad0 100644
--- a/wasm/names_test.go
+++ b/wasm/names_test.go
@@ -14,15 +14,14 @@ func TestCustomNameSection_EncodeData(t *testing.T) {
 	}{
 		{
 			name:  "empty",
-			input: &CustomNameSection{FunctionNames: map[uint32]string{}, LocalNames: map[uint32]map[uint32]string{}},
+			input: &CustomNameSection{},
 		},
 		{
 			name: "only module",
 			// Ex. (module $simple )
 			input: &CustomNameSection{ModuleName: "simple"},
 			expected: []byte{
-				0x00, // Module subsection ID zero
-				0x07, // 7 bytes to follow
+				subsectionIDModuleName, 0x07, // 7 bytes
 				0x06, // the Module name simple is 6 bytes long
 				's', 'i', 'm', 'p', 'l', 'e',
 			},
@@ -38,12 +37,10 @@ func TestCustomNameSection_EncodeData(t *testing.T) {
 				FunctionNames: map[uint32]string{0x00: "hello"},
 			},
 			expected: []byte{
-				0x00, // Module subsection ID zero
-				0x07, // 7 bytes to follow
+				subsectionIDModuleName, 0x07, // 7 bytes
 				0x06, // the Module name simple is 6 bytes long
 				's', 'i', 'm', 'p', 'l', 'e',
-				0x01, // function subsection ID one
-				0x08, // 8 bytes to follow
+				subsectionIDFunctionNames, 0x08, // 8 bytes
 				0x01, // one function name
 				0x00, // the function index is zero
 				0x05, // the function name hello is 5 bytes long
@@ -63,8 +60,7 @@ func TestCustomNameSection_EncodeData(t *testing.T) {
 				},
 			},
 			expected: []byte{
-				0x01, // function subsection ID one
-				0x2b, // 43 bytes to follow
+				subsectionIDFunctionNames, 0x2b, // 43 bytes
 				0x02, // two function names
 				0x00, // the function index is zero
 				0x16, // the function name runtime.args_sizes_get is 22 bytes long
@@ -91,11 +87,11 @@ func TestCustomNameSection_EncodeData(t *testing.T) {
 				},
 			},
 			expected: []byte{
-				0x01, 0x0b, // function subsection ID one (7 bytes)
+				subsectionIDFunctionNames, 0x0b, // 7 bytes
 				0x02,                      // two function names
 				0x00, 0x03, 'm', 'u', 'l', // index 0, size of "mul", "mul"
 				0x01, 0x03, 'a', 'd', 'd', // index 1, size of "add", "add"
-				0x02, 0x11, // local subsection ID two (17 bytes)
+				subsectionIDLocalNames, 0x11, // 17 bytes
 				0x02,       // two functions
 				0x00, 0x02, // index 0 has 2 locals
 				0x00, 0x01, 'x', // index 0, size of "x", "x"
@@ -165,99 +161,149 @@ func TestEncodeSizePrefixed(t *testing.T) {
 	require.Equal(t, []byte{5, 'h', 'e', 'l', 'l', 'o'}, encodeSizePrefixed([]byte("hello")))
 }
 
+// TestDecodeCustomNameSection relies on unit tests for CustomNameSection.EncodeData, specifically that the encoding is
+// both known and correct. This avoids having to copy/paste or share variables to assert against byte arrays.
 func TestDecodeCustomNameSection(t *testing.T) {
 	tests := []struct {
-		name     string
-		input    []byte
-		expected *CustomNameSection
-	}{
+		name  string
+		input *CustomNameSection // round trip test!
+	}{{
+		name:  "empty",
+		input: &CustomNameSection{},
+	},
 		{
-			name: "one function name",
-			//	(module (import "" "Hello" (func $hello)))
-			input: []byte{
-				0x01, // function subsection ID one
-				0x08, // 8 bytes to follow
-				0x01, // one function name
-				0x00, // the function index is zero
-				0x05, // the function name hello is 5 bytes long
-				'h', 'e', 'l', 'l', 'o',
-			},
-			expected: &CustomNameSection{
+			name:  "only module",
+			input: &CustomNameSection{ModuleName: "simple"},
+		},
+		{
+			name: "module and function name",
+			input: &CustomNameSection{
+				ModuleName:    "simple",
 				FunctionNames: map[uint32]string{0x00: "hello"},
 			},
 		},
 		{
-			name: "two function names", // Ex. TinyGo which at one point didn't set a module name
-			//	(module
-			//		(import "wasi_snapshot_preview1" "args_sizes_get" (func $runtime.args_sizes_get (param i32, i32) (result i32)))
-			//		(import "wasi_snapshot_preview1" "fd_write" (func $runtime.fd_write (param i32, i32, i32, i32) (result i32)))
-			//	)
-			input: []byte{
-				0x01, // function subsection ID one
-				0x2b, // 43 bytes to follow
-				0x02, // two function names
-				0x00, // the function index is zero
-				0x16, // the function name runtime.args_sizes_get is 22 bytes long
-				'r', 'u', 'n', 't', 'i', 'm', 'e', '.', 'a', 'r', 'g', 's', '_', 's', 'i', 'z', 'e', 's', '_', 'g', 'e', 't',
-				0x01, // the function index is zero
-				0x10, // the function name runtime.fd_write is 16 bytes long
-				'r', 'u', 'n', 't', 'i', 'm', 'e', '.', 'f', 'd', '_', 'w', 'r', 'i', 't', 'e',
-			},
-			expected: &CustomNameSection{
+			name: "two function names",
+			input: &CustomNameSection{
 				FunctionNames: map[uint32]string{
 					0x00: "runtime.args_sizes_get",
 					0x01: "runtime.fd_write",
 				},
 			},
 		},
+		{
+			name: "function with local names",
+			input: &CustomNameSection{
+				FunctionNames: map[uint32]string{
+					0x00: "mul",
+					0x01: "add",
+				},
+				LocalNames: map[uint32]map[uint32]string{
+					0x00: {0x00: "x", 0x01: "y"},
+					0x01: {0x00: "l", 0x01: "r"},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			ns, err := DecodeCustomNameSection(tc.input)
+			ns, err := DecodeCustomNameSection(tc.input.EncodeData())
 			require.NoError(t, err)
-			require.Equal(t, tc.expected, ns)
+			require.Equal(t, tc.input, ns)
 		})
 	}
 }
 
 func TestDecodeCustomNameSection_Errors(t *testing.T) {
+	// currently, we ignore the size of known subsections
+	ignoredSubsectionSize := byte(50)
 	tests := []struct {
 		name        string
 		input       []byte
 		expectedErr string
 	}{
 		{
-			name:        "empty",
-			input:       []byte{},
-			expectedErr: "failed to read subsection ID: EOF",
+			name:        "EOF after module name subsection ID",
+			input:       []byte{subsectionIDModuleName},
+			expectedErr: "failed to read the size of subsection[0]: EOF",
+		},
+		{
+			name:        "EOF after function names subsection ID",
+			input:       []byte{subsectionIDFunctionNames},
+			expectedErr: "failed to read the size of subsection[1]: EOF",
+		},
+		{
+			name:        "EOF after local names subsection ID",
+			input:       []byte{subsectionIDLocalNames},
+			expectedErr: "failed to read the size of subsection[2]: EOF",
+		},
+		{
+			name:        "EOF after unknown subsection ID",
+			input:       []byte{4},
+			expectedErr: "failed to read the size of subsection[4]: EOF",
 		},
 		{
-			name:        "EOF after subsection ID",
-			input:       []byte{0x01},
-			expectedErr: "failed to read the size of subsection 1: EOF",
+			name:        "EOF after module name subsection size",
+			input:       []byte{subsectionIDModuleName, ignoredSubsectionSize},
+			expectedErr: "failed to read module name size: EOF",
 		},
 		{
-			name:        "EOF after size",
-			input:       []byte{0x01, 0x10},
-			expectedErr: "failed to read the size of name vector: EOF",
+			name:        "EOF after function names subsection size",
+			input:       []byte{subsectionIDFunctionNames, ignoredSubsectionSize},
+			expectedErr: "failed to read the function count of subsection[1]: EOF",
 		},
 		{
-			name:        "EOF after function name vector size",
-			input:       []byte{0x01, 0x10, 0x01},
-			expectedErr: "failed to read function index: EOF",
+			name:        "EOF after local names subsection size",
+			input:       []byte{subsectionIDLocalNames, ignoredSubsectionSize},
+			expectedErr: "failed to read the function count of subsection[2]: EOF",
+		},
+		{
+			name:        "EOF skipping unknown subsection size",
+			input:       []byte{4, 100},
+			expectedErr: "failed to skip subsection[4]: EOF",
+		},
+		{
+			name:        "EOF after module name size",
+			input:       []byte{subsectionIDModuleName, ignoredSubsectionSize, 5},
+			expectedErr: "failed to read module name: EOF",
+		},
+		{
+			name:        "EOF after function name count",
+			input:       []byte{subsectionIDFunctionNames, ignoredSubsectionSize, 2},
+			expectedErr: "failed to read a function index in subsection[1]: EOF",
+		},
+		{
+			name:        "EOF after local names function count",
+			input:       []byte{subsectionIDLocalNames, ignoredSubsectionSize, 2},
+			expectedErr: "failed to read a function index in subsection[2]: EOF",
 		},
 		{
 			name:        "EOF after function name index",
-			input:       []byte{0x01, 0x10, 0x01, 0x00},
-			expectedErr: "failed to read function name size: EOF",
+			input:       []byte{subsectionIDFunctionNames, ignoredSubsectionSize, 2, 0},
+			expectedErr: "failed to read function[0] name size: EOF",
+		},
+		{
+			name:        "EOF after local names function index",
+			input:       []byte{subsectionIDLocalNames, ignoredSubsectionSize, 2, 0},
+			expectedErr: "failed to read the local count for function[0]: EOF",
 		},
 		{
 			name:        "EOF after function name size",
-			input:       []byte{0x01, 0x10, 0x01, 0x00, 0x01},
-			expectedErr: "failed to read function name: EOF",
+			input:       []byte{subsectionIDFunctionNames, ignoredSubsectionSize, 2, 0, 5},
+			expectedErr: "failed to read function[0] name: EOF",
+		},
+		{
+			name:        "EOF after local names count for a function index",
+			input:       []byte{subsectionIDLocalNames, ignoredSubsectionSize, 2, 0, 2},
+			expectedErr: "failed to read a local index of function[0]: EOF",
+		},
+		{
+			name:        "EOF after local name size",
+			input:       []byte{subsectionIDLocalNames, ignoredSubsectionSize, 2, 0, 2, 1},
+			expectedErr: "failed to read function[0] local[1] name size: EOF",
 		},
 	}