Implement Emscripten jmp support in indirect function calls (invoke_xxx)

imports/emscripten/emscripten.go

@@ -104,6 +104,10 @@ func NewFunctionExporterForModule(guest wazero.CompiledModule) (FunctionExporter
 			ret = append(ret, internal.NotifyMemoryGrowth)
 			continue
 		}
+		if importName == internal.FunctionThrowLongjmp {
+			ret = append(ret, internal.ThrowLongjmp)
+			continue
+		}
 		if !strings.HasPrefix(importName, internal.InvokePrefix) {
 			continue // not invoke, and maybe not emscripten
 		}

imports/emscripten/emscripten_test.go

@@ -370,6 +370,8 @@ func TestInstantiateForModule(t *testing.T) {
 			expectedResults: []uint64{42},
 			expectedLog: `--> .call_v_i32(0)
 	==> env.invoke_i(index=0)
+		--> .stackSave()
+		<-- 65536
 		--> .v_i32()
 		<-- 42
 	<== 42
@@ -384,6 +386,8 @@ func TestInstantiateForModule(t *testing.T) {
 			expectedResults: []uint64{42},
 			expectedLog: `--> .call_i32_i32(2,42)
 	==> env.invoke_ii(index=2,a1=42)
+		--> .stackSave()
+		<-- 65536
 		--> .i32_i32(42)
 		<-- 42
 	<== 42
@@ -398,6 +402,8 @@ func TestInstantiateForModule(t *testing.T) {
 			expectedResults: []uint64{3},
 			expectedLog: `--> .call_i32i32_i32(4,1,2)
 	==> env.invoke_iii(index=4,a1=1,a2=2)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32_i32(1,2)
 		<-- 3
 	<== 3
@@ -412,6 +418,8 @@ func TestInstantiateForModule(t *testing.T) {
 			expectedResults: []uint64{7},
 			expectedLog: `--> .call_i32i32i32_i32(6,1,2,4)
 	==> env.invoke_iiii(index=6,a1=1,a2=2,a3=4)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32i32_i32(1,2,4)
 		<-- 7
 	<== 7
@@ -426,6 +434,8 @@ func TestInstantiateForModule(t *testing.T) {
 			expectedResults: []uint64{15},
 			expectedLog: `--> .calli32_i32i32i32i32_i32(8,1,2,4,8)
 	==> env.invoke_iiiii(index=8,a1=1,a2=2,a3=4,a4=8)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32i32i32_i32(1,2,4,8)
 		<-- 15
 	<== 15
@@ -438,6 +448,8 @@ func TestInstantiateForModule(t *testing.T) {
 			tableOffset: 10,
 			expectedLog: `--> .call_v_v(10)
 	==> env.invoke_v(index=10)
+		--> .stackSave()
+		<-- 65536
 		--> .v_v()
 		<--
 	<==
@@ -451,6 +463,8 @@ func TestInstantiateForModule(t *testing.T) {
 			params:      []uint64{42},
 			expectedLog: `--> .call_i32_v(12,42)
 	==> env.invoke_vi(index=12,a1=42)
+		--> .stackSave()
+		<-- 65536
 		--> .i32_v(42)
 		<--
 	<==
@@ -464,6 +478,8 @@ func TestInstantiateForModule(t *testing.T) {
 			params:      []uint64{1, 2},
 			expectedLog: `--> .call_i32i32_v(14,1,2)
 	==> env.invoke_vii(index=14,a1=1,a2=2)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32_v(1,2)
 		<--
 	<==
@@ -477,6 +493,8 @@ func TestInstantiateForModule(t *testing.T) {
 			params:      []uint64{1, 2, 4},
 			expectedLog: `--> .call_i32i32i32_v(16,1,2,4)
 	==> env.invoke_viii(index=16,a1=1,a2=2,a3=4)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32i32_v(1,2,4)
 		<--
 	<==
@@ -490,10 +508,31 @@ func TestInstantiateForModule(t *testing.T) {
 			params:      []uint64{1, 2, 4, 8},
 			expectedLog: `--> .calli32_i32i32i32i32_v(18,1,2,4,8)
 	==> env.invoke_viiii(index=18,a1=1,a2=2,a3=4,a4=8)
+		--> .stackSave()
+		<-- 65536
 		--> .i32i32i32i32_v(1,2,4,8)
 		<--
 	<==
 <--
+`,
+		},
+		{
+			name:        "invoke_v_with_longjmp",
+			funcName:    "call_invoke_v_with_longjmp_throw",
+			tableOffset: 20,
+			params:      []uint64{},
+			expectedLog: `--> .call_invoke_v_with_longjmp_throw(20)
+	==> env.invoke_v(index=20)
+		--> .stackSave()
+		<-- 42
+		--> .call_longjmp_throw()
+			==> env._emscripten_throw_longjmp()
+		--> .stackRestore(42)
+		<--
+		--> .setThrew(1,0)
+		<--
+	<==
+<--
 `,
 		},
 	}

imports/emscripten/testdata/invoke.wat

@@ -10,8 +10,17 @@
   (import "env" "invoke_vii" (func $invoke_vii (param i32 i32 i32)))
   (import "env" "invoke_viii" (func $invoke_viii (param i32 i32 i32 i32)))
   (import "env" "invoke_viiii" (func $invoke_viiii (param i32 i32 i32 i32 i32)))
+  (import "env" "_emscripten_throw_longjmp" (func $_emscripten_throw_longjmp))
 
-  (table 20 20 funcref)
+  (table 22 22 funcref)
+
+  (global $__stack_pointer (mut i32) (i32.const 65536))
+  (func $stackSave (export "stackSave") (result i32)
+    global.get $__stack_pointer)
+  (func $stackRestore (export "stackRestore") (param i32)
+    local.get 0
+    global.set $__stack_pointer)
+  (func $setThrew (export "setThrew") (param i32 i32))
 
   (func $v_i32 (result i32) (i32.const 42))
   (func $v_i32_unreachable (result i32) unreachable)
@@ -112,4 +121,25 @@
   ;; numbers and expect unreachable on 19.
   (func $calli32_i32i32i32i32_v (export "calli32_i32i32i32i32_v") (param i32 i32 i32 i32 i32)
     (call $invoke_viiii (local.get 0) (local.get 1) (local.get 2) (local.get 3) (local.get 4)))
+
+  (func $call_longjmp_throw
+    (call $_emscripten_throw_longjmp)
+    (global.set $__stack_pointer (i32.const 43)))
+
+  (func $call_longjmp_throw_unreachable unreachable)
+
+  (elem (i32.const 20) $call_longjmp_throw $call_longjmp_throw_unreachable)
+
+  ;; $call_invoke_v_with_longjmp_throw should be called with 20 or 21 and
+  ;; expect unreachable on 21. $call_invoke_v_with_longjmp_throw mimics
+  ;; Emscripten by setting the stack pointer to a different value than default.
+  ;; We ensure that the stack pointer was not changed to 43 by $call_longjump.
+  (func $call_invoke_v_with_longjmp_throw (export "call_invoke_v_with_longjmp_throw") (param i32)
+    (global.set $__stack_pointer (i32.const 42))
+    (call $invoke_v (local.get 0))
+    global.get $__stack_pointer
+    i32.const 42
+    i32.ne
+    (if (then unreachable))
+  )
 )

internal/emscripten/emscripten.go

@@ -2,10 +2,12 @@ package emscripten
 
 import (
 	"context"
+	"errors"
 	"strconv"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/wasm"
+	"github.com/tetratelabs/wazero/sys"
 )
 
 const FunctionNotifyMemoryGrowth = "emscripten_notify_memory_growth"
@@ -18,6 +20,25 @@ var NotifyMemoryGrowth = &wasm.HostFunc{
 	Code:       wasm.Code{GoFunc: api.GoModuleFunc(func(context.Context, api.Module, []uint64) {})},
 }
 
+// Emscripten uses this host method to throw an error that can then be caught
+// in the dynamic invoke functions. Emscripten uses this to allow for
+// setjmp/longjmp support. When this error is seen in the invoke handler,
+// it ignores the error and does not act on it.
+const FunctionThrowLongjmp = "_emscripten_throw_longjmp"
+
+var (
+	ThrowLongjmpError = errors.New("_emscripten_throw_longjmp")
+	ThrowLongjmp      = &wasm.HostFunc{
+		ExportName: FunctionThrowLongjmp,
+		Name:       FunctionThrowLongjmp,
+		ParamTypes: []wasm.ValueType{},
+		ParamNames: []string{},
+		Code: wasm.Code{GoFunc: api.GoModuleFunc(func(context.Context, api.Module, []uint64) {
+			panic(ThrowLongjmpError)
+		})},
+	}
+)
+
 // InvokePrefix is the naming convention of Emscripten dynamic functions.
 //
 // All `invoke_` functions have an initial "index" parameter of
@@ -89,8 +110,62 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 		panic(err)
 	}
 
-	err = f.CallWithStack(ctx, stack)
+	// The Go implementation below mimics the Emscripten JS behaviour to support
+	// longjmps from indirect function calls. The implementation of these
+	// indirection function calls in Emscripten JS is like this:
+	//
+	// function invoke_iii(index,a1,a2) {
+	//  var sp = stackSave();
+	//  try {
+	//    return getWasmTableEntry(index)(a1,a2);
+	//  } catch(e) {
+	//    stackRestore(sp);
+	//    if (e !== e+0) throw e;
+	//    _setThrew(1, 0);
+	//  }
+	//}
+
+	// This is the equivalent of "var sp = stackSave();".
+	// We reuse savedStack to save allocations. We allocate with a size of 2
+	// here to accommodate for the input and output of setThrew.
+	var savedStack [2]uint64
+	err = mod.ExportedFunction("stackSave").CallWithStack(ctx, savedStack[:])
 	if err != nil {
 		panic(err)
 	}
+
+	err = f.CallWithStack(ctx, stack)
+	if err != nil {
+		// Module closed: any calls will just fail with the same error.
+		if _, ok := err.(*sys.ExitError); ok {
+			panic(err)
+		}
+
+		// This is the equivalent of "stackRestore(sp);".
+		// Do not overwrite err here to preserve the original error.
+		if err := mod.ExportedFunction("stackRestore").CallWithStack(ctx, savedStack[:]); err != nil {
+			panic(err)
+		}
+
+		// If we encounter ThrowLongjmpError, this means that the C code did a
+		// longjmp, which in turn called _emscripten_throw_longjmp and that is
+		// a host function that panics with ThrowLongjmpError. In that case we
+		// ignore the error because we have restored the stack to what it was
+		// before the indirect function call, so the program can continue.
+		// This is the equivalent of the "if (e !== e+0) throw e;" line in the
+		// JS implementation, which checks if the error is not a number, which
+		// is what the JS implementation throws (Infinity for
+		// _emscripten_throw_longjmp, a memory address for C++ exceptions).
+		if !errors.Is(err, ThrowLongjmpError) {
+			panic(err)
+		}
+
+		// This is the equivalent of "_setThrew(1, 0);".
+		savedStack[0] = 1
+		savedStack[1] = 0
+		err = mod.ExportedFunction("setThrew").CallWithStack(ctx, savedStack[:])
+		if err != nil {
+			panic(err) // setThrew failed
+		}
+	}
 }