Cache iast metrics (#4017)

* Remove IastNamespace * cache IAST metrics * Remove IastMetric WeakMap * Remove not used method * Clean up * use const tags * Refactor * Clean up * Include more tests * One more test * test * Get count only once per metric * Ensure int value * merge more than count metrics * Iast metrics are count metrics only by now. Use same naming for IastMetric methods as for CountMetric * Use track when merging metrics * Remove IastMetric.track and use only IastMetric.inc * Use undefined * Ugaitz suggestion Co-authored-by: Ugaitz Urien <ugaitz.urien@datadoghq.com> * Remove not used metric --------- Co-authored-by: Ugaitz Urien <ugaitz.urien@datadoghq.com>
DataDog · Mar 20, 2024 · 5b7b502 · 5b7b502
1 parent 7d3606c
commit 5b7b502
Show file tree

Hide file tree

Showing 14 changed files with 306 additions and 118 deletions.
diff --git a/packages/dd-trace/src/appsec/iast/iast-plugin.js b/packages/dd-trace/src/appsec/iast/iast-plugin.js
@@ -5,7 +5,8 @@ const { channel } = require('dc-polyfill')
 const iastLog = require('./iast-log')
 const Plugin = require('../../plugins/plugin')
 const iastTelemetry = require('./telemetry')
-const { getInstrumentedMetric, getExecutedMetric, TagKey, EXECUTED_SOURCE } = require('./telemetry/iast-metric')
+const { getInstrumentedMetric, getExecutedMetric, TagKey, EXECUTED_SOURCE, formatTags } =
+  require('./telemetry/iast-metric')
 const { storage } = require('../../../../datadog-core')
 const { getIastContext } = require('./iast-context')
 const instrumentations = require('../../../../datadog-instrumentations/src/helpers/instrumentations')
@@ -20,25 +21,29 @@ const instrumentations = require('../../../../datadog-instrumentations/src/helpe
  * - tagKey can be only SOURCE_TYPE (Source) or VULNERABILITY_TYPE (Sink)
  */
 class IastPluginSubscription {
-  constructor (moduleName, channelName, tag, tagKey = TagKey.VULNERABILITY_TYPE) {
+  constructor (moduleName, channelName, tagValues, tagKey = TagKey.VULNERABILITY_TYPE) {
     this.moduleName = moduleName
     this.channelName = channelName
-    this.tag = tag
-    this.tagKey = tagKey
-    this.executedMetric = getExecutedMetric(this.tagKey)
-    this.instrumentedMetric = getInstrumentedMetric(this.tagKey)
+
+    tagValues = Array.isArray(tagValues) ? tagValues : [tagValues]
+    this.tags = formatTags(tagValues, tagKey)
+
+    this.executedMetric = getExecutedMetric(tagKey)
+    this.instrumentedMetric = getInstrumentedMetric(tagKey)
+
     this.moduleInstrumented = false
   }
 
   increaseInstrumented () {
-    if (this.moduleInstrumented) return
+    if (!this.moduleInstrumented) {
+      this.moduleInstrumented = true
 
-    this.moduleInstrumented = true
-    this.instrumentedMetric.inc(this.tag)
+      this.tags.forEach(tag => this.instrumentedMetric.inc(undefined, tag))
+    }
   }
 
   increaseExecuted (iastContext) {
-    this.executedMetric.inc(this.tag, iastContext)
+    this.tags.forEach(tag => this.executedMetric.inc(iastContext, tag))
   }
 
   matchesModuleInstrumented (name) {
@@ -76,10 +81,16 @@ class IastPlugin extends Plugin {
     }
   }
 
-  _execHandlerAndIncMetric ({ handler, metric, tag, iastContext = getIastContext(storage.getStore()) }) {
+  _execHandlerAndIncMetric ({ handler, metric, tags, iastContext = getIastContext(storage.getStore()) }) {
     try {
       const result = handler()
-      iastTelemetry.isEnabled() && metric.inc(tag, iastContext)
+      if (iastTelemetry.isEnabled()) {
+        if (Array.isArray(tags)) {
+          tags.forEach(tag => metric.inc(iastContext, tag))
+        } else {
+          metric.inc(iastContext, tags)
+        }
+      }
       return result
     } catch (e) {
       iastLog.errorAndPublish(e)

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/operations.js b/packages/dd-trace/src/appsec/iast/taint-tracking/operations.js
@@ -19,7 +19,7 @@ let onRemoveTransaction = (transactionId, iastContext) => {}
 function onRemoveTransactionInformationTelemetry (transactionId, iastContext) {
   const metrics = TaintedUtils.getMetrics(transactionId, iastTelemetry.verbosity)
   if (metrics?.requestCount) {
-    REQUEST_TAINTED.add(metrics.requestCount, null, iastContext)
+    REQUEST_TAINTED.inc(iastContext, metrics.requestCount)
   }
 }
 

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js b/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js
@@ -14,6 +14,10 @@ const {
   HTTP_REQUEST_PATH_PARAM,
   HTTP_REQUEST_URI
 } = require('./source-types')
+const { EXECUTED_SOURCE } = require('../telemetry/iast-metric')
+
+const REQ_HEADER_TAGS = EXECUTED_SOURCE.formatTags(HTTP_REQUEST_HEADER_VALUE, HTTP_REQUEST_HEADER_NAME)
+const REQ_URI_TAGS = EXECUTED_SOURCE.formatTags(HTTP_REQUEST_URI)
 
 class TaintTrackingPlugin extends SourceIastPlugin {
   constructor () {
@@ -97,7 +101,7 @@ class TaintTrackingPlugin extends SourceIastPlugin {
   taintHeaders (headers, iastContext) {
     this.execSource({
       handler: () => taintObject(iastContext, headers, HTTP_REQUEST_HEADER_VALUE, true, HTTP_REQUEST_HEADER_NAME),
-      tag: [HTTP_REQUEST_HEADER_VALUE, HTTP_REQUEST_HEADER_NAME],
+      tags: REQ_HEADER_TAGS,
       iastContext
     })
   }
@@ -107,7 +111,7 @@ class TaintTrackingPlugin extends SourceIastPlugin {
       handler: function () {
         req.url = newTaintedString(iastContext, req.url, HTTP_REQUEST_URI, HTTP_REQUEST_URI)
       },
-      tag: [HTTP_REQUEST_URI],
+      tags: REQ_URI_TAGS,
       iastContext
     })
   }

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter-telemetry.js b/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter-telemetry.js
@@ -14,7 +14,7 @@ const telemetryRewriter = {
 
     const metrics = response.metrics
     if (metrics && metrics.instrumentedPropagation) {
-      INSTRUMENTED_PROPAGATION.add(metrics.instrumentedPropagation)
+      INSTRUMENTED_PROPAGATION.inc(undefined, metrics.instrumentedPropagation)
     }
 
     return response

diff --git a/packages/dd-trace/src/appsec/iast/taint-tracking/taint-tracking-impl.js b/packages/dd-trace/src/appsec/iast/taint-tracking/taint-tracking-impl.js
@@ -40,7 +40,7 @@ function getContextDefault () {
 
 function getContextDebug () {
   const iastContext = getContextDefault()
-  EXECUTED_PROPAGATION.inc(null, iastContext)
+  EXECUTED_PROPAGATION.inc(iastContext)
   return iastContext
 }
 

diff --git a/packages/dd-trace/src/appsec/iast/telemetry/iast-metric.js b/packages/dd-trace/src/appsec/iast/telemetry/iast-metric.js
@@ -19,37 +19,46 @@ const TagKey = {
   PROPAGATION_TYPE: 'propagation_type'
 }
 
+function formatTags (tags, tagKey) {
+  return tags.map(tagValue => tagValue ? [`${tagKey}:${tagValue.toLowerCase()}`] : undefined)
+}
+
+function getNamespace (scope, context) {
+  let namespace = globalNamespace
+
+  if (scope === Scope.REQUEST) {
+    namespace = getNamespaceFromContext(context) || globalNamespace
+  }
+  return namespace
+}
+
 class IastMetric {
   constructor (name, scope, tagKey) {
     this.name = name
     this.scope = scope
     this.tagKey = tagKey
   }
 
-  getNamespace (context) {
-    return getNamespaceFromContext(context) || globalNamespace
+  formatTags (...tags) {
+    return formatTags(tags, this.tagKey)
   }
 
-  getTag (tagValue) {
-    return tagValue ? { [this.tagKey]: tagValue } : undefined
+  inc (context, tags, value = 1) {
+    const namespace = getNamespace(this.scope, context)
+    namespace.count(this.name, tags).inc(value)
   }
+}
 
-  addValue (value, tagValue, context) {
-    this.getNamespace(context)
-      .count(this.name, this.getTag(tagValue))
-      .inc(value)
-  }
+class NoTaggedIastMetric extends IastMetric {
+  constructor (name, scope) {
+    super(name, scope)
 
-  add (value, tagValue, context) {
-    if (Array.isArray(tagValue)) {
-      tagValue.forEach(tag => this.addValue(value, tag, context))
-    } else {
-      this.addValue(value, tagValue, context)
-    }
+    this.tags = []
   }
 
-  inc (tagValue, context) {
-    this.add(1, tagValue, context)
+  inc (context, value = 1) {
+    const namespace = getNamespace(this.scope, context)
+    namespace.count(this.name, this.tags).inc(value)
   }
 }
 
@@ -61,21 +70,18 @@ function getInstrumentedMetric (tagKey) {
   return tagKey === TagKey.VULNERABILITY_TYPE ? INSTRUMENTED_SINK : INSTRUMENTED_SOURCE
 }
 
-const INSTRUMENTED_PROPAGATION = new IastMetric('instrumented.propagation', Scope.GLOBAL)
+const INSTRUMENTED_PROPAGATION = new NoTaggedIastMetric('instrumented.propagation', Scope.GLOBAL)
 const INSTRUMENTED_SOURCE = new IastMetric('instrumented.source', Scope.GLOBAL, TagKey.SOURCE_TYPE)
 const INSTRUMENTED_SINK = new IastMetric('instrumented.sink', Scope.GLOBAL, TagKey.VULNERABILITY_TYPE)
 
 const EXECUTED_SOURCE = new IastMetric('executed.source', Scope.REQUEST, TagKey.SOURCE_TYPE)
 const EXECUTED_SINK = new IastMetric('executed.sink', Scope.REQUEST, TagKey.VULNERABILITY_TYPE)
 
-const REQUEST_TAINTED = new IastMetric('request.tainted', Scope.REQUEST)
+const REQUEST_TAINTED = new NoTaggedIastMetric('request.tainted', Scope.REQUEST)
 
 // DEBUG using metrics
-const EXECUTED_PROPAGATION = new IastMetric('executed.propagation', Scope.REQUEST)
-const EXECUTED_TAINTED = new IastMetric('executed.tainted', Scope.REQUEST)
-
-// DEBUG using distribution endpoint
-const INSTRUMENTATION_TIME = new IastMetric('instrumentation.time', Scope.GLOBAL)
+const EXECUTED_PROPAGATION = new NoTaggedIastMetric('executed.propagation', Scope.REQUEST)
+const EXECUTED_TAINTED = new NoTaggedIastMetric('executed.tainted', Scope.REQUEST)
 
 module.exports = {
   INSTRUMENTED_PROPAGATION,
@@ -89,13 +95,14 @@ module.exports = {
 
   REQUEST_TAINTED,
 
-  INSTRUMENTATION_TIME,
-
   PropagationType,
   TagKey,
 
   IastMetric,
+  NoTaggedIastMetric,
 
   getExecutedMetric,
-  getInstrumentedMetric
+  getInstrumentedMetric,
+
+  formatTags
 }
diff --git a/packages/dd-trace/src/appsec/iast/telemetry/namespaces.js b/packages/dd-trace/src/appsec/iast/telemetry/namespaces.js
@@ -10,13 +10,13 @@ const DD_IAST_METRICS_NAMESPACE = Symbol('_dd.iast.request.metrics.namespace')
 function initRequestNamespace (context) {
   if (!context) return
 
-  const namespace = new Namespace('iast')
+  const namespace = new IastNamespace()
   context[DD_IAST_METRICS_NAMESPACE] = namespace
   return namespace
 }
 
 function getNamespaceFromContext (context) {
-  return context && context[DD_IAST_METRICS_NAMESPACE]
+  return context?.[DD_IAST_METRICS_NAMESPACE]
 }
 
 function finalizeRequestNamespace (context, rootSpan) {
@@ -40,11 +40,14 @@ function finalizeRequestNamespace (context, rootSpan) {
 }
 
 function merge (metrics) {
-  metrics.forEach(metric => metric.points.forEach(point => {
-    globalNamespace
-      .count(metric.metric, getTagsObject(metric.tags))
-      .inc(point[1])
-  }))
+  metrics.forEach(metric => {
+    const { metric: metricName, type, tags, points } = metric
+
+    if (points?.length && type === 'count') {
+      const gMetric = globalNamespace.count(metricName, getTagsObject(tags))
+      points.forEach(point => gMetric.inc(point[1]))
+    }
+  })
 }
 
 function getTagsObject (tags) {
@@ -56,11 +59,34 @@ function getTagsObject (tags) {
 class IastNamespace extends Namespace {
   constructor () {
     super('iast')
+
+    this.iastMetrics = new Map()
   }
 
-  reset () {
-    this.metrics.clear()
-    this.distributions.clear()
+  getIastMetrics (name) {
+    let metrics = this.iastMetrics.get(name)
+    if (!metrics) {
+      metrics = new Map()
+      this.iastMetrics.set(name, metrics)
+    }
+
+    return metrics
+  }
+
+  getMetric (name, tags, type = 'count') {
+    const metrics = this.getIastMetrics(name)
+
+    let metric = metrics.get(tags)
+    if (!metric) {
+      metric = super[type](name, Array.isArray(tags) ? [...tags] : tags)
+      metrics.set(tags, metric)
+    }
+
+    return metric
+  }
+
+  count (name, tags) {
+    return this.getMetric(name, tags, 'count')
   }
 }
 
@@ -72,5 +98,7 @@ module.exports = {
   finalizeRequestNamespace,
   globalNamespace,
 
-  DD_IAST_METRICS_NAMESPACE
+  DD_IAST_METRICS_NAMESPACE,
+
+  IastNamespace
 }
diff --git a/packages/dd-trace/src/appsec/iast/telemetry/span-tags.js b/packages/dd-trace/src/appsec/iast/telemetry/span-tags.js
@@ -1,11 +1,11 @@
 'use strict'
 
 function addMetricsToSpan (rootSpan, metrics, tagPrefix) {
-  if (!rootSpan || !rootSpan.addTags || !metrics) return
+  if (!rootSpan?.addTags || !metrics) return
 
   const flattenMap = new Map()
   metrics
-    .filter(data => data && data.metric)
+    .filter(data => data?.metric)
     .forEach(data => {
       const name = taggedMetricName(data)
       let total = flattenMap.get(name)
@@ -27,19 +27,20 @@ function addMetricsToSpan (rootSpan, metrics, tagPrefix) {
 }
 
 function flatten (metricData) {
-  return metricData.points && metricData.points.map(point => point[1]).reduce((total, value) => total + value, 0)
+  const { points } = metricData
+  return points ? points.map(point => point[1]).reduce((total, value) => total + value, 0) : 0
 }
 
 function taggedMetricName (data) {
   const metric = data.metric
-  const tags = data.tags && filterTags(data.tags)
-  return !tags || !tags.length
+  const tags = filterTags(data.tags)
+  return !tags?.length
     ? metric
     : `${metric}.${processTagValue(tags)}`
 }
 
 function filterTags (tags) {
-  return tags.filter(tag => !tag.startsWith('lib_language') && !tag.startsWith('version'))
+  return tags?.filter(tag => !tag.startsWith('lib_language') && !tag.startsWith('version'))
 }
 
 function processTagValue (tags) {