import { TLogType } from "./../query/logger/logger.types"; import * as ivm from "isolated-vm"; import { VModule } from "./module.class"; import { Plugin } from "./plugin.class"; import { QueryResponse } from "./vm.constants"; import { TLog } from "src/query/logger/logger.types"; import { LoggerService } from "src/query/logger/logger.service"; export class Vm { private memoryLimit: number; private modules: VModule[]; private context: any; private jail: any; private plugins: Plugin[]; private functions: string[]; private isolate: ivm.Isolate; private timeLimit?: bigint; private cpuTimeLimit?: bigint; private sessionId: string | null; constructor(configs: { memoryLimit: number; timeLimit?: bigint; cpuTimeLimit?: bigint; modules: VModule[]; plugins: Plugin[]; functions: string[]; }) { this.memoryLimit = configs.memoryLimit; this.modules = configs.modules; this.plugins = configs.plugins; this.timeLimit = configs.timeLimit; this.cpuTimeLimit = configs.cpuTimeLimit; this.functions = configs.functions; } async init(): Promise { this.isolate = new ivm.Isolate({ memoryLimit: this.memoryLimit }); this.context = await this.isolate.createContext(); this.jail = this.context.global; this.jail.set("global", this.jail.derefInto()); for (const mod of this.modules) { await this.context.eval(mod.getSource()); } for (const fn of this.functions) { await this.context.eval(fn); } for (const plugin of this.plugins) { const pluginName = plugin.getName(); await this.context.evalClosure( "globalThis[$0] = globalThis[$0] || Object.create(null);", [pluginName], { arguments: { copy: true } } ); for (const method of plugin.getMethods()) { const fnRef = new ivm.Reference(async (...args) => { return await plugin[method](...args); }); await this.context.evalClosure( ` globalThis[$0][$1] = (...args) => $2.apply(undefined, args, { arguments: { copy: true }, result: { promise: true, copy: true } }); `, [pluginName, method, fnRef], { arguments: { copy: true } } ); } } return this; } setFunction(name: string, func: (...args) => any) { this.jail.setSync(name, func, { arguments: { copy: true } }); } async runScript( script: string, args: Record, headers: Record, log: TLog ): Promise { let resolvePromise: (value: any) => void; let rejectPromise: (reason?: any) => void; const resultPromise = new Promise((resolve, reject) => { resolvePromise = resolve; rejectPromise = reject; }); this.setFunction("returnResult", (res) => { resolvePromise({ ...res, log }); }); this.setFunction("log", (...args) => { if (!log) { return; } const logType = args.find( (arg) => arg && typeof arg === "object" && arg.type && Object.values(TLogType).includes(arg.type) ); log = LoggerService.log(log, { content: args .map((arg) => (typeof arg === "string" ? arg : JSON.stringify(arg))) .join(" "), type: logType?.type || TLogType.info, timeStamp: new Date().getTime(), }); }); this.setFunction("error", (error: any) => { console.error("Script error:", error); rejectPromise(error); }); const scriptWithResult = ` (async () => { ${script} try { const result = await main(${JSON.stringify( args )}, ${JSON.stringify(headers)}); returnResult(result) } catch (e) { error(e) } })() `; const compiledScript = await this.isolate.compileScript(scriptWithResult); const interval = setInterval(() => { if ( this.isolate.cpuTime > this.cpuTimeLimit || this.isolate.wallTime > this.timeLimit ) { this.isolate.dispose(); rejectPromise(new Error("Script execution timed out")); } }, 500); compiledScript.run(this.context); try { return await resultPromise; } finally { clearInterval(interval); this.onFinish(); } } onFinish() { this.plugins.forEach((plugin) => { plugin.onFinish(); }); } }