feat: accept WebAssembly.Module, Response inputs

`createPlugin` now accepts two additional manifest types (`response` and `module`) as well as `Response` and `WebAssembly.Module`. There are four goals here: 1. Allow us to target the Cloudflare Workers platform. CF Workers only support loading Wasm via `import` statements; these resolve to `WebAssembly.Module` objects, which means we need to allow users to pass `Module`s in addition to our other types. 2. Play nicely with V8's [Wasm caching][1]; in particular V8 will use metadata from the `Response` to build a key for caching the results of Wasm compilation. 3. This sets us up to implement [Wasm linking][2] by allowing us to introspect plugin modules imports and exports before instantiation. 4. And finally, resolving to modules instead of arraybuffers allows us to add [hooks for observe-sdk][3] (especially in advance of adding [thread pooling][4]). Because Bun lacks support for `WebAssembly.compileStreaming` and `Response.clone()`, we provide an alternate implementation for converting a response to a module and its metadata. One caveat is that there's no way to get the source bytes of a `WebAssembly.Module`, so `{module}` cannot be used with `{hash}` in a `Manifest`. Fixes #9 [1]: https://v8.dev/blog/wasm-code-caching#stream [2]: #29 [3]: #3 [4]: #31
extism · Nov 27, 2023 · 07b6aa7 · 07b6aa7
1 parent 8e4fc5a
commit 07b6aa7
Show file tree

Hide file tree

Showing 12 changed files with 201 additions and 54 deletions.
diff --git a/deno.json b/deno.json
@@ -1,6 +1,7 @@
 {
   "imports": {
     "js-sdk:worker-url": "./src/worker-url.ts",
+    "js-sdk:response-to-module": "./src/polyfills/response-to-module.ts",
     "js-sdk:minimatch": "./src/polyfills/deno-minimatch.ts",
     "js-sdk:capabilities": "./src/polyfills/deno-capabilities.ts",
     "js-sdk:wasi": "./src/polyfills/deno-wasi.ts",

diff --git a/justfile b/justfile
@@ -151,6 +151,7 @@ build_node_cjs out='cjs' args='[]':
         "minify": false,
         "alias": {
           "js-sdk:capabilities": "./src/polyfills/node-capabilities.ts",
+          "js-sdk:response-to-module": "./src/polyfills/response-to-module.ts",
           "js-sdk:minimatch": "./src/polyfills/node-minimatch.ts",
           "js-sdk:worker-url": "./dist/worker/node/worker-url.ts",
           "js-sdk:fs": "node:fs/promises",
@@ -182,6 +183,7 @@ build_node_esm out='esm' args='[]':
         "minify": false,
         "alias": {
           "js-sdk:capabilities": "./src/polyfills/node-capabilities.ts",
+          "js-sdk:response-to-module": "./src/polyfills/response-to-module.ts",
           "js-sdk:minimatch": "./src/polyfills/node-minimatch.ts",
           "js-sdk:worker-url": "./dist/worker/node/worker-url.ts",
           "js-sdk:fs": "node:fs/promises",
@@ -202,6 +204,7 @@ build_bun out='bun' args='[]':
         "minify": false,
         "alias": {
           "js-sdk:worker-url": "./src/polyfills/bun-worker-url.ts",
+          "js-sdk:response-to-module": "./src/polyfills/bun-response-to-module.ts",
           "js-sdk:minimatch": "./src/polyfills/node-minimatch.ts",
           "js-sdk:capabilities": "./src/polyfills/bun-capabilities.ts",
           "js-sdk:fs": "node:fs/promises",
@@ -222,6 +225,7 @@ build_browser out='browser' args='[]':
         "format": "esm",
         "alias": {
           "js-sdk:capabilities": "./src/polyfills/browser-capabilities.ts",
+          "js-sdk:response-to-module": "./src/polyfills/response-to-module.ts",
           "js-sdk:minimatch": "./src/polyfills/node-minimatch.ts",
           "node:worker_threads": "./src/polyfills/host-node-worker_threads.ts",
           "js-sdk:fs": "./src/polyfills/browser-fs.ts",

diff --git a/src/background-plugin.ts b/src/background-plugin.ts
@@ -349,7 +349,7 @@ class HttpContext {
 export async function createBackgroundPlugin(
   opts: InternalConfig,
   names: string[],
-  modules: ArrayBuffer[],
+  modules: WebAssembly.Module[],
 ): Promise<BackgroundPlugin> {
   const worker = new Worker(WORKER_URL);
   const context = new CallContext(SharedArrayBuffer, opts.logger, opts.config);
@@ -394,7 +394,7 @@ export async function createBackgroundPlugin(
     });
   });
 
-  worker.postMessage(message, modules);
+  worker.postMessage(message);
   await onready;
 
   return new BackgroundPlugin(worker, sharedData, opts, context);

diff --git a/src/foreground-plugin.ts b/src/foreground-plugin.ts
@@ -4,17 +4,15 @@ import { loadWasi } from 'js-sdk:wasi';
 
 export const EXTISM_ENV = 'extism:host/env';
 
+type InstantiatedModule = { guestType: string; module: WebAssembly.Module; instance: WebAssembly.Instance };
+
 export class ForegroundPlugin {
   #context: CallContext;
-  #modules: { guestType: string; module: WebAssembly.WebAssemblyInstantiatedSource }[];
+  #modules: InstantiatedModule[];
   #names: string[];
   #active: boolean = false;
 
-  constructor(
-    context: CallContext,
-    names: string[],
-    modules: { guestType: string; module: WebAssembly.WebAssemblyInstantiatedSource }[],
-  ) {
+  constructor(context: CallContext, names: string[], modules: InstantiatedModule[]) {
     this.#context = context;
     this.#names = names;
     this.#modules = modules;
@@ -41,7 +39,7 @@ export class ForegroundPlugin {
           ? [this.lookupTarget(search[0]), search[1]]
           : [
               this.#modules.find((guest) => {
-                const exports = WebAssembly.Module.exports(guest.module.module);
+                const exports = WebAssembly.Module.exports(guest.module);
                 return exports.find((item) => {
                   return item.name === search[0] && item.kind === 'function';
                 });
@@ -53,7 +51,7 @@ export class ForegroundPlugin {
         return false;
       }
 
-      const func = target.module.instance.exports[name] as any;
+      const func = target.instance.exports[name] as any;
 
       if (!func) {
         return false;
@@ -74,7 +72,7 @@ export class ForegroundPlugin {
         ? [this.lookupTarget(search[0]), search[1]]
         : [
             this.#modules.find((guest) => {
-              const exports = WebAssembly.Module.exports(guest.module.module);
+              const exports = WebAssembly.Module.exports(guest.module);
               return exports.find((item) => {
                 return item.name === search[0] && item.kind === 'function';
               });
@@ -85,7 +83,7 @@ export class ForegroundPlugin {
     if (!target) {
       throw Error(`Plugin error: target "${search.join('" "')}" does not exist`);
     }
-    const func = target.module.instance.exports[name] as any;
+    const func = target.instance.exports[name] as any;
     if (!func) {
       throw Error(`Plugin error: function "${search.join('" "')}" does not exist`);
     }
@@ -124,7 +122,7 @@ export class ForegroundPlugin {
     return output;
   }
 
-  private lookupTarget(name: any): { guestType: string; module: WebAssembly.WebAssemblyInstantiatedSource } {
+  private lookupTarget(name: any): InstantiatedModule {
     const target = String(name ?? '0');
     const idx = this.#names.findIndex((xs) => xs === target);
     if (idx === -1) {
@@ -134,15 +132,15 @@ export class ForegroundPlugin {
   }
 
   async getExports(name?: string): Promise<WebAssembly.ModuleExportDescriptor[]> {
-    return WebAssembly.Module.exports(this.lookupTarget(name).module.module) || [];
+    return WebAssembly.Module.exports(this.lookupTarget(name).module) || [];
   }
 
   async getImports(name?: string): Promise<WebAssembly.ModuleImportDescriptor[]> {
-    return WebAssembly.Module.imports(this.lookupTarget(name).module.module) || [];
+    return WebAssembly.Module.imports(this.lookupTarget(name).module) || [];
   }
 
   async getInstance(name?: string): Promise<WebAssembly.Instance> {
-    return this.lookupTarget(name).module.instance;
+    return this.lookupTarget(name).instance;
   }
 
   async close(): Promise<void> {
@@ -153,7 +151,7 @@ export class ForegroundPlugin {
 export async function createForegroundPlugin(
   opts: InternalConfig,
   names: string[],
-  sources: ArrayBuffer[],
+  modules: WebAssembly.Module[],
   context: CallContext = new CallContext(ArrayBuffer, opts.logger, opts.config),
 ): Promise<ForegroundPlugin> {
   const wasi = opts.wasiEnabled ? await loadWasi(opts.allowedPaths) : null;
@@ -171,27 +169,27 @@ export async function createForegroundPlugin(
     }
   }
 
-  const modules = await Promise.all(
-    sources.map(async (source) => {
-      const module = await WebAssembly.instantiate(source, imports);
+  const instances = await Promise.all(
+    modules.map(async (module) => {
+      const instance = await WebAssembly.instantiate(module, imports);
       if (wasi) {
-        await wasi?.initialize(module.instance);
+        await wasi?.initialize(instance);
       }
 
-      const guestType = module.instance.exports.hs_init
+      const guestType = instance.exports.hs_init
         ? 'haskell'
-        : module.instance.exports._initialize
+        : instance.exports._initialize
         ? 'reactor'
-        : module.instance.exports._start
+        : instance.exports._start
         ? 'command'
         : 'none';
 
-      const initRuntime: any = module.instance.exports.hs_init ? module.instance.exports.hs_init : () => {};
+      const initRuntime: any = instance.exports.hs_init ? instance.exports.hs_init : () => {};
       initRuntime();
 
-      return { module, guestType };
+      return { module, instance, guestType };
     }),
   );
 
-  return new ForegroundPlugin(context, names, modules);
+  return new ForegroundPlugin(context, names, instances);
 }
diff --git a/src/interfaces.ts b/src/interfaces.ts
@@ -203,16 +203,39 @@ export interface ManifestWasmPath {
 }
 
 /**
- * The WASM to load as bytes, a path, or a url
+ * Represents a WASM module as a response
+ */
+export interface ManifestWasmResponse {
+  response: Response;
+}
+
+/**
+ * Represents a WASM module as a response
+ */
+export interface ManifestWasmModule {
+  module: WebAssembly.Module;
+}
+
+/**
+ * The WASM to load as bytes, a path, a fetch `Response`, a `WebAssembly.Module`, or a url
  *
  * @property name The name of the Wasm module. Used when disambiguating {@link Plugin#call | `Plugin#call`} targets when the
  * plugin embeds multiple Wasm modules.
  *
  * @property hash The expected SHA-256 hash of the associated Wasm module data. {@link createPlugin} validates incoming Wasm against
  * provided hashes. If running on Node v18, `node` must be invoked using the `--experimental-global-webcrypto` flag.
  *
+ * ⚠️ `module` cannot be used in conjunction with `hash`: the Web Platform does not currently provide a way to get source
+ * bytes from a `WebAssembly.Module` in order to hash.
+ *
  */
-export type ManifestWasm = (ManifestWasmUrl | ManifestWasmData | ManifestWasmPath) & {
+export type ManifestWasm = (
+  | ManifestWasmUrl
+  | ManifestWasmData
+  | ManifestWasmPath
+  | ManifestWasmResponse
+  | ManifestWasmModule
+) & {
   name?: string | undefined;
   hash?: string | undefined;
 };
@@ -241,9 +264,9 @@ export interface Manifest {
 /**
  * Any type that can be converted into an Extism {@link Manifest}.
  * - `object` instances that implement {@link Manifest} are validated.
- * - `ArrayBuffer` instances are converted into {@link Manifest}s with a single {@link ManifestWasmData} member.
+ * - `ArrayBuffer` instances are converted into {@link Manifest}s with a single {@link ManifestUint8Array} member.
  * - `URL` instances are fetched and their responses interpreted according to their `content-type` response header. `application/wasm` and `application/octet-stream` items
- *   are treated as {@link ManifestWasmData} items; `application/json` and `text/json` are treated as JSON-encoded {@link Manifest}s.
+ *   are treated as {@link ManifestUint8Array} items; `application/json` and `text/json` are treated as JSON-encoded {@link Manifest}s.
  * - `string` instances that start with `http://`, `https://`, or `file://` are treated as URLs.
  * - `string` instances that start with `{` treated as JSON-encoded {@link Manifest}s.
  * - All other `string` instances are treated as {@link ManifestWasmPath}.
@@ -266,7 +289,7 @@ export interface Manifest {
  * @throws {@link TypeError} when `URL` parameters don't resolve to a known `content-type`
  * @throws {@link TypeError} when the resulting {@link Manifest} does not contain a `wasm` member with valid {@link ManifestWasm} items.
  */
-export type ManifestLike = Manifest | ArrayBuffer | string | URL;
+export type ManifestLike = Manifest | Response | WebAssembly.Module | ArrayBuffer | string | URL;
 
 export interface Capabilities {
   /**

diff --git a/src/manifest.ts b/src/manifest.ts
@@ -1,11 +1,24 @@
-import type { Manifest, ManifestWasmUrl, ManifestWasmData, ManifestWasmPath, ManifestLike } from './interfaces.ts';
+import type {
+  Manifest,
+  ManifestWasmUrl,
+  ManifestWasmData,
+  ManifestWasmPath,
+  ManifestWasmResponse,
+  ManifestWasmModule,
+  ManifestLike,
+} from './interfaces.ts';
 import { readFile } from 'js-sdk:fs';
+import { responseToModule } from 'js-sdk:response-to-module';
 
 async function _populateWasmField(candidate: ManifestLike, _fetch: typeof fetch): Promise<ManifestLike> {
   if (candidate instanceof ArrayBuffer) {
     return { wasm: [{ data: new Uint8Array(candidate as ArrayBuffer) }] };
   }
 
+  if (candidate instanceof WebAssembly.Module) {
+    return { wasm: [{ module: candidate as WebAssembly.Module }] };
+  }
+
   if (typeof candidate === 'string') {
     if (candidate.search(/^\s*\{/g) === 0) {
       return JSON.parse(candidate);
@@ -18,24 +31,28 @@ async function _populateWasmField(candidate: ManifestLike, _fetch: typeof fetch)
     candidate = new URL(candidate);
   }
 
-  if (candidate instanceof URL) {
-    const response = await _fetch(candidate, { redirect: 'follow' });
+  if (candidate?.constructor?.name === 'Response') {
+    const response: Response = candidate as Response;
     const contentType = response.headers.get('content-type') || 'application/octet-stream';
 
     switch (contentType.split(';')[0]) {
       case 'application/octet-stream':
       case 'application/wasm':
-        return _populateWasmField(await response.arrayBuffer(), _fetch);
+        return { wasm: [{ response }] };
       case 'application/json':
       case 'text/json':
         return _populateWasmField(JSON.parse(await response.text()), _fetch);
       default:
         throw new TypeError(
-          `While processing manifest URL "${candidate}"; expected content-type of "text/json", "application/json", "application/octet-stream", or "application/wasm"; got "${contentType}" after stripping off charset.`,
+          `While processing manifest URL "${response.url}"; expected content-type of "text/json", "application/json", "application/octet-stream", or "application/wasm"; got "${contentType}" after stripping off charset.`,
         );
     }
   }
 
+  if (candidate instanceof URL) {
+    return _populateWasmField(await _fetch(candidate, { redirect: 'follow' }), _fetch);
+  }
+
   if (!('wasm' in candidate)) {
     throw new TypeError('Expected "wasm" key in manifest');
   }
@@ -54,42 +71,61 @@ async function _populateWasmField(candidate: ManifestLike, _fetch: typeof fetch)
   return { ...(candidate as Manifest) };
 }
 
-export async function intoManifest(candidate: ManifestLike, _fetch: typeof fetch = fetch): Promise<Manifest> {
+async function intoManifest(candidate: ManifestLike, _fetch: typeof fetch = fetch): Promise<Manifest> {
   const manifest = (await _populateWasmField(candidate, _fetch)) as Manifest;
   manifest.config ??= {};
   return manifest;
 }
 
-export async function toWasmModuleData(manifest: Manifest, _fetch: typeof fetch): Promise<[string[], ArrayBuffer[]]> {
+export async function toWasmModuleData(
+  input: ManifestLike,
+  _fetch: typeof fetch,
+): Promise<[string[], WebAssembly.Module[]]> {
   const names: string[] = [];
 
+  const manifest = await intoManifest(input, _fetch);
+
   const manifestsWasm = await Promise.all(
     manifest.wasm.map(async (item, idx) => {
-      let buffer: ArrayBuffer;
+      let module: WebAssembly.Module;
+      let buffer: ArrayBuffer | undefined;
       if ((item as ManifestWasmData).data) {
         const data = (item as ManifestWasmData).data;
-
-        if ((data as Uint8Array).buffer) {
-          buffer = data.buffer;
-        } else {
-          buffer = data as ArrayBuffer;
-        }
+        buffer = data.buffer ? data.buffer : data;
+        module = await WebAssembly.compile(data);
       } else if ((item as ManifestWasmPath).path) {
         const path = (item as ManifestWasmPath).path;
         const data = await readFile(path);
         buffer = data.buffer as ArrayBuffer;
-      } else {
+        module = await WebAssembly.compile(data);
+      } else if ((item as ManifestWasmUrl).url) {
         const response = await _fetch((item as ManifestWasmUrl).url, {
           headers: {
             accept: 'application/wasm;q=0.9,application/octet-stream;q=0.8',
             'user-agent': 'extism',
           },
         });
-
-        buffer = await response.arrayBuffer();
+        const result = await responseToModule(response, Boolean(item.hash));
+        buffer = result.data;
+        module = result.module;
+      } else if ((item as ManifestWasmResponse).response) {
+        const result = await responseToModule((item as ManifestWasmResponse).response, Boolean(item.hash));
+        buffer = result.data;
+        module = result.module;
+      } else if ((item as ManifestWasmModule).module) {
+        (<any>names[idx]) = item.name ?? String(idx);
+        return (item as ManifestWasmModule).module;
+      } else {
+        throw new Error(
+          `Unrecognized wasm item at index ${idx}. Keys include: "${Object.keys(item).sort().join(',')}"`,
+        );
       }
 
       if (item.hash) {
+        if (!buffer) {
+          throw new Error('Item specified a hash but WebAssembly.Module source data is unavailable for hashing');
+        }
+
         const hashBuffer = new Uint8Array(await crypto.subtle.digest('SHA-256', buffer));
         const checkBuffer = new Uint8Array(32);
         let eq = true;
@@ -108,7 +144,7 @@ export async function toWasmModuleData(manifest: Manifest, _fetch: typeof fetch)
       }
 
       (<any>names[idx]) = item.name ?? String(idx);
-      return buffer;
+      return module;
     }),
   );