Repository: tetratelabs/wazero
Pull request #583 (2022-05-20 07:57:17)

compiler: fixes calling convention violation.
If the call_indirect's offset register is the same as the calling-convention
reserved register, the function call target address resulted in the wrong 
address and therefore segfault.

Fixes partially #526 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

mathetake: will backfill the regression tests as a toehold! 
mathetake: ok fixed arm64 as well as added unit tests!

Patch:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 10900699..e399180d 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -4244,15 +4244,15 @@ func (c *amd64Compiler) allocateRegister(t registerType) (reg asm.Register, err
 // Pass indexReg == asm.NilRegister to indicate that use addr argument as the source of target function's address.
 // Otherwise, the added code tries to read the function address from the register for indexReg argument.
 //
-// Note: this is the counter part for returnFunction, and see the comments there as well
+// Note: this is the counterpart for returnFunction, and see the comments there as well
 // to understand how the function calls are achieved.
-func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressRegister asm.Register, functype *wasm.FunctionType) error {
+func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, functionAddressRegister asm.Register, functype *wasm.FunctionType) error {
 	// Release all the registers as our calling convention requires the caller-save.
 	c.compileReleaseAllRegistersToStack()
 
 	// First, we have to make sure that
-	if !isNilRegister(codeAddressRegister) {
-		c.locationStack.markRegisterUsed(codeAddressRegister)
+	if !isNilRegister(functionAddressRegister) {
+		c.locationStack.markRegisterUsed(functionAddressRegister)
 	}
 
 	// Obtain the temporary registers to be used in the followings.
@@ -4264,7 +4264,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 	c.locationStack.markRegisterUsed(freeRegs...)
 
 	// Alias these free tmp registers for readability.
-	callFrameStackPointerRegister, tmpRegister, targetcodeAddressRegister,
+	callFrameStackPointerRegister, tmpRegister, targetFunctionAddressRegister,
 		callFrameStackTopAddressRegister := freeRegs[0], freeRegs[1], freeRegs[2], freeRegs[3]
 
 	// First, we read the current call frame stack pointer.
@@ -4280,10 +4280,10 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 	jmpIfNotCallFrameStackNeedsGrow := c.assembler.CompileJump(amd64.JNE)
 
 	// Otherwise, we have to make the builtin function call to grow the call frame stack.
-	if !isNilRegister(codeAddressRegister) {
+	if !isNilRegister(functionAddressRegister) {
 		// If we need to get the target funcaddr from register (call_indirect case), we must save it before growing the
 		// call-frame stack, as the register is not saved across function calls.
-		savedOffsetLocation := c.pushRuntimeValueLocationOnRegister(codeAddressRegister, runtimeValueTypeI64)
+		savedOffsetLocation := c.pushRuntimeValueLocationOnRegister(functionAddressRegister, runtimeValueTypeI64)
 		c.compileReleaseRegisterToStack(savedOffsetLocation)
 	}
 
@@ -4293,13 +4293,13 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 	}
 
 	// For call_indirect, we need to push the value back to the register.
-	if !isNilRegister(codeAddressRegister) {
+	if !isNilRegister(functionAddressRegister) {
 		// Since this is right after callGoFunction, we have to initialize the stack base pointer
 		// to properly load the value on memory stack.
 		c.compileReservedStackBasePointerInitialization()
 
 		savedOffsetLocation := c.locationStack.pop()
-		savedOffsetLocation.setRegister(codeAddressRegister)
+		savedOffsetLocation.setRegister(functionAddressRegister)
 		c.compileLoadValueOnStackToRegister(savedOffsetLocation)
 	}
 
@@ -4377,7 +4377,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 
 	// 3) Set rc.next to specify which function is executed on the current call frame (needs to make builtin function calls).
 	{
-		if isNilRegister(codeAddressRegister) {
+		if isNilRegister(functionAddressRegister) {
 			// We must set the target function's address(pointer) of *code into the next call-frame stack.
 			// In the example, this is equivalent to writing the value into "rc.next".
 			//
@@ -4392,17 +4392,18 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 				// Note: FunctionIndex is limited up to 2^27 so this offset never exceeds 32-bit integer.
 				// *8 because the size of *code equals 8 bytes.
 				tmpRegister, int64(index)*8,
-				targetcodeAddressRegister,
+				targetFunctionAddressRegister,
 			)
 		} else {
-			targetcodeAddressRegister = codeAddressRegister
+			targetFunctionAddressRegister = functionAddressRegister
 		}
 		// Finally, we are ready to place the address of the target function's *code into the new call-frame.
 		// In the example, this is equivalent to set "rc.next".
-		c.assembler.CompileRegisterToMemory(amd64.MOVQ, targetcodeAddressRegister, callFrameStackTopAddressRegister, callFrameFunctionOffset)
+		c.assembler.CompileRegisterToMemory(amd64.MOVQ, targetFunctionAddressRegister,
+			callFrameStackTopAddressRegister, callFrameFunctionOffset)
 	}
 
-	// 4) Set ra.1 so that we can return back to this function properly.
+	// 4) Set ra.1 so that we can return to this function properly.
 	//
 	// We have to set the return address for the current call frame (which is "ra.1" in the example).
 	// First, Get the return address into the tmpRegister.
@@ -4420,14 +4421,22 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 	// So we increment the call frame stack pointer.
 	c.assembler.CompileNoneToMemory(amd64.INCQ, amd64ReservedRegisterForCallEngine, callEngineGlobalContextCallFrameStackPointerOffset)
 
+	if amd64CallingConventionModuleInstanceAddressRegister == targetFunctionAddressRegister {
+		// This case we must move the value on targetFunctionAddressRegister to another register, otherwise
+		// the address (jump target below) will be modified and result in segfault.
+		// See #526.
+		c.assembler.CompileRegisterToRegister(amd64.MOVQ, targetFunctionAddressRegister, tmpRegister)
+		targetFunctionAddressRegister = tmpRegister
+	}
+
 	// Also, we have to put the target function's *wasm.ModuleInstance into amd64CallingConventionModuleInstanceAddressRegister.
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ, targetcodeAddressRegister, functionModuleInstanceAddressOffset,
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ, targetFunctionAddressRegister, functionModuleInstanceAddressOffset,
 		amd64CallingConventionModuleInstanceAddressRegister)
 
 	// And jump into the initial address of the target function.
-	c.assembler.CompileJumpToMemory(amd64.JMP, targetcodeAddressRegister, functionCodeInitialAddressOffset)
+	c.assembler.CompileJumpToMemory(amd64.JMP, targetFunctionAddressRegister, functionCodeInitialAddressOffset)
 
-	// All the registers used are temporary so we mark them unused.
+	// All the registers used are temporary, so we mark them unused.
 	c.locationStack.markRegisterUnused(freeRegs...)
 
 	// On the function return, we have to initialize the state.
@@ -4447,11 +4456,11 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, codeAddressReg
 }
 
 // returnFunction adds instructions to return from the current callframe back to the caller's frame.
-// If this is the current one is the origin, we return back to the callEngine.execWasmFunction with the Returned status.
+// If this is the current one is the origin, we return to the callEngine.execWasmFunction with the Returned status.
 // Otherwise, we jump into the callers' return address stored in callFrame.returnAddress while setting
 // up all the necessary change on the callEngine's state.
 //
-// Note: this is the counter part for callFunction, and see the comments there as well
+// Note: this is the counterpart for callFunction, and see the comments there as well
 // to understand how the function calls are achieved.
 func (c *amd64Compiler) compileReturnFunction() error {
 	// Release all the registers as our calling convention requires the caller-save.
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index da87a863..75c3f7f4 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -907,7 +907,7 @@ func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
 }
 
 // compileCallImpl implements compiler.compileCall and compiler.compileCallIndirect for the arm64 architecture.
-func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister asm.Register, functype *wasm.FunctionType) error {
+func (c *arm64Compiler) compileCallImpl(index wasm.Index, targetFunctionAddressRegister asm.Register, functype *wasm.FunctionType) error {
 	// Release all the registers as our calling convention requires the caller-save.
 	if err := c.compileReleaseAllRegistersToStack(); err != nil {
 		return err
@@ -920,7 +920,7 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister as
 	c.markRegisterUsed(freeRegisters...)
 
 	// Alias for readability.
-	callFrameStackPointerRegister, callFrameStackTopAddressRegister, codeRegister, oldStackBasePointer,
+	callFrameStackPointerRegister, callFrameStackTopAddressRegister, targetFunctionRegister, oldStackBasePointer,
 		tmp := freeRegisters[0], freeRegisters[1], freeRegisters[2], freeRegisters[3], freeRegisters[4]
 
 	// First, we have to check if we need to grow the callFrame stack.
@@ -941,10 +941,10 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister as
 
 	// If these values equal, we need to grow the callFrame stack.
 	// For call_indirect, we need to push the value back to the register.
-	if !isNilRegister(codeAddressRegister) {
+	if !isNilRegister(targetFunctionAddressRegister) {
 		// If we need to get the target funcaddr from register (call_indirect case), we must save it before growing the
 		// call-frame stack, as the register is not saved across function calls.
-		savedOffsetLocation := c.pushRuntimeValueLocationOnRegister(codeAddressRegister, runtimeValueTypeI64)
+		savedOffsetLocation := c.pushRuntimeValueLocationOnRegister(targetFunctionAddressRegister, runtimeValueTypeI64)
 		c.compileReleaseRegisterToStack(savedOffsetLocation)
 	}
 
@@ -953,13 +953,13 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister as
 	}
 
 	// For call_indirect, we need to push the value back to the register.
-	if !isNilRegister(codeAddressRegister) {
+	if !isNilRegister(targetFunctionAddressRegister) {
 		// Since this is right after callGoFunction, we have to initialize the stack base pointer
 		// to properly load the value on memory stack.
 		c.compileReservedStackBasePointerRegisterInitialization()
 
 		savedOffsetLocation := c.locationStack.pop()
-		savedOffsetLocation.setRegister(codeAddressRegister)
+		savedOffsetLocation.setRegister(targetFunctionAddressRegister)
 		c.compileLoadValueOnStackToRegister(savedOffsetLocation)
 	}
 
@@ -1025,18 +1025,18 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister as
 
 	// Next, read the index of the target function (= &ce.functions[offset])
 	// into codeIndexRegister.
-	if isNilRegister(codeAddressRegister) {
+	if isNilRegister(targetFunctionAddressRegister) {
 		c.assembler.CompileMemoryToRegister(
 			arm64.MOVD,
 			tmp, int64(index)*8, // * 8 because the size of *function equals 8 bytes.
-			codeRegister)
+			targetFunctionRegister)
 	} else {
-		codeRegister = codeAddressRegister
+		targetFunctionRegister = targetFunctionAddressRegister
 	}
 
 	// Finally, we are ready to write the address of the target function's *function into the new call-frame.
 	c.assembler.CompileRegisterToMemory(arm64.MOVD,
-		codeRegister,
+		targetFunctionRegister,
 		callFrameStackTopAddressRegister, callFrameFunctionOffset)
 
 	// 4) Set ra.current so that we can return back to this function properly.
@@ -1059,15 +1059,23 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, codeAddressRegister as
 		tmp,
 		arm64ReservedRegisterForCallEngine, callEngineGlobalContextCallFrameStackPointerOffset)
 
-	// Also, we have to put the code's moduleinstance address into arm64CallingConventionModuleInstanceAddressRegister.
+	if targetFunctionRegister == arm64CallingConventionModuleInstanceAddressRegister {
+		// This case we must move the value on targetFunctionAddressRegister to another register, otherwise
+		// the address (jump target below) will be modified and result in segfault.
+		// See #526.
+		c.assembler.CompileRegisterToRegister(arm64.MOVD, targetFunctionRegister, tmp)
+		targetFunctionRegister = tmp
+	}
+
+	// Also, we have to put the code's moduleInstance address into arm64CallingConventionModuleInstanceAddressRegister.
 	c.assembler.CompileMemoryToRegister(arm64.MOVD,
-		codeRegister, functionModuleInstanceAddressOffset,
+		targetFunctionRegister, functionModuleInstanceAddressOffset,
 		arm64CallingConventionModuleInstanceAddressRegister,
 	)
 
 	// Then, br into the target function's initial address.
 	c.assembler.CompileMemoryToRegister(arm64.MOVD,
-		codeRegister, functionCodeInitialAddressOffset,
+		targetFunctionRegister, functionCodeInitialAddressOffset,
 		tmp)
 
 	c.assembler.CompileJumpToMemory(arm64.B, tmp)
diff --git a/internal/wazeroir/compiler.go b/internal/wazeroir/compiler.go
index 28ec3b44..401a4624 100644
--- a/internal/wazeroir/compiler.go
+++ b/internal/wazeroir/compiler.go
@@ -228,13 +228,13 @@ func CompileFunctions(_ context.Context, enabledFeatures wasm.Features, module *
 	}
 
 	var ret []*CompilationResult
-	for funcInxdex := range module.FunctionSection {
-		typeID := module.FunctionSection[funcInxdex]
+	for funcIndex := range module.FunctionSection {
+		typeID := module.FunctionSection[funcIndex]
 		sig := module.TypeSection[typeID]
-		code := module.CodeSection[funcInxdex]
+		code := module.CodeSection[funcIndex]
 		r, err := compile(enabledFeatures, sig, code.Body, code.LocalTypes, module.TypeSection, functions, globals)
 		if err != nil {
-			return nil, fmt.Errorf("failed to lower func[%d/%d] to wazeroir: %w", funcInxdex, len(functions)-1, err)
+			return nil, fmt.Errorf("failed to lower func[%d/%d] to wazeroir: %w", funcIndex, len(functions)-1, err)
 		}
 		r.Globals = globals
 		r.Functions = functions


Test patch:
diff --git a/internal/engine/compiler/impl_amd64_test.go b/internal/engine/compiler/impl_amd64_test.go
index 70560ff5..0478c94a 100644
--- a/internal/engine/compiler/impl_amd64_test.go
+++ b/internal/engine/compiler/impl_amd64_test.go
@@ -2,13 +2,73 @@ package compiler
 
 import (
 	"testing"
+	"unsafe"
 
 	"github.com/tetratelabs/wazero/internal/asm"
-	amd64 "github.com/tetratelabs/wazero/internal/asm/amd64"
+	"github.com/tetratelabs/wazero/internal/asm/amd64"
 	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
 
+// TestAmd64Compiler_indirectCallWithTargetOnCallingConvReg is the regression test for #526.
+// In short, the offset register for call_indirect might be the same as amd64CallingConventionModuleInstanceAddressRegister
+// and that must not be a failure.
+func TestAmd64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
+	env := newCompilerEnvironment()
+	table := make([]wasm.Reference, 1)
+	env.addTable(&wasm.TableInstance{References: table})
+	// Ensure that the module instance has the type information for targetOperation.TypeIndex,
+	// and the typeID  matches the table[targetOffset]'s type ID.
+	operation := &wazeroir.OperationCallIndirect{TypeIndex: 0}
+	env.module().TypeIDs = []wasm.FunctionTypeID{0}
+	env.module().Engine = &moduleEngine{functions: []*function{}}
+
+	me := env.moduleEngine()
+	{ // Compiling call target.
+		compiler := env.requireNewCompiler(t, newCompiler, nil)
+		err := compiler.compilePreamble()
+		require.NoError(t, err)
+		err = compiler.compileReturnFunction()
+		require.NoError(t, err)
+
+		c, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		f := &function{
+			parent:                &code{codeSegment: c},
+			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
+			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
+			source:                &wasm.FunctionInstance{TypeID: 0},
+		}
+		me.functions = append(me.functions, f)
+		table[0] = uintptr(unsafe.Pointer(f))
+	}
+
+	compiler := env.requireNewCompiler(t, newCompiler, &wazeroir.CompilationResult{
+		Signature: &wasm.FunctionType{},
+		Types:     []*wasm.FunctionType{{}},
+		HasTable:  true,
+	}).(*amd64Compiler)
+	err := compiler.compilePreamble()
+	require.NoError(t, err)
+
+	// Place the offset into the calling-convention reserved register.
+	offsetLoc := compiler.pushRuntimeValueLocationOnRegister(amd64CallingConventionModuleInstanceAddressRegister,
+		runtimeValueTypeI32)
+	compiler.assembler.CompileConstToRegister(amd64.MOVQ, 0, offsetLoc.register)
+
+	require.NoError(t, compiler.compileCallIndirect(operation))
+
+	err = compiler.compileReturnFunction()
+	require.NoError(t, err)
+
+	// Generate the code under test and run.
+	code, _, _, err := compiler.compile()
+	require.NoError(t, err)
+	env.exec(code)
+}
+
 func TestAmd64Compiler_compile_Mul_Div_Rem(t *testing.T) {
 	for _, kind := range []wazeroir.OperationKind{
 		wazeroir.OperationKindMul,
diff --git a/internal/engine/compiler/impl_arm64_test.go b/internal/engine/compiler/impl_arm64_test.go
index e6dbe938..03361f77 100644
--- a/internal/engine/compiler/impl_arm64_test.go
+++ b/internal/engine/compiler/impl_arm64_test.go
@@ -2,11 +2,72 @@ package compiler
 
 import (
 	"testing"
+	"unsafe"
 
 	arm64 "github.com/tetratelabs/wazero/internal/asm/arm64"
 	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasm"
+	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
 
+// TestArm64Compiler_indirectCallWithTargetOnCallingConvReg is the regression test for #526.
+// In short, the offset register for call_indirect might be the same as arm64CallingConventionModuleInstanceAddressRegister
+// and that must not be a failure.
+func TestArm64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
+	env := newCompilerEnvironment()
+	table := make([]wasm.Reference, 1)
+	env.addTable(&wasm.TableInstance{References: table})
+	// Ensure that the module instance has the type information for targetOperation.TypeIndex,
+	// and the typeID  matches the table[targetOffset]'s type ID.
+	operation := &wazeroir.OperationCallIndirect{TypeIndex: 0}
+	env.module().TypeIDs = []wasm.FunctionTypeID{0}
+	env.module().Engine = &moduleEngine{functions: []*function{}}
+
+	me := env.moduleEngine()
+	{ // Compiling call target.
+		compiler := env.requireNewCompiler(t, newCompiler, nil)
+		err := compiler.compilePreamble()
+		require.NoError(t, err)
+		err = compiler.compileReturnFunction()
+		require.NoError(t, err)
+
+		c, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		f := &function{
+			parent:                &code{codeSegment: c},
+			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
+			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
+			source:                &wasm.FunctionInstance{TypeID: 0},
+		}
+		me.functions = append(me.functions, f)
+		table[0] = uintptr(unsafe.Pointer(f))
+	}
+
+	compiler := env.requireNewCompiler(t, newCompiler, &wazeroir.CompilationResult{
+		Signature: &wasm.FunctionType{},
+		Types:     []*wasm.FunctionType{{}},
+		HasTable:  true,
+	}).(*arm64Compiler)
+	err := compiler.compilePreamble()
+	require.NoError(t, err)
+
+	// Place the offset into the calling-convention reserved register.
+	offsetLoc := compiler.pushRuntimeValueLocationOnRegister(arm64CallingConventionModuleInstanceAddressRegister,
+		runtimeValueTypeI32)
+	compiler.assembler.CompileConstToRegister(arm64.MOVD, 0, offsetLoc.register)
+
+	require.NoError(t, compiler.compileCallIndirect(operation))
+
+	err = compiler.compileReturnFunction()
+	require.NoError(t, err)
+
+	// Generate the code under test and run.
+	code, _, _, err := compiler.compile()
+	require.NoError(t, err)
+	env.exec(code)
+}
+
 func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 	env := newCompilerEnvironment()
 	compiler := env.requireNewCompiler(t, newArm64Compiler, nil).(*arm64Compiler)
diff --git a/internal/integration_test/asm/amd64_debug/impl_test.go b/internal/integration_test/asm/amd64_debug/impl_test.go
index 99ac373a..597fbf1a 100644
--- a/internal/integration_test/asm/amd64_debug/impl_test.go
+++ b/internal/integration_test/asm/amd64_debug/impl_test.go
@@ -8,7 +8,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/asm"
-	amd64 "github.com/tetratelabs/wazero/internal/asm/amd64"
+	"github.com/tetratelabs/wazero/internal/asm/amd64"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )