FastGPT/projects/sandbox/src/worker/runJs.ts

import { RunCodeDto, RunCodeResponse } from 'src/sandbox/dto/create-sandbox.dto';
import { parentPort } from 'worker_threads';
import { workerResponse } from './utils';

// @ts-ignore
const ivm = require('isolated-vm');

parentPort?.on('message', ({ code, variables = {} }: RunCodeDto) => {
  const resolve = (data: RunCodeResponse) => workerResponse({ parentPort, type: 'success', data });
  const reject = (error: any) => workerResponse({ parentPort, type: 'error', data: error });
  try {
    const isolate = new ivm.Isolate({ memoryLimit: 32 });
    const context = isolate.createContextSync();
    const jail = context.global;

    // custom function
    const logData = [];
    const CustomLogStr = 'CUSTOM_LOG';
    code = code.replace(/console\.log/g, `${CustomLogStr}`);
    jail.setSync(CustomLogStr, function (...args) {
      logData.push(
        args
          .map((item) => (typeof item === 'object' ? JSON.stringify(item, null, 2) : item))
          .join(', ')
      );
    });

    jail.setSync('responseData', function (args: any): any {
      if (typeof args === 'object') {
        resolve({
          codeReturn: args,
          log: logData.join('\n')
        });
      } else {
        reject('Not an invalid response, must return an object');
      }
    });

    // Add global variables
    jail.setSync('variables', new ivm.ExternalCopy(variables).copyInto());

    const scriptCode = `
      ${code}
      responseData(main(variables))`;

    context.evalSync(scriptCode, { timeout: 6000 });
  } catch (err) {
    console.log(err);
    reject(err);
  }

  process.exit();
});