Skip to content

Commit

Permalink
add manifest timeout
Browse files Browse the repository at this point in the history
  • Loading branch information
mhmd-azeez committed Jun 26, 2024
1 parent 5f658d4 commit 99618e2
Show file tree
Hide file tree
Showing 6 changed files with 37 additions and 11 deletions.
3 changes: 2 additions & 1 deletion src/background-plugin.ts
Original file line number Diff line number Diff line change
@@ -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 { withTimeout } from './utils';
import { WORKER_URL } from './worker-url.ts';
import { Worker } from 'node:worker_threads';
import { CAPABILITIES } from './polyfills/deno-capabilities.ts';
Expand Down Expand Up @@ -143,7 +144,7 @@ class BackgroundPlugin {
async call(funcName: string, input?: string | Uint8Array): Promise<PluginOutput | null> {
const index = this.#context[STORE](input);

const [errorIdx, outputIdx] = await this.callBlock(funcName, index);
const [errorIdx, outputIdx] = await withTimeout(this.callBlock(funcName, index), this.opts.timeoutMs);

const shouldThrow = errorIdx !== null;
const idx = errorIdx ?? outputIdx;
Expand Down
10 changes: 7 additions & 3 deletions src/foreground-plugin.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { CallContext, RESET, GET_BLOCK, BEGIN, END, ENV, STORE } from './call-context.ts';
import { PluginOutput, type InternalConfig, InternalWasi } from './interfaces.ts';
import { withTimeout } from './utils.ts';
import { loadWasi } from './polyfills/deno-wasi.ts';

export const EXTISM_ENV = 'extism:host/env';
Expand All @@ -11,11 +12,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> {
Expand Down Expand Up @@ -61,7 +64,8 @@ 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 [errorIdx, outputIdx] = await withTimeout(this.callBlock(funcName, inputIdx), this.#opts.timeoutMs);
const shouldThrow = errorIdx !== null;
const idx = errorIdx ?? outputIdx;

Expand Down Expand Up @@ -127,7 +131,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(
Expand Down
13 changes: 13 additions & 0 deletions src/interfaces.ts
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,13 @@ export interface ExtismPluginOptions {
sharedArrayBufferSize?: number;
}

export type ManifestOptions = {
allowedPaths?: { [key: string]: string } | undefined;
allowedHosts?: string[] | undefined;
config?: PluginConfigLike;
timeoutMs?: number | undefined;
};

export interface InternalConfig {
logger: Console;
allowedHosts: string[];
Expand All @@ -203,6 +210,7 @@ export interface InternalConfig {
wasiEnabled: boolean;
config: PluginConfig;
sharedArrayBufferSize: number;
timeoutMs: number;
}

export interface InternalWasi {
Expand Down Expand Up @@ -307,6 +315,11 @@ export interface Manifest {
* ```
*/
allowed_hosts?: string[] | undefined;

/**
* Plugin call timeout in milliseconds
*/
timeout_ms?: number | undefined;
}

/**
Expand Down
8 changes: 2 additions & 6 deletions src/manifest.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,11 @@ import type {
ManifestWasmModule,
ManifestLike,
PluginConfigLike,
ManifestOptions,
} from './interfaces.ts';
import { readFile } from './polyfills/node-fs.ts';
import { responseToModule } from './polyfills/response-to-module.ts';

export type ManifestOptions = {
allowedPaths?: { [key: string]: string } | undefined;
allowedHosts?: string[] | undefined;
config?: PluginConfigLike;
};

async function _populateWasmField(candidate: ManifestLike, _fetch: typeof fetch): Promise<ManifestLike> {
if (candidate instanceof ArrayBuffer) {
return { wasm: [{ data: new Uint8Array(candidate as ArrayBuffer) }] };
Expand Down Expand Up @@ -98,6 +93,7 @@ export async function toWasmModuleData(
allowedPaths: manifest.allowed_paths,
allowedHosts: manifest.allowed_hosts,
config: manifest.config,
timeoutMs: manifest.timeout_ms,
};

const manifestsWasm = await Promise.all(
Expand Down
3 changes: 2 additions & 1 deletion src/mod.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { CAPABILITIES } from './polyfills/deno-capabilities.ts';

import type { Manifest, ManifestLike, InternalConfig, ExtismPluginOptions, Plugin } from './interfaces.ts';
import type { ManifestLike, InternalConfig, ExtismPluginOptions, Plugin } from './interfaces.ts';

import { toWasmModuleData as _toWasmModuleData } from './manifest.ts';

Expand Down Expand Up @@ -105,6 +105,7 @@ export async function createPlugin(
config: opts.config,
enableWasiOutput: opts.enableWasiOutput,
sharedArrayBufferSize: Number(opts.sharedArrayBufferSize) || 1 << 16,
timeoutMs: manifestOpts.timeoutMs ?? Number.MAX_SAFE_INTEGER,
};

return (opts.runInWorker ? _createBackgroundPlugin : _createForegroundPlugin)(ic, names, moduleData);
Expand Down
11 changes: 11 additions & 0 deletions src/utils.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
export function withTimeout<T>(promise: Promise<T>, ms: number): Promise<T> {
return new Promise((resolve, reject) => {
const timeout = setTimeout(() => {
reject(new Error('Function call timed out'));
}, ms);

promise.then(resolve, reject).finally(() => {
clearTimeout(timeout);
});
});
}

0 comments on commit 99618e2

Please sign in to comment.