few-line-engine/src/vm/vm.class.ts

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<Vm> {
    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<string, any>,
    headers: Record<string, any>,
    log: TLog
  ): Promise<QueryResponse> {
    let resolvePromise: (value: any) => void;
    let rejectPromise: (reason?: any) => void;

    const resultPromise = new Promise<any>((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();
    });
  }
}