artillery/dist/lib/commons/engine_util.js

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.js";
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;
try {
    xmlCapture = require('artillery-xml-capture');
}
catch (_e) {
    xmlCapture = null;
}
export { createThink, createLoopWithCount, createParallel, isProbableEnough, template, captureOrMatch, evil, ensurePropertyIsAList, renderVariables as _renderVariables };
function createThink(requestSpec, opts) {
    opts = opts || {};
    const thinkspec = requestSpec.think;
    const f = function think(context, callback) {
        const templatedThink = template(thinkspec, context);
        let thinktime = 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);
    };
    return f;
}
// "count" can be an integer (negative or positive) or a string defining a range
// like "1-15"
function createLoopWithCount(count, steps, opts = {}) {
    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 = null;
        let loopValue = 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;
                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) {
            const zero = (cb2) => 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) {
                        shouldContinue = b;
                        return cb(err, context2);
                    });
                }
                else {
                    return cb(err, context2);
                }
            });
        }, (err, finalContext) => {
            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, opts) {
    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) => {
            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) {
    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, context, inPlace) {
    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;
                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, context) {
    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, vars) {
    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)[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, code) {
    const context = vm.createContext(sandbox);
    const script = new vm.Script(code);
    try {
        return script.runInContext(context);
    }
    catch (_e) {
        return null;
    }
}
function parseLoopCount(countSpec) {
    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, defaultStrict) {
    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, prop) {
    if (Array.isArray(obj[prop])) {
        return obj;
    }
    obj[prop] = [].concat(typeof obj[prop] === 'undefined' ? [] : obj[prop]);
    return obj;
}
function captureOrMatch(params, 
// WS/Socket.IO engines pass body-only faux responses:
response, context, done) {
    if ((!params.capture || params.capture.length === 0) &&
        (!params.match || params.match.length === 0)) {
        return done(null, null);
    }
    const result = {
        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, 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, doc) => {
            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, response) {
    let parser;
    let extractor;
    let expr;
    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, callback) {
    let r = null;
    let err = null;
    try {
        if (typeof body === 'string') {
            r = JSON.parse(body);
        }
        else {
            r = body;
        }
    }
    catch (e) {
        err = e;
    }
    return callback(err, r);
}
function dummyParser(body, callback) {
    return callback(null, body);
}
// doc is a JSON object
function extractJSONPath(doc, expr, opts) {
    // 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, expr, opts) {
    const group = opts.group;
    const flags = opts.flags;
    let str;
    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, headerName) {
    return headers[headerName] || '';
}
function extractCheerio(doc, expr, opts) {
    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) {
    // 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;
    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) {
    const headers = res.headers;
    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, high) {
    return Math.floor(Math.random() * (high - low + 1) + low);
}
function sanitiseValue(value) {
    if (value === 0 || value === false || value === null || value === undefined)
        return value;
    return value ? value : '';
}