add support to api security sampling (#4755)

* add support to api security sampling

* fix express plugin schema extraction

* use priority simpler to get span priority

* use lru cache package

* use route path instead of url

* use route.path or url to generate the key

* use ttlcache

* Fix standalone integration test

* Increase test timeout

* simplify force sample

* avoid checking is sampled twice

* use route.path or url to generate the key

* remove unnecessary tests of sample delay

* fix non experimental options test

* remove unused isSampled

* always sample request if delay is 0

---------

Co-authored-by: Igor Unanua <igor.unanua@datadoghq.com>
Co-authored-by: simon-id <simon.id@datadoghq.com>

LICENSE-3rdparty.csv

@@ -8,6 +8,7 @@ require,@datadog/pprof,Apache license 2.0,Copyright 2019 Google Inc.
 require,@datadog/sketches-js,Apache license 2.0,Copyright 2020 Datadog Inc.
 require,@opentelemetry/api,Apache license 2.0,Copyright OpenTelemetry Authors
 require,@opentelemetry/core,Apache license 2.0,Copyright OpenTelemetry Authors
+require,@isaacs/ttlcache,ISC,Copyright (c) 2022-2023 - Isaac Z. Schlueter and Contributors
 require,crypto-randomuuid,MIT,Copyright 2021 Node.js Foundation and contributors
 require,dc-polyfill,MIT,Copyright 2023 Datadog Inc.
 require,ignore,MIT,Copyright 2013 Kael Zhang and contributors

docs/test.ts

@@ -115,7 +115,6 @@ tracer.init({
     },
     apiSecurity: {
       enabled: true,
-      requestSampling: 1.0
     },
     rasp: {
       enabled: true

index.d.ts

@@ -662,19 +662,13 @@ declare namespace tracer {
         mode?: 'safe' | 'extended' | 'disabled'
       },
       /**
-       * Configuration for Api Security sampling
+       * Configuration for Api Security
        */
       apiSecurity?: {
         /** Whether to enable Api Security.
-         * @default false
+         * @default true
          */
         enabled?: boolean,
-
-        /** Controls the request sampling rate (between 0 and 1) in which Api Security is triggered.
-         * The value will be coerced back if it's outside of the 0-1 range.
-         * @default 0.1
-         */
-        requestSampling?: number
       },
       /**
        * Configuration for RASP

integration-tests/standalone-asm.spec.js

@@ -81,33 +81,42 @@ describe('Standalone ASM', () => {
       })
     })
 
-    it('should keep second req because RateLimiter allows 1 req/min and discard the next', async () => {
-      // 1st req kept because waf init
-      // 2nd req kept because it's the first one hitting RateLimiter
-      // next in the first minute are dropped
-      await doWarmupRequests(proc)
-
-      return curlAndAssertMessage(agent, proc, ({ headers, payload }) => {
+    it('should keep fifth req because RateLimiter allows 1 req/min', async () => {
+      const promise = curlAndAssertMessage(agent, proc, ({ headers, payload }) => {
         assert.propertyVal(headers, 'datadog-client-computed-stats', 'yes')
         assert.isArray(payload)
-        assert.strictEqual(payload.length, 4)
-
-        const secondReq = payload[1]
-        assert.isArray(secondReq)
-        assert.strictEqual(secondReq.length, 5)
+        if (payload.length === 4) {
+          assertKeep(payload[0][0])
+          assertDrop(payload[1][0])
+          assertDrop(payload[2][0])
+          assertDrop(payload[3][0])
+
+          // req after a minute
+        } else {
+          const fifthReq = payload[0]
+          assert.isArray(fifthReq)
+          assert.strictEqual(fifthReq.length, 5)
+
+          const { meta, metrics } = fifthReq[0]
+          assert.notProperty(meta, 'manual.keep')
+          assert.notProperty(meta, '_dd.p.appsec')
+
+          assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_KEEP)
+          assert.propertyVal(metrics, '_dd.apm.enabled', 0)
+        }
+      }, 70000, 2)
 
-        const { meta, metrics } = secondReq[0]
-        assert.notProperty(meta, 'manual.keep')
-        assert.notProperty(meta, '_dd.p.appsec')
+      // 1st req kept because waf init
+      // next in the first minute are dropped
+      // 5nd req kept because RateLimiter allows 1 req/min
+      await doWarmupRequests(proc)
 
-        assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_KEEP)
-        assert.propertyVal(metrics, '_dd.apm.enabled', 0)
+      await new Promise(resolve => setTimeout(resolve, 60000))
 
-        assertDrop(payload[2][0])
+      await curl(proc)
 
-        assertDrop(payload[3][0])
-      })
-    })
+      return promise
+    }).timeout(70000)
 
     it('should keep attack requests', async () => {
       await doWarmupRequests(proc)

package.json

@@ -89,6 +89,7 @@
     "@datadog/native-metrics": "^3.0.1",
     "@datadog/pprof": "5.4.1",
     "@datadog/sketches-js": "^2.1.0",
+    "@isaacs/ttlcache": "^1.4.1",
     "@opentelemetry/api": ">=1.0.0 <1.9.0",
     "@opentelemetry/core": "^1.14.0",
     "crypto-randomuuid": "^1.0.0",

packages/datadog-instrumentations/src/express.js

@@ -29,7 +29,7 @@ function wrapResponseJson (json) {
         obj = arguments[1]
       }
 
-      responseJsonChannel.publish({ req: this.req, body: obj })
+      responseJsonChannel.publish({ req: this.req, res: this, body: obj })
     }
 
     return json.apply(this, arguments)

packages/dd-trace/src/appsec/api_security_sampler.js

@@ -1,61 +1,84 @@
 'use strict'
 
+const TTLCache = require('@isaacs/ttlcache')
+const web = require('../plugins/util/web')
 const log = require('../log')
+const { AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
+
+const MAX_SIZE = 4096
 
 let enabled
-let requestSampling
+let sampledRequests
 
-const sampledRequests = new WeakSet()
+class NoopTTLCache {
+  clear () { }
+  set (key) { return undefined }
+  has (key) { return false }
+}
 
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
-  setRequestSampling(apiSecurity.requestSampling)
+  sampledRequests = apiSecurity.sampleDelay === 0
+    ? new NoopTTLCache()
+    : new TTLCache({ max: MAX_SIZE, ttl: apiSecurity.sampleDelay * 1000 })
 }
 
 function disable () {
   enabled = false
+  sampledRequests?.clear()
 }
 
-function setRequestSampling (sampling) {
-  requestSampling = parseRequestSampling(sampling)
-}
+function sampleRequest (req, res, force = false) {
+  if (!enabled) return false
 
-function parseRequestSampling (requestSampling) {
-  let parsed = parseFloat(requestSampling)
+  const key = computeKey(req, res)
+  if (!key || isSampled(key)) return false
 
-  if (isNaN(parsed)) {
-    log.warn(`Incorrect API Security request sampling value: ${requestSampling}`)
+  const rootSpan = web.root(req)
+  if (!rootSpan) return false
 
-    parsed = 0
-  } else {
-    parsed = Math.min(1, Math.max(0, parsed))
+  let priority = getSpanPriority(rootSpan)
+  if (!priority) {
+    rootSpan._prioritySampler?.sample(rootSpan)
+    priority = getSpanPriority(rootSpan)
   }
 
-  return parsed
-}
-
-function sampleRequest (req) {
-  if (!enabled || !requestSampling) {
+  if (priority === AUTO_REJECT || priority === USER_REJECT) {
     return false
   }
 
-  const shouldSample = Math.random() <= requestSampling
-
-  if (shouldSample) {
-    sampledRequests.add(req)
+  if (force) {
+    sampledRequests.set(key)
   }
 
-  return shouldSample
+  return true
+}
+
+function isSampled (key) {
+  return sampledRequests.has(key)
+}
+
+function computeKey (req, res) {
+  const route = web.getContext(req)?.paths?.join('') || ''
+  const method = req.method
+  const status = res.statusCode
+
+  if (!method || !status) {
+    log.warn('Unsupported groupkey for API security')
+    return null
+  }
+  return method + route + status
 }
 
-function isSampled (req) {
-  return sampledRequests.has(req)
+function getSpanPriority (span) {
+  const spanContext = span.context?.()
+  return spanContext._sampling?.priority
 }
 
 module.exports = {
   configure,
   disable,
-  setRequestSampling,
   sampleRequest,
-  isSampled
+  isSampled,
+  computeKey
 }

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

@@ -145,10 +145,6 @@ function incomingHttpStartTranslator ({ req, res, abortController }) {
     persistent[addresses.HTTP_CLIENT_IP] = clientIp
   }
 
-  if (apiSecuritySampler.sampleRequest(req)) {
-    persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
-  }
-
   const actions = waf.run({ persistent }, req)
 
   handleResults(actions, req, res, rootSpan, abortController)
@@ -172,6 +168,10 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
+  if (apiSecuritySampler.sampleRequest(req, res, true)) {
+    persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
+  }
+
   if (Object.keys(persistent).length) {
     waf.run({ persistent }, req)
   }
@@ -228,9 +228,9 @@ function onRequestProcessParams ({ req, res, abortController, params }) {
   handleResults(results, req, res, rootSpan, abortController)
 }
 
-function onResponseBody ({ req, body }) {
+function onResponseBody ({ req, res, body }) {
   if (!body || typeof body !== 'object') return
-  if (!apiSecuritySampler.isSampled(req)) return
+  if (!apiSecuritySampler.sampleRequest(req, res)) return
 
   // we don't support blocking at this point, so no results needed
   waf.run({

packages/dd-trace/src/appsec/remote_config/capabilities.js

@@ -11,7 +11,7 @@ module.exports = {
   ASM_CUSTOM_RULES: 1n << 8n,
   ASM_CUSTOM_BLOCKING_RESPONSE: 1n << 9n,
   ASM_TRUSTED_IPS: 1n << 10n,
-  ASM_API_SECURITY_SAMPLE_RATE: 1n << 11n,
+  ASM_API_SECURITY_SAMPLE_RATE: 1n << 11n, // deprecated
   APM_TRACING_SAMPLE_RATE: 1n << 12n,
   APM_TRACING_LOGS_INJECTION: 1n << 13n,
   APM_TRACING_HTTP_HEADER_TAGS: 1n << 14n,

packages/dd-trace/src/appsec/remote_config/index.js

@@ -4,7 +4,6 @@ const Activation = require('../activation')
 
 const RemoteConfigManager = require('./manager')
 const RemoteConfigCapabilities = require('./capabilities')
-const apiSecuritySampler = require('../api_security_sampler')
 
 let rc
 
@@ -24,18 +23,12 @@ function enable (config, appsec) {
       rc.updateCapabilities(RemoteConfigCapabilities.ASM_ACTIVATION, true)
     }
 
-    if (config.appsec.apiSecurity?.enabled) {
-      rc.updateCapabilities(RemoteConfigCapabilities.ASM_API_SECURITY_SAMPLE_RATE, true)
-    }
-
     rc.setProductHandler('ASM_FEATURES', (action, rcConfig) => {
       if (!rcConfig) return
 
       if (activation === Activation.ONECLICK) {
         enableOrDisableAppsec(action, rcConfig, config, appsec)
       }
-
-      apiSecuritySampler.setRequestSampling(rcConfig.api_security?.request_sample_rate)
     })
   }
 

packages/dd-trace/src/config.js

@@ -444,7 +444,7 @@ class Config {
     const defaults = setHiddenProperty(this, '_defaults', {})
 
     this._setValue(defaults, 'appsec.apiSecurity.enabled', true)
-    this._setValue(defaults, 'appsec.apiSecurity.requestSampling', 0.1)
+    this._setValue(defaults, 'appsec.apiSecurity.sampleDelay', 30)
     this._setValue(defaults, 'appsec.blockedTemplateGraphql', undefined)
     this._setValue(defaults, 'appsec.blockedTemplateHtml', undefined)
     this._setValue(defaults, 'appsec.blockedTemplateJson', undefined)
@@ -571,7 +571,7 @@ class Config {
       AWS_LAMBDA_FUNCTION_NAME,
       DD_AGENT_HOST,
       DD_API_SECURITY_ENABLED,
-      DD_API_SECURITY_REQUEST_SAMPLE_RATE,
+      DD_API_SECURITY_SAMPLE_DELAY,
       DD_APPSEC_AUTOMATED_USER_EVENTS_TRACKING,
       DD_APPSEC_ENABLED,
       DD_APPSEC_GRAPHQL_BLOCKED_TEMPLATE_JSON,
@@ -700,7 +700,7 @@ class Config {
       DD_API_SECURITY_ENABLED && isTrue(DD_API_SECURITY_ENABLED),
       DD_EXPERIMENTAL_API_SECURITY_ENABLED && isTrue(DD_EXPERIMENTAL_API_SECURITY_ENABLED)
     ))
-    this._setUnit(env, 'appsec.apiSecurity.requestSampling', DD_API_SECURITY_REQUEST_SAMPLE_RATE)
+    this._setValue(env, 'appsec.apiSecurity.sampleDelay', maybeFloat(DD_API_SECURITY_SAMPLE_DELAY))
     this._setValue(env, 'appsec.blockedTemplateGraphql', maybeFile(DD_APPSEC_GRAPHQL_BLOCKED_TEMPLATE_JSON))
     this._setValue(env, 'appsec.blockedTemplateHtml', maybeFile(DD_APPSEC_HTTP_BLOCKED_TEMPLATE_HTML))
     this._envUnprocessed['appsec.blockedTemplateHtml'] = DD_APPSEC_HTTP_BLOCKED_TEMPLATE_HTML
@@ -880,7 +880,6 @@ class Config {
     tagger.add(tags, options.tags)
 
     this._setBoolean(opts, 'appsec.apiSecurity.enabled', options.appsec.apiSecurity?.enabled)
-    this._setUnit(opts, 'appsec.apiSecurity.requestSampling', options.appsec.apiSecurity?.requestSampling)
     this._setValue(opts, 'appsec.blockedTemplateGraphql', maybeFile(options.appsec.blockedTemplateGraphql))
     this._setValue(opts, 'appsec.blockedTemplateHtml', maybeFile(options.appsec.blockedTemplateHtml))
     this._optsUnprocessed['appsec.blockedTemplateHtml'] = options.appsec.blockedTemplateHtml