1 files changed, 247 insertions, 0 deletions
diff --git a/packages/integrations/cloudflare/src/utils/cloudflare-module-loader.ts b/packages/integrations/cloudflare/src/utils/cloudflare-module-loader.ts
new file mode 100644
index 000000000..33eb40d74
--- /dev/null
+++ b/packages/integrations/cloudflare/src/utils/cloudflare-module-loader.ts
@@ -0,0 +1,247 @@
+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>;
+}
+
+export type ModuleType = 'CompiledWasm' | 'Text' | 'Data';
+
+/**
+ * Enables support for various non-standard extensions in module imports that cloudflare workers supports.
+ *
+ * See https://developers.cloudflare.com/pages/functions/module-support/ for reference
+ *
+ * This adds supports for imports in the following formats:
+ * - .wasm
+ * - .wasm?module
+ * - .bin
+ * - .txt
+ *
+ * @param enabled - if true, will load all cloudflare pages supported types
+ * @returns Vite plugin with additional extension method to hook into astro build
+ */
+export function cloudflareModuleLoader(
+	enabled: boolean
+): PluginOption & CloudflareModulePluginExtra {
+	/**
+	 * It's likely that eventually cloudflare will add support for custom extensions, like they do in vanilla cloudflare workers,
+	 * by adding rules to your wrangler.tome
+	 * https://developers.cloudflare.com/workers/wrangler/bundling/
+	 */
+	const adaptersByExtension: Record<string, ModuleType> = enabled ? { ...defaultAdapters } : {};
+
+	const extensions = Object.keys(adaptersByExtension);
+
+	let isDev = false;
+	const MAGIC_STRING = '__CLOUDFLARE_ASSET__';
+	const replacements: Replacement[] = [];
+
+	return {
+		name: 'vite:cf-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: extensions.map((x) => `**/*${x}`),
+				build: {
+					rollupOptions: {
+						// mark the wasm files as external so that they are not bundled and instead are loaded from the files
+						external: extensions.map(
+							(x) => new RegExp(`^${MAGIC_STRING}.+${escapeRegExp(x)}.mjs$`, 'i')
+						),
+					},
+				},
+			};
+		},
+
+		async load(id, _) {
+			const maybeExtension = extensions.find((x) => id.endsWith(x));
+			const moduleType: ModuleType | undefined =
+				(maybeExtension && adaptersByExtension[maybeExtension]) || undefined;
+			if (!moduleType || !maybeExtension) {
+				return;
+			}
+			if (!enabled) {
+				throw new Error(
+					`Cloudflare module loading is experimental. The ${maybeExtension} module cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.`
+				);
+			}
+
+			const moduleLoader = renderers[moduleType];
+
+			const filePath = id.replace(/\?\w+$/, '');
+			const extension = maybeExtension.replace(/\?\w+$/, '');
+
+			const data = await fs.readFile(filePath);
+			const base64 = data.toString('base64');
+
+			const inlineModule = moduleLoader(data);
+
+			if (isDev) {
+				// no need to wire up the assets in dev mode, just rewrite
+				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}${extension}`;
+			this.emitFile({
+				type: 'asset',
+				// 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,
+			});
+
+			// 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: inlineModule,
+			});
+
+			return `import module from "${MAGIC_STRING}${chunkId}${extension}.mjs";export default module;`;
+		},
+
+		// output original wasm file relative to the chunk now that chunking has been achieved
+		renderChunk(code, chunk, _) {
+			if (isDev) return;
+
+			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 ext of extensions) {
+				const extension = ext.replace(/\?\w+$/, '');
+				// chunk id can be many things, (alpha numeric, dollars, or underscores, maybe more)
+				replaced = replaced.replaceAll(
+					new RegExp(`${MAGIC_STRING}([^\\s]+?)${escapeRegExp(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}`;
+					}
+				);
+			}
+
+			return { code: replaced };
+		},
+
+		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);
+			}
+			for (const chunk of Object.values(bundle)) {
+				const repls = chunk.name && replacementsByChunkName.get(chunk.name);
+				for (const replacement of repls || []) {
+					replacement.fileName = chunk.fileName;
+				}
+			}
+		},
+
+		/**
+		 * 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 = updated.replaceAll(replacement.nodejsImport, replacement.cloudflareImport);
+				}
+				await fs.writeFile(filepath, updated, 'utf-8');
+			}
+		},
+	};
+}
+
+interface Replacement {
+	fileName?: string;
+	chunkName: string;
+	// desired import for cloudflare
+	cloudflareImport: string;
+	// nodejs import that simulates a wasm module
+	nodejsImport: string;
+}
+
+const renderers: Record<ModuleType, (fileContents: Buffer) => string> = {
+	CompiledWasm(fileContents: Buffer) {
+		const base64 = fileContents.toString('base64');
+		return `const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));export default wasmModule;`;
+	},
+	Data(fileContents: Buffer) {
+		const base64 = fileContents.toString('base64');
+		return `const binModule = Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)).buffer;export default binModule;`;
+	},
+	Text(fileContents: Buffer) {
+		const escaped = JSON.stringify(fileContents.toString('utf-8'));
+		return `const stringModule = ${escaped};export default stringModule;`;
+	},
+};
+
+const defaultAdapters: Record<string, ModuleType> = {
+	// 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
+	'.wasm?module': 'CompiledWasm',
+	// treats the module as a WASM module
+	'.wasm': 'CompiledWasm',
+	'.bin': 'Data',
+	'.txt': 'Text',
+};
+
+/**
+ * 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);
+}
+
+function escapeRegExp(string: string) {
+	return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string
+}