From f8e55b104ef51a6586a98b6d36cd46c6357178eb Mon Sep 17 00:00:00 2001
From: Takeshi Yoneda <takeshi@tetrate.io>
Date: Wed, 31 Aug 2022 10:01:17 +0900
Subject: [PATCH] fix

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>
---
 internal/engine/compiler/engine.go      | 35 ++++++++++++++-----------
 internal/engine/compiler/engine_test.go |  6 +++--
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index a43930df9b3..b197edc4b29 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -103,6 +103,9 @@ type (
 	moduleContext struct {
 		// moduleInstanceAddress is the address of module instance from which we initialize
 		// the following fields. This is set whenever we enter a function or return from function calls.
+		//
+		// On the entry to the native code, this must be initialized to zero to let native code preamble know
+		// that this is the initial function call (which leads to moduleContext initialization pass).
 		moduleInstanceAddress uintptr //lint:ignore U1000 This is only used by Compiler code.
 
 		// globalElement0Address is the address of the first element in the global slice,
@@ -562,12 +565,7 @@ func (ce *callEngine) Call(ctx context.Context, callCtx *wasm.CallContext, param
 		}
 		// TODO: ^^ Will not fail if the function was imported from a closed module.
 
-		if v := recover(); v != nil {
-			err = ce.recoverOnCall(v)
-		}
-
-		// Allows the reuse of CallEngine.
-		ce.stackBasePointerInBytes, ce.stackPointer, ce.callFrameStackPointer = 0, 0, 0
+		err = ce.deferredOnCall(recover())
 	}()
 
 	for _, v := range params {
@@ -578,16 +576,23 @@ func (ce *callEngine) Call(ctx context.Context, callCtx *wasm.CallContext, param
 	return
 }
 
-// recoverOnCall takes the recovered value `recoverOnCall`, and wraps it
-// with the call frame stack traces. Also, reset the state of callEngine
-// so that it can be used for the subsequent calls.
-func (ce *callEngine) recoverOnCall(v interface{}) (err error) {
-	builder := wasmdebug.NewErrorBuilder()
-	for i := uint64(0); i < ce.callFrameStackPointer; i++ {
-		def := ce.callFrameStack[ce.callFrameStackPointer-1-i].function.source.FunctionDefinition
-		builder.AddFrame(def.DebugName(), def.ParamTypes(), def.ResultTypes())
+// deferredOnCall takes the recovered value `recovered`, and wraps it
+// with the call frame stack traces if v != nil. Also, reset
+// the state of callEngine so that it can be used for the subsequent calls.
+//
+// This is defined for testability.
+func (ce *callEngine) deferredOnCall(recovered interface{}) (err error) {
+	if recovered != nil {
+		builder := wasmdebug.NewErrorBuilder()
+		for i := uint64(0); i < ce.callFrameStackPointer; i++ {
+			def := ce.callFrameStack[ce.callFrameStackPointer-1-i].function.source.FunctionDefinition
+			builder.AddFrame(def.DebugName(), def.ParamTypes(), def.ResultTypes())
+		}
+		err = builder.FromRecovered(recovered)
 	}
-	err = builder.FromRecovered(v)
+
+	// Allows the reuse of CallEngine.
+	ce.stackBasePointerInBytes, ce.stackPointer, ce.callFrameStackPointer, ce.moduleInstanceAddress = 0, 0, 0, 0
 	return
 }
 
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index efc85a74633..80773726d94 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -305,7 +305,7 @@ func TestCallEngine_builtinFunctionTableGrow(t *testing.T) {
 	require.Equal(t, uintptr(0xff), table.References[0])
 }
 
-func TestCallEngine_recoverOnCall(t *testing.T) {
+func TestCallEngine_deferredOnCall(t *testing.T) {
 	ce := &callEngine{
 		valueStack:        make([]uint64, 100),
 		valueStackContext: valueStackContext{stackPointer: 3},
@@ -317,11 +317,12 @@ func TestCallEngine_recoverOnCall(t *testing.T) {
 			{function: &function{source: &wasm.FunctionInstance{FunctionDefinition: newMockFunctionDefinition("4")}}},
 			{function: &function{source: &wasm.FunctionInstance{FunctionDefinition: newMockFunctionDefinition("5")}}},
 		},
+		moduleContext: moduleContext{moduleInstanceAddress: 0xdeafbeaf},
 	}
 
 	beforeRecoverValueStack, beforeRecoverCallFrameStack := ce.valueStack, ce.callFrameStack
 
-	err := ce.recoverOnCall(errors.New("some error"))
+	err := ce.deferredOnCall(errors.New("some error"))
 	require.EqualError(t, err, `some error (recovered by wazero)
 wasm stack trace:
 	5()
@@ -334,6 +335,7 @@ wasm stack trace:
 	// for the subsequent calls.
 	require.Equal(t, uint64(0), ce.stackBasePointerInBytes)
 	require.Equal(t, uint64(0), ce.callFrameStackPointer)
+	require.Equal(t, uintptr(0), ce.moduleInstanceAddress)
 	require.Equal(t, beforeRecoverValueStack, ce.valueStack)
 	require.Equal(t, beforeRecoverCallFrameStack, ce.callFrameStack)
 }