artillery/lib/commons/engine_util.ts

Cloud-scale load testing. https://www.artillery.io

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

import { createRequire } from 'node:module';
import vm from 'node:vm';
import async from 'async';
import * as cheerio from 'cheerio';
import createDebug from 'debug';
import deepForEach from 'deep-for-each';
import * as espree from 'espree';
import { JSONPath as jsonpath } from 'jsonpath-plus';
import L from 'lodash';
import ms from 'ms';
import { jitter } from './jitter.ts';

const debug = createDebug('engine_util');
const A = async;

const require = createRequire(import.meta.url);

// Optional dependency - loaded lazily and kept synchronous (no
// top-level await: the package must stay require()-able)
let xmlCapture: {
  parseXML: ParserFn;
  extractXPath: ExtractorFn;
} | null;
try {
  xmlCapture = require('artillery-xml-capture');
} catch (_e) {
  xmlCapture = null;
}

// TODO Write tests

// Pragmatic typing: user scripts and processor hooks feed dynamic
// values through these utilities, so specs/params and VU variables
// are typed Record<string, any> until canonical script input types
// exist (modernization plan, phase 2).

export interface TemplateContext {
  vars: Record<string, any>;
  funcs?: Record<string, (...args: any[]) => any>;
  [key: string]: any;
}

type StepCallback = (err: Error | null | undefined, context?: any) => void;
type StepFunction = (context: any, callback: StepCallback) => void;

type ParserFn = (
  body: unknown,
  callback: (err: unknown, doc?: any) => void
) => unknown;
type ExtractorFn = (doc: any, expr: string, opts: Record<string, any>) => any;

interface CaptureEntry {
  value?: any;
  error?: unknown;
  strict?: boolean;
  failed?: boolean;
}

interface MatchEntry {
  success?: boolean;
  expected?: any;
  got?: any;
  expression?: string;
  error?: unknown;
  strict?: boolean;
}

export interface CaptureOrMatchResult {
  captures: Record<string, CaptureEntry>;
  matches: Record<string, MatchEntry>;
  failedCaptures: boolean;
}

export {
  createThink,
  createLoopWithCount,
  createParallel,
  isProbableEnough,
  template,
  captureOrMatch,
  evil,
  ensurePropertyIsAList,
  renderVariables as _renderVariables
};

function createThink(
  requestSpec: Record<string, any>,
  opts?: Record<string, any>
): StepFunction {
  opts = opts || {};

  const thinkspec = requestSpec.think;

  const f: StepFunction = function think(context, callback) {
    const templatedThink = template(thinkspec, context);
    let thinktime: string | number = Number.isInteger(
      L.toNumber(templatedThink)
    )
      ? ms(`${templatedThink}s`)
      : ms(templatedThink);

    if (typeof thinktime === 'undefined') {
      throw new Error(`Invalid think time: ${templatedThink || thinkspec}`);
    }

    if (requestSpec.jitter || opts?.jitter) {
      thinktime = jitter(`${thinktime}:${requestSpec.jitter || opts?.jitter}`);
    }
    debug(
      'think %s, %s, %s -> %s',
      requestSpec.think,
      requestSpec.jitter,
      opts?.jitter,
      thinktime
    );
    setTimeout(() => {
      callback(null, context);
    }, thinktime as number);
  };

  return f;
}

// "count" can be an integer (negative or positive) or a string defining a range
// like "1-15"
function createLoopWithCount(
  count: number | string,
  steps: StepFunction[],
  opts: Record<string, any> = {}
): StepFunction {
  return function aLoop(context, callback) {
    let count2 = count;
    if (typeof count === 'string') {
      count2 = template(count, context);
    }

    const from = parseLoopCount(count2).from;
    const to = parseLoopCount(count2).to;

    let i = from;
    let newContext = context;
    const loopIndexVar = opts?.loopValue || '$loopCount';
    const loopElementVar = opts?.loopElement || '$loopElement';
    // Should we stop early because the value of "over" is not an array
    let abortEarly = false;

    let overValues: any[] | null = null;
    let loopValue: any = i; // default to the current iteration of the loop, ie same as $loopCount
    if (typeof opts.overValues !== 'undefined') {
      if (opts.overValues && typeof opts.overValues === 'object') {
        const ov = opts.overValues as any[];
        overValues = ov;
        loopValue = ov[i];
      } else if (opts.overValues && typeof opts.overValues === 'string') {
        overValues = L.get(context.vars, opts.overValues);
        if (L.isArray(overValues)) {
          loopValue = overValues[i];
        } else {
          abortEarly = true;
        }
      }
    }

    newContext.vars[loopElementVar] = loopValue;
    newContext.vars[loopIndexVar] = i;

    let shouldContinue = true;

    A.whilst(
      function test() {
        if (abortEarly) {
          return false;
        }
        if (opts.whileTrue) {
          return shouldContinue;
        }
        if (overValues !== null) {
          return i !== overValues.length;
        } else {
          return i < to || to === -1;
        }
      },
      function repeated(cb: (err?: Error | null, context?: any) => void) {
        const zero = (cb2: (err: Error | null, context: any) => void) =>
          cb2(null, newContext);
        const steps2 = L.flatten([zero, steps]);

        A.waterfall(steps2, (err, context2) => {
          if (err) {
            return cb(err, context2);
          }
          i++;
          newContext = context2;

          newContext.vars[loopIndexVar]++;
          if (overValues !== null) {
            newContext.vars[loopElementVar] = overValues[i];
          }

          if (opts.whileTrue) {
            opts.whileTrue(context2, function done(b: boolean) {
              shouldContinue = b;
              return cb(err, context2);
            });
          } else {
            return cb(err, context2);
          }
        });
      },
      (err: Error | null | undefined, finalContext: any) => {
        if (typeof finalContext === 'undefined') {
          // this happens if test() returns false immediately, e.g. with
          // nested loops where one of the inner loops goes over an
          // empty array
          return callback(err, newContext);
        }
        return callback(err, finalContext);
      }
    );
  };
}

function createParallel(
  steps: StepFunction[],
  opts?: Record<string, any>
): StepFunction {
  const limit = opts?.limitValue || 100;

  return function aParallel(context, callback) {
    const newContext = context;
    const newCallback = callback;

    // Remap the steps array to pass the context into each step.
    const newSteps = L.map(steps, (step) => (callback: StepCallback) => {
      step(newContext, callback);
    });

    // Run each of the steps in parallel.
    A.parallelLimit(newSteps, limit, (err, _finalContext) => {
      // We don't need to do anything with the array of contexts returned from each step at the moment.
      return newCallback(err, newContext);
    });
  };
}

function isProbableEnough(obj: Record<string, any>): boolean {
  if (typeof obj.probability === 'undefined') {
    return true;
  }

  let probability = Number(obj.probability) || 0;
  if (probability > 100) {
    probability = 100;
  }

  const r = L.random(100);
  return r < probability;
}

// Returns the same kind of value it was given - dynamic by design.
function template(o: any, context: TemplateContext, inPlace?: boolean): any {
  let result;

  if (typeof o === 'undefined') {
    return undefined;
  }

  if (o && (o.constructor === Object || o.constructor === Array)) {
    if (!inPlace) {
      result = L.cloneDeep(o);
    } else {
      result = o;
    }
    templateObjectOrArray(result, context);
  } else if (typeof o === 'string') {
    if (!/{{/.test(o)) {
      return o;
    }
    const funcCallRegex =
      /{{\s*(\$[A-Za-z0-9_]+\s*\(\s*[A-Za-z0-9_,\s]*\s*\))\s*}}/;
    const match = o.match(funcCallRegex);
    if (match) {
      // This looks like it could be a function call:
      const syntax = espree.parse(match[1]);
      // TODO: Use a proper schema for what we expect here
      if (
        syntax.body &&
        syntax.body.length === 1 &&
        syntax.body[0].type === 'ExpressionStatement'
      ) {
        const funcName = syntax.body[0].expression.callee.name;
        const args = L.map(
          syntax.body[0].expression.arguments,
          (arg) => arg.value
        );
        // NOTE: pre-existing behavior: throws when the script uses a
        // function-call template without any context funcs defined.
        const funcs = context.funcs as NonNullable<TemplateContext['funcs']>;
        if (funcName in funcs) {
          return template(
            o.replace(funcCallRegex, funcs[funcName].apply(null, args)),
            context
          );
        }
      }
    } else {
      if (!o.match(/{{/)) {
        return o;
      }

      result = renderVariables(o, context.vars);
    }
  } else {
    return o;
  }

  return result;
}

// Mutates the object in place
function templateObjectOrArray(o: any, context: TemplateContext): void {
  deepForEach(o, (value, key, subj, path) => {
    const newPath = template(path, context, true);

    let newValue;
    if (value && value.constructor !== Object && value.constructor !== Array) {
      newValue = template(value, context, true);
    } else {
      newValue = value;
    }

    debug(
      `path = ${path} ; value = ${JSON.stringify(
        value
      )} (${typeof value}) ; (subj type: ${
        subj.length ? 'list' : 'hash'
      }) ; newValue = ${JSON.stringify(newValue)} ; newPath = ${newPath}`
    );

    // If path has changed, we need to unset the original path and
    // explicitly walk down the new subtree from this path:
    if (path !== newPath) {
      L.unset(o, path);
      newValue = template(value, context, true);
    }

    if (newPath.endsWith(key)) {
      const keyIndex = newPath.lastIndexOf(key);
      const prefix = newPath.substr(0, keyIndex - 1);
      L.set(o, `${prefix}["${key}"]`, newValue);
    } else {
      L.set(o, newPath, newValue);
    }
  });
}

function renderVariables(str: string, vars: Record<string, any>): any {
  const RX = /{{{?[\s$\w.[\]'"-]+}}}?/g;
  let _rxmatch;
  let result = str.substring(0, str.length);

  // Special case for handling integer/boolean/object substitution:
  //
  // Does the template string contain one variable and nothing else?
  // e.g.: "{{ myvar }" or "{{    myvar }", but NOT " {{ myvar }"
  // If so, we treat it as a special case.
  const matches = str.match(RX);
  if (matches && matches.length === 1) {
    if (matches[0] === str) {
      // there's nothing else in the template but the variable
      const varName = str.replace(/{/g, '').replace(/}/g, '').trim();
      return sanitiseValue(L.get(vars, varName));
    }
  }

  while (result.search(RX) > -1) {
    const templateStr = (result.match(RX) as RegExpMatchArray)[0];
    const varName = templateStr.replace(/{/g, '').replace(/}/g, '').trim();

    let varValue = L.get(vars, varName);

    if (typeof varValue === 'object') {
      varValue = JSON.stringify(varValue);
    }
    result = result.replace(templateStr, varValue);
  }

  return result;
}

// Presume code is valid JS code (i.e. that it has been checked elsewhere)
function evil(sandbox: Record<string, any>, code: string): unknown {
  const context = vm.createContext(sandbox);
  const script = new vm.Script(code);
  try {
    return script.runInContext(context);
  } catch (_e) {
    return null;
  }
}

function parseLoopCount(countSpec: number | string): {
  from: number;
  to: number;
} {
  let from = 0;
  let to = 0;

  if (typeof countSpec === 'number') {
    from = 0;
    to = countSpec;
  } else if (typeof countSpec === 'string') {
    if (Number.isNaN(Number(countSpec))) {
      if (/\d-\d/.test(countSpec)) {
        from = Number(countSpec.split('-')[0]);
        to = Number(countSpec.split('-')[1]);
      } else {
        to = 0;
      }
    } else {
      to = Number(countSpec);
    }
  } else {
    to = 0;
  }

  return { from: from, to: to };
}

function isCaptureFailed(
  v: Record<string, any>,
  defaultStrict?: boolean
): boolean {
  const noValue =
    typeof v.value === 'undefined' ||
    v.value === '' ||
    typeof v.error !== 'undefined';

  if (!noValue) {
    return false;
  }

  return !(
    (typeof defaultStrict === 'undefined' && v.strict === false) ||
    (defaultStrict === true && v.strict === false) ||
    (defaultStrict === false && typeof v.strict === 'undefined') ||
    (defaultStrict === false && v.strict === false)
  );
}

// Helper function to wrap an object's property in a list if it's
// defined, or set it to an empty list if not.
function ensurePropertyIsAList(
  obj: Record<string, any>,
  prop: string
): Record<string, any> {
  if (Array.isArray(obj[prop])) {
    return obj;
  }

  obj[prop] = ([] as any[]).concat(
    typeof obj[prop] === 'undefined' ? [] : obj[prop]
  );
  return obj;
}

function captureOrMatch(
  params: Record<string, any>,
  // WS/Socket.IO engines pass body-only faux responses:
  response: { body?: unknown; headers?: Record<string, any> },
  context: TemplateContext,
  done: (err: Error | null, result: CaptureOrMatchResult | null) => void
) {
  if (
    (!params.capture || params.capture.length === 0) &&
    (!params.match || params.match.length === 0)
  ) {
    return done(null, null);
  }

  const result: CaptureOrMatchResult = {
    captures: {},
    matches: {},
    failedCaptures: false
  };

  // Objects updated in place the first time this runs:
  ensurePropertyIsAList(params, 'capture');
  ensurePropertyIsAList(params, 'match');

  const specs = params.capture.concat(params.match);

  async.eachSeries(
    specs,
    (spec: Record<string, any>, next) => {
      const parsedSpec = parseSpec(spec, response);
      const parser = parsedSpec.parser;
      const extractor = parsedSpec.extractor;
      const expr = parsedSpec.expr;

      // are we looking at body or headers:
      var content = response.body;
      if (spec.header) {
        content = response.headers;
      }

      parser(content, (err: unknown, doc: any) => {
        if (err) {
          if (spec.as) {
            result.captures[spec.as] = {
              error: err,
              strict: spec.strict
            };
            result.captures[spec.as].failed = isCaptureFailed(
              result.captures[spec.as],
              context._defaultStrictCapture
            );
          } else {
            result.matches[spec.expr] = {
              error: err,
              strict: spec.strict
            };
          }
          return next(null);
        }

        const extractedValue = extractor(doc, template(expr, context), spec);

        if (spec.value !== undefined) {
          // this is a match spec
          const expected = template(spec.value, context);
          debug(
            'match: %s, expected: %s, got: %s',
            expr,
            expected,
            extractedValue
          );
          if (extractedValue !== expected) {
            result.matches[expr] = {
              success: false,
              expected: expected,
              got: extractedValue,
              expression: expr,
              strict: spec.strict
            };
          } else {
            result.matches.expr = {
              success: true,
              expected: expected,
              expression: expr
            };
          }
          return next(null);
        }

        if (spec.as) {
          // this is a capture
          debug('capture: %s = %s', spec.as, extractedValue);
          result.captures[spec.as] = {
            value: extractedValue,
            strict: spec.strict
          };

          result.captures[spec.as].failed = isCaptureFailed(
            result.captures[spec.as],
            context._defaultStrictCapture
          );
        }

        return next(null);
      });
    },
    (err) => {
      if (err) {
        return done(err, null);
      } else {
        return done(null, result);
      }
    }
  );
}

function parseSpec(
  spec: Record<string, any>,
  response: { body?: unknown; headers?: Record<string, any> }
): { parser: ParserFn; extractor: ExtractorFn; expr: string } {
  let parser: ParserFn;
  let extractor: ExtractorFn;
  let expr: string;

  if (spec.json) {
    parser = parseJSON;
    extractor = extractJSONPath;
    expr = spec.json;
  } else if (xmlCapture && spec.xpath) {
    parser = xmlCapture.parseXML;
    extractor = xmlCapture.extractXPath;
    expr = spec.xpath;
  } else if (spec.regexp) {
    parser = dummyParser;
    extractor = extractRegExp;
    expr = spec.regexp;
  } else if (spec.header) {
    parser = dummyParser;
    extractor = extractHeader;
    expr = spec.header;
  } else if (spec.selector) {
    parser = dummyParser;
    extractor = extractCheerio;
    expr = spec.selector;
  } else {
    if (isJSON(response)) {
      parser = parseJSON;
      extractor = extractJSONPath;
      expr = spec.json;
    } else if (xmlCapture && isXML(response)) {
      parser = xmlCapture.parseXML;
      extractor = xmlCapture.extractXPath;
      expr = spec.xpath;
    } else {
      // We really don't know what to do here.
      parser = dummyParser;
      extractor = dummyExtractor;
      expr = '';
    }
  }

  return { parser: parser, extractor: extractor, expr: expr };
}

/*
 * Wrap JSON.parse in a callback
 */
function parseJSON(body: unknown, callback: (err: unknown, doc?: any) => void) {
  let r: any = null;
  let err: unknown = null;

  try {
    if (typeof body === 'string') {
      r = JSON.parse(body);
    } else {
      r = body;
    }
  } catch (e) {
    err = e;
  }

  return callback(err, r);
}

function dummyParser(
  body: unknown,
  callback: (err: unknown, doc?: any) => void
) {
  return callback(null, body);
}

// doc is a JSON object
function extractJSONPath(
  doc: any,
  expr: string,
  opts: Record<string, any>
): any {
  // typeof null is 'object' hence the explicit check here
  if (typeof doc !== 'object' || doc === null) {
    return '';
  }

  let results;

  try {
    results = jsonpath({ path: expr, json: doc, wrap: opts.multiple ?? true });
  } catch (queryErr) {
    debug(queryErr);
  }

  if (!results) {
    return '';
  }

  if (opts.multiple === false) {
    return results;
  }

  if (results.length > 1) {
    return results[randomInt(0, results.length - 1)];
  } else {
    return results[0];
  }
}

// doc is a string or an object (body parsed by Request when headers indicate JSON)
function extractRegExp(
  doc: unknown,
  expr: string,
  opts: Record<string, any>
): string | undefined {
  const group = opts.group;
  const flags = opts.flags;
  let str: string;
  if (typeof doc === 'string') {
    str = doc;
  } else {
    str = JSON.stringify(doc); // FIXME: not the same string as the one we got from the server
  }
  let rx;
  if (flags) {
    rx = new RegExp(expr, flags);
  } else {
    rx = new RegExp(expr);
  }
  const match = rx.exec(str);
  if (!match) {
    return '';
  }

  // Captures named group (https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Regular_expressions/Named_capturing_group)
  if (group && match.groups) {
    return match.groups[group];
    // Captures integer index defined group since those don't show up in match.groups
  } else if (group && match[group]) {
    return match[group];
    // Defaults to first match if found and no group defined
  } else if (match[0]) {
    return match[0];
    // If no match returns empty string
  } else {
    return '';
  }
}

function extractHeader(headers: any, headerName: string): any {
  return headers[headerName] || '';
}

function extractCheerio(
  doc: any,
  expr: string,
  opts: Record<string, any>
): string | undefined {
  const $ = cheerio.load(doc);
  const els = $(expr);
  let i = 0;
  if (typeof opts.index !== 'undefined') {
    if (opts.index === 'random') {
      i = Math.ceil(Math.random() * els.get().length - 1);
    } else if (opts.index === 'last') {
      i = els.get().length - 1;
    } else if (typeof Number(opts.index) === 'number') {
      i = Number(opts.index);
    }
  }
  return els.slice(i, i + 1).attr(opts.attr);
}

function dummyExtractor() {
  return '';
}

/*
 * Given a response object determine if it's JSON
 */
function isJSON(res: { headers?: Record<string, any> }): boolean {
  // NOTE: pre-existing behavior: throws when a body-only faux response
  // reaches the content-type fallback (no capture type and no headers).
  const headers = res.headers as Record<string, any>;
  debug('isJSON: content-type = %s', headers['content-type']);
  return Boolean(
    headers['content-type'] &&
      /^application\/json/.test(headers['content-type'])
  );
}

/*
 * Given a response object determine if it's some kind of XML
 */
function isXML(res: { headers?: Record<string, any> }): boolean {
  const headers = res.headers as Record<string, any>;
  return Boolean(
    headers['content-type'] &&
      (/^[a-zA-Z]+\/xml/.test(headers['content-type']) ||
        /^[a-zA-Z]+\/[a-zA-Z]+\+xml/.test(headers['content-type']))
  );
}

function randomInt(low: number, high: number): number {
  return Math.floor(Math.random() * (high - low + 1) + low);
}

function sanitiseValue(value: any): any {
  if (value === 0 || value === false || value === null || value === undefined)
    return value;
  return value ? value : '';
}