Api Security sampling

docs/test.ts

@@ -115,7 +115,7 @@ tracer.init({
     },
     apiSecurity: {
       enabled: true,
-      requestSampling: 1.0
+      sampleDelay: 30
     }
   }
 });

index.d.ts

@@ -678,11 +678,10 @@ declare namespace tracer {
          */
         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
+        /** Specifies the Api Security samples collection interval in seconds.
+         * @default 30.0
          */
-        requestSampling?: number
+        sampleDelay?: number
       }
     };
 

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,51 @@
 'use strict'
 
-const log = require('../log')
+const LRU = require('lru-cache')
 
 let enabled
-let requestSampling
-
-const sampledRequests = new WeakSet()
+let sampledCache
 
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
-  setRequestSampling(apiSecurity.requestSampling)
-}
-
-function disable () {
-  enabled = false
-}
-
-function setRequestSampling (sampling) {
-  requestSampling = parseRequestSampling(sampling)
-}
 
-function parseRequestSampling (requestSampling) {
-  let parsed = parseFloat(requestSampling)
+  if (enabled) {
+    const max = apiSecurity.sampleCacheSize || 4096 // for testing purposes only
+    const ttl = 1000 * apiSecurity.sampleDelay
 
-  if (isNaN(parsed)) {
-    log.warn(`Incorrect API Security request sampling value: ${requestSampling}`)
-
-    parsed = 0
-  } else {
-    parsed = Math.min(1, Math.max(0, parsed))
+    sampledCache = new LRU({ max, ttl })
   }
+}
 
-  return parsed
+function disable () {
+  enabled = false
 }
 
-function sampleRequest (req) {
-  if (!enabled || !requestSampling) {
+function sampleRequest (req, res) {
+  if (!enabled) {
     return false
   }
 
-  const shouldSample = Math.random() <= requestSampling
-
+  const key = getKey(req, res)
+  const shouldSample = !sampledCache.has(key)
   if (shouldSample) {
-    sampledRequests.add(req)
+    sampledCache.set(key)
   }
 
   return shouldSample
 }
 
-function isSampled (req) {
-  return sampledRequests.has(req)
+// rfc mentions using a hash
+function getKey (req = {}, res = {}) {
+  return `${req.method}-${req.url}-${res.statusCode}`
+}
+
+function has (req, res) {
+  return sampledCache.has(getKey(req, res))
 }
 
 module.exports = {
   configure,
   disable,
-  setRequestSampling,
   sampleRequest,
-  isSampled
+  has
 }

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

@@ -21,12 +21,14 @@ const appsecTelemetry = require('./telemetry')
 const apiSecuritySampler = require('./api_security_sampler')
 const web = require('../plugins/util/web')
 const { extractIp } = require('../plugins/util/ip_extractor')
-const { HTTP_CLIENT_IP } = require('../../../../ext/tags')
+const { HTTP_CLIENT_IP, MANUAL_KEEP } = require('../../../../ext/tags')
 const { block, setTemplates } = require('./blocking')
 const { passportTrackEvent } = require('./passport')
 const { storage } = require('../../../datadog-core')
 const graphql = require('./graphql')
 
+const EXTRACT_SCHEMA = { 'extract-schema': true }
+
 let isEnabled = false
 let config
 
@@ -95,10 +97,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)
@@ -134,6 +132,13 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
+  if (apiSecuritySampler.sampleRequest(req, res)) {
+    persistent[addresses.WAF_CONTEXT_PROCESSOR] = EXTRACT_SCHEMA
+
+    const rootSpan = web.root(req)
+    rootSpan?.setTag(MANUAL_KEEP, 'true')
+  }
+
   waf.run({ persistent }, req)
 
   waf.disposeContext(req)
@@ -196,16 +201,24 @@ function onRequestCookieParser ({ req, res, abortController, cookies }) {
   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
 
-  // we don't support blocking at this point, so no results needed
-  waf.run({
-    persistent: {
-      [addresses.HTTP_OUTGOING_BODY]: body
-    }
-  }, req)
+  if (apiSecuritySampler.sampleRequest(req, res)) {
+    const rootSpan = web.root(req)
+    if (!rootSpan) return
+
+    // we don't support blocking at this point, so no results needed
+    waf.run({
+      persistent: {
+        [addresses.WAF_CONTEXT_PROCESSOR]: EXTRACT_SCHEMA,
+        [addresses.HTTP_OUTGOING_BODY]: body
+      }
+    }, req)
+
+    // set manual.keep to force agent to keep the trace. should we delay it until finishRequest?
+    rootSpan.setTag(MANUAL_KEEP, 'true')
+  }
 }
 
 function onPassportVerify ({ credentials, user }) {

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
 
@@ -28,13 +27,9 @@ function enable (config, appsec) {
     }
 
     rc.on('ASM_FEATURES', (action, rcConfig) => {
-      if (!rcConfig) return
-
-      if (activation === Activation.ONECLICK) {
+      if (rcConfig && activation === Activation.ONECLICK) {
         enableOrDisableAppsec(action, rcConfig, config, appsec)
       }
-
-      apiSecuritySampler.setRequestSampling(rcConfig.api_security?.request_sample_rate)
     })
   }
 

packages/dd-trace/src/config.js

@@ -262,10 +262,10 @@ class Config {
       process.env.DD_EXPERIMENTAL_API_SECURITY_ENABLED && isTrue(process.env.DD_EXPERIMENTAL_API_SECURITY_ENABLED),
       true
     )
-    const DD_API_SECURITY_REQUEST_SAMPLE_RATE = coalesce(
-      options.appsec?.apiSecurity?.requestSampling,
-      parseFloat(process.env.DD_API_SECURITY_REQUEST_SAMPLE_RATE),
-      0.1
+    const DD_API_SECURITY_SAMPLE_DELAY = coalesce(
+      maybeFloat(options.appsec?.apiSecurity?.sampleDelay),
+      maybeFloat(process.env.DD_API_SECURITY_SAMPLE_DELAY),
+      30.0
     )
 
     // 0: disabled, 1: logging, 2: garbage collection + logging
@@ -332,8 +332,7 @@ class Config {
       },
       apiSecurity: {
         enabled: DD_API_SECURITY_ENABLED,
-        // Coerce value between 0 and 1
-        requestSampling: Math.min(1, Math.max(0, DD_API_SECURITY_REQUEST_SAMPLE_RATE))
+        sampleDelay: DD_API_SECURITY_SAMPLE_DELAY
       }
     }
 

packages/dd-trace/test/appsec/api_security_sampler.spec.js

@@ -4,68 +4,79 @@ const apiSecuritySampler = require('../../src/appsec/api_security_sampler')
 
 describe('Api Security Sampler', () => {
   let config
+  const req = {
+    url: '/test',
+    method: 'GET'
+  }
+  const res = {
+    statusCode: 200
+  }
 
   beforeEach(() => {
     config = {
       apiSecurity: {
         enabled: true,
-        requestSampling: 1
+        sampleDelay: 30,
+        sampleCacheSize: 2
       }
     }
-
-    sinon.stub(Math, 'random').returns(0.3)
   })
 
   afterEach(sinon.restore)
 
   describe('sampleRequest', () => {
-    it('should sample request if enabled and sampling 1', () => {
+    it('should sample unknown request', () => {
       apiSecuritySampler.configure(config)
 
-      expect(apiSecuritySampler.sampleRequest({})).to.true
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
     })
 
-    it('should not sample request if enabled and sampling 0', () => {
-      config.apiSecurity.requestSampling = 0
+    it('should sample multiple requests based on url, method and statusCode', () => {
       apiSecuritySampler.configure(config)
 
-      expect(apiSecuritySampler.sampleRequest({})).to.false
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
+      expect(apiSecuritySampler.sampleRequest(req, { statusCode: 500 })).to.true
     })
 
-    it('should sample request if enabled and sampling greater than random', () => {
-      config.apiSecurity.requestSampling = 0.5
-
+    it('should sample multiple requests based on url, method and statusCode II', () => {
       apiSecuritySampler.configure(config)
 
-      expect(apiSecuritySampler.sampleRequest({})).to.true
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
+      expect(apiSecuritySampler.sampleRequest({ url: '/otherTest', method: 'POST' }, res)).to.true
     })
 
-    it('should not sample request if enabled and sampling less than random', () => {
-      config.apiSecurity.requestSampling = 0.1
-
+    it('should not sample repeated requests in the configured interval', () => {
       apiSecuritySampler.configure(config)
 
-      expect(apiSecuritySampler.sampleRequest()).to.false
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.false
     })
 
-    it('should not sample request if incorrect config value', () => {
-      config.apiSecurity.requestSampling = NaN
+    it('should sample repeated request after interval', done => {
+      config.apiSecurity.sampleDelay = 0.1
 
       apiSecuritySampler.configure(config)
 
-      expect(apiSecuritySampler.sampleRequest()).to.false
-    })
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
 
-    it('should sample request according to the config', () => {
-      config.apiSecurity.requestSampling = 1
+      setTimeout(() => {
+        expect(apiSecuritySampler.has(req, res)).to.false
 
-      apiSecuritySampler.configure(config)
+        expect(apiSecuritySampler.sampleRequest(req, res)).to.true
+
+        done()
+      }, 200)
+    })
 
-      expect(apiSecuritySampler.sampleRequest({})).to.true
+    it('should mantain a max size cache', () => {
+      apiSecuritySampler.configure(config)
 
-      apiSecuritySampler.setRequestSampling(0)
+      expect(apiSecuritySampler.sampleRequest(req, res)).to.true
+      expect(apiSecuritySampler.sampleRequest(req, { statusCode: 500 })).to.true
+      expect(apiSecuritySampler.sampleRequest(req, { statusCode: 404 })).to.true
 
-      expect(apiSecuritySampler.sampleRequest()).to.false
+      expect(apiSecuritySampler.has(req, res)).to.false
+      expect(apiSecuritySampler.has(req, { statusCode: 500 })).to.true
     })
   })
 })