dd-trace/packages/dd-trace/src/aiguard/index.js

Datadog APM tracing client for JavaScript

'use strict'

const log = require('../log')
const { incomingHttpRequestStart, aiguardChannel } = require('./channels')
const AIGuard = require('./sdk')
const { SOURCE_AUTO, INTEGRATION_NONE } = require('./tags')

let isEnabled = false
let aiguard
let block

function onIncomingHttpRequestStart () {
  // No-op: subscribing ensures the HTTP plugin spreads req onto the store
}

function enable (tracer, config) {
  if (isEnabled) return

  try {
    aiguard = new AIGuard(tracer, config)
    block = config.experimental?.aiguard?.block !== false

    incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
    aiguardChannel.subscribe(onEvaluate)

    isEnabled = true
  } catch (err) {
    log.error('AIGuard: unexpected error during initialization: %s', err.message)
    disable()
  }
}

function disable () {
  if (!isEnabled) return

  incomingHttpRequestStart.unsubscribe(onIncomingHttpRequestStart)
  aiguardChannel.unsubscribe(onEvaluate)

  aiguard = undefined
  isEnabled = false
  block = false
}

/**
 * Handles channel messages with pre-converted messages.
 *
 * @param {{messages: Array<object>, resolve: Function, reject: Function}} ctx
 */
function onEvaluate (ctx) {
  if (!ctx.messages?.length) {
    ctx.resolve()
    return
  }

  const opts = { block, source: SOURCE_AUTO, integration: ctx.integration || INTEGRATION_NONE }
  aiguard.evaluate(ctx.messages, opts)
    .then(() => {
      ctx.resolve()
    })
    .catch(err => {
      if (err.name === 'AIGuardAbortError') {
        ctx.reject(err)
      } else {
        log.error('AIGuard: unexpected error during evaluation: %s', err.message)
        ctx.resolve()
      }
    })
}

module.exports = { enable, disable }