fake-agent.js

'use strict'

const { createHash } = require('crypto')
const { EventEmitter, once } = require('events')
const http = require('http')
const express = require('express')
const bodyParser = require('body-parser')
const msgpack = require('@msgpack/msgpack')
const upload = require('multer')()

const noop = () => {}

/**
 * @typedef {object} RemoteConfigFile
 * @property {number} orgId
 * @property {string} product
 * @property {string} id
 * @property {string} name
 * @property {string} config
 * @property {string} path
 * @property {string} fileHash
 * @property {object} meta
 * @property {object} meta.custom
 * @property {number} meta.custom.v
 * @property {object} meta.hashes
 * @property {string} meta.hashes.sha256
 * @property {number} meta.length
 */

module.exports = class FakeAgent extends EventEmitter {
  port = 0
  advertiseDebuggerV2IntakeSupport = true
  debuggerV2IntakeStatusCode = 202
  /** @type {Set<import('net').Socket>} */
  #sockets = new Set()
  /** @type {Record<string, RemoteConfigFile>} */
  _rcFiles = {}
  _rcTargetsVersion = 0
  /** @type {Set<string>} */
  _rcSeenStates = new Set()

  constructor (port = 0, options = {}) {
    // Redirect rejections to the error event
    super({ captureRejections: true })
    this.port = port

    if (options.advertiseDebuggerV2IntakeSupport !== undefined) {
      this.advertiseDebuggerV2IntakeSupport = options.advertiseDebuggerV2IntakeSupport
    }
    if (options.debuggerV2IntakeStatusCode !== undefined) {
      this.debuggerV2IntakeStatusCode = options.debuggerV2IntakeStatusCode
    }
  }

  /**
   * Start the fake agent.
   * @returns {Promise<FakeAgent>} A promise that resolves when the agent has started up.
   */
  start () {
    return new Promise((resolve, reject) => {
      const timeoutObj = setTimeout(() => {
        reject(new Error('Agent timed out starting up'))
      }, 10_000)
      this.server = http.createServer(buildExpressServer(this))
      this.server.on('error', reject)

      // Track connections to force close them later
      this.server.on('connection', (socket) => {
        this.#sockets.add(socket)
        socket.on('close', () => {
          this.#sockets.delete(socket)
        })
      })

      this.server.listen(this.port, () => {
        this.port = (/** @type {import('net').AddressInfo} */ (
          (/** @type {import('http').Server} */ (this.server)).address()
        )).port
        clearTimeout(timeoutObj)
        resolve(this)
      })
    })
  }

  stop () {
    if (!this.server?.listening) return

    for (const socket of this.#sockets) {
      socket.destroy()
    }
    this.#sockets.clear()
    this.server.close()

    return once(this.server, 'close')
  }

  /**
   * Add a config object to be returned by the fake Remote Config endpoint.
   * @param {object} config - Object containing the Remote Config "file" and metadata
   * @param {number} [config.orgId] - The Datadog organization ID. Defaults to 2.
   * @param {string} config.product - The Remote Config product name
   * @param {string} config.id - The Remote Config config ID
   * @param {string} [config.name] - The Remote Config "name". Defaults to the sha256 hash of `config.id`
   * @param {object} config.config - The Remote Config "file" object
   */
  addRemoteConfig (config) {
    const orgId = config.orgId || 2
    const name = config.name || createHash('sha256').update(config.id).digest('hex')
    const configStr = JSON.stringify(config.config)
    const path = `datadog/${orgId}/${config.product}/${config.id}/${name}`
    const fileHash = createHash('sha256').update(configStr).digest('hex')
    const meta = {
      custom: { v: 1 },
      hashes: { sha256: fileHash },
      length: configStr.length,
    }

    this._rcFiles[config.id] = {
      orgId,
      product: config.product,
      id: config.id,
      name,
      config: configStr,
      path,
      fileHash,
      meta,
    }
    this._rcTargetsVersion++
  }

  /**
   * Update an existing config object
   * @param {string} id - The Remote Config config ID
   * @param {object} config - The Remote Config "file" object
   */
  updateRemoteConfig (id, config) {
    config = JSON.stringify(config)
    config = Object.assign(
      this._rcFiles[id],
      {
        config,
        fileHash: createHash('sha256').update(config).digest('hex'),
      }
    )
    config.meta.custom.v++
    config.meta.hashes.sha256 = config.fileHash
    config.meta.length = config.config.length
    this._rcTargetsVersion++
  }

  /**
   * Remove a specific config object
   * @param {string} id - The ID of the config object that should be removed
   */
  removeRemoteConfig (id) {
    delete this._rcFiles[id]
    this._rcTargetsVersion++
  }

  /**
   * Reset any existing Remote Config state. Usefull in `before` and `beforeEach` blocks.
   */
  resetRemoteConfig () {
    this._rcFiles = {}
    this._rcTargetsVersion = 0
    this._rcSeenStates = new Set()
  }

  // **resolveAtFirstSuccess** - specific use case for Next.js (or any other future libraries)
  // where multiple payloads are generated, and only one is expected to have the proper span (ie next.request),
  // but it't not guaranteed to be the last one (so, expectedMessageCount would not be helpful).
  // It can still fail if it takes longer than `timeout` duration or if none pass the assertions (timeout still called)
  assertMessageReceived (fn, timeout, expectedMessageCount = 1, resolveAtFirstSuccess = true) {
    timeout = timeout || 30000
    let resultResolve
    let resultReject
    let msgCount = 0
    const errors = []

    const timeoutObj = setTimeout(() => {
      this.removeListener('message', messageHandler)
      const errorsMsg = errors.length === 0 ? '' : `, additionally:\n${errors.map(e => e.stack).join('\n')}\n===\n`
      resultReject(new Error(`timeout${errorsMsg}`, { cause: { errors } }))
    }, timeout)

    const resultPromise = /** @type {Promise<void>} */ (new Promise((resolve, reject) => {
      resultResolve = () => {
        clearTimeout(timeoutObj)
        resolve()
      }
      resultReject = (e) => {
        clearTimeout(timeoutObj)
        reject(e)
      }
    }))

    const messageHandler = msg => {
      try {
        msgCount += 1
        fn(msg)
        if (resolveAtFirstSuccess || msgCount === expectedMessageCount) {
          resultResolve()
          this.removeListener('message', messageHandler)
        }
      } catch (e) {
        errors.push(e)
      }
    }
    this.on('message', messageHandler)

    return resultPromise
  }

  /**
   * Assert that a telemetry message is received.
   *
   * @param {object} options
   * @param {Function} [options.fn] - Function called with each matching telemetry message. If it throws,
   *     the error is collected and the listener stays alive for the next message. Defaults to a no-op.
   * @param {string} options.requestType - The telemetry request type to match.
   * @param {number} [options.timeout] - Timeout in milliseconds before the promise rejects.
   * @param {number} [options.expectedMessageCount] - Number of matching messages to wait for.
   * @param {boolean} [options.resolveAtFirstSuccess] - Resolve as soon as `fn` first runs without throwing.
   * @param {string} [options.namespace] - If set, only consider messages whose payload namespace equals this value.
   * @returns {Promise<void>} A promise that resolves when the expected telemetry messages are received and `fn` has
   *     run successfully. If `fn` throws on every matching message, the promise rejects once `timeout` is reached.
   */
  assertTelemetryReceived ({
    fn = noop,
    requestType,
    timeout = 30_000,
    expectedMessageCount = 1,
    resolveAtFirstSuccess = false,
    namespace,
  }) {
    let resultResolve
    let resultReject
    let msgCount = 0
    const errors = []

    const timeoutObj = setTimeout(() => {
      const errorsMsg = errors.length === 0 ? '' : `, additionally:\n${errors.map(e => e.stack).join('\n')}\n===\n`
      resultReject(new Error(`timeout${errorsMsg}`, { cause: { errors } }))
    }, timeout)

    const resultPromise = /** @type {Promise<void>} */ (new Promise((resolve, reject) => {
      resultResolve = () => {
        clearTimeout(timeoutObj)
        resolve()
      }
      resultReject = (e) => {
        clearTimeout(timeoutObj)
        reject(e)
      }
    }))

    const messageHandler = msg => {
      if (msg.payload.request_type !== requestType) return
      if (namespace !== undefined && msg.payload.payload?.namespace !== namespace) return
      msgCount += 1
      try {
        fn(msg)
        if (resolveAtFirstSuccess || msgCount === expectedMessageCount) {
          resultResolve()
        }

        if (resolveAtFirstSuccess) {
          this.removeListener('telemetry', messageHandler)
        }
      } catch (e) {
        errors.push(e)
      }
      if (!resolveAtFirstSuccess && msgCount === expectedMessageCount) {
        this.removeListener('telemetry', messageHandler)
      }
    }
    this.on('telemetry', messageHandler)

    return resultPromise
  }

  /**
   * Collect span groups matching a predicate after firing a trigger.
   *
   * Attaches the listener before invoking `trigger`, so events emitted between request
   * dispatch and listener registration are not missed.
   *
   * @param {object} options
   * @param {() => Promise<unknown>} options.trigger Fired once the listener is in place.
   * @param {(group: object[]) => boolean} options.predicate Group-level filter.
   * @param {number} [options.expectedCount] Resolve after this many matching groups arrive.
   * @param {number} [options.timeout] Timeout in milliseconds before the promise rejects.
   * @returns {Promise<object[][]>} The matching groups, in arrival order.
   */
  collectGroups ({ trigger, predicate, expectedCount = 1, timeout = 30_000 }) {
    const groups = []
    let resolveResult
    let rejectResult

    const handler = ({ payload }) => {
      for (const group of payload) {
        if (predicate(group)) groups.push(group)
      }
      if (groups.length >= expectedCount) resolveResult(groups)
    }

    const timeoutObj = setTimeout(() => {
      rejectResult(new Error(
        `timed out waiting for ${expectedCount} matching span groups, got ${groups.length}`
      ))
    }, timeout)

    const result = /** @type {Promise<object[][]>} */ (new Promise((resolve, reject) => {
      resolveResult = (value) => {
        clearTimeout(timeoutObj)
        this.removeListener('message', handler)
        resolve(value)
      }
      rejectResult = (error) => {
        clearTimeout(timeoutObj)
        this.removeListener('message', handler)
        reject(error)
      }
    }))

    this.on('message', handler)
    trigger().catch(rejectResult)
    return result
  }

  assertLlmObsPayloadReceived (fn, timeout, expectedMessageCount = 1, resolveAtFirstSuccess) {
    timeout = timeout || 30000
    let resultResolve
    let resultReject
    let msgCount = 0
    const errors = []

    const timeoutObj = setTimeout(() => {
      const errorsMsg = errors.length === 0 ? '' : `, additionally:\n${errors.map(e => e.stack).join('\n')}\n===\n`
      resultReject(new Error(`timeout${errorsMsg}`, { cause: { errors } }))
    }, timeout)

    const resultPromise = /** @type {Promise<void>} */ (new Promise((resolve, reject) => {
      resultResolve = () => {
        clearTimeout(timeoutObj)
        resolve()
      }
      resultReject = (e) => {
        clearTimeout(timeoutObj)
        reject(e)
      }
    }))

    const messageHandler = msg => {
      try {
        msgCount += 1
        fn(msg)
        if (resolveAtFirstSuccess || msgCount === expectedMessageCount) {
          resultResolve()
          this.removeListener('llmobs', messageHandler)
        }
      } catch (e) {
        errors.push(e)
      }
    }
    this.on('llmobs', messageHandler)

    return resultPromise
  }
}

function buildExpressServer (agent) {
  const app = express()

  app.use(bodyParser.raw({ limit: Infinity, type: 'application/msgpack' }))
  app.use(bodyParser.json({ limit: Infinity, type: 'application/json' }))

  app.get('/info', (req, res) => {
    const endpoints = ['/evp_proxy/v2', '/debugger/v1/input']
    if (agent.advertiseDebuggerV2IntakeSupport) {
      endpoints.push('/debugger/v2/input')
    }
    res.json({ endpoints })
  })

  app.put('/v0.4/traces', (req, res) => {
    if (req.body.length === 0) return res.status(200).send()
    res.status(200).send({ rate_by_service: { 'service:,env:': 1 } })
    agent.emit('message', {
      headers: req.headers,
      payload: msgpack.decode(req.body, { useBigInt64: true }),
    })
  })

  app.post('/v0.7/config', (req, res) => {
    const {
      client: { products, state },
      cached_target_files: cachedTargetFiles,
    } = req.body

    // Emit the remote config request payload for testing
    agent.emit('remote-config-request', req.body)

    if (state.has_error) {
      // Print the error sent by the client in case it's useful in debugging tests
      console.error(state.error) // eslint-disable-line no-console
    }

    for (const cs of state.config_states) {
      const uniqueState = `${cs.id}-${cs.version}-${cs.apply_state}`
      if (!agent._rcSeenStates.has(uniqueState)) {
        agent._rcSeenStates.add(uniqueState)
        agent.emit('remote-config-ack-update', cs.id, cs.version, cs.apply_state, cs.apply_error)
      }

      if (cs.apply_error) {
        // Print the error sent by the client in case it's useful in debugging tests
        console.error(cs.apply_error) // eslint-disable-line no-console
      }
    }

    res.on('close', () => {
      agent.emit('remote-config-responded')
    })

    if (agent._rcTargetsVersion === state.targets_version) {
      // If the state hasn't changed since the last time the client asked, just return an empty result
      res.json({})
      return
    }

    if (Object.keys(agent._rcFiles).length === 0) {
      // All config files have been removed, but the client has not yet been informed.
      // Return this custom result to let the client know.
      res.json({ client_configs: [] })
      return
    }

    // The actual targets object is much more complicated,
    // but the Node.js tracer currently only cares about the following properties.
    const targets = {
      signed: {
        custom: { opaque_backend_state: 'foo' },
        targets: {},
        version: agent._rcTargetsVersion,
      },
    }
    const targetFiles = []
    const clientConfigs = []

    const files = Object.values(agent._rcFiles).filter(({ product }) => products.includes(product))

    for (const { path, fileHash, meta, config } of files) {
      clientConfigs.push(path)
      targets.signed.targets[path] = meta

      // skip files already cached by the client so we don't send them more than once
      if (cachedTargetFiles.some((cached) =>
        path === cached.path &&
        fileHash === cached.hashes.find((e) => e.algorithm === 'sha256').hash
      )) continue

      targetFiles.push({ path, raw: base64(config) })
    }

    // The real response object also contains a `roots` property which has been omitted here since it's not currently
    // used by the Node.js tracer.
    res.json({
      targets: clientConfigs.length === 0 ? undefined : base64(targets),
      target_files: targetFiles,
      client_configs: clientConfigs,
    })
  })

  app.post('/debugger/v1/input', (req, res) => {
    res.status(202).send()
    agent.emit('debugger-input', {
      headers: req.headers,
      query: req.query,
      payload: req.body,
    })
    agent.emit('debugger-input-v1', {
      headers: req.headers,
      query: req.query,
      payload: req.body,
    })
  })

  app.post('/debugger/v2/input', (req, res) => {
    res.status(agent.debuggerV2IntakeStatusCode).send()
    if (agent.debuggerV2IntakeStatusCode === 404) {
      agent.emit('debugger-input-v2-404')
      return
    }
    agent.emit('debugger-input', {
      headers: req.headers,
      query: req.query,
      payload: req.body,
    })
    agent.emit('debugger-input-v2', {
      headers: req.headers,
      query: req.query,
      payload: req.body,
    })
  })

  app.post('/debugger/v1/diagnostics', upload.any(), (req, res) => {
    res.status(200).send() // TODO: Should we send a 202 here instead?
    // The diagnostics endpoint can receive both probe results and status messages
    // Emit the appropriate events based on payload structure
    const isProbeResult = req.body[0]?.debugger?.snapshot !== undefined
    if (isProbeResult) {
      const event = {
        headers: req.headers,
        payload: req.body,
      }
      agent.emit('debugger-input', event)
      agent.emit('debugger-diagnostics-input', event)
    } else {
      agent.emit('debugger-diagnostics', {
        headers: req.headers,
        payload: JSON.parse((/** @type {Array<{ buffer: Buffer }>} */ (req.files))[0].buffer.toString()),
      })
    }
  })

  app.post('/profiling/v1/input', upload.any(), (req, res) => {
    res.status(200).send()
    agent.emit('message', {
      headers: req.headers,
      payload: req.body,
      files: req.files,
    })
  })

  app.post('/telemetry/proxy/api/v2/apmtelemetry', (req, res) => {
    res.status(200).send()
    agent.emit('telemetry', {
      headers: req.headers,
      payload: req.body,
    })
  })

  app.post('/evp_proxy/v2/api/v2/llmobs', (req, res) => {
    res.status(200).send()
    agent.emit('llmobs', {
      headers: req.headers,
      payload: req.body,
    })
  })

  app.post('/v1/traces', (req, res) => {
    res.status(200).send()
    agent.emit('otlp-traces', {
      headers: req.headers,
      payload: req.body,
    })
  })

  app.post('/v1/logs', (req, res) => {
    res.status(200).send()
    agent.emit('otlp-logs', {
      headers: req.headers,
      payload: req.body,
    })
  })

  app.post('/evp_proxy/v2/api/v2/exposures', (req, res) => {
    res.status(200).send()
    agent.emit('exposures', {
      headers: req.headers,
      payload: req.body,
    })
  })

  // Ensure that any failure inside of Express isn't swallowed and returned as a 500, but instead crashes the test
  app.use((err, req, res, next) => {
    if (!err) next()
    process.nextTick(() => {
      throw err
    })
  })

  return app
}

function base64 (strOrObj) {
  const str = typeof strOrObj === 'string' ? strOrObj : JSON.stringify(strOrObj)
  return Buffer.from(str).toString('base64')
}