[wasm] Fix typing of stack-switching wrappers

- Suspender.suspendOnReturnedPromise expects a function with type
[ti*]->[externref] and returns a function with the same type.
- Suspender.returnPromiseOnSuspend expects a function with type
[ti*]->[to] and returns a function with type [ti*]->[externref].

Changes:
- Check the wrapped function's return types
- Skip type checking of return types when importing a wrapper (and
assert that the return type is externref)
- Add special case for WebAssembly.Function.type of a
WasmExportedFunction: it currently returns the signature declared by
the module. Change the return type to externref if this is a
stack-switching export.

Bug: v8:12191
Change-Id: I6619c306e9613825ad1b021cb3400d73cd684656
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3435190
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#78973}

src/compiler/types.cc

@@ -270,13 +270,14 @@ Type::bitset BitsetType::Lub(const MapRefLike& map) {
     case JS_TEMPORAL_ZONED_DATE_TIME_TYPE:
 #if V8_ENABLE_WEBASSEMBLY
     case WASM_ARRAY_TYPE:
-    case WASM_TAG_OBJECT_TYPE:
     case WASM_GLOBAL_OBJECT_TYPE:
     case WASM_INSTANCE_OBJECT_TYPE:
     case WASM_MEMORY_OBJECT_TYPE:
     case WASM_MODULE_OBJECT_TYPE:
     case WASM_STRUCT_TYPE:
+    case WASM_SUSPENDER_OBJECT_TYPE:
     case WASM_TABLE_OBJECT_TYPE:
+    case WASM_TAG_OBJECT_TYPE:
     case WASM_VALUE_OBJECT_TYPE:
 #endif  // V8_ENABLE_WEBASSEMBLY
     case WEAK_CELL_TYPE:

src/compiler/wasm-compiler.cc

@@ -7574,10 +7574,13 @@ WasmImportData ResolveWasmImportCall(
   Handle<HeapObject> suspender = isolate->factory()->undefined_value();
   if (WasmJSFunction::IsWasmJSFunction(*callable)) {
     auto js_function = Handle<WasmJSFunction>::cast(callable);
-    if (!js_function->MatchesSignature(expected_sig)) {
+    suspender = handle(js_function->GetSuspender(), isolate);
+    if ((!suspender->IsUndefined() &&
+         !js_function->MatchesSignatureForSuspend(expected_sig)) ||
+        (suspender->IsUndefined() &&
+         !js_function->MatchesSignature(expected_sig))) {
       return {WasmImportCallKind::kLinkError, callable, no_suspender};
     }
-    suspender = handle(js_function->GetSuspender(), isolate);
     // Resolve the short-cut to the underlying callable and continue.
     callable = handle(js_function->GetCallable(), isolate);
   }

src/objects/fixed-array.h

@@ -594,6 +594,13 @@ class PodArray : public ByteArray {
     return memcmp(GetDataStartAddress(), buffer, length * sizeof(T)) == 0;
   }
 
+  bool matches(int offset, const T* buffer, int length) {
+    DCHECK_LE(offset, this->length());
+    DCHECK_LE(offset + length, this->length());
+    return memcmp(GetDataStartAddress() + sizeof(T) * offset, buffer,
+                  length * sizeof(T)) == 0;
+  }
+
   T get(int index) {
     T result;
     copy_out(index, &result, 1);

src/wasm/wasm-js.cc

@@ -1912,7 +1912,23 @@ void WebAssemblyFunctionType(const v8::FunctionCallbackInfo<v8::Value>& args) {
   i::Zone zone(i_isolate->allocator(), ZONE_NAME);
   i::Handle<i::Object> arg0 = Utils::OpenHandle(*args[0]);
   if (i::WasmExportedFunction::IsWasmExportedFunction(*arg0)) {
-    sig = i::Handle<i::WasmExportedFunction>::cast(arg0)->sig();
+    auto wasm_exported_function =
+        i::Handle<i::WasmExportedFunction>::cast(arg0);
+    auto sfi = handle(wasm_exported_function->shared(), i_isolate);
+    i::Handle<i::WasmExportedFunctionData> data =
+        handle(sfi->wasm_exported_function_data(), i_isolate);
+    sig = wasm_exported_function->sig();
+    if (!data->suspender().IsUndefined()) {
+      // If this export is wrapped by a Suspender, the function returns a
+      // promise as an externref instead of the original return type.
+      size_t param_count = sig->parameter_count();
+      i::wasm::FunctionSig::Builder builder(&zone, 1, param_count);
+      for (size_t i = 0; i < param_count; ++i) {
+        builder.AddParam(sig->GetParam(0));
+      }
+      builder.AddReturn(i::wasm::kWasmExternRef);
+      sig = builder.Build();
+    }
   } else if (i::WasmJSFunction::IsWasmJSFunction(*arg0)) {
     sig = i::Handle<i::WasmJSFunction>::cast(arg0)->GetSignature(&zone);
   } else {
@@ -2571,6 +2587,10 @@ void WebAssemblySuspenderReturnPromiseOnSuspend(
     thrower.TypeError("Argument 0 must be a wasm function");
   }
   i::WasmExportedFunctionData data = sfi.wasm_exported_function_data();
+  if (data.sig()->return_count() != 1) {
+    thrower.TypeError(
+        "Expected a WebAssembly.Function with exactly one return type");
+  }
   int index = data.function_index();
   i::Handle<i::WasmInstanceObject> instance(
       i::WasmInstanceObject::cast(data.internal().ref()), i_isolate);
@@ -2612,6 +2632,11 @@ void WebAssemblySuspenderSuspendOnReturnedPromise(
     return;
   }
   sig = i::Handle<i::WasmJSFunction>::cast(arg0)->GetSignature(&zone);
+  if (sig->return_count() != 1 ||
+      sig->GetReturn(0) != i::wasm::kWasmExternRef) {
+    thrower.TypeError(
+        "Expected a WebAssembly.Function with return type externref");
+  }
 
   auto callable = handle(
       i::Handle<i::WasmJSFunction>::cast(arg0)->GetCallable(), i_isolate);

src/wasm/wasm-objects.cc

@@ -2150,12 +2150,33 @@ const wasm::FunctionSig* WasmJSFunction::GetSignature(Zone* zone) {
   return zone->New<wasm::FunctionSig>(return_count, parameter_count, types);
 }
 
+bool WasmJSFunction::MatchesSignatureForSuspend(const wasm::FunctionSig* sig) {
+  DCHECK_LE(sig->all().size(), kMaxInt);
+  int sig_size = static_cast<int>(sig->all().size());
+  int parameter_count = static_cast<int>(sig->parameter_count());
+  int return_count = static_cast<int>(sig->return_count());
+  DisallowHeapAllocation no_alloc;
+  WasmJSFunctionData function_data = shared().wasm_js_function_data();
+  if (parameter_count != function_data.serialized_parameter_count()) {
+    return false;
+  }
+  if (sig_size == 0) return true;  // Prevent undefined behavior.
+  // This function is only called for functions wrapped by a
+  // WebAssembly.Suspender object, so the return type has to be externref.
+  CHECK_EQ(function_data.serialized_return_count(), 1);
+  CHECK_EQ(function_data.serialized_signature().get(0), wasm::kWasmExternRef);
+  const wasm::ValueType* expected = sig->all().begin();
+  return function_data.serialized_signature().matches(
+      1, expected + return_count, parameter_count);
+}
+
 // TODO(9495): Update this if function type variance is introduced.
 bool WasmJSFunction::MatchesSignature(const wasm::FunctionSig* sig) {
   DCHECK_LE(sig->all().size(), kMaxInt);
   int sig_size = static_cast<int>(sig->all().size());
   int return_count = static_cast<int>(sig->return_count());
   int parameter_count = static_cast<int>(sig->parameter_count());
+  DisallowHeapAllocation no_alloc;
   WasmJSFunctionData function_data = shared().wasm_js_function_data();
   if (return_count != function_data.serialized_return_count() ||
       parameter_count != function_data.serialized_parameter_count()) {

src/wasm/wasm-objects.h

@@ -631,6 +631,8 @@ class WasmJSFunction : public JSFunction {
   // that lifetime of the signature is hence directly coupled to the zone.
   const wasm::FunctionSig* GetSignature(Zone* zone);
   bool MatchesSignature(const wasm::FunctionSig* sig);
+  // Special typing rule for imports wrapped by a Suspender.
+  bool MatchesSignatureForSuspend(const wasm::FunctionSig* sig);
 
   DECL_CAST(WasmJSFunction)
   OBJECT_CONSTRUCTORS(WasmJSFunction, JSFunction);

test/mjsunit/wasm/stack-switching.js

@@ -15,12 +15,68 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
       /WebAssembly.Suspender must be invoked with 'new'/);
 })();
 
+(function TestSuspenderTypes() {
+  print(arguments.callee.name);
+  let builder = new WasmModuleBuilder();
+  builder.addImport('m', 'import', kSig_v_i);
+  builder.addFunction("export", kSig_i_i)
+      .addBody([kExprLocalGet, 0]).exportFunc();
+  builder.addFunction("wrong1", kSig_ii_v)
+      .addBody([kExprI32Const, 0, kExprI32Const, 0]).exportFunc();
+  builder.addFunction("wrong2", kSig_v_v)
+      .addBody([]).exportFunc();
+  let suspender = new WebAssembly.Suspender();
+  function js_import(i) {
+    return new Promise((resolve) => { resolve(42); });
+  }
+
+  // Wrap the import, instantiate the module, and wrap the export.
+  let wasm_js_import = new WebAssembly.Function(
+      {parameters: ['i32'], results: ['externref']}, js_import);
+  let import_wrapper = suspender.suspendOnReturnedPromise(wasm_js_import);
+  let instance = builder.instantiate({'m': {'import': import_wrapper}});
+  let export_wrapper =
+      suspender.returnPromiseOnSuspend(instance.exports.export);
+
+  // Check type errors.
+  wasm_js_import = new WebAssembly.Function(
+      {parameters: [], results: ['i32']}, js_import);
+  assertThrows(() => suspender.suspendOnReturnedPromise(wasm_js_import),
+      TypeError, /Expected a WebAssembly.Function with return type externref/);
+  assertThrows(() => suspender.returnPromiseOnSuspend(instance.exports.wrong1),
+      TypeError,
+      /Expected a WebAssembly.Function with exactly one return type/);
+  assertThrows(() => suspender.returnPromiseOnSuspend(instance.exports.wrong2),
+      TypeError,
+      /Expected a WebAssembly.Function with exactly one return type/);
+  // Signature mismatch (link error).
+  let wrong_import = new WebAssembly.Function(
+      {parameters: ['f32'], results: ['externref']}, () => {});
+  wrong_import = suspender.suspendOnReturnedPromise(wrong_import);
+  assertThrows(() => builder.instantiate({'m': {'import': wrong_import}}),
+      WebAssembly.LinkError,
+      /imported function does not match the expected type/);
+
+  // Check the wrapped export's signature.
+  let export_sig = WebAssembly.Function.type(export_wrapper);
+  assertEquals(['i32'], export_sig.parameters);
+  assertEquals(['externref'], export_sig.results);
+
+  // Check the wrapped import's signature.
+  let import_sig = WebAssembly.Function.type(import_wrapper);
+  assertEquals(['i32'], import_sig.parameters);
+  assertEquals(['externref'], import_sig.results);
+})();
+
 (function TestStackSwitchNoSuspend() {
   print(arguments.callee.name);
   let builder = new WasmModuleBuilder();
   builder.addGlobal(kWasmI32, true).exportAs('g');
-  builder.addFunction("test", kSig_v_v)
-      .addBody([kExprI32Const, 42, kExprGlobalSet, 0]).exportFunc();
+  builder.addFunction("test", kSig_i_v)
+      .addBody([
+          kExprI32Const, 42,
+          kExprGlobalSet, 0,
+          kExprI32Const, 0]).exportFunc();
   let instance = builder.instantiate();
   let suspender = new WebAssembly.Suspender();
   let wrapper = suspender.returnPromiseOnSuspend(instance.exports.test);
@@ -33,17 +89,21 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   let builder = new WasmModuleBuilder();
   builder.addGlobal(kWasmI32, true).exportAs('g');
   import_index = builder.addImport('m', 'import', kSig_i_v);
-  builder.addFunction("test", kSig_v_v)
+  builder.addFunction("test", kSig_i_v)
       .addBody([
           kExprCallFunction, import_index, // suspend
-          kExprGlobalSet, 0 // resume
+          kExprGlobalSet, 0, // resume
+          kExprI32Const, 0,
       ]).exportFunc();
   let suspender = new WebAssembly.Suspender();
   function js_import() {
     return new Promise((resolve) => { resolve(42); });
   };
-  let wasm_js_import = new WebAssembly.Function({parameters: [], results: ['i32']}, js_import);
-  let suspending_wasm_js_import = suspender.suspendOnReturnedPromise(wasm_js_import);
+  let wasm_js_import = new WebAssembly.Function(
+      {parameters: [], results: ['externref']}, js_import);
+  let suspending_wasm_js_import =
+      suspender.suspendOnReturnedPromise(wasm_js_import);
+
   let instance = builder.instantiate({m: {import: suspending_wasm_js_import}});
   let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
   let combined_promise = wrapped_export();
@@ -61,7 +121,7 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   // for (i = 0; i < 5; ++i) {
   //   g = g + import();
   // }
-  builder.addFunction("test", kSig_v_v)
+  builder.addFunction("test", kSig_i_v)
       .addLocals(kWasmI32, 1)
       .addBody([
           kExprI32Const, 5,
@@ -77,6 +137,7 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
             kExprLocalTee, 0,
             kExprBrIf, 0,
           kExprEnd,
+          kExprI32Const, 0,
       ]).exportFunc();
   let suspender = new WebAssembly.Suspender();
   let i = 0;
@@ -84,8 +145,10 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   function js_import() {
     return new Promise((resolve) => { resolve(++i); });
   };
-  let wasm_js_import = new WebAssembly.Function({parameters: [], results: ['i32']}, js_import);
-  let suspending_wasm_js_import = suspender.suspendOnReturnedPromise(wasm_js_import);
+  let wasm_js_import = new WebAssembly.Function(
+      {parameters: [], results: ['externref']}, js_import);
+  let suspending_wasm_js_import =
+      suspender.suspendOnReturnedPromise(wasm_js_import);
   let instance = builder.instantiate({m: {import: suspending_wasm_js_import}});
   let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
   let chained_promise = wrapped_export();
@@ -97,8 +160,8 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   print(arguments.callee.name);
   let builder = new WasmModuleBuilder();
   let gc_index = builder.addImport('m', 'gc', kSig_v_v);
-  builder.addFunction("test", kSig_v_v)
-      .addBody([kExprCallFunction, gc_index]).exportFunc();
+  builder.addFunction("test", kSig_i_v)
+      .addBody([kExprCallFunction, gc_index, kExprI32Const, 0]).exportFunc();
   let instance = builder.instantiate({'m': {'gc': gc}});
   let suspender = new WebAssembly.Suspender();
   let wrapper = suspender.returnPromiseOnSuspend(instance.exports.test);