Merge pull request #73 from extism/fix/manifest-options

fix: fallback to manifest values if they're not specified in ExtismPluginOptions

src/background-plugin.ts

@@ -1,6 +1,7 @@
 /*eslint-disable no-empty*/
-import { CallContext, RESET, IMPORT_STATE, EXPORT_STATE, STORE, GET_BLOCK } from './call-context.ts';
-import { PluginOutput, SAB_BASE_OFFSET, SharedArrayBufferSection, type InternalConfig } from './interfaces.ts';
+import { CallContext, RESET, IMPORT_STATE, EXPORT_STATE, STORE, GET_BLOCK, ENV } from './call-context.ts';
+import { MemoryOptions, PluginOutput, SAB_BASE_OFFSET, SharedArrayBufferSection, type InternalConfig } from './interfaces.ts';
+import { readBodyUpTo } from './utils.ts';
 import { WORKER_URL } from './worker-url.ts';
 import { Worker } from 'node:worker_threads';
 import { CAPABILITIES } from './polyfills/deno-capabilities.ts';
@@ -36,25 +37,47 @@ const AtomicsWaitAsync =
   })();
 
 class BackgroundPlugin {
-  worker: Worker;
   sharedData: SharedArrayBuffer;
   sharedDataView: DataView;
   hostFlag: Int32Array;
   opts: InternalConfig;
+  worker?: Worker | undefined;
+  modules: WebAssembly.Module[];
+  names: string[];
 
   #context: CallContext;
   #request: [(result: any[]) => void, (result: any[]) => void] | null = null;
 
-  constructor(worker: Worker, sharedData: SharedArrayBuffer, opts: InternalConfig, context: CallContext) {
-    this.worker = worker;
+  constructor(worker: Worker, sharedData: SharedArrayBuffer, names: string[], modules: WebAssembly.Module[], opts: InternalConfig, context: CallContext) {
     this.sharedData = sharedData;
     this.sharedDataView = new DataView(sharedData);
     this.hostFlag = new Int32Array(sharedData);
     this.opts = opts;
+    this.names = names;
+    this.modules = modules;
     this.#context = context;
 
     this.hostFlag[0] = SAB_BASE_OFFSET;
+    this.setWorker(worker);
+  }
+
+  async restartWorker() {
+    await this.close();
+
+    const worker = await createWorker(this.opts, this.names, this.modules, this.sharedData);
+    this.setWorker(worker);
+  }
+
+  setWorker(worker: Worker) {
+
+    this.#context[RESET]();
 
+    if (this.#request) {
+      this.#request[1]([new Error('Call canceled due to call to restartWorker()')])
+    }
+    this.#request = null;
+
+    this.worker = worker;
     this.worker.on('message', (ev) => this.#handleMessage(ev));
   }
 
@@ -126,6 +149,10 @@ class BackgroundPlugin {
 
     this.#request = [resolve as any, reject as any];
 
+    if (!this.worker) {
+      throw new Error('worker not initialized');
+    }
+
     this.worker.postMessage({
       type: 'invoke',
       handler,
@@ -212,7 +239,7 @@ class BackgroundPlugin {
     //
     // - https://github.com/nodejs/node/pull/44409
     // - https://github.com/denoland/deno/issues/14786
-    const timer = setInterval(() => {}, 0);
+    const timer = setInterval(() => { }, 0);
     try {
       if (!func) {
         throw Error(`Plugin error: host function "${ev.namespace}" "${ev.func}" does not exist`);
@@ -337,7 +364,7 @@ class RingBufferWriter {
 
   signal() {
     const old = Atomics.load(this.flag, 0);
-    while (Atomics.compareExchange(this.flag, 0, old, this.outputOffset) === old) {}
+    while (Atomics.compareExchange(this.flag, 0, old, this.outputOffset) === old) { }
     Atomics.notify(this.flag, 0, 1);
   }
 
@@ -412,11 +439,13 @@ class HttpContext {
   fetch: typeof fetch;
   lastStatusCode: number;
   allowedHosts: string[];
+  memoryOptions: MemoryOptions;
 
-  constructor(_fetch: typeof fetch, allowedHosts: string[]) {
+  constructor(_fetch: typeof fetch, allowedHosts: string[], memoryOptions: MemoryOptions) {
     this.fetch = _fetch;
     this.allowedHosts = allowedHosts;
     this.lastStatusCode = 0;
+    this.memoryOptions = memoryOptions;
   }
 
   contribute(functions: Record<string, Record<string, any>>) {
@@ -453,9 +482,23 @@ class HttpContext {
     });
 
     this.lastStatusCode = response.status;
-    const result = callContext.store(new Uint8Array(await response.arrayBuffer()));
 
-    return result;
+    try {
+      let bytes = this.memoryOptions.maxHttpResponseBytes ?
+        await readBodyUpTo(response, this.memoryOptions.maxHttpResponseBytes) :
+        new Uint8Array(await response.arrayBuffer());
+
+      const result = callContext.store(bytes);
+
+      return result;
+    } catch (err) {
+      if (err instanceof Error) {
+        const ptr = callContext.store(new TextEncoder().encode(err.message));
+        callContext[ENV].log_error(ptr);
+        return 0n;
+      }
+      return 0n;
+    }
   }
 }
 
@@ -464,11 +507,28 @@ export async function createBackgroundPlugin(
   names: string[],
   modules: WebAssembly.Module[],
 ): Promise<BackgroundPlugin> {
-  const worker = new Worker(WORKER_URL);
-  const context = new CallContext(SharedArrayBuffer, opts.logger, opts.config);
-  const httpContext = new HttpContext(opts.fetch, opts.allowedHosts);
+  const context = new CallContext(SharedArrayBuffer, opts.logger, opts.config, opts.memory);
+  const httpContext = new HttpContext(opts.fetch, opts.allowedHosts, opts.memory);
   httpContext.contribute(opts.functions);
 
+  // NB(chrisdickinson): We *have* to create the SharedArrayBuffer in
+  // the parent context because -- for whatever reason! -- chromium does
+  // not allow the creation of shared buffers in worker contexts, but firefox
+  // and webkit do.
+  const sharedData = new (SharedArrayBuffer as any)(opts.sharedArrayBufferSize);
+  new Uint8Array(sharedData).subarray(8).fill(0xfe);
+
+  const worker = await createWorker(opts, names, modules, sharedData);
+  return new BackgroundPlugin(worker, sharedData, names, modules, opts, context);
+}
+
+async function createWorker(
+  opts: InternalConfig,
+  names: string[],
+  modules: WebAssembly.Module[],
+  sharedData: SharedArrayBuffer): Promise<Worker> {
+  const worker = new Worker(WORKER_URL);
+
   await new Promise((resolve, reject) => {
     worker.on('message', function handler(ev) {
       if (ev?.type !== 'initialized') {
@@ -480,13 +540,16 @@ export async function createBackgroundPlugin(
     });
   });
 
-  // NB(chrisdickinson): We *have* to create the SharedArrayBuffer in
-  // the parent context because -- for whatever reason! -- chromium does
-  // not allow the creation of shared buffers in worker contexts, but firefox
-  // and webkit do.
-  const sharedData = new (SharedArrayBuffer as any)(opts.sharedArrayBufferSize);
+  const onready = new Promise((resolve, reject) => {
+    worker.on('message', function handler(ev) {
+      if (ev?.type !== 'ready') {
+        reject(new Error(`received unexpected message (type=${ev?.type})`));
+      }
 
-  new Uint8Array(sharedData).subarray(8).fill(0xfe);
+      worker.removeListener('message', handler);
+      resolve(null);
+    });
+  });
 
   const { fetch: _, logger: __, ...rest } = opts;
   const message = {
@@ -498,19 +561,8 @@ export async function createBackgroundPlugin(
     sharedData,
   };
 
-  const onready = new Promise((resolve, reject) => {
-    worker.on('message', function handler(ev) {
-      if (ev?.type !== 'ready') {
-        reject(new Error(`received unexpected message (type=${ev?.type})`));
-      }
-
-      worker.removeListener('message', handler);
-      resolve(null);
-    });
-  });
-
   worker.postMessage(message);
   await onready;
 
-  return new BackgroundPlugin(worker, sharedData, opts, context);
-}
+  return worker;
+}

src/call-context.ts

@@ -1,4 +1,4 @@
-import { type PluginConfig, PluginOutput } from './interfaces.ts';
+import { type PluginConfig, PluginOutput, MemoryOptions } from './interfaces.ts';
 import { CAPABILITIES } from './polyfills/deno-capabilities.ts';
 
 export const BEGIN = Symbol('begin');
@@ -49,16 +49,18 @@ export class CallContext {
   #logger: Console;
   #decoder: TextDecoder;
   #encoder: TextEncoder;
-  #arrayBufferType: { new (size: number): ArrayBufferLike };
+  #arrayBufferType: { new(size: number): ArrayBufferLike };
   #config: PluginConfig;
   #vars: Map<string, number> = new Map();
+  #memoryOptions: MemoryOptions;
 
   /** @hidden */
-  constructor(type: { new (size: number): ArrayBufferLike }, logger: Console, config: PluginConfig) {
+  constructor(type: { new(size: number): ArrayBufferLike }, logger: Console, config: PluginConfig, memoryOptions: MemoryOptions) {
     this.#arrayBufferType = type;
     this.#logger = logger;
     this.#decoder = new TextDecoder();
     this.#encoder = new TextEncoder();
+    this.#memoryOptions = memoryOptions;
 
     this.#stack = [];
 
@@ -76,6 +78,18 @@ export class CallContext {
     const block = new Block(new this.#arrayBufferType(Number(size)), true);
     const index = this.#blocks.length;
     this.#blocks.push(block);
+
+    if (this.#memoryOptions.maxPages) {
+      const pageSize = 64 * 1024;
+      const totalBytes = this.#blocks.reduce((acc, block) => acc + (block?.buffer.byteLength ?? 0), 0)
+      const totalPages = Math.ceil(totalBytes / pageSize);
+
+      if (totalPages > this.#memoryOptions.maxPages) {
+        this.#logger.error(`memory limit exceeded: ${totalPages} pages requested, ${this.#memoryOptions.maxPages} allowed`);
+        return 0n;
+      }
+    }
+
     return Block.indexToAddress(index);
   }
 
@@ -282,6 +296,15 @@ export class CallContext {
         return 0n;
       }
 
+      const valueBlock = this.#blocks[Block.addressToIndex(valueaddr)];
+      if (this.#memoryOptions.maxVarBytes) {
+        const currentBytes = [...this.#vars.values()].map(idx => this.#blocks[idx]?.byteLength ?? 0).reduce((acc, length) => acc + length, 0)
+        const totalBytes = currentBytes + (valueBlock?.byteLength ?? 0);
+        if (totalBytes > this.#memoryOptions.maxVarBytes) {
+          throw Error(`var memory limit exceeded: ${totalBytes} bytes requested, ${this.#memoryOptions.maxVarBytes} allowed`);
+        }
+      }
+
       this.#vars.set(key, Block.addressToIndex(valueaddr));
     },
 

src/foreground-plugin.ts

@@ -11,11 +11,13 @@ export class ForegroundPlugin {
   #instancePair: InstantiatedModule;
   #active: boolean = false;
   #wasi: InternalWasi[];
+  #opts: InternalConfig;
 
-  constructor(context: CallContext, instancePair: InstantiatedModule, wasi: InternalWasi[]) {
+  constructor(opts: InternalConfig, context: CallContext, instancePair: InstantiatedModule, wasi: InternalWasi[]) {
     this.#context = context;
     this.#instancePair = instancePair;
     this.#wasi = wasi;
+    this.#opts = opts;
   }
 
   async reset(): Promise<boolean> {
@@ -61,6 +63,7 @@ export class ForegroundPlugin {
 
   async call(funcName: string, input?: string | Uint8Array): Promise<PluginOutput | null> {
     const inputIdx = this.#context[STORE](input);
+
     const [errorIdx, outputIdx] = await this.callBlock(funcName, inputIdx);
     const shouldThrow = errorIdx !== null;
     const idx = errorIdx ?? outputIdx;
@@ -103,7 +106,7 @@ export async function createForegroundPlugin(
   opts: InternalConfig,
   names: string[],
   modules: WebAssembly.Module[],
-  context: CallContext = new CallContext(ArrayBuffer, opts.logger, opts.config),
+  context: CallContext = new CallContext(ArrayBuffer, opts.logger, opts.config, opts.memory),
 ): Promise<ForegroundPlugin> {
   const imports: Record<string, Record<string, any>> = {
     [EXTISM_ENV]: context[ENV],
@@ -127,7 +130,7 @@ export async function createForegroundPlugin(
 
   const instance = await instantiateModule(['main'], modules[mainIndex], imports, opts, wasiList, names, modules, seen);
 
-  return new ForegroundPlugin(context, [modules[mainIndex], instance], wasiList);
+  return new ForegroundPlugin(opts, context, [modules[mainIndex], instance], wasiList);
 }
 
 async function instantiateModule(
@@ -211,9 +214,9 @@ async function instantiateModule(
       const instance = providerExports.find((xs) => xs.name === '_start')
         ? await instantiateModule([...current, module], provider, imports, opts, wasiList, names, modules, new Map())
         : !linked.has(provider)
-        ? (await instantiateModule([...current, module], provider, imports, opts, wasiList, names, modules, linked),
-          linked.get(provider))
-        : linked.get(provider);
+          ? (await instantiateModule([...current, module], provider, imports, opts, wasiList, names, modules, linked),
+            linked.get(provider))
+          : linked.get(provider);
 
       if (!instance) {
         // circular import, either make a trampoline or bail
@@ -254,10 +257,10 @@ async function instantiateModule(
   const guestType = instance.exports.hs_init
     ? 'haskell'
     : instance.exports._initialize
-    ? 'reactor'
-    : instance.exports._start
-    ? 'command'
-    : 'none';
+      ? 'reactor'
+      : instance.exports._start
+        ? 'command'
+        : 'none';
 
   if (wasi) {
     await wasi?.initialize(instance);