Repository: tetratelabs/wazero (go)

Issue:
Executing a Function from another Module 
Executing a host function from a function called from another module returns a pointer to original module. Memory address referenced from the other module is not accessible, as the pointer is from the original module.

```go
//go:wasm-module host
//export printModuleName()
func printModuleName()


//export function1
func function1() {
	printModuleName()
}
```

```go
//go:wasm-module module1
//export function1
func function1()

func function2() {
	function1()
}
```

```go
func printModuleName(ctx, module) error {
	fmt.Println(module.Name())
}

````

calling function2 from module1
```go
module2.ExportedFunction("function2").Call()
```

prints


```
module2
```

Thus memory addressed in module1 are not accessible from function1 being called by function2. 

Not sure if this is intended or not. 

Pull request description:
Pass correct api.Module to host functions
Fixes #1211 


Previously, host functions are getting `api.Module` for the "originating" module, which is the module for api.Function currently invoked, except that the api.Module is modified withMemory for the caller's memory instance, therefore there haven't been no problem. The only issues were the methods except `Memory()` of api.Module, and this commit fixes them.


Patch:
diff --git a/internal/engine/compiler/arch_arm64.s b/internal/engine/compiler/arch_arm64.s
index 0a423ce4..bec94728 100644
--- a/internal/engine/compiler/arch_arm64.s
+++ b/internal/engine/compiler/arch_arm64.s
@@ -8,8 +8,8 @@ TEXT ·nativecall(SB), NOSPLIT|NOFRAME, $0-24
 
 	// In arm64, return address is stored in R30 after jumping into the code.
 	// We save the return address value into archContext.compilerReturnAddress in Engine.
-	// Note that the const 136 drifts after editting Engine or archContext struct. See TestArchContextOffsetInEngine.
-	MOVD R30, 136(R0)
+	// Note that the const 144 drifts after editting Engine or archContext struct. See TestArchContextOffsetInEngine.
+	MOVD R30, 144(R0)
 
 	// Load the address of *wasm.ModuleInstance into arm64CallingConventionModuleInstanceAddressRegister.
 	MOVD moduleInstanceAddress+16(FP), R29
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index a3269811..398334f1 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -229,6 +229,9 @@ type (
 		// returnAddress is the return address which the engine jumps into
 		// after executing a builtin function or host function.
 		returnAddress uintptr
+
+		// callerFunctionInstance holds the caller's wasm.FunctionInstance, and is only valid if currently executing a host function.
+		callerFunctionInstance *wasm.FunctionInstance
 	}
 
 	// callFrame holds the information to which the caller function can return.
@@ -331,6 +334,7 @@ const (
 	callEngineExitContextNativeCallStatusCodeOffset     = 120
 	callEngineExitContextBuiltinFunctionCallIndexOffset = 124
 	callEngineExitContextReturnAddressOffset            = 128
+	callEngineExitContextCallerFunctionInstanceOffset   = 136
 
 	// Offsets for function.
 	functionCodeInitialAddressOffset    = 0
@@ -925,7 +929,7 @@ entry:
 			fn := calleeHostFunction.parent.goFunc
 			switch fn := fn.(type) {
 			case api.GoModuleFunction:
-				fn.Call(ce.ctx, callCtx.WithMemory(ce.memoryInstance), stack)
+				fn.Call(ce.ctx, ce.callerFunctionInstance.Module.CallCtx, stack)
 			case api.GoFunction:
 				fn.Call(ce.ctx, stack)
 			}
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 6032486f..54b3bcef 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -191,6 +191,27 @@ func (c *amd64Compiler) compileGoDefinedHostFunction() error {
 		}
 	}
 
+	// Host function needs access to the caller's Function Instance, and the caller's information is stored in the stack
+	// (as described in the doc of callEngine.stack). Here, we get the caller's *wasm.FunctionInstance from the stack,
+	// and save it in callEngine.exitContext.callerFunctionInstance so we can pass it to the host function
+	// without sacrificing the performance.
+	c.compileReservedStackBasePointerInitialization()
+	// Alias for readability.
+	tmp := amd64.RegAX
+	// Get the location of the callerFunction (*function) in the stack, which depends on the signature.
+	_, _, callerFunction := c.locationStack.getCallFrameLocations(c.ir.Signature)
+	// Load the value into the tmp register: tmp = &function{..}
+	callerFunction.setRegister(tmp)
+	c.compileLoadValueOnStackToRegister(callerFunction)
+	// tmp = *(tmp+functionSourceOffset) = &wasm.FunctionInstance{...}
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ, tmp, functionSourceOffset, tmp)
+	// Load it onto callEngine.exitContext.callerFunctionInstance.
+	c.assembler.CompileRegisterToMemory(amd64.MOVQ,
+		tmp,
+		amd64ReservedRegisterForCallEngine, callEngineExitContextCallerFunctionInstanceOffset)
+	// Reset the state of callerFunction value location so that we won't mess up subsequent code generation below.
+	c.locationStack.releaseRegister(callerFunction)
+
 	if err := c.compileCallGoHostFunction(); err != nil {
 		return err
 	}
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index 140d7c4d..97381540 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -85,11 +85,11 @@ var arm64CallingConventionModuleInstanceAddressRegister = arm64.RegR29
 
 const (
 	// arm64CallEngineArchContextCompilerCallReturnAddressOffset is the offset of archContext.nativeCallReturnAddress in callEngine.
-	arm64CallEngineArchContextCompilerCallReturnAddressOffset = 136
+	arm64CallEngineArchContextCompilerCallReturnAddressOffset = 144
 	// arm64CallEngineArchContextMinimum32BitSignedIntOffset is the offset of archContext.minimum32BitSignedIntAddress in callEngine.
-	arm64CallEngineArchContextMinimum32BitSignedIntOffset = 144
+	arm64CallEngineArchContextMinimum32BitSignedIntOffset = 152
 	// arm64CallEngineArchContextMinimum64BitSignedIntOffset is the offset of archContext.minimum64BitSignedIntAddress in callEngine.
-	arm64CallEngineArchContextMinimum64BitSignedIntOffset = 152
+	arm64CallEngineArchContextMinimum64BitSignedIntOffset = 160
 )
 
 func isZeroRegister(r asm.Register) bool {
@@ -392,6 +392,27 @@ func (c *arm64Compiler) compileGoDefinedHostFunction() error {
 		}
 	}
 
+	// Host function needs access to the caller's Function Instance, and the caller's information is stored in the stack
+	// (as described in the doc of callEngine.stack). Here, we get the caller's *wasm.FunctionInstance from the stack,
+	// and save it in callEngine.exitContext.callerFunctionInstance so we can pass it to the host function
+	// without sacrificing the performance.
+	c.compileReservedStackBasePointerRegisterInitialization()
+	// Alias for readability.
+	tmp := arm64CallingConventionModuleInstanceAddressRegister
+	// Get the location of the callerFunction (*function) in the stack, which depends on the signature.
+	_, _, callerFunction := c.locationStack.getCallFrameLocations(c.ir.Signature)
+	// Load the value into the tmp register: tmp = &function{..}
+	callerFunction.setRegister(tmp)
+	c.compileLoadValueOnStackToRegister(callerFunction)
+	// tmp = *(tmp+functionSourceOffset) = &wasm.FunctionInstance{...}
+	c.assembler.CompileMemoryToRegister(arm64.LDRD, tmp, functionSourceOffset, tmp)
+	// Load it onto callEngine.exitContext.callerFunctionInstance.
+	c.assembler.CompileRegisterToMemory(arm64.STRD,
+		tmp,
+		arm64ReservedRegisterForCallEngine, callEngineExitContextCallerFunctionInstanceOffset)
+	// Reset the state of callerFunction value location so that we won't mess up subsequent code generation below.
+	c.locationStack.releaseRegister(callerFunction)
+
 	if err := c.compileCallGoFunction(nativeCallStatusCodeCallGoHostFunction, 0); err != nil {
 		return err
 	}
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index bdcaaf69..24248336 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -859,7 +859,6 @@ func (ce *callEngine) callFunction(ctx context.Context, callCtx *wasm.CallContex
 
 func (ce *callEngine) callGoFunc(ctx context.Context, callCtx *wasm.CallContext, f *function, stack []uint64) {
 	lsn := f.parent.listener
-	callCtx = callCtx.WithMemory(ce.callerMemory())
 	if lsn != nil {
 		params := stack[:f.source.Type.ParamNumInUint64]
 		ctx = lsn.Before(ctx, callCtx, f.source.Definition, params)
@@ -934,7 +933,7 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, callCtx *wasm.CallCont
 				frame.pc = op.us[0]
 			}
 		case wazeroir.OperationKindCall:
-			ce.callFunction(ctx, callCtx, &functions[op.us[0]])
+			ce.callFunction(ctx, f.source.Module.CallCtx, &functions[op.us[0]])
 			frame.pc++
 		case wazeroir.OperationKindCallIndirect:
 			offset := ce.popValue()
@@ -952,7 +951,7 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, callCtx *wasm.CallCont
 				panic(wasmruntime.ErrRuntimeIndirectCallTypeMismatch)
 			}
 
-			ce.callFunction(ctx, callCtx, tf)
+			ce.callFunction(ctx, f.source.Module.CallCtx, tf)
 			frame.pc++
 		case wazeroir.OperationKindDrop:
 			ce.drop(op.rs[0])


Test patch:
diff --git a/internal/engine/compiler/compiler_controlflow_test.go b/internal/engine/compiler/compiler_controlflow_test.go
index ebd16072..ef806f88 100644
--- a/internal/engine/compiler/compiler_controlflow_test.go
+++ b/internal/engine/compiler/compiler_controlflow_test.go
@@ -17,13 +17,23 @@ func TestCompiler_compileHostFunction(t *testing.T) {
 	err := compiler.compileGoDefinedHostFunction()
 	require.NoError(t, err)
 
-	// Generate and run the code under test.
+	// Get the location of caller function's location stored in the stack, which depends on the type.
+	// In this test, the host function has empty sig.
+	_, _, callerFuncLoc := compiler.runtimeValueLocationStack().getCallFrameLocations(&wasm.FunctionType{})
+
+	// Generate the machine code for the test.
 	code, _, err := compiler.compile()
 	require.NoError(t, err)
+
+	// Set the caller's function which always exists in the real usecase.
+	f := &function{source: &wasm.FunctionInstance{}}
+	env.stack()[callerFuncLoc.stackPointer] = uint64(uintptr(unsafe.Pointer(f)))
 	env.exec(code)
 
 	// On the return, the code must exit with the host call status.
 	require.Equal(t, nativeCallStatusCodeCallGoHostFunction, env.compilerStatus())
+	// Plus, the exitContext holds the caller's wasm.FunctionInstance.
+	require.Equal(t, f.source, env.ce.exitContext.callerFunctionInstance)
 
 	// Re-enter the return address.
 	require.NotEqual(t, uintptr(0), uintptr(env.ce.returnAddress))
diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go
index d1903e96..cd09eeb0 100644
--- a/internal/engine/compiler/compiler_test.go
+++ b/internal/engine/compiler/compiler_test.go
@@ -58,6 +58,7 @@ func init() {
 	requireEqual(int(unsafe.Offsetof(ce.statusCode)), callEngineExitContextNativeCallStatusCodeOffset, "callEngineExitContextNativeCallStatusCodeOffset")
 	requireEqual(int(unsafe.Offsetof(ce.builtinFunctionCallIndex)), callEngineExitContextBuiltinFunctionCallIndexOffset, "callEngineExitContextBuiltinFunctionCallIndexOffset")
 	requireEqual(int(unsafe.Offsetof(ce.returnAddress)), callEngineExitContextReturnAddressOffset, "callEngineExitContextReturnAddressOffset")
+	requireEqual(int(unsafe.Offsetof(ce.callerFunctionInstance)), callEngineExitContextCallerFunctionInstanceOffset, "callEngineExitContextCallerFunctionInstanceOffset")
 
 	// Size and offsets for callFrame.
 	var frame callFrame
diff --git a/internal/integration_test/engine/adhoc_test.go b/internal/integration_test/engine/adhoc_test.go
index df8b1314..bb8f89eb 100644
--- a/internal/integration_test/engine/adhoc_test.go
+++ b/internal/integration_test/engine/adhoc_test.go
@@ -47,6 +47,7 @@ var tests = map[string]func(t *testing.T, r wazero.Runtime){
 	"un-signed extend global":                           testGlobalExtend,
 	"user-defined primitive in host func":               testUserDefinedPrimitiveHostFunc,
 	"ensures invocations terminate on module close":     testEnsureTerminationOnClose,
+	"call host function indirectly":                     callHostFunctionIndirect,
 }
 
 func TestEngineCompiler(t *testing.T) {
@@ -479,6 +480,62 @@ func testHostFunctionNumericParameter(t *testing.T, r wazero.Runtime) {
 	}
 }
 
+func callHostFunctionIndirect(t *testing.T, r wazero.Runtime) {
+	// With the following call graph,
+	//  originWasmModule -- call --> importingWasmModule -- call --> hostModule
+	// this ensures that hostModule's hostFn only has access importingWasmModule, not originWasmModule.
+
+	const hostModule, importingWasmModule, originWasmModule = "host", "importing", "origin"
+	const hostFn, importingWasmModuleFn, originModuleFn = "host_fn", "call_host_func", "origin"
+	importingModule := &wasm.Module{
+		TypeSection:     []*wasm.FunctionType{{Params: []wasm.ValueType{}, Results: []wasm.ValueType{}}},
+		ImportSection:   []*wasm.Import{{Module: hostModule, Name: hostFn, Type: wasm.ExternTypeFunc, DescFunc: 0}},
+		FunctionSection: []wasm.Index{0},
+		ExportSection:   []*wasm.Export{{Name: importingWasmModuleFn, Type: wasm.ExternTypeFunc, Index: 1}},
+		CodeSection:     []*wasm.Code{{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}},
+		NameSection:     &wasm.NameSection{ModuleName: importingWasmModule},
+	}
+
+	originModule := &wasm.Module{
+		TypeSection:     []*wasm.FunctionType{{Params: []wasm.ValueType{}, Results: []wasm.ValueType{}}},
+		ImportSection:   []*wasm.Import{{Module: importingWasmModule, Name: importingWasmModuleFn, Type: wasm.ExternTypeFunc, DescFunc: 0}},
+		FunctionSection: []wasm.Index{0},
+		ExportSection:   []*wasm.Export{{Name: "origin", Type: wasm.ExternTypeFunc, Index: 1}},
+		CodeSection:     []*wasm.Code{{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}},
+		NameSection:     &wasm.NameSection{ModuleName: originWasmModule},
+	}
+
+	require.NoError(t, importingModule.Validate(api.CoreFeaturesV2))
+	require.NoError(t, originModule.Validate(api.CoreFeaturesV2))
+	importingModuleBytes := binaryencoding.EncodeModule(importingModule)
+	originModuleBytes := binaryencoding.EncodeModule(originModule)
+
+	var originInst, importingInst api.Module
+	_, err := r.NewHostModuleBuilder(hostModule).
+		NewFunctionBuilder().
+		WithFunc(func(ctx context.Context, mod api.Module) {
+			// Module must be the caller (importing module), not the origin.
+			require.Equal(t, mod, importingInst)
+			require.NotEqual(t, mod, originInst)
+			// Name must be the caller, not origin.
+			require.Equal(t, importingWasmModule, mod.Name())
+		}).
+		Export(hostFn).
+		Instantiate(testCtx)
+	require.NoError(t, err)
+
+	importingInst, err = r.Instantiate(testCtx, importingModuleBytes)
+	require.NoError(t, err)
+	originInst, err = r.Instantiate(testCtx, originModuleBytes)
+	require.NoError(t, err)
+
+	originFn := originInst.ExportedFunction(originModuleFn)
+	require.NotNil(t, originFn)
+
+	_, err = originFn.Call(testCtx)
+	require.NoError(t, err)
+}
+
 func callReturnImportWasm(t *testing.T, importedModule, importingModule string, vt wasm.ValueType) []byte {
 	// test an imported function by re-exporting it
 	module := &wasm.Module{
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 8a05d802..111aa95a 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -287,11 +287,6 @@ func runTestModuleEngine_Call_HostFn_Mem(t *testing.T, et EngineTester, readMem
 			fn:       &importing.Functions[importing.Exports[callImportReadMemName].Index],
 			expected: importingMemoryVal,
 		},
-		{
-			name:     callImportCallReadMemName,
-			fn:       &importing.Functions[importing.Exports[callImportCallReadMemName].Index],
-			expected: importingMemoryVal,
-		},
 	}
 	for _, tt := range tests {
 		tc := tt
@@ -620,10 +615,8 @@ var (
 )
 
 const (
-	readMemName               = "read_mem"
-	callReadMemName           = "call->read_mem"
-	callImportReadMemName     = "call_import->read_mem"
-	callImportCallReadMemName = "call_import->call->read_mem"
+	readMemName           = "read_mem"
+	callImportReadMemName = "call_import->read_mem"
 )
 
 func readMemGo(_ context.Context, m api.Module) uint64 {
@@ -754,28 +747,16 @@ func setupCallTests(t *testing.T, e wasm.Engine, divBy *wasm.Code, fnlf experime
 func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf experimental.FunctionListenerFactory) (*wasm.ModuleInstance, *wasm.ModuleInstance, func()) {
 	ft := &wasm.FunctionType{Results: []wasm.ValueType{i64}, ResultNumInUint64: 1}
 
-	callReadMem := &wasm.Code{ // shows indirect calls still use the same memory
-		IsHostFunction: true,
-		Body: []byte{
-			wasm.OpcodeCall, 1,
-			// On the return from the another host function,
-			// we should still be able to access the memory.
-			wasm.OpcodeI32Const, 0,
-			wasm.OpcodeI32Load, 0x2, 0x0,
-			wasm.OpcodeEnd,
-		},
-	}
 	hostModule := &wasm.Module{
 		TypeSection:     []*wasm.FunctionType{ft},
-		FunctionSection: []wasm.Index{0, 0},
-		CodeSection:     []*wasm.Code{callReadMem, readMem},
+		FunctionSection: []wasm.Index{0},
+		CodeSection:     []*wasm.Code{readMem},
 		ExportSection: []*wasm.Export{
-			{Name: callReadMemName, Type: wasm.ExternTypeFunc, Index: 0},
-			{Name: readMemName, Type: wasm.ExternTypeFunc, Index: 1},
+			{Name: readMemName, Type: wasm.ExternTypeFunc, Index: 0},
 		},
 		NameSection: &wasm.NameSection{
 			ModuleName:    "host",
-			FunctionNames: wasm.NameMap{{Index: 0, Name: readMemName}, {Index: 1, Name: callReadMemName}},
+			FunctionNames: wasm.NameMap{{Index: 0, Name: readMemName}},
 		},
 		ID: wasm.ModuleID{0},
 	}
@@ -786,7 +767,6 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 	host.Functions = host.BuildFunctions(hostModule, nil)
 	host.BuildExports(hostModule.ExportSection)
 	readMemFn := &host.Functions[host.Exports[readMemName].Index]
-	callReadMemFn := &host.Functions[host.Exports[callReadMemName].Index]
 
 	hostME, err := e.NewModuleEngine(host.Name, hostModule, host.Functions)
 	require.NoError(t, err)
@@ -797,22 +777,18 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 		ImportSection: []*wasm.Import{
 			// Placeholder for two import functions from `importedModule`.
 			{Type: wasm.ExternTypeFunc, DescFunc: 0},
-			{Type: wasm.ExternTypeFunc, DescFunc: 0},
 		},
-		FunctionSection: []wasm.Index{0, 0},
+		FunctionSection: []wasm.Index{0},
 		ExportSection: []*wasm.Export{
-			{Name: callImportReadMemName, Type: wasm.ExternTypeFunc, Index: 2},
-			{Name: callImportCallReadMemName, Type: wasm.ExternTypeFunc, Index: 3},
+			{Name: callImportReadMemName, Type: wasm.ExternTypeFunc, Index: 1},
 		},
 		CodeSection: []*wasm.Code{
-			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Calling the index 0 = callReadMemFn.
-			{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}}, // Calling the index 1 = readMemFn.
+			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Calling the index 1 = readMemFn.
 		},
 		NameSection: &wasm.NameSection{
 			ModuleName: "importing",
 			FunctionNames: wasm.NameMap{
 				{Index: 2, Name: callImportReadMemName},
-				{Index: 3, Name: callImportCallReadMemName},
 			},
 		},
 		// Indicates that this module has a memory so that compilers are able to assembe memory-related initialization.
@@ -825,7 +801,7 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 
 	// Add the exported function.
 	importing := &wasm.ModuleInstance{Name: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
-	importingFunctions := importing.BuildFunctions(importingModule, []*wasm.FunctionInstance{readMemFn, callReadMemFn})
+	importingFunctions := importing.BuildFunctions(importingModule, []*wasm.FunctionInstance{readMemFn})
 	// Note: adds imported functions readMemFn and callReadMemFn at index 0 and 1.
 	importing.Functions = importingFunctions
 	importing.BuildExports(importingModule.ExportSection)
diff --git a/runtime_test.go b/runtime_test.go
index 55d8a40a..27517323 100644
--- a/runtime_test.go
+++ b/runtime_test.go
@@ -549,69 +549,79 @@ func TestRuntime_CloseWithExitCode(t *testing.T) {
 }
 
 func TestHostFunctionWithCustomContext(t *testing.T) {
-	const fistString = "hello"
-	const secondString = "hello call"
-	hostCtx := &HostContext{fistString}
-	r := NewRuntime(hostCtx)
-	defer r.Close(hostCtx)
-
-	// Define a function that will be set as the start function
-	var calledStart bool
-	var calledCall bool
-	start := func(ctx context.Context, module api.Module) {
-		hts, ok := ctx.(*HostContext)
-		if !ok {
-			t.Fatal("decorate call context could effect host ctx cast failed, please consider it.")
-		}
-		calledStart = true
-		require.NotNil(t, hts)
-		require.Equal(t, fistString, hts.Content)
-	}
+	for _, tc := range []struct {
+		name   string
+		config RuntimeConfig
+	}{
+		{name: "compiler", config: NewRuntimeConfigCompiler()},
+		{name: "interpreter", config: NewRuntimeConfigInterpreter()},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			const fistString = "hello"
+			const secondString = "hello call"
+			hostCtx := &HostContext{fistString}
+			r := NewRuntimeWithConfig(hostCtx, tc.config)
+			defer r.Close(hostCtx)
+
+			// Define a function that will be set as the start function
+			var calledStart bool
+			var calledCall bool
+			start := func(ctx context.Context, module api.Module) {
+				hts, ok := ctx.(*HostContext)
+				if !ok {
+					t.Fatal("decorate call context could effect host ctx cast failed, please consider it.")
+				}
+				calledStart = true
+				require.NotNil(t, hts)
+				require.Equal(t, fistString, hts.Content)
+			}
 
-	callFunc := func(ctx context.Context, module api.Module) {
-		hts, ok := ctx.(*HostContext)
-		if !ok {
-			t.Fatal("decorate call context could effect host ctx cast failed, please consider it.")
-		}
-		calledCall = true
-		require.NotNil(t, hts)
-		require.Equal(t, secondString, hts.Content)
-	}
+			callFunc := func(ctx context.Context, module api.Module) {
+				hts, ok := ctx.(*HostContext)
+				if !ok {
+					t.Fatal("decorate call context could effect host ctx cast failed, please consider it.")
+				}
+				calledCall = true
+				require.NotNil(t, hts)
+				require.Equal(t, secondString, hts.Content)
+			}
 
-	_, err := r.NewHostModuleBuilder("env").
-		NewFunctionBuilder().WithFunc(start).Export("host").
-		NewFunctionBuilder().WithFunc(callFunc).Export("host2").
-		Instantiate(hostCtx)
-	require.NoError(t, err)
+			_, err := r.NewHostModuleBuilder("env").
+				NewFunctionBuilder().WithFunc(start).Export("host").
+				NewFunctionBuilder().WithFunc(callFunc).Export("host2").
+				Instantiate(hostCtx)
+			require.NoError(t, err)
 
-	one := uint32(0)
-	binary := binaryencoding.EncodeModule(&wasm.Module{
-		TypeSection: []*wasm.FunctionType{{}, {}},
-		ImportSection: []*wasm.Import{
-			{Module: "env", Name: "host", Type: wasm.ExternTypeFunc, DescFunc: 0},
-			{Module: "env", Name: "host2", Type: wasm.ExternTypeFunc, DescFunc: 0},
-		},
-		FunctionSection: []wasm.Index{0, 1},
-		CodeSection: []*wasm.Code{
-			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Call the imported env.host.
-			{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}}, // Call the imported env.host.
-		},
-		ExportSection: []*wasm.Export{
-			{Type: api.ExternTypeFunc, Name: "callHost", Index: uint32(3)},
-		},
-		StartSection: &one,
-	})
+			startFnIndex := uint32(2)
+			binary := binaryencoding.EncodeModule(&wasm.Module{
+				TypeSection: []*wasm.FunctionType{{}},
+				ImportSection: []*wasm.Import{
+					{Module: "env", Name: "host", Type: wasm.ExternTypeFunc, DescFunc: 0},
+					{Module: "env", Name: "host2", Type: wasm.ExternTypeFunc, DescFunc: 0},
+				},
+				FunctionSection: []wasm.Index{0, 0},
+				CodeSection: []*wasm.Code{
+					{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Call the imported env.host.
+					{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}}, // Call the imported env.host.
+				},
+				ExportSection: []*wasm.Export{
+					{Type: api.ExternTypeFunc, Name: "callHost", Index: uint32(3)},
+				},
+				StartSection: &startFnIndex,
+			})
 
-	// Instantiate the module, which calls the start function. This will fail if the context wasn't as intended.
-	ins, err := r.Instantiate(hostCtx, binary)
-	require.NoError(t, err)
-	require.True(t, calledStart)
+			// Instantiate the module, which calls the start function. This will fail if the context wasn't as intended.
+			ins, err := r.Instantiate(hostCtx, binary)
+			require.NoError(t, err)
+			require.True(t, calledStart)
 
-	// add the new context content for call with used in host function
-	hostCtx.Content = secondString
-	_, err = ins.ExportedFunction("callHost").Call(hostCtx)
-	require.NoError(t, err)
-	require.True(t, calledCall)
+			// add the new context content for call with used in host function
+			hostCtx.Content = secondString
+			_, err = ins.ExportedFunction("callHost").Call(hostCtx)
+			require.NoError(t, err)
+			require.True(t, calledCall)
+		})
+	}
 }
 
 func TestRuntime_Close_ClosesCompiledModules(t *testing.T) {