Forward extra arguments and guard buffer registration

Author: andreiborza

packages/wasm/src/patchWebAssembly.ts

@@ -68,7 +68,7 @@ function registerFromBufferSource(
 ): void {
   const url = getWasmSourceUrl(source);
   if (url) {
-    registerModule(module, url);
+    registerSafely(registerModule, module, url);
   }
 }
 
@@ -82,34 +82,26 @@ function patchNonStreamingWebAssembly(registerModule: RegisterModuleCallback): v
 
   nonStreamingPatched = true;
 
-  const origInstantiate = WebAssembly.instantiate;
-  WebAssembly.instantiate = function instantiate(
-    source: BufferSource | WebAssembly.Module,
-    importObject?: WebAssembly.Imports,
-  ) {
+  // Double-cast, because the overloaded native signature (buffer vs. module
+  // first argument) cannot be widened to a pass-through shape in one step.
+  const origInstantiate = WebAssembly.instantiate as unknown as (
+    source: unknown,
+    ...rest: unknown[]
+  ) => Promise<WebAssembly.WebAssemblyInstantiatedSource>;
+  WebAssembly.instantiate = function instantiate(source: BufferSource | WebAssembly.Module, ...rest: unknown[]) {
     if (source instanceof WebAssembly.Module) {
-      return (
-        origInstantiate as (
-          moduleObject: WebAssembly.Module,
-          importObject?: WebAssembly.Imports,
-        ) => Promise<WebAssembly.Instance>
-      )(source, importObject);
+      return origInstantiate(source, ...rest);
     }
 
-    return (
-      origInstantiate as (
-        bytes: BufferSource,
-        importObject?: WebAssembly.Imports,
-      ) => Promise<WebAssembly.WebAssemblyInstantiatedSource>
-    )(source, importObject).then(result => {
+    return origInstantiate(source, ...rest).then(result => {
       registerFromBufferSource(registerModule, result.module, source);
       return result;
     });
   } as typeof WebAssembly.instantiate;
 
-  const origCompile = WebAssembly.compile;
-  WebAssembly.compile = function compile(source: BufferSource): Promise<WebAssembly.Module> {
-    return origCompile(source).then(module => {
+  const origCompile = WebAssembly.compile as (source: unknown, ...rest: unknown[]) => Promise<WebAssembly.Module>;
+  WebAssembly.compile = function compile(source: BufferSource, ...rest: unknown[]): Promise<WebAssembly.Module> {
+    return origCompile(source, ...rest).then(module => {
       registerFromBufferSource(registerModule, module, source);
       return module;
     });

packages/wasm/test/patchWebAssembly.test.ts

@@ -154,3 +154,51 @@ describe('patchWebAssembly() non-streaming registration', () => {
     expect(IMAGES).toHaveLength(0);
   });
 });
+
+describe('patchWebAssembly() non-streaming argument forwarding', () => {
+  const savedGlobals = saveWasmGlobals();
+
+  afterEach(() => {
+    restoreWasmGlobals(savedGlobals);
+  });
+
+  it('forwards every argument to instantiate', async () => {
+    const orig = vi.fn().mockResolvedValue({ module: MODULE, instance: {} });
+    WebAssembly.instantiate = orig as unknown as typeof WebAssembly.instantiate;
+
+    patchWebAssembly(registerModule);
+
+    const bytes = new Uint8Array(8);
+    const compileOptions = { builtins: ['js-string'] };
+    await (WebAssembly.instantiate as unknown as (...args: unknown[]) => Promise<unknown>)(
+      bytes,
+      WASM_IMPORTS,
+      compileOptions,
+    );
+
+    expect(orig).toHaveBeenCalledWith(bytes, WASM_IMPORTS, compileOptions);
+  });
+
+  it('forwards every argument to compile', async () => {
+    const orig = vi.fn().mockResolvedValue(MODULE);
+    WebAssembly.compile = orig as unknown as typeof WebAssembly.compile;
+
+    patchWebAssembly(registerModule);
+
+    const bytes = new Uint8Array(8);
+    const compileOptions = { builtins: ['js-string'] };
+    await (WebAssembly.compile as unknown as (...args: unknown[]) => Promise<unknown>)(bytes, compileOptions);
+
+    expect(orig).toHaveBeenCalledWith(bytes, compileOptions);
+  });
+
+  it('resolves the original result even if registration throws', async () => {
+    patchWebAssembly(() => {
+      throw new Error('registration failed');
+    });
+
+    const buffer = await fetchWasmBytes();
+
+    await expect(WebAssembly.compile(buffer)).resolves.toBeInstanceOf(WebAssembly.Module);
+  });
+});