/**
Copyright: 2024 Contrast Security, Inc
  Contact: support@contrastsecurity.com
  License: Commercial

  NOTICE: This Software and the patented inventions embodied within may only be
  used as part of Contrast Security’s commercial offerings. Even though it is
  made available through public repositories, use of this Software is subject to
  the applicable End User Licensing Agreement found at
  https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
  between Contrast Security and the End User. The Software may not be reverse
  engineered, modified, repackaged, sold, redistributed or otherwise used in a
  way not consistent with the End User License Agreement.
*/
/**
 * @module lib/hooks/patcher
 */
'use strict';

// Disable the prefer-rest-params rule, which disallows use of `arguments`.
// With the way we're monkey patching, it's more sementically correct to use
// the arguments keyword instead of providing a `...args` rest param.
/* eslint-disable prefer-rest-params */

const logger = require('../core/logger')('contrast:hooks');
const agent = require('../agent');
const perfLoggingEnabled =
  agent.config && agent.config.agent.node.req_perf_logging;
const perfLogger = perfLoggingEnabled
  ? require('../core/logger/perf-logger')(agent)
  : undefined;
const tracker = require('../tracker.js');
const { AsyncStorage } = require('../core/async-storage');
const {
  Scopes,
  SCOPE_NAMES: { NO_PROPAGATION }
} = require('../core/async-storage/scopes');
const promisifyCustom = Symbol.for('nodejs.util.promisify.custom');
const some = require('../util/some');

// When a pre-hook returns the result of this function, the value passed will
// be used in place of the original function's return value, and the original
// function will NOT be executed.
function preempt(value) {
  return new PatcherPreempt(value);
}

function isFunction(fn) {
  return fn instanceof Function || typeof fn === 'function';
}

function PatcherPreempt(v) {
  this.value = v;
}

/**
 * function => {fn, ...}
 * allows us to check if a function is hooked without pinning random properties
 * to the function, and lets us reference the original function any time we want
 */
const hookedFunctions = new WeakMap();

/*
 * mapping of hooked func -> array of pre/post funcs
 */
let hooks;

/**
 * clears instrumentation for all patched functions
 *
 * used in testing for better isolation between tests
 */
function resetInstrumentation() {
  hooks = new WeakMap();
}
resetInstrumentation();

const functionToString = Function.prototype.toString;
Function.prototype.toString = function() {
  if (hookedFunctions.has(this)) {
    return Reflect.apply(
      functionToString,
      hookedFunctions.get(this).fn,
      arguments
    );
  }
  return Reflect.apply(functionToString, this, arguments);
};

Object.defineProperty(Function.prototype, '__isContrastHooked', {
  get() {
    return hookedFunctions.has(this);
  }
});

Object.defineProperty(Function.prototype, '_contrast_toString', {
  value: function _contrast_toString() {
    return Reflect.apply(functionToString, this, arguments);
  }
});

function runHooks(type, data, thisTarget, fnHooks) {
  fnHooks.forEach((hook, key) => {
    if (hook[type]) {
      hook[type].apply(thisTarget, [data]);
    }
  });
}

/**
 * Checks if any of the provided arguments are tracked.
 * @param {any[]} args
 * @return {boolean}
 */
function argsTracked(args) {
  return some(args, (arg) => arg && tracker.getData(arg));
}

/**
 * Whether to skip the running of instrumentation in registered pre/post hooks.
 * When this returns `false` (don't skip), the pre/post hooks will execute in
 * a NO_INSTRUMENTATION scope, and the original function will be invoked in
 * the scope provided by the options to `patcher.patch()`.
 * @param {Object} options patcher options
 */
function skipInstrumentation(options) {
  // certain instrumentation must always run(i.e. source instrumentation
  // and protect sensors)
  if (options.alwaysRun) {
    return false;
  }

  // we are in our own code(i.e tsReporter), or a deadzoned module(i.e bcrypt)
  // or we are in a pre/post function of instrumentation
  if (Scopes.isNoInstrumentationScoped()) {
    return true;
  }

  // we only want to run the top-most propagator instrumentation
  // so if we are already in a no propagation scope, skip further instrumentation
  // by libraries inside a propagator
  if (options.scope === NO_PROPAGATION && Scopes.isNoPropagationScoped()) {
    return true;
  }

  // If we do not have contrast context no reason to run instrumentation.
  // Right now this means in an active request but in future could mean
  // more. For example, actively processing data from a message queue in
  // a microservice.
  return !AsyncStorage.getContext();
}

/**
 * @this  {Function}        the hooked function
 * @param {Function} fn     the original function
 * @param {any}      target the constructor being called (if called with new)
 * @param {Object}   data
 */
function getResult(fn, data, target) {
  if (data.result instanceof PatcherPreempt) {
    // If the pre-hook returned patcher.preempt(value) just use the passed
    // value and don't execute the original function.
    return data.result.value;
  }

  return runOriginalFunction.call(this, fn, data, target);
}

function runOriginalFunction(fn, { args, name }, target) {
  // Invoke the original function in the current context with the
  // passed-in arguments, and store its result.
  if (target) {
    // NODE-818 need to unwrap Object because it does not
    // properly reflect
    if (name === 'Object') {
      target = unwrap(target);
    }

    return Reflect.construct(fn, args, target);
  }

  return fn.apply(this, args);
}

const runWrapper =
  !process.hrtime.bigint || !perfLoggingEnabled
    ? runWithoutRecordingTime
    : runAndRecordTime;

function runWithoutRecordingTime(func) {
  return () => func();
}

function runAndRecordTime(func, funcKey, type) {
  return () => {
    const start = process.hrtime.bigint();
    const rv = func();
    perfLogger[type](funcKey, Number(process.hrtime.bigint() - start));
    return rv;
  };
}

/**
 *
 * @param {Function} fn original function being hooked
 * @param {Object}   options
 * @param {string}   options.funcKey - identifier for this patch patchType:funcName
 * @param {string}   options.patchType - one of patcher.PATCH_TYPES
 * @param {string}   options.name - name of the function (usually signature)
 * @param {boolean}  options.alwaysRun - flag to say always run instrumentation for given function
 * @param {string}   options.scope - name of scope to run original function in
 * @returns {Function} the hooked function
 */
function hookFunction(fn, options) {
  if (hookedFunctions.has(fn)) {
    // f has already been hooked
    return fn;
  }

  const { funcKey } = options;
  logger.trace(`hook ${funcKey}`);

  // eslint-disable-next-line complexity
  function hooked(...args) {
    const target = new.target;

    if (skipInstrumentation(options)) {
      return runOriginalFunction.call(
        this,
        fn,
        { args, name: options.name },
        target
      );
    }

    // short-circuit optimization for add
    if (fn.name === '__add' && !argsTracked(args)) {
      return args[0] + args[1];
    }

    const data = {
      hooked,
      funcKey,
      obj: new.target || this,
      name: options.name,
      args,
      result: undefined
    };

    if (perfLoggingEnabled) {
      perfLogger.logEvent(funcKey);
    }

    let hasPreHook, hasPostHook;
    const fnHooks = hooks && hooks.get(hooked);

    // If there's a pre event run it in a no instrumentation scope
    // this will prevent other instrumentation from running
    // within the pre function
    if (fnHooks) {
      for (const storedHooks of fnHooks.values()) {
        hasPreHook = Boolean(hasPreHook || (storedHooks && storedHooks.pre));
        hasPostHook = Boolean(hasPostHook || (storedHooks && storedHooks.post));
      }
    }

    if (hasPreHook) {
      Scopes.runInNoInstrumentationScope(
        runWrapper(() => runHooks('pre', data, this, fnHooks), funcKey, 'pre'),
        `${funcKey} pre`
      );
    }

    // Run original function in scope that was passed in
    // or just run the original function when the passed scope is undefined
    if (options.scope) {
      data.result = Scopes.runIn(
        options.scope,
        runWrapper(
          () => getResult.call(this, fn, data, target),
          funcKey,
          'orig'
        ),
        funcKey
      );
    } else {
      data.result = runWrapper(
        () => getResult.call(this, fn, data, target),
        funcKey,
        'orig'
      )();
    }

    // If there's a post event run it in a no instrumentation scope
    // this will prevent other instrumentation from running
    // within the post function
    if (hasPostHook) {
      Scopes.runInNoInstrumentationScope(
        runWrapper(
          () => runHooks('post', data, this, fnHooks),
          funcKey,
          'post'
        ),
        `${funcKey} post`
      );
    }

    return data.result;
  }

  // copy over all properties if there are properties on the original function.
  // an example of this is the version of buffer you get from require('buffer')
  // see NODE-335 for further background.
  const descriptors = Object.getOwnPropertyDescriptors(fn);
  for (const key of [
    ...Object.getOwnPropertySymbols(descriptors),
    ...Object.getOwnPropertyNames(descriptors)
  ]) {
    if (
      key === promisifyCustom &&
      typeof descriptors[key].value === 'function'
    ) {
      // We must assume that custom promisified function values are true aliases
      // of each of our hook targets. As such we hook these also.
      Object.defineProperty(hooked, key, {
        ...descriptors[key],
        value: patch(descriptors[key].value, options)
      });
    } else {
      Object.defineProperty(hooked, key, descriptors[key]);
    }
  }

  hookedFunctions.set(hooked, { fn });
  if (!fn.prototype) hooked.prototype = undefined;

  return hooked;
}

function noPropsPassed(props) {
  return props == null || (Array.isArray(props) && props.length === 0);
}

function hookableMethods(obj, props, callback) {
  if (typeof props === 'string') {
    // A single prop string was passed. Create an array.
    props = [props];
  } else if (noPropsPassed(props)) {
    if (typeof obj === 'function') {
      // No props were passed because the module export is a single function.
      // This needs to be special cased.
      return [''];
    }
    props = Object.keys(obj);
  }

  // Execute callback for every method in the props array.
  let i = props.length;
  while (i--) {
    // If the property isn't a function...
    if (
      !isFunction(obj[props[i]]) ||
      // ...or the callback returns false...
      // (if used)
      (callback && callback(obj, props[i]) === false)
    ) {
      // ...remove it from the props array to be returned.
      props.splice(i, 1);
    }
  }

  // Return an array of method names for which the callback didn't fail.
  return props;
}

/**
 * Returns true if a property is unwritable.
 *
 * @param {Object} obj
 * @param {string} prop
 * @param {Object} [desc] Property descriptor of obj[prop]
 * @returns {boolean}
 */
function unwritable(obj, prop, desc) {
  if (!obj) {
    return true;
  }

  if (!desc) {
    desc = getDescriptor(obj, prop);
  }

  return desc && !(desc.writable || desc.set);
}

/**
 * Attempts to find a descriptor for a property on an object or its parents
 *
 * @param {Object} obj
 * @param {string} prop
 * @return {Object|undefined} the descriptor, or undefined if it couldn't get one
 */
function getDescriptor(obj, prop) {
  if (!obj) {
    return;
  }

  // if it isn't its own property, but a parent's, return writability for the parent
  if (obj && obj[prop] && !Object.getOwnPropertyDescriptor(obj, prop)) {
    const proto = Object.getPrototypeOf(obj);
    return getDescriptor(proto, prop);
  }

  return Object.getOwnPropertyDescriptor(obj, prop);
}

/**
 * @param {Object} obj
 * @param {string} prop
 * @param {Object} opts options for hooking
 * @returns {Function} The hooked function
 */
function hook(obj, prop, opts) {
  // indicates object is a function and not an object with methods
  if (prop === '') {
    return hookFunction(obj, opts);
  }

  const desc = getDescriptor(obj, prop);

  if (!unwritable(obj, prop, desc)) {
    obj[prop] = hookFunction(obj[prop], opts);
    return obj[prop];
  }

  try {
    logger.debug(
      `${opts.name}:${prop} may be unwritable. hooking by defineProperty`
    );
    if (desc.configurable) {
      if (desc.get) {
        const f = hookFunction(obj[prop], opts);
        Object.defineProperty(obj, prop, {
          get() {
            return f;
          },
          set: desc.set
        });
      } else if (desc.value) {
        Object.defineProperty(obj, prop, {
          enumerable: desc.enumerable,
          writable: desc.writable,
          configurable: desc.configurable,
          value: hookFunction(obj[prop], opts)
        });
      } else {
        throw new Error('unknown object descriptor setup');
      }
    } else {
      logger.info(
        `unable to patch unconfigurable property ${opts.name}:${prop}`
      );
    }
  } catch (err) {
    logger.info(
      `unable to patch unconfigurable property ${opts.name}:${prop}`,
      err
    );
  }

  return obj[prop];
}

/**
 * Notes:
 * - Use `alwaysRun` to run pre/post hooks independent of whether the hooked
 *   function is being called in a scope prohibiting instrumentation.
 * - Use the `ALLOW_ALL` scope to run the original function in a scope that will
 *   allow instrumentation regardless of current scope.
 * - In pre hooks, if the intention of the hook is to bind a callback(s) to the
 *   Contrast CLS, please make use of `Scopes.runInAllowAllScope()` in the bound
 *   function. Otherwise the callback will be bound to the pre/post
 *   NO_INSTRUMENTATION scope, which is likely not the intention.
 * - Return type can be used to return a patched function from a require hook
 *   on a single function export.
 */
function patch(obj, props, options) {
  // If the props argument was omitted, shuffle the arguments.
  if (options == null) {
    options = props;
    props = null;
  }

  if (!options.patchType || !options.name) {
    logger.error('patchType and name are required options');
    return;
  }

  options.funcKey = `${options.patchType}:${options.name}`;

  if (obj === undefined) {
    logger.error(
      `attempt to hook property ${props} in nonexistent object ${options.name}`
    );
    return null;
  }

  const hookedMethods = hookableMethods(obj, props).map(function(prop) {
    const opts = typeof options === 'object' ? { ...options } : options;

    // staticName means do not append methods to final name
    // only used for Function(aka global.ContrastFunction)
    if (prop.length > 0 && !opts.staticName) {
      // There is a function name
      opts.funcKey += `.${prop}`;
    } else {
      // Modules with a single function export
      // e.g: module.exports = function() {...}
    }

    const fn = hook(obj, prop, opts);

    const fnHooks = hooks.get(fn) || new Map();
    if (fnHooks.get(opts.funcKey)) {
      logger.info(`${opts.funcKey} already patched ignoring`);
    } else {
      fnHooks.set(opts.funcKey, { pre: opts.pre, post: opts.post });
    }
    hooks.set(fn, fnHooks);

    return fn;
  });

  /*
   * special case - if a function was being directly hooked, not a property
   * on an object (i.e. module export set to a function) return the patched
   * function
   */
  if (typeof obj === 'function' && !props) {
    return hookedMethods[0];
  }
  return obj;
}

/**
 * Remove a handler for an event given an id (rule name, key, etc) and an
 * event name
 *
 * @param {string} id
 * @param {string|Array} events
 * @param {string} [emitter='hook']
 */
function unpatch(fn, funcName, patchType) {
  const fnHooks = hooks.get(fn);
  if (!fnHooks) {
    return;
  }

  fnHooks.delete(`${patchType}:${funcName}`);
}

/**
 * WeakMap lookup to get original, unwrapped version of a function.
 */
function unwrap(fn) {
  return hookedFunctions.has(fn) ? hookedFunctions.get(fn).fn : fn;
}

module.exports = {
  patch,
  preempt,
  resetInstrumentation,
  unpatch,
  unwrap
};
