appium/packages/execute-driver-plugin/lib/execute-child.ts

import {fileURLToPath} from 'node:url';
import {promisify} from 'node:util';
import vm from 'node:vm';

import {logger, util} from '@appium/support';

import type {DriverScriptMessageEvent, RunScriptResult, ScriptResult} from './types.js';
import {wrapHostBindingForVmContext} from './vm-host-binding.js';

const log = logger.getLogger('ExecuteDriver Child');
let send: (res: ScriptResult) => Promise<void>;

// duplicate defining these keys here so we don't need to re-load a huge appium
// dependency tree into memory just to run a wdio script
export const W3C_ELEMENT_KEY = util.W3C_WEB_ELEMENT_IDENTIFIER;
export const MJSONWP_ELEMENT_KEY = 'ELEMENT';

/**
 * Run the script in a VM.
 * @param eventParams - The driver script message event parameters
 * @returns Promise resolving to the run script result
 * @throws {TypeError}
 */
async function runScript(eventParams: DriverScriptMessageEvent): Promise<RunScriptResult> {
  const {driverOpts, script, timeoutMs} = eventParams;
  if (typeof timeoutMs !== 'number' || Number.isNaN(timeoutMs)) {
    throw new TypeError('Timeout parameter must be a number');
  }

  /**
   * set up fake logger
   */
  const logs: {error: any[]; warn: any[]; log: any[]} = {
    error: [],
    warn: [],
    log: [],
  };
  const consoleFns: {
    error: (...args: any[]) => void;
    warn: (...args: any[]) => void;
    log: (...args: any[]) => void;
  } = {
    error: wrapHostBindingForVmContext((...logMsgs) => logs.error.push(...logMsgs)),
    warn: wrapHostBindingForVmContext((...logMsgs) => logs.warn.push(...logMsgs)),
    log: wrapHostBindingForVmContext((...logMsgs) => logs.log.push(...logMsgs)),
  };

  const {attach} = await import('webdriverio');

  const driver = await attach(driverOpts);
  const sandboxDriver = wrapHostBindingForVmContext(driver);
  const sandboxConsole = wrapHostBindingForVmContext(consoleFns);
  const sandboxSetTimeout = wrapHostBindingForVmContext(setTimeout);
  const sandboxClearTimeout = wrapHostBindingForVmContext(clearTimeout);

  const fullScript = `(async () => {${script}})();`;

  log.info('Running driver script in Node vm');

  // run the driver script, giving user access to the driver object, a fake console logger,
  // and standard setTimeout/clearTimeout functions. Each host value is proxied so
  // main-realm prototype metadata cannot be used to obtain the host `Function` (VM escape).
  let result = await vm.runInNewContext(
    fullScript,
    {
      driver: sandboxDriver,
      console: sandboxConsole,
      setTimeout: sandboxSetTimeout,
      clearTimeout: sandboxClearTimeout,
    },
    {timeout: timeoutMs, breakOnSigint: true},
  );

  result = coerceScriptResult(result);
  log.info('Successfully ensured driver script result is appropriate type for return');
  return {result, logs};
}

/**
 * We can get any manner of crazy thing back from a vm executing untrusted
 * code. We might also get WebdriverIO objects that aren't suitable for JSON
 * response. So make sure we convert the things we know about to their
 * appropriate response format, and squash other weird things.
 *
 * @param obj - object to convert and sanitize
 * @return safely converted object
 */
function coerceScriptResult(obj: any): any {
  // first ensure obj is of a type that can be JSON encoded safely. This will
  // get rid of custom objects, functions, etc... and turn them into POJOs
  try {
    obj = JSON.parse(JSON.stringify(obj));
  } catch {
    log.warn(
      'Could not convert executeDriverScript to safe response!' +
        `Result was: ${JSON.stringify(obj)}. Will make it null`,
    );
    return null;
  }

  let res: any;

  // now we begin our recursive case options
  if (util.isPlainObject(obj)) {
    // if we have an object, it's either an element object or something else
    // webdriverio has no monadic object types other than element and driver,
    // and we don't want to allow special casing return of driver
    res = {};

    if (obj[MJSONWP_ELEMENT_KEY] || obj[W3C_ELEMENT_KEY]) {
      // if it's an element object, clear out anything that's not the key, and then return the
      // object. Note that if the element object contains only one element key type (MJSONWP or
      // W3C), keep it that way in our response. But if the element contains both key types
      // (because the client is being backwards compatible) then keep both keys in our response as
      // well.
      if (obj[MJSONWP_ELEMENT_KEY]) {
        res[MJSONWP_ELEMENT_KEY] = obj[MJSONWP_ELEMENT_KEY];
      }

      if (obj[W3C_ELEMENT_KEY]) {
        res[W3C_ELEMENT_KEY] = obj[W3C_ELEMENT_KEY];
      }
      return res;
    }

    // otherwise, recurse into the object
    for (const key of Object.keys(obj)) {
      res[key] = coerceScriptResult(obj[key]);
    }
    return res;
  }

  // also handle arrays
  if (Array.isArray(obj)) {
    return obj.map(coerceScriptResult);
  }

  // base case, if it's not an object or array, return straightaway
  return obj;
}

/**
 * Entry point to runScript
 * @param eventParams - The driver script message event parameters
 */
async function main(eventParams: DriverScriptMessageEvent): Promise<void> {
  /**
   * keep the response of runScript
   */
  let res: ScriptResult;
  log.info('Parameters received from parent process');
  try {
    res = {success: await runScript(eventParams)};
    log.info('runScript success');
  } catch (error: any) {
    log.info('runScript error');
    res = {error: {message: error.message, stack: error.stack}};
  }
  await send(res);
}

// ensure we're running this script in IPC mode
if (process.argv[1] && fileURLToPath(import.meta.url) === process.argv[1] && typeof process.send === 'function') {
  send = promisify(process.send).bind(process) as (res: ScriptResult) => Promise<void>;
  log.info('Running driver execution in child process');
  process.on('message', main);
}