feat: add vm2 sandbox support as vm runtime backup

Author: moonrailgun

pnpm-lock.yaml

@@ -1,7 +1,5 @@
 import { MonitorProvider } from './type.js';
-import ivm from 'isolated-vm';
-import { buildSandbox, environmentScript } from '../../../utils/sandbox.js';
-import { env } from '../../../utils/env.js';
+import { runCodeInVM } from '../../../utils/vm/index.js';
 
 export const custom: MonitorProvider<{
   code: string;
@@ -22,45 +20,3 @@ export const custom: MonitorProvider<{
     return result;
   },
 };
-
-export async function runCodeInVM(_code: string) {
-  const start = Date.now();
-  const isolate = new ivm.Isolate({ memoryLimit: env.sandboxMemoryLimit });
-
-  // avoid end comment with line break
-  const code = `${environmentScript}
-
-;(async () => {
-  ${_code}
-})()`;
-
-  const [context, script] = await Promise.all([
-    isolate.createContext(),
-    isolate.compileScript(code),
-  ]);
-
-  const logger: any[][] = [];
-
-  buildSandbox(context, {
-    console: {
-      log: (...args: any[]) => {
-        logger.push(['log', Date.now(), ...args]);
-      },
-      warn: (...args: any[]) => {
-        logger.push(['warn', Date.now(), ...args]);
-      },
-      error: (...args: any[]) => {
-        logger.push(['error', Date.now(), ...args]);
-      },
-    },
-  });
-
-  const res = await script.run(context, {
-    promise: true,
-  });
-
-  context.release();
-  script.release();
-
-  return { logger, result: res, usage: Date.now() - start };
-}

src/server/model/monitor/provider/custom.ts

@@ -77,6 +77,7 @@
     "trpc-to-openapi": "^2.1.3",
     "uuid": "^9.0.1",
     "vite-express": "^0.13.0",
+    "vm2": "^3.9.19",
     "winston": "^3.17.0",
     "yup": "^1.6.1",
     "zeromq": "^6.3.0",

src/server/package.json

@@ -26,14 +26,14 @@ import {
   MonitorModelSchema,
   MonitorStatusPageModelSchema,
 } from '../../prisma/zod/index.js';
-import { runCodeInVM } from '../../model/monitor/provider/custom.js';
 import { createAuditLog } from '../../model/auditLog.js';
 import {
   MonitorInfoWithNotificationIds,
   monitorPublicInfoSchema,
 } from '../../model/_schema/monitor.js';
 import { monitorPageManager } from '../../model/monitor/page/manager.js';
 import { token } from '../../model/notification/token/index.js';
+import { runCodeInVM } from '../../utils/vm/index.js';
 
 export const monitorRouter = router({
   all: workspaceProcedure

src/server/trpc/routers/monitor.ts

@@ -69,9 +69,12 @@ export const env = {
   allowRegister: checkEnvTrusty(process.env.ALLOW_REGISTER),
   allowOpenapi: checkEnvTrusty(process.env.ALLOW_OPENAPI ?? 'true'),
   websiteId: process.env.WEBSITE_ID,
-  sandboxMemoryLimit: process.env.SANDBOX_MEMORY_LIMIT
-    ? Number(process.env.SANDBOX_MEMORY_LIMIT)
-    : 16, // unit: MB
+  sandbox: {
+    useVM2: checkEnvTrusty(process.env.USE_VM2),
+    memoryLimit: process.env.SANDBOX_MEMORY_LIMIT
+      ? Number(process.env.SANDBOX_MEMORY_LIMIT)
+      : 16, // unit: MB
+  },
   puppeteerExecutablePath: process.env.PUPPETEER_EXECUTABLE_PATH,
   dbDebug: checkEnvTrusty(process.env.DB_DEBUG),
   amapToken: process.env.AMAP_TOKEN,

src/server/utils/env.ts

@@ -0,0 +1,82 @@
+import ivm from 'isolated-vm';
+import { buildSandbox, environmentScript } from './sandbox.js';
+import { env } from '../env.js';
+import { runCodeInVM2 } from './sandbox-vm2.js';
+
+export async function runCodeInVM(_code: string): Promise<{
+  logger: any[][];
+  result: any;
+  usage: number;
+}> {
+  const code = `(async () => {${_code}})();`;
+
+  try {
+    // Try to use VM2 first if enabled via environment variable
+    if (env.sandbox.useVM2) {
+      try {
+        const ret = await runCodeInVM2(code);
+        return ret;
+      } catch (err) {
+        console.error(
+          '[Monitor] VM2 execution failed, falling back to isolated-vm:',
+          err
+        );
+        throw err;
+      }
+    } else {
+      // Use isolated-vm
+      try {
+        const ret = await runCodeInIVM(code);
+        return ret;
+      } catch (err) {
+        console.error('[Monitor] isolated-vm execution failed:', err);
+        throw err;
+      }
+    }
+  } catch (err) {
+    console.error('[Monitor] Code execution failed:', err);
+    throw err;
+  }
+}
+
+async function runCodeInIVM(_code: string) {
+  const start = Date.now();
+  const isolate = new ivm.Isolate({ memoryLimit: env.sandbox.memoryLimit });
+
+  // avoid end comment with line break
+  const code = `${environmentScript}
+
+;(async () => {
+  ${_code}
+})()`;
+
+  const [context, script] = await Promise.all([
+    isolate.createContext(),
+    isolate.compileScript(code),
+  ]);
+
+  const logger: any[][] = [];
+
+  buildSandbox(context, {
+    console: {
+      log: (...args: any[]) => {
+        logger.push(['log', Date.now(), ...args]);
+      },
+      warn: (...args: any[]) => {
+        logger.push(['warn', Date.now(), ...args]);
+      },
+      error: (...args: any[]) => {
+        logger.push(['error', Date.now(), ...args]);
+      },
+    },
+  });
+
+  const res = await script.run(context, {
+    promise: true,
+  });
+
+  context.release();
+  script.release();
+
+  return { logger, result: res, usage: Date.now() - start };
+}

src/server/utils/vm/index.ts

@@ -0,0 +1,144 @@
+import { Worker } from 'worker_threads';
+import { env } from '../env.js';
+import path from 'path';
+import fs from 'fs';
+import { nanoid } from 'nanoid';
+import os from 'os';
+
+// Function to create a temporary worker file
+async function createTempWorkerFile(
+  code: string,
+  memoryLimitMB: number
+): Promise<string> {
+  const tempDir = os.tmpdir();
+  const id = nanoid();
+  const filePath = path.join(tempDir, `worker-${id}.js`);
+
+  // The worker script that will use VM2 to execute the code
+  const workerScript = `
+    const { parentPort } = require('worker_threads');
+    const { VM } = require('vm2');
+
+    // Set up memory limit
+    const memoryLimitBytes = ${memoryLimitMB} * 1024 * 1024;
+
+    // Create a sandbox with console logging
+    const logger = [];
+
+    (async () => {
+      try {
+        const vm = new VM({
+          timeout: 5000, // 5 seconds timeout
+          sandbox: {
+            console: {
+              log: (...args) => {
+                logger.push(['log', Date.now(), ...args]);
+              },
+              warn: (...args) => {
+                logger.push(['warn', Date.now(), ...args]);
+              },
+              error: (...args) => {
+                logger.push(['error', Date.now(), ...args]);
+              }
+            },
+            request: async (config) => {
+              // This is a simplified version of the request function
+              // In a real implementation, you would use axios or another HTTP client
+              const axios = require('axios');
+              try {
+                const result = await axios.request(config);
+                return {
+                  headers: { ...result.headers },
+                  data: result.data,
+                  status: result.status
+                };
+              } catch (error) {
+                throw new Error(error.message);
+              }
+            }
+          },
+          compiler: 'javascript',
+          eval: false,
+          wasm: false
+        });
+
+        // Execute the code
+        const start = Date.now();
+        const result = await vm.run(${JSON.stringify(code)});
+        const usage = Date.now() - start;
+
+        // Send the result back to the parent
+        parentPort.postMessage({ success: true, result, logger, usage });
+      } catch (error) {
+        // Send the error back to the parent
+        parentPort.postMessage({
+          success: false,
+          error: error.message,
+          logger
+        });
+      }
+    })();
+  `;
+
+  await fs.promises.writeFile(filePath, workerScript);
+  return filePath;
+}
+
+// Function to run code in VM2 via worker threads
+export async function runCodeInVM2(
+  code: string,
+  memoryLimitMB = env.sandbox.memoryLimit
+): Promise<{ logger: any[][]; result: any; usage: number }> {
+  const start = Date.now();
+
+  // Create a temporary worker file
+  const workerFilePath = await createTempWorkerFile(code, memoryLimitMB);
+
+  return new Promise((resolve, reject) => {
+    // Create a worker
+    const worker = new Worker(workerFilePath);
+
+    // Set a timeout to kill the worker if it takes too long
+    const timeout = setTimeout(() => {
+      worker.terminate();
+      fs.unlinkSync(workerFilePath);
+      reject(new Error('Execution timed out'));
+    }, 10000); // 10 seconds timeout
+
+    // Listen for messages from the worker
+    worker.on('message', (message) => {
+      clearTimeout(timeout);
+
+      // Clean up the temporary worker file
+      fs.unlinkSync(workerFilePath);
+
+      if (message.success) {
+        resolve({
+          logger: message.logger || [],
+          result: message.result,
+          usage: message.usage || Date.now() - start,
+        });
+      } else {
+        const error = new Error(message.error || 'Unknown error');
+        // error.logger = message.logger || [];
+        reject(error);
+      }
+    });
+
+    // Handle worker errors
+    worker.on('error', (error) => {
+      clearTimeout(timeout);
+      fs.unlinkSync(workerFilePath);
+      reject(error);
+    });
+
+    // Handle worker exit
+    worker.on('exit', (code) => {
+      clearTimeout(timeout);
+      if (code !== 0) {
+        fs.unlinkSync(workerFilePath);
+        reject(new Error(`Worker stopped with exit code ${code}`));
+      }
+    });
+  });
+}