fix(cloudflare): "shared" wasm module imports (#249)

.changeset/lovely-mails-change.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/cloudflare': patch
+---
+
+Fixes build errors when wasm modules are imported from a file that is shared in both prerendered static pages and server side rendered pages

packages/cloudflare/src/index.ts

@@ -1,5 +1,7 @@
 import type { AstroConfig, AstroIntegration, RouteData } from 'astro';
 import type { OutputChunk, ProgramNode } from 'rollup';
+import type { PluginOption } from 'vite';
+import type { CloudflareModulePluginExtra } from './utils/wasm-module-loader.js';
 
 import { createReadStream } from 'node:fs';
 import { appendFile, rename, stat, unlink } from 'node:fs/promises';
@@ -18,7 +20,7 @@ import { createRoutesFile, getParts } from './utils/generate-routes-json.js';
 import { setImageConfig } from './utils/image-config.js';
 import { mutateDynamicPageImportsInPlace, mutatePageMapInPlace } from './utils/index.js';
 import { NonServerChunkDetector } from './utils/non-server-chunk-detector.js';
-import { wasmModuleLoader } from './utils/wasm-module-loader.js';
+import { cloudflareModuleLoader } from './utils/wasm-module-loader.js';
 
 export type { Runtime } from './entrypoints/server.js';
 
@@ -62,13 +64,20 @@ export type Options = {
 		/** Configuration persistence settings. Default '.wrangler/state/v3' */
 		persist?: boolean | { path: string };
 	};
-	/** Enable WebAssembly support */
+	/**
+	 * Allow bundling cloudflare worker specific file types
+	 * https://developers.cloudflare.com/workers/wrangler/bundling/
+	 */
 	wasmModuleImports?: boolean;
 };
 
 export default function createIntegration(args?: Options): AstroIntegration {
 	let _config: AstroConfig;
 
+	const cloudflareModulePlugin: PluginOption & CloudflareModulePluginExtra = cloudflareModuleLoader(
+		args?.wasmModuleImports ?? false
+	);
+
 	// Initialize the unused chunk analyzer as a shared state between hooks.
 	// The analyzer is used on earlier hooks to collect information about used hooks on a Vite plugin
 	// and then later after the full build to clean up unused chunks, so it has to be shared between them.
@@ -91,9 +100,7 @@ export default function createIntegration(args?: Options): AstroIntegration {
 					vite: {
 						// load .wasm files as WebAssembly modules
 						plugins: [
-							wasmModuleLoader({
-								disabled: !args?.wasmModuleImports,
-							}),
+							cloudflareModulePlugin,
 							chunkAnalyzer.getPlugin(),
 							{
 								name: 'dynamic-imports-analyzer',
@@ -274,6 +281,7 @@ export default function createIntegration(args?: Options): AstroIntegration {
 				}
 			},
 			'astro:build:done': async ({ pages, routes, dir, logger }) => {
+				await cloudflareModulePlugin.afterBuildCompleted(_config);
 				const PLATFORM_FILES = ['_headers', '_redirects', '_routes.json'];
 				if (_config.base !== '/') {
 					for (const file of PLATFORM_FILES) {

packages/cloudflare/src/utils/wasm-module-loader.ts

@@ -1,23 +1,29 @@
-import * as fs from 'node:fs';
+import * as fs from 'node:fs/promises';
 import * as path from 'node:path';
+import * as url from 'node:url';
 import type { AstroConfig } from 'astro';
+import type { OutputBundle } from 'rollup';
+import type { PluginOption } from 'vite';
 
+export interface CloudflareModulePluginExtra {
+	afterBuildCompleted(config: AstroConfig): Promise<void>;
+}
 /**
- * Loads '*.wasm?module' imports as WebAssembly modules, which is the only way to load WASM in cloudflare workers.
+ * Enables support for wasm modules within cloudflare pages functions
+ *
+ * Loads '*.wasm?module' and `*.wasm` 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'
+ * @param enabled - if true, load '.wasm' imports as Uint8Arrays, otherwise will throw errors when encountered to clarify that it must be enabled
  * @returns Vite plugin to load WASM tagged with '?module' as a WASM modules
  */
-export function wasmModuleLoader({
-	disabled,
-}: {
-	disabled: boolean;
-}): NonNullable<AstroConfig['vite']['plugins']>[number] {
-	const postfix = '.wasm?module';
+export function cloudflareModuleLoader(
+	enabled: boolean
+): PluginOption & CloudflareModulePluginExtra {
+	const enabledAdapters = cloudflareImportAdapters.filter((x) => enabled);
 	let isDev = false;
+	const MAGIC_STRING = '__CLOUDFLARE_ASSET__';
+	const replacements: Replacement[] = [];
 
 	return {
 		name: 'vite:wasm-module-loader',
@@ -28,46 +34,51 @@ export function wasmModuleLoader({
 		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'],
+				assetsInclude: enabledAdapters.map((x) => `**/*.${x.qualifiedExtension}`),
 				build: {
 					rollupOptions: {
 						// mark the wasm files as external so that they are not bundled and instead are loaded from the files
-						external: [/^__WASM_ASSET__.+\.wasm$/i, /^__WASM_ASSET__.+\.wasm.mjs$/i],
+						external: enabledAdapters.map(
+							(x) => new RegExp(`^${MAGIC_STRING}.+\\.${x.extension}.mjs$`, 'i')
+						),
 					},
 				},
 			};
 		},
 
-		load(id, _) {
-			if (!id.endsWith(postfix)) {
+		async load(id, _) {
+			const importAdapter = cloudflareImportAdapters.find((x) => id.endsWith(x.qualifiedExtension));
+			if (!importAdapter) {
 				return;
 			}
-			if (disabled) {
+			if (!enabled) {
 				throw new Error(
-					`WASM module's cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.`
+					`Cloudflare module loading is experimental. The ${importAdapter.qualifiedExtension} module cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.`
 				);
 			}
 
-			const filePath = id.slice(0, -1 * '?module'.length);
+			const filePath = id.replace(/\?module$/, '');
 
-			const data = fs.readFileSync(filePath);
+			const data = await fs.readFile(filePath);
 			const base64 = data.toString('base64');
 
-			const base64Module = `const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));export default wasmModule;`;
+			const inlineModule = importAdapter.asNodeModule(data);
 
 			if (isDev) {
 				// no need to wire up the assets in dev mode, just rewrite
-				return base64Module;
+				return inlineModule;
 			}
 			// just some shared ID
 			const 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`;
+			const assetName = `${path.basename(filePath).split('.')[0]}.${hash}.${
+				importAdapter.extension
+			}`;
 			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
+				// emit the data explicitly as an esset with `fileName` rather than `name` so that
+				// vite doesn't give it a random hash-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: assetName,
 				source: data,
@@ -77,51 +88,122 @@ export function wasmModuleLoader({
 			const chunkId = this.emitFile({
 				type: 'prebuilt-chunk',
 				fileName: `${assetName}.mjs`,
-				code: base64Module,
+				code: inlineModule,
 			});
 
-			return `import wasmModule from "__WASM_ASSET__${chunkId}.wasm.mjs";export default wasmModule;`;
+			return `import module from "${MAGIC_STRING}${chunkId}.${importAdapter.extension}.mjs";export default module;`;
 		},
 
-		// output original wasm file relative to the chunk
+		// output original wasm file relative to the chunk now that chunking has been achieved
 		renderChunk(code, chunk, _) {
 			if (isDev) return;
 
-			if (!/__WASM_ASSET__/g.test(code)) return;
-
-			const isPrerendered = Object.keys(chunk.modules).some(
-				(moduleId) => this.getModuleInfo(moduleId)?.meta?.astro?.pageOptions?.prerender === true
-			);
+			if (!code.includes(MAGIC_STRING)) return;
+
+			// SSR will need the .mjs suffix removed from the import before this works in cloudflare, but this is done as a final step
+			// so as to support prerendering from nodejs runtime
+			let replaced = code;
+			for (const loader of enabledAdapters) {
+				replaced = replaced.replaceAll(
+					// chunk id can be many things, (alpha numeric, dollars, or underscores, maybe more)
+					new RegExp(`${MAGIC_STRING}([^\\s]+?)\\.${loader.extension}\\.mjs`, 'g'),
+					(s, assetId) => {
+						const fileName = this.getFileName(assetId);
+						const relativePath = path
+							.relative(path.dirname(chunk.fileName), fileName)
+							.replaceAll('\\', '/'); // fix windows paths for import
+
+						// record this replacement for later, to adjust it to import the unbundled asset
+						replacements.push({
+							chunkName: chunk.name,
+							cloudflareImport: relativePath.replace(/\.mjs$/, ''),
+							nodejsImport: relativePath,
+						});
+						return `./${relativePath}`;
+					}
+				);
+			}
+			if (replaced.includes(MAGIC_STRING)) {
+				console.error('failed to replace', replaced);
+			}
 
-			let final = code;
+			return { code: replaced };
+		},
 
-			// SSR
-			if (!isPrerendered) {
-				final = code.replaceAll(/__WASM_ASSET__([A-Za-z\d]+).wasm.mjs/g, (s, assetId) => {
-					const fileName = this.getFileName(assetId).replace(/\.mjs$/, '');
-					const relativePath = path
-						.relative(path.dirname(chunk.fileName), fileName)
-						.replaceAll('\\', '/'); // fix windows paths for import
-					return `./${relativePath}`;
-				});
+		generateBundle(_, bundle: OutputBundle) {
+			// associate the chunk name to the final file name. After the prerendering is done, we can use this to replace the imports in the _worker.js
+			// in a targetted way
+			const replacementsByChunkName = new Map<string, Replacement[]>();
+			for (const replacement of replacements) {
+				const repls = replacementsByChunkName.get(replacement.chunkName) || [];
+				if (!repls.length) {
+					replacementsByChunkName.set(replacement.chunkName, repls);
+				}
+				repls.push(replacement);
 			}
-
-			// SSG
-			if (isPrerendered) {
-				final = code.replaceAll(/__WASM_ASSET__([A-Za-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}`;
-				});
+			for (const chunk of Object.values(bundle)) {
+				const repls = chunk.name && replacementsByChunkName.get(chunk.name);
+				for (const replacement of repls || []) {
+					replacement.fileName = chunk.fileName;
+				}
 			}
+		},
 
-			return { code: final };
+		/**
+		 * Once prerendering is complete, restore the imports in the _worker.js to cloudflare compatible ones, removing the .mjs suffix.
+		 */
+		async afterBuildCompleted(config: AstroConfig) {
+			const baseDir = url.fileURLToPath(config.outDir);
+			const replacementsByFileName = new Map<string, Replacement[]>();
+			for (const replacement of replacements) {
+				if (!replacement.fileName) continue;
+				const repls = replacementsByFileName.get(replacement.fileName) || [];
+				if (!repls.length) {
+					replacementsByFileName.set(replacement.fileName, repls);
+				}
+				repls.push(replacement);
+			}
+			for (const [fileName, repls] of replacementsByFileName.entries()) {
+				const filepath = path.join(baseDir, '_worker.js', fileName);
+				const contents = await fs.readFile(filepath, 'utf-8');
+				let updated = contents;
+				for (const replacement of repls) {
+					updated = contents.replaceAll(replacement.nodejsImport, replacement.cloudflareImport);
+				}
+				await fs.writeFile(filepath, updated, 'utf-8');
+			}
 		},
 	};
 }
 
+export type ImportType = 'wasm';
+
+interface Replacement {
+	fileName?: string;
+	chunkName: string;
+	// desired import for cloudflare
+	cloudflareImport: string;
+	// nodejs import that simulates a wasm module
+	nodejsImport: string;
+}
+
+interface ModuleImportAdapter {
+	extension: ImportType;
+	qualifiedExtension: string;
+	asNodeModule(fileContents: Buffer): string;
+}
+
+const wasmImportAdapter: ModuleImportAdapter = {
+	extension: 'wasm',
+	qualifiedExtension: 'wasm?module',
+	asNodeModule(fileContents: Buffer) {
+		const base64 = fileContents.toString('base64');
+		return `const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));export default wasmModule;`;
+	},
+};
+
+const cloudflareImportAdapters = [wasmImportAdapter];
+
 /**
  * Returns a deterministic 32 bit hash code from a string
  */

packages/cloudflare/test/fixtures/wasm/src/pages/hybridshared.ts

@@ -0,0 +1,16 @@
+import { type APIContext } from 'astro';
+
+import {add} from '../util/add';
+
+export const prerender = true
+
+export async function GET(
+  context: APIContext
+): Promise<Response> {
+  return new Response(JSON.stringify({ answer: add(20, 1) }), {
+    status: 200,
+    headers: {
+      'Content-Type': 'application/json',
+    },
+  });
+}

packages/cloudflare/test/fixtures/wasm/src/pages/shared/[a]/[b].ts

@@ -0,0 +1,18 @@
+import { type APIContext } from 'astro';
+import { add } from '../../../util/add';
+
+
+export const prerender = false;
+
+export async function GET(
+  context: APIContext
+): Promise<Response> {
+	const a = Number.parseInt(context.params.a ?? "0");
+	const b = Number.parseInt(context.params.b ?? "0");
+  return new Response(JSON.stringify({ answer: add(a, b) }), {
+    status: 200,
+    headers: {
+      'Content-Type': 'application/json',
+    },
+  });
+}

packages/cloudflare/test/fixtures/wasm/src/util/add.ts

@@ -0,0 +1,8 @@
+import mod from './add.wasm?module';
+
+
+const addModule: any = new WebAssembly.Instance(mod);
+
+export function add(a, b) {
+    return addModule.exports.add(a, b);
+}

packages/cloudflare/test/wasm.test.js

@@ -32,4 +32,25 @@ describe('WasmImport', () => {
 		const json = await res.json();
 		assert.deepEqual(json, { answer: 42 });
 	});
+
+	it('can render static', async () => {
+		const res = await fetch('http://127.0.0.1:8788/hybrid');
+		assert.equal(res.status, 200);
+		const json = await res.json();
+		assert.deepEqual(json, { answer: 21 });
+	});
+
+	it('can render shared', async () => {
+		const res = await fetch('http://127.0.0.1:8788/shared/40/2');
+		assert.equal(res.status, 200);
+		const json = await res.json();
+		assert.deepEqual(json, { answer: 42 });
+	});
+
+	it('can render static shared', async () => {
+		const res = await fetch('http://127.0.0.1:8788/hybridshared');
+		assert.equal(res.status, 200);
+		const json = await res.json();
+		assert.deepEqual(json, { answer: 21 });
+	});
 });