feat(wasm): add SSR support for .wasm?init

docs/guide/features.md

@@ -691,6 +691,12 @@ In the production build, `.wasm` files smaller than `assetInlineLimit` will be i
 Use [`vite-plugin-wasm`](https://github.com/Menci/vite-plugin-wasm) or other community plugins to handle this.
 :::
 
+::: warning For SSR build, Node.js compatible runtimes are only supported
+
+Due to the lack of a universal way to load a file, the internal implementation for `.wasm?init` relies on `node:fs` module. This means that this feature will only work in Node.js compatible runtimes for SSR builds.
+
+:::
+
 ### Accessing the WebAssembly Module
 
 If you need access to the `Module` object, e.g. to instantiate it multiple times, use an [explicit URL import](./assets#explicit-url-imports) to resolve the asset, and then perform the instantiation:
@@ -710,31 +716,6 @@ const main = async () => {
 main()
 ```
 
-### Fetching the module in Node.js
-
-In SSR, the `fetch()` happening as part of the `?init` import, may fail with `TypeError: Invalid URL`.
-See the issue [Support wasm in SSR](https://github.com/vitejs/vite/issues/8882).
-
-Here is an alternative, assuming the project base is the current directory:
-
-```js twoslash
-import 'vite/client'
-// ---cut---
-import wasmUrl from 'foo.wasm?url'
-import { readFile } from 'node:fs/promises'
-
-const main = async () => {
-  const resolvedUrl = (await import('./test/boot.test.wasm?url')).default
-  const buffer = await readFile('.' + resolvedUrl)
-  const { instance } = await WebAssembly.instantiate(buffer, {
-    /* ... */
-  })
-  /* ... */
-}
-
-main()
-```
-
 ## Web Workers
 
 ### Import with Constructors

packages/vite/src/node/plugins/asset.ts

@@ -1,6 +1,7 @@
 import path from 'node:path'
 import fsp from 'node:fs/promises'
 import { Buffer } from 'node:buffer'
+import { pathToFileURL } from 'node:url'
 import * as mrmime from 'mrmime'
 import type {
   NormalizedOutputOptions,
@@ -318,10 +319,11 @@ export function assetPlugin(config: ResolvedConfig): Plugin {
 export async function fileToUrl(
   pluginContext: PluginContext,
   id: string,
+  asFileUrl = false,
 ): Promise<string> {
   const { environment } = pluginContext
   if (!environment.config.isBundled) {
-    return fileToDevUrl(environment, id)
+    return fileToDevUrl(environment, id, asFileUrl)
   } else {
     return fileToBuiltUrl(pluginContext, id)
   }
@@ -330,7 +332,7 @@ export async function fileToUrl(
 export async function fileToDevUrl(
   environment: Environment,
   id: string,
-  skipBase = false,
+  asFileUrl = false,
 ): Promise<string> {
   const config = environment.getTopLevelConfig()
   const publicFile = checkPublicFile(id, config)
@@ -353,6 +355,10 @@ export async function fileToDevUrl(
     }
   }
 
+  if (asFileUrl) {
+    return pathToFileURL(cleanedId).href
+  }
+
   let rtn: string
   if (publicFile) {
     // in public dir during dev, keep the url as-is
@@ -365,9 +371,6 @@ export async function fileToDevUrl(
     // (this is special handled by the serve static middleware
     rtn = path.posix.join(FS_PREFIX, id)
   }
-  if (skipBase) {
-    return rtn
-  }
   const base = joinUrlSegments(config.server.origin ?? '', config.decodedBase)
   return joinUrlSegments(base, removeLeadingSlash(rtn))
 }

packages/vite/src/node/plugins/wasm.ts

@@ -1,11 +1,17 @@
+import MagicString from 'magic-string'
 import { exactRegex } from 'rolldown/filter'
-import type { Plugin } from '../plugin'
-import { fileToUrl } from './asset'
+import type { BindingMagicString } from 'rolldown'
+import { createToImportMetaURLBasedRelativeRuntime } from '../build'
+import { type Plugin, perEnvironmentPlugin } from '../plugin'
+import { cleanUrl } from '../../shared/utils'
+import { assetUrlRE, fileToUrl } from './asset'
 
 const wasmHelperId = '\0vite/wasm-helper.js'
 
 const wasmInitRE = /(?<![?#].*)\.wasm\?init/
 
+const wasmInitUrlRE: RegExp = /__VITE_WASM_INIT__([\w$]+)__/g
+
 const wasmHelper = async (opts = {}, url: string) => {
   let result
   if (url.startsWith('data:')) {
@@ -26,53 +32,129 @@ const wasmHelper = async (opts = {}, url: string) => {
     }
     result = await WebAssembly.instantiate(bytes, opts)
   } else {
-    // https://github.com/mdn/webassembly-examples/issues/5
-    // WebAssembly.instantiateStreaming requires the server to provide the
-    // correct MIME type for .wasm files, which unfortunately doesn't work for
-    // a lot of static file servers, so we just work around it by getting the
-    // raw buffer.
-    const response = await fetch(url)
-    const contentType = response.headers.get('Content-Type') || ''
-    if (
-      'instantiateStreaming' in WebAssembly &&
-      contentType.startsWith('application/wasm')
-    ) {
-      result = await WebAssembly.instantiateStreaming(response, opts)
-    } else {
-      const buffer = await response.arrayBuffer()
-      result = await WebAssembly.instantiate(buffer, opts)
-    }
+    result = await instantiateFromUrl(url, opts)
   }
   return result.instance
 }
 
 const wasmHelperCode = wasmHelper.toString()
 
+const instantiateFromUrl = async (url: string, opts?: WebAssembly.Imports) => {
+  // https://github.com/mdn/webassembly-examples/issues/5
+  // WebAssembly.instantiateStreaming requires the server to provide the
+  // correct MIME type for .wasm files, which unfortunately doesn't work for
+  // a lot of static file servers, so we just work around it by getting the
+  // raw buffer.
+  const response = await fetch(url)
+  const contentType = response.headers.get('Content-Type') || ''
+  if (
+    'instantiateStreaming' in WebAssembly &&
+    contentType.startsWith('application/wasm')
+  ) {
+    return WebAssembly.instantiateStreaming(response, opts)
+  } else {
+    const buffer = await response.arrayBuffer()
+    return WebAssembly.instantiate(buffer, opts)
+  }
+}
+
+const instantiateFromUrlCode = instantiateFromUrl.toString()
+
+const instantiateFromFile = async (
+  fileUrlString: string,
+  opts?: WebAssembly.Imports,
+) => {
+  const { readFile } = await import('node:fs/promises')
+  const fileUrl = new URL(fileUrlString, /** #__KEEP__ */ import.meta.url)
+  const buffer = await readFile(fileUrl)
+  return WebAssembly.instantiate(buffer, opts)
+}
+
+const instantiateFromFileCode = instantiateFromFile.toString()
+
 export const wasmHelperPlugin = (): Plugin => {
-  return {
-    name: 'vite:wasm-helper',
+  return perEnvironmentPlugin('vite:wasm-helper', (env) => {
+    return {
+      name: 'vite:wasm-helper',
 
-    resolveId: {
-      filter: { id: exactRegex(wasmHelperId) },
-      handler(id) {
-        return id
+      resolveId: {
+        filter: { id: exactRegex(wasmHelperId) },
+        handler(id) {
+          return id
+        },
       },
-    },
 
-    load: {
-      filter: { id: [exactRegex(wasmHelperId), wasmInitRE] },
-      async handler(id) {
-        if (id === wasmHelperId) {
-          return `export default ${wasmHelperCode}`
-        }
+      load: {
+        filter: { id: [exactRegex(wasmHelperId), wasmInitRE] },
+        async handler(id) {
+          const ssr = this.environment.config.consumer === 'server'
 
-        const url = await fileToUrl(this, id)
+          if (id === wasmHelperId) {
+            return `
+const instantiateFromUrl = ${ssr ? instantiateFromFileCode : instantiateFromUrlCode}
+export default ${wasmHelperCode}
+`
+          }
 
-        return `
+          id = id.split('?')[0]
+          let url = await fileToUrl(this, id, ssr)
+          if (ssr && assetUrlRE.test(url)) {
+            url = url.replace('__VITE_ASSET__', '__VITE_WASM_INIT__')
+          }
+          return `
   import initWasm from "${wasmHelperId}"
   export default opts => initWasm(opts, ${JSON.stringify(url)})
   `
+        },
       },
-    },
-  }
+
+      renderChunk:
+        env.config.consumer === 'server'
+          ? {
+              filter: { code: wasmInitUrlRE },
+              async handler(code, chunk, opts, meta) {
+                const toRelativeRuntime =
+                  createToImportMetaURLBasedRelativeRuntime(
+                    opts.format,
+                    this.environment.config.isWorker,
+                  )
+
+                let match: RegExpExecArray | null
+                let s: BindingMagicString | MagicString | undefined
+
+                wasmInitUrlRE.lastIndex = 0
+                while ((match = wasmInitUrlRE.exec(code))) {
+                  const [full, referenceId] = match
+                  const file = this.getFileName(referenceId)
+                  chunk.viteMetadata!.importedAssets.add(cleanUrl(file))
+                  const { runtime } = toRelativeRuntime(file, chunk.fileName)
+
+                  s ??= meta.magicString ?? new MagicString(code)
+
+                  s.update(
+                    match.index,
+                    match.index + full.length,
+                    `"+${runtime}+"`,
+                  )
+                }
+
+                if (!s) return null
+
+                return meta.magicString
+                  ? {
+                      code: s as BindingMagicString,
+                    }
+                  : {
+                      code: s.toString(),
+                      map: this.environment.config.build.sourcemap
+                        ? (s as MagicString).generateMap({
+                            hires: 'boundary',
+                          })
+                        : null,
+                    }
+              },
+            }
+          : undefined,
+    }
+  })
 }

playground/ssr-wasm/__tests__/serve.ts

@@ -0,0 +1,42 @@
+// this is automatically detected by playground/vitestSetup.ts and will replace
+// the default e2e test serve behavior
+
+import path from 'node:path'
+import kill from 'kill-port'
+import { build } from 'vite'
+import { hmrPorts, isBuild, ports, rootDir } from '~utils'
+
+export const port = ports['ssr-wasm']
+
+export async function preServe() {
+  if (isBuild) {
+    await build({ root: rootDir, logLevel: 'silent' })
+  }
+}
+
+export async function serve(): Promise<{ close(): Promise<void> }> {
+  await kill(port)
+
+  const { createServer } = await import(path.resolve(rootDir, 'server.js'))
+  const { app, vite } = await createServer(rootDir, hmrPorts['ssr-wasm'])
+
+  return new Promise((resolve, reject) => {
+    try {
+      const server = app.listen(port, () => {
+        resolve({
+          // for test teardown
+          async close() {
+            await new Promise((resolve) => {
+              server.close(resolve)
+            })
+            if (vite) {
+              await vite.close()
+            }
+          },
+        })
+      })
+    } catch (e) {
+      reject(e)
+    }
+  })
+}

playground/ssr-wasm/__tests__/ssr-wasm.spec.ts

@@ -0,0 +1,37 @@
+import { expect, test } from 'vitest'
+import { port } from './serve'
+import { findAssetFile, isBuild, listAssets, page } from '~utils'
+
+const url = `http://localhost:${port}`
+
+test('should work when inlined', async () => {
+  await page.goto(`${url}/static-light`)
+  expect(await page.textContent('.static-light')).toMatch('42')
+})
+
+test('should work when output', async () => {
+  await page.goto(`${url}/static-heavy`)
+  expect(await page.textContent('.static-heavy')).toMatch('24')
+})
+
+test.runIf(isBuild)('should not contain wasm file when inlined', async () => {
+  const assets = await listAssets()
+  const lightWasm = assets.find((f) => /light-.+\.wasm$/.test(f))
+  expect(lightWasm).toBeUndefined()
+
+  const staticLight = await findAssetFile(/^static-light-.+\.js$/)
+  expect(staticLight).toContain('data:application/wasm;base64,')
+})
+
+test.runIf(isBuild)(
+  'should contain and reference wasm file when output',
+  async () => {
+    const assets = await listAssets()
+    const heavyWasm = assets.find((f) => /heavy-.+\.wasm$/.test(f))
+    expect(heavyWasm).toBeDefined()
+
+    const staticHeavy = await findAssetFile(/^static-heavy-.+\.js$/)
+    expect(staticHeavy).toContain(heavyWasm)
+    expect(staticHeavy).not.toContain('data:application/wasm;base64,')
+  },
+)

playground/ssr-wasm/package.json

@@ -0,0 +1,15 @@
+{
+  "name": "@vitejs/test-ssr-wasm",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "dev": "node server",
+    "build": "vite build",
+    "preview": "NODE_ENV=production node server"
+  },
+  "dependencies": {},
+  "devDependencies": {
+    "express": "^5.2.1"
+  }
+}