compiler: force moduleContext initialization after Go function calls (#988)

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

internal/engine/compiler/engine_test.go

@@ -58,6 +58,11 @@ func TestCompiler_Engine_NewModuleEngine(t *testing.T) {
 	enginetest.RunTestEngine_NewModuleEngine(t, et)
 }
 
+func TestCompiler_MemoryGrowInRecursiveCall(t *testing.T) {
+	defer functionLog.Reset()
+	enginetest.RunTestEngine_MemoryGrowInRecursiveCall(t, et)
+}
+
 func TestCompiler_ModuleEngine_LookupFunction(t *testing.T) {
 	defer functionLog.Reset()
 	enginetest.RunTestModuleEngine_LookupFunction(t, et)

internal/engine/compiler/impl_amd64.go

@@ -186,6 +186,12 @@ func (c *amd64Compiler) compileGoDefinedHostFunction() error {
 
 	// Initializes the reserved stack base pointer which is used to retrieve the call frame stack.
 	c.compileReservedStackBasePointerInitialization()
+
+	// Go function can change the module state in arbitrary way, so we have to force
+	// the callEngine.moduleContext initialization on the function return. To do so,
+	// we zero-out callEngine.moduleInstanceAddress.
+	c.assembler.CompileConstToMemory(amd64.MOVQ,
+		0, amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset)
 	return c.compileReturnFunction()
 }
 

internal/engine/compiler/impl_arm64.go

@@ -398,6 +398,14 @@ func (c *arm64Compiler) compileGoDefinedHostFunction() error {
 
 	// Initializes the reserved stack base pointer which is used to retrieve the call frame stack.
 	c.compileReservedStackBasePointerRegisterInitialization()
+
+	// Go function can change the module state in arbitrary way, so we have to force
+	// the callEngine.moduleContext initialization on the function return. To do so,
+	// we zero-out callEngine.moduleInstanceAddress.
+	c.assembler.CompileRegisterToMemory(arm64.STRD,
+		arm64.RegRZR,
+		arm64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset)
+
 	return c.compileReturnFunction()
 }
 

internal/engine/interpreter/interpreter_test.go

@@ -95,6 +95,11 @@ func (e engineTester) CompiledFunctionPointerValue(me wasm.ModuleEngine, funcInd
 	return uint64(uintptr(unsafe.Pointer(internal.functions[funcIndex])))
 }
 
+func TestInterpreter_MemoryGrowInRecursiveCall(t *testing.T) {
+	defer functionLog.Reset()
+	enginetest.RunTestEngine_MemoryGrowInRecursiveCall(t, et)
+}
+
 func TestInterpreter_Engine_NewModuleEngine(t *testing.T) {
 	enginetest.RunTestEngine_NewModuleEngine(t, et)
 }

internal/gojs/misc_test.go

@@ -45,7 +45,6 @@ func Test_stdio(t *testing.T) {
 }
 
 func Test_stdio_large(t *testing.T) {
-	t.Skip("TODO: #980 memory out of bounds when run with compiler")
 	t.Parallel()
 
 	size := 2 * 1024 * 1024 // 2MB

internal/testing/enginetest/enginetest.go

@@ -55,6 +55,64 @@ type EngineTester interface {
 	CompiledFunctionPointerValue(tme wasm.ModuleEngine, funcIndex wasm.Index) uint64
 }
 
+// RunTestEngine_MemoryGrowInRecursiveCall ensures that it's safe to grow memory in the recursive Wasm calls.
+func RunTestEngine_MemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
+	enabledFeatures := api.CoreFeaturesV1
+	e := et.NewEngine(enabledFeatures)
+	s, ns := wasm.NewStore(enabledFeatures, e)
+
+	const hostModuleName = "env"
+	const hostFnName = "grow_memory"
+	var growFn api.Function
+	hm, err := wasm.NewHostModule(hostModuleName, map[string]interface{}{hostFnName: func() {
+		// Does the recursive call into Wasm, which grows memory.
+		_, err := growFn.Call(context.Background())
+		require.NoError(t, err)
+	}}, map[string]*wasm.HostFuncNames{hostFnName: {}}, enabledFeatures)
+	require.NoError(t, err)
+
+	err = s.Engine.CompileModule(testCtx, hm, nil)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(testCtx, ns, hm, hostModuleName, nil)
+	require.NoError(t, err)
+
+	m := &wasm.Module{
+		TypeSection:     []*wasm.FunctionType{{Params: []wasm.ValueType{}, Results: []wasm.ValueType{}}},
+		FunctionSection: []wasm.Index{0, 0},
+		CodeSection: []*wasm.Code{
+			{
+				Body: []byte{
+					// Calls the imported host function, which in turn calls the next in-Wasm function recursively.
+					wasm.OpcodeCall, 0,
+					// Access the memory and this should succeed as we already had memory grown at this point.
+					wasm.OpcodeI32Const, 0,
+					wasm.OpcodeI32Load, 0x2, 0x0,
+					wasm.OpcodeDrop,
+					wasm.OpcodeEnd,
+				},
+			},
+			{
+				// Grows memory by 1 page.
+				Body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeMemoryGrow, wasm.OpcodeDrop, wasm.OpcodeEnd},
+			},
+		},
+		MemorySection: &wasm.Memory{Max: 1000},
+		ImportSection: []*wasm.Import{{Module: hostModuleName, Name: hostFnName, DescFunc: 0}},
+	}
+	m.BuildFunctionDefinitions()
+
+	err = s.Engine.CompileModule(testCtx, m, nil)
+	require.NoError(t, err)
+
+	inst, err := s.Instantiate(testCtx, ns, m, t.Name(), nil)
+	require.NoError(t, err)
+
+	growFn = inst.Function(2)
+	_, err = inst.Function(1).Call(context.Background())
+	require.NoError(t, err)
+}
+
 func RunTestEngine_NewModuleEngine(t *testing.T, et EngineTester) {
 	e := et.NewEngine(api.CoreFeaturesV1)