Merge pull request #12986 from Budibase/isolated-vm-wrapper

Isolated vm wrapper

packages/backend-core/src/context/types.ts

@@ -1,5 +1,4 @@
-import { IdentityContext } from "@budibase/types"
-import { Isolate, Context, Module } from "isolated-vm"
+import { IdentityContext, VM } from "@budibase/types"
 
 // keep this out of Budibase types, don't want to expose context info
 export type ContextMap = {
@@ -10,9 +9,5 @@ export type ContextMap = {
   isScim?: boolean
   automationId?: string
   isMigrating?: boolean
-  isolateRefs?: {
-    jsIsolate: Isolate
-    jsContext: Context
-    helpersModule: Module
-  }
+  vm?: VM
 }

packages/server/src/api/routes/tests/row.spec.ts

@@ -2031,7 +2031,7 @@ describe.each([
 
   describe("Formula JS protection", () => {
     it("should time out JS execution if a single cell takes too long", async () => {
-      await config.withEnv({ JS_PER_EXECUTION_TIME_LIMIT_MS: 20 }, async () => {
+      await config.withEnv({ JS_PER_INVOCATION_TIMEOUT_MS: 20 }, async () => {
         const js = Buffer.from(
           `
               let i = 0;
@@ -2071,8 +2071,8 @@ describe.each([
     it("should time out JS execution if a multiple cells take too long", async () => {
       await config.withEnv(
         {
-          JS_PER_EXECUTION_TIME_LIMIT_MS: 20,
-          JS_PER_REQUEST_TIME_LIMIT_MS: 40,
+          JS_PER_INVOCATION_TIMEOUT_MS: 20,
+          JS_PER_REQUEST_TIMEOUT_MS: 40,
         },
         async () => {
           const js = Buffer.from(

packages/server/src/environment.ts

@@ -71,10 +71,10 @@ const environment = {
   SELF_HOSTED: process.env.SELF_HOSTED,
   HTTP_MB_LIMIT: process.env.HTTP_MB_LIMIT,
   FORKED_PROCESS_NAME: process.env.FORKED_PROCESS_NAME || "main",
-  JS_PER_EXECUTION_TIME_LIMIT_MS:
-    parseIntSafe(process.env.JS_PER_EXECUTION_TIME_LIMIT_MS) || 1000,
-  JS_PER_REQUEST_TIME_LIMIT_MS: parseIntSafe(
-    process.env.JS_PER_REQUEST_TIME_LIMIT_MS
+  JS_PER_INVOCATION_TIMEOUT_MS:
+    parseIntSafe(process.env.JS_PER_INVOCATION_TIMEOUT_MS) || 1000,
+  JS_PER_REQUEST_TIMEOUT_MS: parseIntSafe(
+    process.env.JS_PER_REQUEST_TIMEOUT_MS
   ),
   // old
   CLIENT_ID: process.env.CLIENT_ID,

packages/server/src/jsRunner/index.ts

@@ -1,155 +1,33 @@
-import ivm from "isolated-vm"
 import env from "../environment"
 import { setJSRunner, JsErrorTimeout } from "@budibase/string-templates"
-import { context } from "@budibase/backend-core"
 import tracer from "dd-trace"
-import url from "url"
-import crypto from "crypto"
-import querystring from "querystring"
-import { BundleType, loadBundle } from "./bundles"
 
-class ExecutionTimeoutError extends Error {
-  constructor(message: string) {
-    super(message)
-    this.name = "ExecutionTimeoutError"
-  }
-}
+import { IsolatedVM } from "./vm"
+import { context } from "@budibase/backend-core"
 
 export function init() {
-  const helpersSource = loadBundle(BundleType.HELPERS)
   setJSRunner((js: string, ctx: Record<string, any>) => {
     return tracer.trace("runJS", {}, span => {
       try {
         const bbCtx = context.getCurrentContext()!
 
-        const isolateRefs = bbCtx.isolateRefs
-        if (!isolateRefs) {
-          const jsIsolate = new ivm.Isolate({
+        let { vm } = bbCtx
+        if (!vm) {
+          // Can't copy the native helpers into the isolate. We just ignore them as they are handled properly from the helpersSource
+          const { helpers, ...ctxToPass } = ctx
+
+          vm = new IsolatedVM({
             memoryLimit: env.JS_RUNNER_MEMORY_LIMIT,
+            invocationTimeout: env.JS_PER_INVOCATION_TIMEOUT_MS,
+            isolateAccumulatedTimeout: env.JS_PER_REQUEST_TIMEOUT_MS,
           })
-          const jsContext = jsIsolate.createContextSync()
+            .withContext(ctxToPass)
+            .withHelpers()
 
-          const injectedRequire = `
-            const require = function(val){
-              switch (val) {
-                case "url": 
-                  return {
-                    resolve: (...params) => urlResolveCb(...params),
-                    parse: (...params) => urlParseCb(...params),
-                  }
-                case "querystring":
-                  return {
-                    escape: (...params) => querystringEscapeCb(...params),
-                  }
-              }
-            };`
-
-          const global = jsContext.global
-          global.setSync(
-            "urlResolveCb",
-            new ivm.Callback((...params: Parameters<typeof url.resolve>) =>
-              url.resolve(...params)
-            )
-          )
-
-          global.setSync(
-            "urlParseCb",
-            new ivm.Callback((...params: Parameters<typeof url.parse>) =>
-              url.parse(...params)
-            )
-          )
-
-          global.setSync(
-            "querystringEscapeCb",
-            new ivm.Callback(
-              (...params: Parameters<typeof querystring.escape>) =>
-                querystring.escape(...params)
-            )
-          )
-
-          global.setSync(
-            "helpersStripProtocol",
-            new ivm.Callback((str: string) => {
-              var parsed = url.parse(str) as any
-              parsed.protocol = ""
-              return parsed.format()
-            })
-          )
-
-          const helpersModule = jsIsolate.compileModuleSync(
-            `${injectedRequire};${helpersSource}`
-          )
-
-          const cryptoModule = jsIsolate.compileModuleSync(
-            `export default { randomUUID: cryptoRandomUUIDCb }`
-          )
-          cryptoModule.instantiateSync(jsContext, specifier => {
-            throw new Error(`No imports allowed. Required: ${specifier}`)
-          })
-
-          global.setSync(
-            "cryptoRandomUUIDCb",
-            new ivm.Callback(
-              (...params: Parameters<typeof crypto.randomUUID>) => {
-                return crypto.randomUUID(...params)
-              }
-            )
-          )
-
-          helpersModule.instantiateSync(jsContext, specifier => {
-            if (specifier === "crypto") {
-              return cryptoModule
-            }
-            throw new Error(`No imports allowed. Required: ${specifier}`)
-          })
-
-          for (const [key, value] of Object.entries(ctx)) {
-            if (key === "helpers") {
-              // Can't copy the native helpers into the isolate. We just ignore them as they are handled properly from the helpersSource
-              continue
-            }
-            global.setSync(key, value)
-          }
-
-          bbCtx.isolateRefs = { jsContext, jsIsolate, helpersModule }
+          bbCtx.vm = vm
         }
 
-        let { jsIsolate, jsContext, helpersModule } = bbCtx.isolateRefs!
-
-        const perRequestLimit = env.JS_PER_REQUEST_TIME_LIMIT_MS
-        if (perRequestLimit) {
-          const cpuMs = Number(jsIsolate.cpuTime) / 1e6
-          if (cpuMs > perRequestLimit) {
-            throw new ExecutionTimeoutError(
-              `CPU time limit exceeded (${cpuMs}ms > ${perRequestLimit}ms)`
-            )
-          }
-        }
-
-        const script = jsIsolate.compileModuleSync(
-          `import helpers from "compiled_module";const result=${js};cb(result)`,
-          {}
-        )
-
-        script.instantiateSync(jsContext, specifier => {
-          if (specifier === "compiled_module") {
-            return helpersModule
-          }
-
-          throw new Error(`"${specifier}" import not allowed`)
-        })
-
-        let result
-        jsContext.global.setSync(
-          "cb",
-          new ivm.Callback((value: any) => {
-            result = value
-          })
-        )
-
-        script.evaluateSync({
-          timeout: env.JS_PER_EXECUTION_TIME_LIMIT_MS,
-        })
+        const result = vm.execute(js)
 
         return result
       } catch (error: any) {

packages/server/src/jsRunner/vm/index.ts

@@ -0,0 +1,202 @@
+import ivm from "isolated-vm"
+
+import url from "url"
+import crypto from "crypto"
+import querystring from "querystring"
+
+import { BundleType, loadBundle } from "../bundles"
+import { VM } from "@budibase/types"
+
+class ExecutionTimeoutError extends Error {
+  constructor(message: string) {
+    super(message)
+    this.name = "ExecutionTimeoutError"
+  }
+}
+
+class ModuleHandler {
+  private modules: {
+    import: string
+    moduleKey: string
+    module: ivm.Module
+  }[] = []
+
+  private generateRandomKey = () => `i${crypto.randomUUID().replace(/-/g, "")}`
+
+  registerModule(module: ivm.Module, imports: string) {
+    this.modules.push({
+      moduleKey: this.generateRandomKey(),
+      import: imports,
+      module: module,
+    })
+  }
+
+  generateImports() {
+    return this.modules
+      .map(m => `import ${m.import} from "${m.moduleKey}"`)
+      .join(";")
+  }
+
+  getModule(key: string) {
+    const module = this.modules.find(m => m.moduleKey === key)
+    return module?.module
+  }
+}
+
+export class IsolatedVM implements VM {
+  private isolate: ivm.Isolate
+  private vm: ivm.Context
+  private jail: ivm.Reference
+  private invocationTimeout: number
+  private isolateAccumulatedTimeout?: number
+
+  private moduleHandler = new ModuleHandler()
+
+  private readonly resultKey = "results"
+
+  constructor({
+    memoryLimit,
+    invocationTimeout,
+    isolateAccumulatedTimeout,
+  }: {
+    memoryLimit: number
+    invocationTimeout: number
+    isolateAccumulatedTimeout?: number
+  }) {
+    this.isolate = new ivm.Isolate({ memoryLimit })
+    this.vm = this.isolate.createContextSync()
+    this.jail = this.vm.global
+    this.jail.setSync("global", this.jail.derefInto())
+
+    this.addToContext({
+      [this.resultKey]: { out: "" },
+    })
+
+    this.invocationTimeout = invocationTimeout
+    this.isolateAccumulatedTimeout = isolateAccumulatedTimeout
+  }
+
+  withHelpers() {
+    const urlModule = this.registerCallbacks({
+      resolve: url.resolve,
+      parse: url.parse,
+    })
+
+    const querystringModule = this.registerCallbacks({
+      escape: querystring.escape,
+    })
+
+    this.addToContext({
+      helpersStripProtocol: new ivm.Callback((str: string) => {
+        var parsed = url.parse(str) as any
+        parsed.protocol = ""
+        return parsed.format()
+      }),
+    })
+
+    const injectedRequire = `const require=function req(val) {
+        switch (val) {
+            case "url": return ${urlModule};
+            case "querystring": return ${querystringModule};
+        }
+      }`
+    const helpersSource = loadBundle(BundleType.HELPERS)
+    const helpersModule = this.isolate.compileModuleSync(
+      `${injectedRequire};${helpersSource}`
+    )
+
+    helpersModule.instantiateSync(this.vm, specifier => {
+      if (specifier === "crypto") {
+        const cryptoModule = this.registerCallbacks({
+          randomUUID: crypto.randomUUID,
+        })
+        const module = this.isolate.compileModuleSync(
+          `export default ${cryptoModule}`
+        )
+        module.instantiateSync(this.vm, specifier => {
+          throw new Error(`No imports allowed. Required: ${specifier}`)
+        })
+        return module
+      }
+      throw new Error(`No imports allowed. Required: ${specifier}`)
+    })
+
+    this.moduleHandler.registerModule(helpersModule, "helpers")
+    return this
+  }
+
+  withContext(context: Record<string, any>) {
+    this.addToContext(context)
+    return this
+  }
+
+  execute(code: string): any {
+    if (this.isolateAccumulatedTimeout) {
+      const cpuMs = Number(this.isolate.cpuTime) / 1e6
+      if (cpuMs > this.isolateAccumulatedTimeout) {
+        throw new ExecutionTimeoutError(
+          `CPU time limit exceeded (${cpuMs}ms > ${this.isolateAccumulatedTimeout}ms)`
+        )
+      }
+    }
+
+    code = `${this.moduleHandler.generateImports()};results.out=${code};`
+
+    const script = this.isolate.compileModuleSync(code)
+
+    script.instantiateSync(this.vm, specifier => {
+      const module = this.moduleHandler.getModule(specifier)
+      if (module) {
+        return module
+      }
+
+      throw new Error(`"${specifier}" import not allowed`)
+    })
+
+    script.evaluateSync({ timeout: this.invocationTimeout })
+
+    const result = this.getResult()
+    return result
+  }
+
+  private registerCallbacks(functions: Record<string, any>) {
+    const libId = crypto.randomUUID().replace(/-/g, "")
+
+    const x: Record<string, string> = {}
+    for (const [funcName, func] of Object.entries(functions)) {
+      const key = `f${libId}${funcName}cb`
+      x[funcName] = key
+
+      this.addToContext({
+        [key]: new ivm.Callback((...params: any[]) => (func as any)(...params)),
+      })
+    }
+
+    const mod =
+      `{` +
+      Object.entries(x)
+        .map(([key, func]) => `${key}: ${func}`)
+        .join() +
+      "}"
+    return mod
+  }
+
+  private addToContext(context: Record<string, any>) {
+    for (let key in context) {
+      const value = context[key]
+      this.jail.setSync(
+        key,
+        typeof value === "function"
+          ? value
+          : new ivm.ExternalCopy(value).copyInto({ release: true })
+      )
+    }
+  }
+
+  private getResult() {
+    const ref = this.vm.global.getSync(this.resultKey, { reference: true })
+    const result = ref.copySync()
+    ref.release()
+    return result.out
+  }
+}

packages/server/src/utilities/scriptRunner.ts

@@ -1,63 +1,23 @@
-import ivm from "isolated-vm"
+import env from "../environment"
+import { IsolatedVM } from "../jsRunner/vm"
 
 const JS_TIMEOUT_MS = 1000
-
 class ScriptRunner {
-  vm: IsolatedVM
+  private code
+  private vm
 
   constructor(script: string, context: any) {
-    const code = `let fn = () => {\n${script}\n}; results.out = fn();`
-    this.vm = new IsolatedVM({ memoryLimit: 8 })
-    this.vm.context = {
-      ...context,
-      results: { out: "" },
-    }
-    this.vm.code = code
+    this.code = `(() => {${script}})();`
+    this.vm = new IsolatedVM({
+      memoryLimit: env.JS_RUNNER_MEMORY_LIMIT,
+      invocationTimeout: JS_TIMEOUT_MS,
+    }).withContext(context)
   }
 
   execute() {
-    this.vm.runScript()
-    const results = this.vm.getValue("results")
-    return results.out
-  }
-}
-
-class IsolatedVM {
-  isolate: ivm.Isolate
-  vm: ivm.Context
-  #jail: ivm.Reference
-  script: any
-
-  constructor({ memoryLimit }: { memoryLimit: number }) {
-    this.isolate = new ivm.Isolate({ memoryLimit })
-    this.vm = this.isolate.createContextSync()
-    this.#jail = this.vm.global
-    this.#jail.setSync("global", this.#jail.derefInto())
-  }
-
-  getValue(key: string) {
-    const ref = this.vm.global.getSync(key, { reference: true })
-    const result = ref.copySync()
-    ref.release()
+    const result = this.vm.execute(this.code)
     return result
   }
-
-  set context(context: Record<string, any>) {
-    for (let key in context) {
-      this.#jail.setSync(key, this.copyRefToVm(context[key]))
-    }
-  }
-
-  set code(code: string) {
-    this.script = this.isolate.compileScriptSync(code)
-  }
-
-  runScript() {
-    this.script.runSync(this.vm, { timeout: JS_TIMEOUT_MS })
-  }
-
-  copyRefToVm(value: Object): ivm.Copy<Object> {
-    return new ivm.ExternalCopy(value).copyInto({ release: true })
-  }
 }
+
 export default ScriptRunner

packages/types/src/sdk/index.ts

@@ -20,3 +20,4 @@ export * from "./cli"
 export * from "./websocket"
 export * from "./permissions"
 export * from "./row"
+export * from "./vm"

packages/types/src/sdk/vm.ts

@@ -0,0 +1,3 @@
+export interface VM {
+  execute(code: string): any
+}