feat(@astrojs/cloudflare): Add support for wasm module imports (#8542)

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

2 files changed, 234 insertions, 31 deletions

diff --git a/packages/integrations/cloudflare/src/index.ts b/packages/integrations/cloudflare/src/index.ts
index 24c22d8f1..4ae43a110 100644
--- a/packages/integrations/cloudflare/src/index.ts
+++ b/packages/integrations/cloudflare/src/index.ts
@@ -9,10 +9,11 @@ import { AstroError } from 'astro/errors';
 import esbuild from 'esbuild';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
-import { sep } from 'node:path';
+import { basename, dirname, relative, sep } from 'node:path';
 import { fileURLToPath, pathToFileURL } from 'node:url';
 import glob from 'tiny-glob';
 import { getEnvVars } from './parser.js';
+import { wasmModuleLoader } from './wasm-module-loader.js';
 
 export type { AdvancedRuntime } from './server.advanced.js';
 export type { DirectoryRuntime } from './server.directory.js';
@@ -26,11 +27,13 @@ type Options = {
 	 * 'remote': use a dynamic real-live req.cf object, and env vars defined in wrangler.toml & .dev.vars (astro dev is enough)
 	 */
 	runtime?: 'off' | 'local' | 'remote';
+	wasmModuleImports?: boolean;
 };
 
 interface BuildConfig {
 	server: URL;
 	client: URL;
+	assets: string;
 	serverEntry: string;
 	split?: boolean;
 }
@@ -189,6 +192,15 @@ export default function createIntegration(args?: Options): AstroIntegration {
 						serverEntry: '_worker.mjs',
 						redirects: false,
 					},
+					vite: {
+						// load .wasm files as WebAssembly modules
+						plugins: [
+							wasmModuleLoader({
+								disabled: !args?.wasmModuleImports,
+								assetsDirectory: config.build.assets,
+							}),
+						],
+					},
 				});
 			},
 			'astro:config:done': ({ setAdapter, config }) => {
@@ -280,6 +292,7 @@ export default function createIntegration(args?: Options): AstroIntegration {
 			},
 			'astro:build:done': async ({ pages, routes, dir }) => {
 				const functionsUrl = new URL('functions/', _config.root);
+				const assetsUrl = new URL(_buildConfig.assets, _buildConfig.client);
 
 				if (isModeDirectory) {
 					await fs.promises.mkdir(functionsUrl, { recursive: true });
@@ -291,36 +304,71 @@ export default function createIntegration(args?: Options): AstroIntegration {
 					const entryPaths = entryPointsURL.map((entry) => fileURLToPath(entry));
 					const outputUrl = new URL('$astro', _buildConfig.server);
 					const outputDir = fileURLToPath(outputUrl);
-
-					await esbuild.build({
-						target: 'es2020',
-						platform: 'browser',
-						conditions: ['workerd', 'worker', 'browser'],
-						external: [
-							'node:assert',
-							'node:async_hooks',
-							'node:buffer',
-							'node:diagnostics_channel',
-							'node:events',
-							'node:path',
-							'node:process',
-							'node:stream',
-							'node:string_decoder',
-							'node:util',
-						],
-						entryPoints: entryPaths,
-						outdir: outputDir,
-						allowOverwrite: true,
-						format: 'esm',
-						bundle: true,
-						minify: _config.vite?.build?.minify !== false,
-						banner: {
-							js: SHIM,
-						},
-						logOverride: {
-							'ignored-bare-import': 'silent',
-						},
-					});
+					//
+					// Sadly, when wasmModuleImports is enabled, this needs to build esbuild for each depth of routes/entrypoints
+					// independently so that relative import paths to the assets are the correct depth of '../' traversals
+					// This is inefficient, so wasmModuleImports is opt-in. This could potentially be improved in the future by
+					// taking advantage of the esbuild "onEnd" hook to rewrite import code per entry point relative to where the final
+					// destination of the entrypoint is
+					const entryPathsGroupedByDepth = !args.wasmModuleImports
+						? [entryPaths]
+						: entryPaths
+								.reduce((sum, thisPath) => {
+									const depthFromRoot = thisPath.split(sep).length;
+									sum.set(depthFromRoot, (sum.get(depthFromRoot) || []).concat(thisPath));
+									return sum;
+								}, new Map<number, string[]>())
+								.values();
+
+					for (const pathsGroup of entryPathsGroupedByDepth) {
+						// for some reason this exports to "entry.pages" on windows instead of "pages" on unix environments.
+						// This deduces the name of the "pages" build directory
+						const pagesDirname = relative(fileURLToPath(_buildConfig.server), pathsGroup[0]).split(
+							sep
+						)[0];
+						const absolutePagesDirname = fileURLToPath(new URL(pagesDirname, _buildConfig.server));
+						const urlWithinFunctions = new URL(
+							relative(absolutePagesDirname, pathsGroup[0]),
+							functionsUrl
+						);
+						const relativePathToAssets = relative(
+							dirname(fileURLToPath(urlWithinFunctions)),
+							fileURLToPath(assetsUrl)
+						);
+						await esbuild.build({
+							target: 'es2020',
+							platform: 'browser',
+							conditions: ['workerd', 'worker', 'browser'],
+							external: [
+								'node:assert',
+								'node:async_hooks',
+								'node:buffer',
+								'node:diagnostics_channel',
+								'node:events',
+								'node:path',
+								'node:process',
+								'node:stream',
+								'node:string_decoder',
+								'node:util',
+							],	
+							entryPoints: pathsGroup,
+							outbase: absolutePagesDirname,
+							outdir: outputDir,
+							allowOverwrite: true,
+							format: 'esm',
+							bundle: true,
+							minify: _config.vite?.build?.minify !== false,
+							banner: {
+								js: SHIM,
+							},
+							logOverride: {
+								'ignored-bare-import': 'silent',
+							},
+							plugins: !args?.wasmModuleImports
+								? []
+								: [rewriteWasmImportPath({ relativePathToAssets })],
+						});
+					}
 
 					const outputFiles: Array<string> = await glob(`**/*`, {
 						cwd: outputDir,
@@ -393,6 +441,15 @@ export default function createIntegration(args?: Options): AstroIntegration {
 						logOverride: {
 							'ignored-bare-import': 'silent',
 						},
+						plugins: !args?.wasmModuleImports
+							? []
+							: [
+									rewriteWasmImportPath({
+										relativePathToAssets: isModeDirectory
+											? relative(fileURLToPath(functionsUrl), fileURLToPath(assetsUrl))
+											: relative(fileURLToPath(_buildConfig.client), fileURLToPath(assetsUrl)),
+									}),
+							  ],
 					});
 
 					// Rename to worker.js
@@ -602,3 +659,30 @@ function deduplicatePatterns(patterns: string[]) {
 			return true;
 		});
 }
+
+/**
+ *
+ * @param relativePathToAssets - relative path from the final location for the current esbuild output bundle, to the assets directory.
+ */
+function rewriteWasmImportPath({
+	relativePathToAssets,
+}: {
+	relativePathToAssets: string;
+}): esbuild.Plugin {
+	return {
+		name: 'wasm-loader',
+		setup(build) {
+			build.onResolve({ filter: /.*\.wasm.mjs$/ }, (args) => {
+				const updatedPath = [
+					relativePathToAssets.replaceAll('\\', '/'),
+					basename(args.path).replace(/\.mjs$/, ''),
+				].join('/');
+
+				return {
+					path: updatedPath, // change the reference to the changed module
+					external: true, // mark it as external in the bundle
+				};
+			});
+		},
+	};
+}
diff --git a/packages/integrations/cloudflare/src/wasm-module-loader.ts b/packages/integrations/cloudflare/src/wasm-module-loader.ts
new file mode 100644
index 000000000..7d34d48c3
--- /dev/null
+++ b/packages/integrations/cloudflare/src/wasm-module-loader.ts
@@ -0,0 +1,119 @@
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import { type Plugin } from 'vite';
+
+/**
+ * Loads '*.wasm?module' imports as WebAssembly modules, which is the only way to load WASM in cloudflare workers.
+ * Current proposal for WASM modules: https://github.com/WebAssembly/esm-integration/tree/main/proposals/esm-integration
+ * Cloudflare worker WASM from javascript support: https://developers.cloudflare.com/workers/runtime-apis/webassembly/javascript/
+ * @param disabled - if true throws a helpful error message if wasm is encountered and wasm imports are not enabled,
+ * 								otherwise it will error obscurely in the esbuild and vite builds
+ * @param assetsDirectory - the folder name for the assets directory in the build directory. Usually '_astro'
+ * @returns Vite plugin to load WASM tagged with '?module' as a WASM modules
+ */
+export function wasmModuleLoader({
+	disabled,
+	assetsDirectory,
+}: {
+	disabled: boolean;
+	assetsDirectory: string;
+}): Plugin {
+	const postfix = '.wasm?module';
+	let isDev = false;
+
+	return {
+		name: 'vite:wasm-module-loader',
+		enforce: 'pre',
+		configResolved(config) {
+			isDev = config.command === 'serve';
+		},
+		config(_, __) {
+			// let vite know that file format and the magic import string is intentional, and will be handled in this plugin
+			return {
+				assetsInclude: ['**/*.wasm?module'],
+				build: { rollupOptions: { external: /^__WASM_ASSET__.+\.wasm\.mjs$/i } },
+			};
+		},
+
+		load(id, _) {
+			if (!id.endsWith(postfix)) {
+				return;
+			}
+			if (disabled) {
+				throw new Error(
+					`WASM module's cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.`
+				);
+			}
+
+			const filePath = id.slice(0, -1 * '?module'.length);
+
+			const data = fs.readFileSync(filePath);
+			const base64 = data.toString('base64');
+
+			const base64Module = `
+const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));
+export default wasmModule
+`;
+			if (isDev) {
+				// no need to wire up the assets in dev mode, just rewrite
+				return base64Module;
+			} else {
+				// just some shared ID
+				let hash = hashString(base64);
+				// emit the wasm binary as an asset file, to be picked up later by the esbuild bundle for the worker.
+				// give it a shared deterministic name to make things easy for esbuild to switch on later
+				const assetName = path.basename(filePath).split('.')[0] + '.' + hash + '.wasm';
+				this.emitFile({
+					type: 'asset',
+					// put it explicitly in the _astro assets directory with `fileName` rather than `name` so that
+					// vite doesn't give it a random id in its name. We need to be able to easily rewrite from
+					// the .mjs loader and the actual wasm asset later in the ESbuild for the worker
+					fileName: path.join(assetsDirectory, assetName),
+					source: fs.readFileSync(filePath),
+				});
+
+				// however, by default, the SSG generator cannot import the .wasm as a module, so embed as a base64 string
+				const chunkId = this.emitFile({
+					type: 'prebuilt-chunk',
+					fileName: assetName + '.mjs',
+					code: base64Module,
+				});
+
+				return `
+import wasmModule from "__WASM_ASSET__${chunkId}.wasm.mjs";
+export default wasmModule;
+	`;
+			}
+		},
+
+		// output original wasm file relative to the chunk
+		renderChunk(code, chunk, _) {
+			if (isDev) return;
+
+			if (!/__WASM_ASSET__/g.test(code)) return;
+
+			const final = code.replaceAll(/__WASM_ASSET__([a-z\d]+).wasm.mjs/g, (s, assetId) => {
+				const fileName = this.getFileName(assetId);
+				const relativePath = path
+					.relative(path.dirname(chunk.fileName), fileName)
+					.replaceAll('\\', '/'); // fix windows paths for import
+				return `./${relativePath}`;
+			});
+
+			return { code: final };
+		},
+	};
+}
+
+/**
+ * Returns a deterministic 32 bit hash code from a string
+ */
+function hashString(str: string): string {
+	let hash = 0;
+	for (let i = 0; i < str.length; i++) {
+		const char = str.charCodeAt(i);
+		hash = (hash << 5) - hash + char;
+		hash &= hash; // Convert to 32bit integer
+	}
+	return new Uint32Array([hash])[0].toString(36);
+}