Repository: tetratelabs/wazero (go)

Issue:
Host functions called from another host function cannot know the right MemoryInstance
In the current implementation, host functions try to get the `MemoryInstance` of `HostFunctionCallContext` from `<function>.ModuleInstance.Memory` of the parent stack's (in interpreter) or of the function itself (in jit).
However, a host function doesn't have associated `MemoryInstance` nor wasm modules in the first place. So, a host function only finds `nil` as `MemoryInstance` when it's called from another host function. 

FYI, it looks wasmtime also doesn't provide the access to memory instances when a host function is called from another host function because there's no caller context in that case, but let me check.

https://github.com/tetratelabs/wazero/blob/a6cd1549ba093e35d09fe07cd1a049b70e798185/wasm/interpreter/interpreter.go#L516
https://github.com/tetratelabs/wazero/blob/a6cd1549ba093e35d09fe07cd1a049b70e798185/wasm/jit/engine.go#L390

This also prevents the WASI implementation from testing APIs by "import-and-exported-back" style, where tests call each WASI API by having a wasm module import a WASI API and re-export it to the host.



Pull request:
Uses importing module memory for host function calls
Fixes #195

Patch:
diff --git a/wasi/types.go b/wasi/types.go
index fd9fabbf..d0ec0e78 100644
--- a/wasi/types.go
+++ b/wasi/types.go
@@ -260,70 +260,78 @@ var errnoToString = [...]string{
 const (
 	// FunctionArgsGet reads command-line argument data.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "args_get"
-	//		(func $wasi.args_get (param $argv i32) (param $argv_buf i32) (result (;errno;) i32)))
-	//
+	// See ImportArgsGet
 	// See API.ArgsGet
 	// See FunctionArgsSizesGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-args_getargv-pointerpointeru8-argv_buf-pointeru8---errno
 	FunctionArgsGet = "args_get"
 
+	// ImportArgsGet is the WebAssembly 1.0 (MVP) Text format import of FunctionArgsGet
+	ImportArgsGet = `(import "wasi_snapshot_preview1" "args_get"
+    (func $wasi.args_get (param $argv i32) (param $argv_buf i32) (result (;errno;) i32)))`
+
 	// FunctionArgsSizesGet returns command-line argument data sizes.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "args_sizes_get"
-	//		(func $wasi.args_sizes_get (param $result.argc i32) (param $result.argv_buf_size i32) (result (;errno;) i32)))
-	//
+	// See ImportArgsSizesGet
 	// See API.ArgsSizesGet
 	// See FunctionArgsGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-args_sizes_get---errno-size-size
 	FunctionArgsSizesGet = "args_sizes_get"
 
+	// ImportArgsSizesGet is the WebAssembly 1.0 (MVP) Text format import of FunctionArgsSizesGet
+	ImportArgsSizesGet = `(import "wasi_snapshot_preview1" "args_sizes_get"
+    (func $wasi.args_sizes_get (param $result.argc i32) (param $result.argv_buf_size i32) (result (;errno;) i32)))`
+
 	// FunctionEnvironGet reads environment variable data.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "environ_get"
-	//		(func $wasi.environ_get (param $environ i32) (param $environ_buf i32) (result (;errno;) i32)))
-	//
+	// See ImportEnvironGet
 	// See API.EnvironGet
 	// See FunctionEnvironSizesGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-environ_getenviron-pointerpointeru8-environ_buf-pointeru8---errno
 	FunctionEnvironGet = "environ_get"
 
+	// ImportEnvironGet is the WebAssembly 1.0 (MVP) Text format import of FunctionEnvironGet
+	ImportEnvironGet = `(import "wasi_snapshot_preview1" "environ_get"
+    (func $wasi.environ_get (param $environ i32) (param $environ_buf i32) (result (;errno;) i32)))`
+
 	// FunctionEnvironSizesGet returns environment variable data sizes.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "environ_sizes_get"
-	//		(func $wasi.environ_sizes_get (param $result.environc i32) (param $result.environBufSize i32) (result (;errno;) i32)))
-	//
+	// See ImportEnvironSizesGet
 	// See API.EnvironSizesGet
 	// See FunctionEnvironGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-environ_sizes_get---errno-size-size
 	FunctionEnvironSizesGet = "environ_sizes_get"
 
+	// ImportEnvironSizesGet is the WebAssembly 1.0 (MVP) Text format import of FunctionEnvironSizesGet
+	ImportEnvironSizesGet = `
+(import "wasi_snapshot_preview1" "environ_sizes_get"
+    (func $wasi.environ_sizes_get (param $result.environc i32) (param $result.environBufSize i32) (result (;errno;) i32)))`
+
 	// FunctionClockResGet returns the resolution of a clock.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "clock_res_get"
-	//		(func $wasi.clock_res_get (param $id i32) (param $result.resolution i32) (result (;errno;) i32)))
-	//
+	// See ImportClockResGet
 	// See API.ClockResGet
 	// See FunctionClockTimeGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-clock_res_getid-clockid---errno-timestamp
 	FunctionClockResGet = "clock_res_get"
 
+	// ImportClockResGet is the WebAssembly 1.0 (MVP) Text format import of FunctionClockResGet
+	ImportClockResGet = `
+(import "wasi_snapshot_preview1" "clock_res_get"
+    (func $wasi.clock_res_get (param $id i32) (param $result.resolution i32) (result (;errno;) i32)))`
+
 	// FunctionClockTimeGet returns the time value of a clock.
 	//
-	// In WebAssembly 1.0 (MVP) Text format, this signature is:
-	//	(import "wasi_snapshot_preview1" "clock_time_get"
-	//		(func $wasi.clock_time_get (param $id i32) (param $precision i64) (param $result.timestamp i32) (result (;errno;) i32)))
-	//
+	// See ImportClockTimeGet
 	// See API.ClockResGet
 	// See FunctionClockResGet
 	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/phases/snapshot/docs.md#-clock_time_getid-clockid-precision-timestamp---errno-timestamp
 	FunctionClockTimeGet = "clock_time_get"
 
+	// ImportClockTimeGet is the WebAssembly 1.0 (MVP) Text format import of FunctionClockTimeGet
+	ImportClockTimeGet = `(import "wasi_snapshot_preview1" "clock_time_get"
+    (func $wasi.clock_time_get (param $id i32) (param $precision i64) (param $result.timestamp i32) (result (;errno;) i32)))`
+
 	FunctionFdAdvise             = "fd_advise"
 	FunctionFdAllocate           = "fd_allocate"
 	FunctionFdClose              = "fd_close"
diff --git a/wasm/interpreter/interpreter.go b/wasm/interpreter/interpreter.go
index 3f196271..128887ca 100644
--- a/wasm/interpreter/interpreter.go
+++ b/wasm/interpreter/interpreter.go
@@ -517,6 +517,8 @@ func (it *interpreter) callHostFunc(ctx context.Context, f *interpreterFunction)
 	var memory *wasm.MemoryInstance
 	if len(it.frames) > 0 {
 		memory = it.frames[len(it.frames)-1].f.funcInstance.ModuleInstance.Memory
+	} else {
+		memory = f.funcInstance.ModuleInstance.Memory
 	}
 	val.Set(reflect.ValueOf(wasm.NewHostFunctionCallContext(ctx, memory)))
 	in[0] = val
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index e9884cdd..9efb37fe 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -4070,62 +4070,38 @@ func (c *amd64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
 
 // setupMemoryOffset pops the top value from the stack (called "base"), and returns the result of addition with
 // base and offsetArg, which we call "offset". The returned offsetRegister is the register number with the offset calculation value.
-// targetSizeInByte is the original memory operation's target size in byte. For example, 4 = 32 / 8 for Load32 operation.
+// targetSizeInBytes is the original memory operation's target size in byte. For example, 4 = 32 / 8 for Load32 operation.
 // This is used for all Store* and Load* instructions.
 //
 // Note that this also emits the instructions to check the out of bounds memory access. That means
-// if the base+offsetArg+targetSizeInByte exceeds the memory size, we exit this function with
-// jitCallStatusCodeMemoryOutOfBounds status code since we read memory as [base+offsetArg: base+offsetArg+targetSizeInByte].
-func (c *amd64Compiler) setupMemoryOffset(offsetArg uint32, targetSizeInByte int64) (offsetRegister int16, err error) {
+// if the base+offsetArg+targetSizeInBytes exceeds the memory size, we exit this function with
+// jitCallStatusCodeMemoryOutOfBounds status code since we read memory as [base+offsetArg: base+offsetArg+targetSizeInBytes].
+func (c *amd64Compiler) setupMemoryOffset(offsetArg uint32, targetSizeInBytes int64) (int16, error) {
 	base := c.locationStack.pop()
-	if err = c.ensureOnGeneralPurposeRegister(base); err != nil {
+	if err := c.ensureOnGeneralPurposeRegister(base); err != nil {
 		return 0, err
 	}
 
-	// First, we calculate the offset on the memory region.
-	addOffsetToBase := c.newProg()
-	addOffsetToBase.As = x86.AADDL // 32-bit!
-	addOffsetToBase.To.Type = obj.TYPE_REG
-	addOffsetToBase.To.Reg = base.register
-	addOffsetToBase.From.Type = obj.TYPE_CONST
-	addOffsetToBase.From.Offset = int64(offsetArg)
-	c.addInstruction(addOffsetToBase)
-
-	// If the base+offset already overflows from uint32 range, we exit with the out of boundary status.
-	overflowJmp := c.newProg()
-	overflowJmp.As = x86.AJCS
-	overflowJmp.To.Type = obj.TYPE_BRANCH
-	c.addInstruction(overflowJmp)
-
-	// Otherwise, we calculate base+offset+targetSizeInByte and check if it is within memory boundary.
-	tmpReg, err := c.allocateRegister(generalPurposeRegisterTypeInt)
-	if err != nil {
-		return 0, err
+	result := base.register
+	if offsetConst := int64(offsetArg) + targetSizeInBytes; offsetConst <= math.MaxUint32 {
+		addOffsetToBase := c.newProg()
+		addOffsetToBase.As = x86.AADDQ
+		addOffsetToBase.To.Type = obj.TYPE_REG
+		addOffsetToBase.To.Reg = result
+		addOffsetToBase.From.Type = obj.TYPE_CONST
+		addOffsetToBase.From.Offset = offsetConst
+		c.addInstruction(addOffsetToBase)
+	} else {
+		// If the offset const is too large, we exit with jitCallStatusCodeMemoryOutOfBounds.
+		c.exit(jitCallStatusCodeMemoryOutOfBounds)
+		return result, nil
 	}
 
-	// Copy the 32-bit base+offset as to the temporary register as 64-bit integer.
-	copyOffset := c.newProg()
-	copyOffset.As = x86.AMOVLQZX // Zero extend
-	copyOffset.To.Type = obj.TYPE_REG
-	copyOffset.To.Reg = tmpReg
-	copyOffset.From.Type = obj.TYPE_REG
-	copyOffset.From.Reg = base.register
-	c.addInstruction(copyOffset)
-
-	// Adds targetSizeInByte to base+offset stored in the temporary register.
-	addTargetSize := c.newProg()
-	addTargetSize.As = x86.AADDQ
-	addTargetSize.To.Type = obj.TYPE_REG
-	addTargetSize.To.Reg = tmpReg
-	addTargetSize.From.Type = obj.TYPE_CONST
-	addTargetSize.From.Offset = targetSizeInByte
-	c.addInstruction(addTargetSize)
-
 	// Now we compare the value with the memory length which is held by engine.
 	cmp := c.newProg()
 	cmp.As = x86.ACMPQ
 	cmp.To.Type = obj.TYPE_REG
-	cmp.To.Reg = tmpReg
+	cmp.To.Reg = result
 	cmp.From.Type = obj.TYPE_MEM
 	cmp.From.Reg = reservedRegisterForEngine
 	cmp.From.Offset = engineModuleContextMemorySliceLenOffset
@@ -4138,13 +4114,22 @@ func (c *amd64Compiler) setupMemoryOffset(offsetArg uint32, targetSizeInByte int
 	c.addInstruction(okJmp)
 
 	// Otherwise, we exit the function with out of bounds status code.
-	c.addSetJmpOrigins(overflowJmp)
 	c.exit(jitCallStatusCodeMemoryOutOfBounds)
 
 	c.addSetJmpOrigins(okJmp)
 
-	c.locationStack.markRegisterUnused(base.register)
-	return base.register, nil
+	// Finally we have to subtract targetSizeInBytes from the result register
+	// because we access memory with [offset: offset+targetSizeInBytes]
+	sub := c.newProg()
+	sub.As = x86.ASUBL
+	sub.To.Type = obj.TYPE_REG
+	sub.To.Reg = result
+	sub.From.Type = obj.TYPE_CONST
+	sub.From.Offset = targetSizeInBytes
+	c.addInstruction(sub)
+
+	c.locationStack.markRegisterUnused(result)
+	return result, nil
 }
 
 // compileStore implements compiler.compileStore for the amd64 architecture.
diff --git a/wasm/store.go b/wasm/store.go
index 3771303c..d5a5b30f 100644
--- a/wasm/store.go
+++ b/wasm/store.go
@@ -617,14 +617,14 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 			}
 		}
 
-		typeInstace, err := s.getTypeInstance(module.TypeSection[typeIndex])
+		typeInstance, err := s.getTypeInstance(module.TypeSection[typeIndex])
 		if err != nil {
 			return rollbackFuncs, err
 		}
 
 		f := &FunctionInstance{
 			Name:           name,
-			FunctionType:   typeInstace,
+			FunctionType:   typeInstance,
 			Body:           module.CodeSection[codeIndex].Body,
 			LocalTypes:     module.CodeSection[codeIndex].LocalTypes,
 			ModuleInstance: target,
@@ -785,6 +785,12 @@ func (s *Store) buildExportInstances(module *Module, target *ModuleInstance) (ro
 				return nil, fmt.Errorf("unknown function for export[%s]", name)
 			}
 			ei = &ExportInstance{Kind: exp.Kind, Function: target.Functions[index]}
+			// The module instance of the host function is a fake that only includes the function and its types.
+			// We need to assign the ModuleInstance when re-exporting so that any memory defined in the target is
+			// available to the HostFunctionCallContext.
+			if ei.Function.HostFunction != nil {
+				ei.Function.ModuleInstance = target
+			}
 		case ExportKindGlobal:
 			if index >= uint32(len(target.Globals)) {
 				return nil, fmt.Errorf("unknown global for export[%s]", name)
@@ -887,6 +893,11 @@ func (c *HostFunctionCallContext) Context() context.Context {
 	return c.ctx
 }
 
+// AddHostFunction exports a function so that it can be imported under the given module and name. If a function already
+// exists for this module and name it is ignored rather than overwritten.
+//
+// Note: The HostFunctionCallContext.Memory of the fn will be from the importing module.
+// Note: The ModuleInstance of this host function is lazy created and only includes exported functions and their types.
 func (s *Store) AddHostFunction(moduleName, funcName string, fn reflect.Value) error {
 	getTypeOf := func(kind reflect.Kind) (ValueType, error) {
 		switch kind {
@@ -932,7 +943,7 @@ func (s *Store) AddHostFunction(moduleName, funcName string, fn reflect.Value) e
 		return fmt.Errorf("invalid signature: %w", err)
 	}
 
-	typeInstace, err := s.getTypeInstance(sig)
+	typeInstance, err := s.getTypeInstance(sig)
 	if err != nil {
 		return err
 	}
@@ -940,7 +951,7 @@ func (s *Store) AddHostFunction(moduleName, funcName string, fn reflect.Value) e
 	f := &FunctionInstance{
 		Name:           fmt.Sprintf("%s.%s", moduleName, funcName),
 		HostFunction:   &fn,
-		FunctionType:   typeInstace,
+		FunctionType:   typeInstance,
 		ModuleInstance: m,
 	}
 
@@ -950,8 +961,10 @@ func (s *Store) AddHostFunction(moduleName, funcName string, fn reflect.Value) e
 	if err = s.addFunctionInstance(f); err != nil {
 		return err
 	}
+	// TODO: This races on adding export. A future design may be able to eliminate this race, possibly via a HostModule
+	// type.
 	if err = m.addExport(funcName, &ExportInstance{Kind: ExportKindFunc, Function: f}); err != nil {
-		s.Functions = s.Functions[:len(s.Functions)-1] // revert the add on conflict
+		s.Functions = s.Functions[:len(s.Functions)-1] // lost race: revert the add on conflict
 		return err
 	}
 	return nil


Test patch:
diff --git a/tests/engine/adhoc_test.go b/tests/engine/adhoc_test.go
index d998ae79..436e48b5 100644
--- a/tests/engine/adhoc_test.go
+++ b/tests/engine/adhoc_test.go
@@ -201,31 +201,41 @@ func recursiveEntry(t *testing.T, newEngine func() wasm.Engine) {
 }
 
 // importedAndExportedFunc fails if the engine cannot call an "imported-and-then-exported-back" function
+// Notably, this uses memory, which ensures wasm.HostFunctionCallContext is valid in both interpreter and JIT engines.
 func importedAndExportedFunc(t *testing.T, newEngine func() wasm.Engine) {
 	ctx := context.Background()
-	mod, err := text.DecodeModule([]byte(`(module
-		;; arbitrary function with params
-		(import "env" "add_int" (func $add_int (param i32 i32) (result i32)))
-		;; add_int is imported from the environment, but it's also exported back to the environment
-		(export "add_int" (func $add_int))
+	mod, err := text.DecodeModule([]byte(`(module $test
+		(import "" "store_int"
+			(func $store_int (param $offset i32) (param $val i64) (result (;errno;) i32)))
+		(memory $memory 1 1)
+		(export "memory" (memory $memory))
+		;; store_int is imported from the environment, but it's also exported back to the environment
+		(export "store_int" (func $store_int))
 		)`))
 	require.NoError(t, err)
 
 	store := wasm.NewStore(newEngine())
 
-	addInt := func(ctx *wasm.HostFunctionCallContext, x int32, y int32) int32 {
-		return x + y
+	storeInt := func(ctx *wasm.HostFunctionCallContext, offset uint32, val uint64) uint32 {
+		if !ctx.Memory.PutUint64(offset, val) {
+			return 1
+		}
+		return 0
 	}
-	err = store.AddHostFunction("env", "add_int", reflect.ValueOf(addInt))
+
+	err = store.AddHostFunction("", "store_int", reflect.ValueOf(storeInt))
 	require.NoError(t, err)
 
 	err = store.Instantiate(mod, "test")
 	require.NoError(t, err)
 
-	// We should be able to call the exported add_int and it should add two ints.
-	results, _, err := store.CallFunction(ctx, "test", "add_int", uint64(12), uint64(30))
+	// Call store_int and ensure it didn't return an error code.
+	results, _, err := store.CallFunction(ctx, "test", "store_int", 1, math.MaxUint64)
 	require.NoError(t, err)
-	require.Equal(t, uint64(42), results[0])
+	require.Equal(t, uint64(0), results[0])
+
+	// Since offset=1 and val=math.MaxUint64, we expect to have written exactly 8 bytes, with all bits set, at index 1.
+	require.Equal(t, []byte{0x0, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x0}, store.Memories[0].Buffer[0:10])
 }
 
 //  hostFuncWithFloatParam fails if a float parameter corrupts a host function value
@@ -237,7 +247,7 @@ func hostFuncWithFloatParam(t *testing.T, newEngine func() wasm.Engine) {
 		;; 'identity_float' just returns the given float value as is.
 		(import "test" "identity_float" (func $test.identity_float (param {{ . }}) (result {{ . }})))
 
-		;; 'call->test.identity_float' proxies 'test.identity_float' via call in order to test floats aren't corrupted 
+		;; 'call->test.identity_float' proxies 'test.identity_float' via call in order to test floats aren't corrupted
 		;; when the float values are passed through OpCodeCall
 		(func $call->test.identity_float (param {{ . }}) (result {{ . }})
 			local.get 0
diff --git a/wasi/testdata/.gitignore b/wasi/testdata/.gitignore
deleted file mode 100644
index 33fb246a..00000000
--- a/wasi/testdata/.gitignore
+++ /dev/null
@@ -1,2 +0,0 @@
-# don't accidentally check-in binaries when a developer uses wat2wasm to test manually
-*.wasm
diff --git a/wasi/testdata/args.wat b/wasi/testdata/args.wat
deleted file mode 100644
index fcfa914d..00000000
--- a/wasi/testdata/args.wat
+++ /dev/null
@@ -1,32 +0,0 @@
-;; This is a wat file to just export args WASI API to the host environment for testing the APIs.
-;; This is currently separated as a wat file and pre-compiled because our text parser doesn't
-;; implement 'memory' yet. After it supports 'memory', we can remove this file and embed this
-;; wat file in the Go test code.
-;;
-;; Note: Although this is a raw wat file which should be moved under /tests/wasi in principle,
-;; this file is put here for now, because this is a temporary file until the parser supports
-;; the enough syntax, and this file will be embedded in unit test codes after that.
-(module
-  (import "wasi_snapshot_preview1" "args_sizes_get"
-    (func $wasi.args_sizes_get (param $result.argc i32) (param $result.argv_buf_size i32) (result (;errno;) i32)))
-  (import "wasi_snapshot_preview1" "args_get"
-    (func $wasi.args_get (param $argv i32) (param $argv_buf i32) (result (;errno;) i32)))
-  (memory 1)  ;; just an arbitrary size big enough for tests
-  (export "memory" (memory 0))
-  ;; Define wrapper functions instead of just exporting the imported WASI APIS for now
-  ;; because wazero's interpreter has a bug that it crashes when an imported-and-exported host function
-  ;; is called from the host environment, which will be fixed soon.
-  ;; After it's fixed, these wrapper functions are no longer necessary.
-  (func $args_sizes_get (param i32 i32) (result i32)
-        local.get 0
-        local.get 1
-        call $wasi.args_sizes_get
-        )
-  (func $args_get (param i32 i32) (result i32)
-        local.get 0
-        local.get 1
-        call $wasi.args_get
-        )
-  (export "args_sizes_get" (func $args_sizes_get))
-  (export "args_get" (func $args_get))
-  )
diff --git a/wasi/testdata/clock.wat b/wasi/testdata/clock.wat
deleted file mode 100644
index 36ac423f..00000000
--- a/wasi/testdata/clock.wat
+++ /dev/null
@@ -1,25 +0,0 @@
-;; This is a wat file to just export clock WASI API to the host environment for testing the APIs.
-;; This is currently separated as a wat file and pre-compiled because our text parser doesn't
-;; implement 'memory' yet. After it supports 'memory', we can remove this file and embed this
-;; wat file in the Go test code.
-;;
-;; Note: Although this is a raw wat file which should be moved under /tests/wasi in principle,
-;; this file is put here for now, because this is a temporary file until the parser supports
-;; the enough syntax, and this file will be embedded in unit test codes after that.
-(module
-  (import "wasi_snapshot_preview1" "clock_time_get"
-    (func $wasi.clock_time_get (param $id i32) (param $precision i64) (param $result.timestamp i32) (result (;errno;) i32)))
-  (memory 1)  ;; just an arbitrary size big enough for tests
-  (export "memory" (memory 0))
-  ;; Define wrapper functions instead of just exporting the imported WASI APIS for now
-  ;; because wazero's interpreter has a bug that it crashes when an imported-and-exported host function
-  ;; is called from the host environment, which will be fixed soon.
-  ;; After it's fixed, these wrapper functions are no longer necessary.
-  (func $clock_time_get (param i32 i64 i32) (result i32)
-        local.get 0
-        local.get 1
-        local.get 2
-        call $wasi.clock_time_get
-        )
-  (export "clock_time_get" (func $clock_time_get))
-  )
diff --git a/wasi/wasi_test.go b/wasi/wasi_test.go
index 1cbbef32..56b96005 100644
--- a/wasi/wasi_test.go
+++ b/wasi/wasi_test.go
@@ -3,6 +3,7 @@ package wasi
 import (
 	"context"
 	_ "embed"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -31,10 +32,6 @@ func TestNewAPI_Args(t *testing.T) {
 	})
 }
 
-// argsWat is a wasm module to call args_get and args_sizes_get.
-//go:embed testdata/args.wat
-var argsWat []byte
-
 func TestAPI_ArgsGet(t *testing.T) {
 	args, err := Args("a", "bc")
 	require.NoError(t, err)
@@ -49,7 +46,7 @@ func TestAPI_ArgsGet(t *testing.T) {
 		3, 0, 0, 0, // little endian-encoded offset of "bc"
 		'?', // stopped after encoding
 	} // tr
-	store, wasiAPI := instantiateWasmStore(t, argsWat, "test", args)
+	store, wasiAPI := instantiateWasmStore(t, FunctionArgsGet, ImportArgsGet, "test", args)
 
 	t.Run("API.ArgsGet", func(t *testing.T) {
 		maskMemory(store, maskLength)
@@ -75,7 +72,7 @@ func TestAPI_ArgsGet(t *testing.T) {
 func TestAPI_ArgsGet_Errors(t *testing.T) {
 	args, err := Args("a", "bc")
 	require.NoError(t, err)
-	store, wasiAPI := instantiateWasmStore(t, argsWat, "test", args)
+	store, wasiAPI := instantiateWasmStore(t, FunctionArgsGet, ImportArgsGet, "test", args)
 
 	memorySize := uint32(len(store.Memories[0].Buffer))
 	validAddress := uint32(0) // arbitrary valid address as arguments to args_get. We chose 0 here.
@@ -132,7 +129,7 @@ func TestAPI_ArgsSizesGet(t *testing.T) {
 		0x5, 0x0, 0x0, 0x0, // little endian-encoded size of null terminated strings
 		'?', // stopped after encoding
 	} // tr
-	store, wasiAPI := instantiateWasmStore(t, argsWat, "test", args)
+	store, wasiAPI := instantiateWasmStore(t, FunctionArgsSizesGet, ImportArgsSizesGet, "test", args)
 
 	t.Run("API.ArgsSizesGet", func(t *testing.T) {
 		maskMemory(store, maskLength)
@@ -158,7 +155,7 @@ func TestAPI_ArgsSizesGet(t *testing.T) {
 func TestAPI_ArgsSizesGet_Errors(t *testing.T) {
 	args, err := Args("a", "bc")
 	require.NoError(t, err)
-	store, _ := instantiateWasmStore(t, argsWat, "test", args)
+	store, _ := instantiateWasmStore(t, FunctionArgsSizesGet, ImportArgsSizesGet, "test", args)
 
 	memorySize := uint32(len(store.Memories[0].Buffer))
 	validAddress := uint32(0) // arbitrary valid address as arguments to args_sizes_get. We chose 0 here.
@@ -205,10 +202,6 @@ func TestAPI_ArgsSizesGet_Errors(t *testing.T) {
 // TODO TestAPI_EnvironSizesGet TestAPI_EnvironSizesGet_Errors
 // TODO TestAPI_ClockResGet TestAPI_ClockResGet_Errors
 
-// clockWat is a wasm module to call clock_time_get.
-//go:embed testdata/clock.wat
-var clockWat []byte
-
 func TestAPI_ClockTimeGet(t *testing.T) {
 	epochNanos := uint64(1640995200000000000) // midnight UTC 2022-01-01
 	resultTimestamp := uint32(1)              // arbitrary offset
@@ -219,7 +212,7 @@ func TestAPI_ClockTimeGet(t *testing.T) {
 		'?', // stopped after encoding
 	} // tr
 
-	store, wasiAPI := instantiateWasmStore(t, clockWat, "test")
+	store, wasiAPI := instantiateWasmStore(t, FunctionClockTimeGet, ImportClockTimeGet, "test")
 	wasiAPI.(*api).timeNowUnixNano = func() uint64 { return epochNanos }
 
 	t.Run("API.ClockTimeGet", func(t *testing.T) {
@@ -245,7 +238,7 @@ func TestAPI_ClockTimeGet(t *testing.T) {
 
 func TestAPI_ClockTimeGet_Errors(t *testing.T) {
 	epochNanos := uint64(1640995200000000000) // midnight UTC 2022-01-01
-	store, wasiAPI := instantiateWasmStore(t, clockWat, "test")
+	store, wasiAPI := instantiateWasmStore(t, FunctionClockTimeGet, ImportClockTimeGet, "test")
 	wasiAPI.(*api).timeNowUnixNano = func() uint64 { return epochNanos }
 
 	memorySize := uint32(len(store.Memories[0].Buffer))
@@ -317,8 +310,13 @@ func TestAPI_ClockTimeGet_Errors(t *testing.T) {
 // TODO: TestAPI_SockSend TestAPI_SockSend_Errors
 // TODO: TestAPI_SockShutdown TestAPI_SockShutdown_Errors
 
-func instantiateWasmStore(t *testing.T, wat []byte, moduleName string, opts ...Option) (*wasm.Store, API) {
-	mod, err := text.DecodeModule(wat)
+func instantiateWasmStore(t *testing.T, wasiFunction, wasiImport string, moduleName string, opts ...Option) (*wasm.Store, API) {
+	mod, err := text.DecodeModule([]byte(fmt.Sprintf(`(module
+  %[2]s
+  (memory 1)  ;; just an arbitrary size big enough for tests
+  (export "memory" (memory 0))
+  (export "%[1]s" (func $wasi.%[1]s))
+)`, wasiFunction, wasiImport)))
 	require.NoError(t, err)
 
 	store := wasm.NewStore(interpreter.NewEngine())
diff --git a/wasm/interpreter/interpreter_test.go b/wasm/interpreter/interpreter_test.go
index 5bd9f24d..774f1e9a 100644
--- a/wasm/interpreter/interpreter_test.go
+++ b/wasm/interpreter/interpreter_test.go
@@ -1,10 +1,13 @@
 package interpreter
 
 import (
+	"context"
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/tetratelabs/wazero/wasm"
 	"github.com/tetratelabs/wazero/wasm/buildoptions"
 )
 
@@ -38,3 +41,29 @@ func TestInterpreter_PushFrame_StackOverflow(t *testing.T) {
 	it.pushFrame(f3)
 	require.Panics(t, func() { it.pushFrame(f4) })
 }
+
+func TestInterpreter_CallHostFunc(t *testing.T) {
+	t.Run("defaults to module memory when call stack empty", func(t *testing.T) {
+		memory := &wasm.MemoryInstance{}
+		var ctxMemory *wasm.MemoryInstance
+		hostFn := reflect.ValueOf(func(ctx *wasm.HostFunctionCallContext) {
+			ctxMemory = ctx.Memory
+		})
+		it := interpreter{functions: map[wasm.FunctionAddress]*interpreterFunction{
+			0: {hostFn: &hostFn, funcInstance: &wasm.FunctionInstance{
+				FunctionType: &wasm.TypeInstance{
+					Type: &wasm.FunctionType{
+						Params:  []wasm.ValueType{},
+						Results: []wasm.ValueType{},
+					},
+				},
+				ModuleInstance: &wasm.ModuleInstance{Memory: memory},
+			},
+			},
+		}}
+
+		// When calling a host func directly, there may be no stack. This ensures the module's memory is used.
+		it.callHostFunc(context.Background(), it.functions[0])
+		require.Same(t, memory, ctxMemory)
+	})
+}
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 60e7e99f..bb321822 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -4803,24 +4803,25 @@ func TestAmd64Compiler_setupMemoryOffset(t *testing.T) {
 
 					compiler.locationStack.pushValueLocationOnRegister(reg)
 
-					// Generate the code under test.
 					err = compiler.releaseAllRegistersToStack()
 					require.NoError(t, err)
 					compiler.exit(jitCallStatusCodeReturned)
+
+					// Generate the code under test and run.
 					code, _, _, err := compiler.compile()
 					require.NoError(t, err)
-
-					// Run code.
 					env.exec(code)
 
-					// If the memory offset exceeds the length of memory, we must exit the function
-					// with jitCallStatusCodeMemoryOutOfBounds status code.
 					mem := env.memory()
-					baseOffset := int(base) + int(offset)
-					if baseOffset >= math.MaxUint32 || len(mem) < baseOffset+int(targetSizeInByte) {
+					baseOffset := int64(base) + int64(offset)
+					if ceil := baseOffset + int64(targetSizeInByte); int64(len(mem)) < ceil {
+						// If the targe memory region's ceil exceeds the length of memory, we must exit the function
+						// with jitCallStatusCodeMemoryOutOfBounds status code.
 						require.Equal(t, jitCallStatusCodeMemoryOutOfBounds, env.jitStatus())
 					} else {
 						require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+						require.Equal(t, uint64(1), env.stackPointer())
+						require.Equal(t, uint64(baseOffset), env.stackTopAsUint64())
 					}
 				})
 			}
diff --git a/wasm/store_test.go b/wasm/store_test.go
index 2b59ce46..1f4b8345 100644
--- a/wasm/store_test.go
+++ b/wasm/store_test.go
@@ -109,6 +109,33 @@ func TestStore_AddHostFunction(t *testing.T) {
 	require.Equal(t, map[string]*ExportInstance{"fn": exp}, m.Exports)
 }
 
+func TestStore_ExportImportedHostFunction(t *testing.T) {
+	s := NewStore(nopEngineInstance)
+	hostFunction := func(_ *HostFunctionCallContext) {
+	}
+
+	err := s.AddHostFunction("", "host_fn", reflect.ValueOf(hostFunction))
+	require.NoError(t, err)
+
+	t.Run("ModuleInstance is the importing module", func(t *testing.T) {
+		err = s.Instantiate(&Module{
+			TypeSection:   []*FunctionType{{}},
+			ImportSection: []*Import{{Kind: ImportKindFunc, Name: "host_fn", DescFunc: 0}},
+			MemorySection: []*MemoryType{{1, nil}},
+			ExportSection: map[string]*Export{"host.fn": {Kind: ExportKindFunc, Name: "host.fn", Index: 0}},
+		}, "test")
+		require.NoError(t, err)
+
+		ei, err := s.getExport("test", "host.fn", ExportKindFunc)
+		require.NoError(t, err)
+
+		// We expect the host function to be called in context of the importing module.
+		// Otherwise, it would be the pseudo-module of the host, which only includes types and function definitions.
+		// Notably, this ensures the host function call context has the correct memory (from the importing module).
+		require.Equal(t, s.ModuleInstances["test"], ei.Function.ModuleInstance)
+	})
+}
+
 func TestStore_BuildFunctionInstances_FunctionNames(t *testing.T) {
 	name := "test"