manual keep trace on asm standalone and api security (#5649)

* manual keep trace on asm standalone and api security

* set priority to user reject from auto

* add test for sample request

* keep auto reject

Author: IlyasShabi

integration-tests/standalone-asm.spec.js

@@ -25,13 +25,35 @@ describe('Standalone ASM', () => {
     await sandbox.remove()
   })
 
+  function assertKeep ({ meta, metrics }) {
+    assert.propertyVal(meta, '_dd.p.ts', '02')
+
+    assert.propertyVal(metrics, '_sampling_priority_v1', USER_KEEP)
+    assert.propertyVal(metrics, '_dd.apm.enabled', 0)
+  }
+
+  function assertDrop ({ meta, metrics }) {
+    assert.notProperty(meta, '_dd.p.ts')
+
+    assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_REJECT)
+    assert.propertyVal(metrics, '_dd.apm.enabled', 0)
+  }
+
+  async function doWarmupRequests (procOrUrl, number = 3) {
+    for (let i = number; i > 0; i--) {
+      await curl(procOrUrl)
+    }
+  }
+
   describe('enabled', () => {
     beforeEach(async () => {
       agent = await new FakeAgent().start()
 
       env = {
         AGENT_PORT: agent.port,
-        DD_EXPERIMENTAL_APPSEC_STANDALONE_ENABLED: 'true'
+        DD_APM_TRACING_ENABLED: 'false',
+        DD_APPSEC_ENABLED: 'true',
+        DD_API_SECURITY_ENABLED: 'false'
       }
 
       const execArgv = []
@@ -44,28 +66,6 @@ describe('Standalone ASM', () => {
       await agent.stop()
     })
 
-    function assertKeep (payload) {
-      const { meta, metrics } = payload
-
-      assert.propertyVal(meta, '_dd.p.ts', '02')
-
-      assert.propertyVal(metrics, '_sampling_priority_v1', USER_KEEP)
-      assert.propertyVal(metrics, '_dd.apm.enabled', 0)
-    }
-
-    function assertDrop (payload) {
-      const { metrics } = payload
-      assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_REJECT)
-      assert.propertyVal(metrics, '_dd.apm.enabled', 0)
-      assert.notProperty(metrics, '_dd.p.ts')
-    }
-
-    async function doWarmupRequests (procOrUrl, number = 3) {
-      for (let i = number; i > 0; i--) {
-        await curl(procOrUrl)
-      }
-    }
-
     // first req initializes the waf and reports the first appsec event adding manual.keep tag
     it('should send correct headers and tags on first req', async () => {
       return curlAndAssertMessage(agent, proc, ({ headers, payload }) => {
@@ -254,6 +254,58 @@ describe('Standalone ASM', () => {
     })
   })
 
+  describe('With API Security enabled', () => {
+    beforeEach(async () => {
+      agent = await new FakeAgent().start()
+
+      env = {
+        AGENT_PORT: agent.port,
+        DD_APM_TRACING_ENABLED: 'false',
+        DD_APPSEC_ENABLED: 'true'
+      }
+
+      const execArgv = []
+
+      proc = await spawnProc(startupTestFile, { cwd, env, execArgv })
+    })
+
+    afterEach(async () => {
+      proc.kill()
+      await agent.stop()
+    })
+
+    it('should keep fifth req because of api security sampler', async () => {
+      const promise = curlAndAssertMessage(agent, proc, ({ headers, payload }) => {
+        assert.propertyVal(headers, 'datadog-client-computed-stats', 'yes')
+        assert.isArray(payload)
+        if (payload.length === 4) {
+          assertKeep(payload[0][0])
+          assertDrop(payload[1][0])
+          assertDrop(payload[2][0])
+          assertDrop(payload[3][0])
+
+          // req after 30s
+        } else {
+          const fifthReq = payload[0]
+          assert.isArray(fifthReq)
+          assert.strictEqual(fifthReq.length, 5)
+          assertKeep(fifthReq[0])
+        }
+      }, 40000, 2)
+
+      // 1st req kept because waf init
+      // next in the 30s are dropped
+      // 5nd req manual kept because of api security sampler
+      await doWarmupRequests(proc)
+
+      await new Promise(resolve => setTimeout(resolve, 30000))
+
+      await curl(proc)
+
+      return promise
+    }).timeout(40000)
+  })
+
   describe('disabled', () => {
     beforeEach(async () => {
       agent = await new FakeAgent().start()

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

@@ -4,10 +4,13 @@ const TTLCache = require('@isaacs/ttlcache')
 const web = require('../plugins/util/web')
 const log = require('../log')
 const { AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
+const { keepTrace } = require('../priority_sampler')
+const { ASM } = require('../standalone/product')
 
 const MAX_SIZE = 4096
 
 let enabled
+let asmStandaloneEnabled
 
 /**
  * @type {TTLCache}
@@ -20,11 +23,12 @@ class NoopTTLCache {
   has (_key) { return false }
 }
 
-function configure ({ apiSecurity }) {
-  enabled = apiSecurity.enabled
-  sampledRequests = apiSecurity.sampleDelay === 0
+function configure ({ appsec, apmTracingEnabled }) {
+  enabled = appsec.apiSecurity.enabled
+  asmStandaloneEnabled = apmTracingEnabled === false
+  sampledRequests = appsec.apiSecurity.sampleDelay === 0
     ? new NoopTTLCache()
-    : new TTLCache({ max: MAX_SIZE, ttl: apiSecurity.sampleDelay * 1000 })
+    : new TTLCache({ max: MAX_SIZE, ttl: appsec.apiSecurity.sampleDelay * 1000 })
 }
 
 function disable () {
@@ -41,14 +45,18 @@ function sampleRequest (req, res, force = false) {
   const rootSpan = web.root(req)
   if (!rootSpan) return false
 
-  let priority = getSpanPriority(rootSpan)
-  if (!priority) {
-    rootSpan._prioritySampler?.sample(rootSpan)
-    priority = getSpanPriority(rootSpan)
-  }
-
-  if (priority === AUTO_REJECT || priority === USER_REJECT) {
-    return false
+  if (asmStandaloneEnabled) {
+    keepTrace(rootSpan, ASM)
+  } else {
+    let priority = getSpanPriority(rootSpan)
+    if (!priority) {
+      rootSpan._prioritySampler?.sample(rootSpan)
+      priority = getSpanPriority(rootSpan)
+    }
+
+    if (priority === AUTO_REJECT || priority === USER_REJECT) {
+      return false
+    }
   }
 
   if (force) {

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

@@ -59,7 +59,7 @@ function enable (_config) {
 
     Reporter.setRateLimit(_config.appsec.rateLimit)
 
-    apiSecuritySampler.configure(_config.appsec)
+    apiSecuritySampler.configure(_config)
 
     UserTracking.setCollectionMode(_config.appsec.eventTracking.mode, false)
 

packages/dd-trace/src/standalone/product.js

@@ -3,14 +3,12 @@
 const { SAMPLING_MECHANISM_APPSEC } = require('../constants')
 const RateLimiter = require('../rate_limiter')
 
-const dropAll = new RateLimiter(0)
-const onePerMinute = new RateLimiter(1, 'minute')
-
 function getProductRateLimiter (config) {
   if (config?.appsec?.enabled || config?.iast?.enabled) {
-    return onePerMinute
+    return new RateLimiter(1, 'minute') // onePerMinute
   }
-  return dropAll
+
+  return new RateLimiter(0) // dropAll
 }
 
 const PRODUCTS = {

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

@@ -69,7 +69,7 @@ describe('API Security Sampler', () => {
 
   describe('with TTLCache', () => {
     beforeEach(() => {
-      apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } })
+      apiSecuritySampler.configure({ appsec: { apiSecurity: { enabled: true, sampleDelay: 30 } } })
     })
 
     it('should not sample before 30 seconds', () => {
@@ -152,7 +152,7 @@ describe('API Security Sampler', () => {
 
   describe('with NoopTTLCache', () => {
     beforeEach(() => {
-      apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 0 } })
+      apiSecuritySampler.configure({ appsec: { apiSecurity: { enabled: true, sampleDelay: 0 } } })
     })
 
     it('should always return true for sampleRequest', () => {
@@ -197,4 +197,43 @@ describe('API Security Sampler', () => {
       assert.isTrue(apiSecuritySampler.sampleRequest(req, res, true))
     })
   })
+
+  describe('ASM Standalone', () => {
+    let keepTraceStub
+
+    beforeEach(() => {
+      keepTraceStub = sinon.stub()
+      apiSecuritySampler = proxyquire('../../src/appsec/api_security_sampler', {
+        '../plugins/util/web': webStub,
+        '../priority_sampler': {
+          keepTrace: keepTraceStub
+        },
+        '../standalone/product': {
+          ASM: 'asm'
+        }
+      })
+      apiSecuritySampler.configure({
+        appsec: {
+          apiSecurity: {
+            enabled: true,
+            sampleDelay: 30
+          }
+        },
+        apmTracingEnabled: false
+      })
+    })
+
+    it('should keep trace with ASM product when in standalone mode', () => {
+      assert.isTrue(apiSecuritySampler.sampleRequest(req, res, true))
+      assert.isFalse(apiSecuritySampler.sampleRequest(req, res, true))
+      assert.isTrue(keepTraceStub.calledOnceWith(span, 'asm'))
+    })
+
+    it('should not check priority sampling in standalone mode', () => {
+      span.context.returns({ _sampling: { priority: AUTO_REJECT } })
+      assert.isTrue(apiSecuritySampler.sampleRequest(req, res, true))
+      assert.isFalse(apiSecuritySampler.sampleRequest(req, res, true))
+      assert.isTrue(keepTraceStub.calledOnceWith(span, 'asm'))
+    })
+  })
 })