Repository: tetratelabs/wazero
Pull request #702 (2022-07-15 02:25:05)

validation: correct expected types on loop with params
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/wasm/func_validation.go b/internal/wasm/func_validation.go
index 60616d56..390073d6 100644
--- a/internal/wasm/func_validation.go
+++ b/internal/wasm/func_validation.go
@@ -393,11 +393,11 @@ func (m *Module) validateFunctionWithMaxStackValues(
 			pc += num - 1
 			// Check type soundness.
 			target := controlBlockStack[len(controlBlockStack)-int(index)-1]
-			targetResultType := target.blockType.Results
+			var targetResultType []ValueType
 			if target.op == OpcodeLoop {
-				// Loop operation doesn't require results since the continuation is
-				// the beginning of the loop.
-				targetResultType = []ValueType{}
+				targetResultType = target.blockType.Params
+			} else {
+				targetResultType = target.blockType.Results
 			}
 			if err = valueTypeStack.popResults(op, targetResultType, false); err != nil {
 				return err
@@ -411,7 +411,7 @@ func (m *Module) validateFunctionWithMaxStackValues(
 				return fmt.Errorf("read immediate: %v", err)
 			} else if int(index) >= len(controlBlockStack) {
 				return fmt.Errorf(
-					"invalid ln param given for %s: index=%d with %d for the current lable stack length",
+					"invalid ln param given for %s: index=%d with %d for the current label stack length",
 					OpcodeBrIfName, index, len(controlBlockStack))
 			}
 			pc += num - 1
@@ -420,11 +420,11 @@ func (m *Module) validateFunctionWithMaxStackValues(
 			}
 			// Check type soundness.
 			target := controlBlockStack[len(controlBlockStack)-int(index)-1]
-			targetResultType := target.blockType.Results
+			var targetResultType []ValueType
 			if target.op == OpcodeLoop {
-				// Loop operation doesn't require results since the continuation is
-				// the beginning of the loop.
-				targetResultType = []ValueType{}
+				targetResultType = target.blockType.Params
+			} else {
+				targetResultType = target.blockType.Results
 			}
 			if err := valueTypeStack.popResults(op, targetResultType, false); err != nil {
 				return err
@@ -455,7 +455,7 @@ func (m *Module) validateFunctionWithMaxStackValues(
 				return fmt.Errorf("read immediate: %w", err)
 			} else if int(ln) >= len(controlBlockStack) {
 				return fmt.Errorf(
-					"invalid ln param given for %s: ln=%d with %d for the current lable stack length",
+					"invalid ln param given for %s: ln=%d with %d for the current label stack length",
 					OpcodeBrTableName, ln, len(controlBlockStack))
 			}
 			pc += n + num - 1
@@ -464,36 +464,38 @@ func (m *Module) validateFunctionWithMaxStackValues(
 				return fmt.Errorf("cannot pop the required operand for %s", OpcodeBrTableName)
 			}
 			lnLabel := controlBlockStack[len(controlBlockStack)-1-int(ln)]
-			var expTypes []ValueType
+			var defaultLabelType []ValueType
 			if lnLabel.op != OpcodeLoop { // Loop operation doesn't require results since the continuation is the beginning of the loop.
-				expTypes = make([]ValueType, len(lnLabel.blockType.Results))
+				defaultLabelType = make([]ValueType, len(lnLabel.blockType.Results))
 				// Below, we might modify the slice in case of unreachable. Therefore,
 				// we have to copy the content of block result types, otherwise the original
 				// function type might result in invalid value types if the block is the outermost label
 				// which equals the function's type.
-				copy(expTypes, lnLabel.blockType.Results)
+				copy(defaultLabelType, lnLabel.blockType.Results)
+			} else {
+				defaultLabelType = lnLabel.blockType.Params
 			}
 
 			if enabledFeatures.Get(FeatureReferenceTypes) {
 				// As of reference-types proposal, br_table on unreachable state
 				// can choose unknown types for expected parameter types for each label.
 				// https://github.com/WebAssembly/reference-types/pull/116
-				for i := range expTypes {
-					index := len(expTypes) - 1 - i
-					exp := expTypes[index]
+				for i := range defaultLabelType {
+					index := len(defaultLabelType) - 1 - i
+					exp := defaultLabelType[index]
 					actual, err := valueTypeStack.pop()
 					if err != nil {
 						return err
 					}
 					if actual == valueTypeUnknown {
 						// Re-assign the expected type to unknown.
-						expTypes[index] = valueTypeUnknown
+						defaultLabelType[index] = valueTypeUnknown
 					} else if actual != exp {
 						return typeMismatchError(true, OpcodeBrTableName, actual, exp, i)
 					}
 				}
 			} else {
-				if err = valueTypeStack.popResults(op, expTypes, false); err != nil {
+				if err = valueTypeStack.popResults(op, defaultLabelType, false); err != nil {
 					return err
 				}
 			}
@@ -503,17 +505,17 @@ func (m *Module) validateFunctionWithMaxStackValues(
 					return fmt.Errorf("invalid l param given for %s", OpcodeBrTableName)
 				}
 				label := controlBlockStack[len(controlBlockStack)-1-int(l)]
-				expType2 := label.blockType.Results
-				if label.op == OpcodeLoop {
-					// Loop operation doesn't require results since the continuation is
-					// the beginning of the loop.
-					expType2 = []ValueType{}
+				var tableLabelType []ValueType
+				if label.op != OpcodeLoop {
+					tableLabelType = label.blockType.Results
+				} else {
+					tableLabelType = label.blockType.Params
 				}
-				if len(expTypes) != len(expType2) {
-					return fmt.Errorf("incosistent block type length for %s at %d; %v (ln=%d) != %v (l=%d)", OpcodeBrTableName, l, expTypes, ln, expType2, l)
+				if len(defaultLabelType) != len(tableLabelType) {
+					return fmt.Errorf("inconsistent block type length for %s at %d; %v (ln=%d) != %v (l=%d)", OpcodeBrTableName, l, defaultLabelType, ln, tableLabelType, l)
 				}
-				for i := range expTypes {
-					if expTypes[i] != valueTypeUnknown && expTypes[i] != expType2[i] {
+				for i := range defaultLabelType {
+					if defaultLabelType[i] != valueTypeUnknown && defaultLabelType[i] != tableLabelType[i] {
 						return fmt.Errorf("incosistent block type for %s at %d", OpcodeBrTableName, l)
 					}
 				}


Test patch:
diff --git a/internal/wasm/func_validation_test.go b/internal/wasm/func_validation_test.go
index 7dfdfc2e..d2c472e4 100644
--- a/internal/wasm/func_validation_test.go
+++ b/internal/wasm/func_validation_test.go
@@ -3368,7 +3368,7 @@ func TestFuncValidation_UnreachableBrTable_NotModifyTypes(t *testing.T) {
 				// Setting the destination as lalbels of index 0 which
 				// is the function return. Therefore, the expected value types
 				// for each label can be "any".
-				OpcodeBrTable, 0, 0, 0,
+				OpcodeBrTable, 2, 0, 0, 0,
 				OpcodeEnd,
 			}},
 		},
@@ -3381,3 +3381,128 @@ func TestFuncValidation_UnreachableBrTable_NotModifyTypes(t *testing.T) {
 	require.Equal(t, i32, afterValidationResultTypes[0])
 	require.Equal(t, i64, afterValidationResultTypes[1])
 }
+
+func TestModule_funcValidation_loopWithParams(t *testing.T) {
+	tests := []struct {
+		name   string
+		body   []byte
+		expErr string
+	}{
+		{
+			name: "br",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeBr, 0,
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+		},
+		{
+			name: "br_if",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeI32Const, 1, // operand for br_if
+				OpcodeBrIf, 0,
+				OpcodeUnreachable,
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+		},
+		{
+			name: "br_table",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeI32Const, 4, // Operand for br_table.
+				OpcodeBrTable, 2, 0, 0, 0, 0, // Jump into the loop header anyway.
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+		},
+		{
+			name: "br_table - nested",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeI32Const, 4, // Operand for br_table.
+				OpcodeBrTable, 2, 0, 1, 0, // Jump into the loop header anyway.
+				OpcodeEnd,
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+		},
+		{
+			name: "br / mismatch",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				OpcodeDrop,
+				OpcodeBr, 0, // trying to jump the loop head after dropping the value which causes the type mismatch.
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+			expErr: `not enough results in br block
+	have ()
+	want (i32)`,
+		},
+		{
+			name: "br_if / mismatch",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				// Use up the param for the br_if, therefore at the time of jumping, we don't have any param to the header.
+				OpcodeBrIf, 0, // trying to jump the loop head after dropping the value which causes the type mismatch.
+				OpcodeUnreachable,
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+			expErr: `not enough results in br_if block
+	have ()
+	want (i32)`,
+		},
+		{
+			name: "br_table",
+			body: []byte{
+				OpcodeI32Const, 1,
+				OpcodeLoop, 1, // loop (param i32)
+				// Use up the param for the br_table, therefore at the time of jumping, we don't have any param to the header.
+				OpcodeBrTable, 2, 0, 0, 0, // Jump into the loop header anyway.
+				OpcodeEnd,
+				OpcodeUnreachable,
+				OpcodeEnd,
+			},
+			expErr: `not enough results in br_table block
+	have ()
+	want (i32)`,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+		t.Run(tc.name, func(t *testing.T) {
+			m := &Module{
+				TypeSection: []*FunctionType{
+					v_i32,
+					i32_v,
+				},
+				FunctionSection: []Index{0},
+				CodeSection:     []*Code{{Body: tc.body}},
+			}
+			err := m.validateFunction(FeatureMultiValue, 0, []Index{0}, nil, nil, nil, nil)
+			if tc.expErr != "" {
+				require.EqualError(t, err, tc.expErr)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}