Replace manual.keep tag usage with an specific method to keep the tra…

…ce (#4739)

* Allow to set sampling priority

* Span.keep method

* Replace manual.keep tag usage with Span.keep()

* Update standalone integration tests

* use PrioritySampler.keepTrace

* Lint

* PrioritySampler.keepTrace test

integration-tests/standalone-asm.spec.js

@@ -10,6 +10,7 @@ const {
   curlAndAssertMessage,
   curl
 } = require('./helpers')
+const { USER_KEEP, AUTO_REJECT, AUTO_KEEP } = require('../ext/priority')
 
 describe('Standalone ASM', () => {
   let sandbox, cwd, startupTestFile, agent, proc, env
@@ -43,22 +44,18 @@ describe('Standalone ASM', () => {
       await agent.stop()
     })
 
-    function assertKeep (payload, manual = true) {
+    function assertKeep (payload) {
       const { meta, metrics } = payload
-      if (manual) {
-        assert.propertyVal(meta, 'manual.keep', 'true')
-      } else {
-        assert.notProperty(meta, 'manual.keep')
-      }
+
       assert.propertyVal(meta, '_dd.p.appsec', '1')
 
-      assert.propertyVal(metrics, '_sampling_priority_v1', 2)
+      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', 0)
+      assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_REJECT)
       assert.propertyVal(metrics, '_dd.apm.enabled', 0)
       assert.notProperty(metrics, '_dd.p.appsec')
     }
@@ -103,7 +100,7 @@ describe('Standalone ASM', () => {
         assert.notProperty(meta, 'manual.keep')
         assert.notProperty(meta, '_dd.p.appsec')
 
-        assert.propertyVal(metrics, '_sampling_priority_v1', 1)
+        assert.propertyVal(metrics, '_sampling_priority_v1', AUTO_KEEP)
         assert.propertyVal(metrics, '_dd.apm.enabled', 0)
 
         assertDrop(payload[2][0])
@@ -213,7 +210,7 @@ describe('Standalone ASM', () => {
 
             const innerReq = payload.find(p => p[0].resource === 'GET /down')
             assert.notStrictEqual(innerReq, undefined)
-            assertKeep(innerReq[0], false)
+            assertKeep(innerReq[0])
           }, undefined, undefined, true)
         })
 

packages/dd-trace/src/appsec/iast/vulnerability-reporter.js

@@ -1,10 +1,11 @@
 'use strict'
 
-const { MANUAL_KEEP } = require('../../../../../ext/tags')
 const LRU = require('lru-cache')
 const vulnerabilitiesFormatter = require('./vulnerabilities-formatter')
 const { IAST_ENABLED_TAG_KEY, IAST_JSON_TAG_KEY } = require('./tags')
 const standalone = require('../standalone')
+const { SAMPLING_MECHANISM_APPSEC } = require('../../constants')
+const { keepTrace } = require('../../priority_sampler')
 
 const VULNERABILITIES_KEY = 'vulnerabilities'
 const VULNERABILITY_HASHES_MAX_SIZE = 1000
@@ -56,9 +57,10 @@ function sendVulnerabilities (vulnerabilities, rootSpan) {
         const tags = {}
         // TODO: Store this outside of the span and set the tag in the exporter.
         tags[IAST_JSON_TAG_KEY] = JSON.stringify(jsonToSend)
-        tags[MANUAL_KEEP] = 'true'
         span.addTags(tags)
 
+        keepTrace(span, SAMPLING_MECHANISM_APPSEC)
+
         standalone.sample(span)
 
         if (!rootSpan) span.finish()

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

@@ -13,8 +13,9 @@ const {
   getRequestMetrics
 } = require('./telemetry')
 const zlib = require('zlib')
-const { MANUAL_KEEP } = require('../../../../ext/tags')
 const standalone = require('./standalone')
+const { SAMPLING_MECHANISM_APPSEC } = require('../constants')
+const { keepTrace } = require('../priority_sampler')
 
 // default limiter, configurable with setRateLimit()
 let limiter = new Limiter(100)
@@ -96,8 +97,6 @@ function reportWafInit (wafVersion, rulesVersion, diagnosticsRules = {}) {
     metricsQueue.set('_dd.appsec.event_rules.errors', JSON.stringify(diagnosticsRules.errors))
   }
 
-  metricsQueue.set(MANUAL_KEEP, 'true')
-
   incrementWafInitMetric(wafVersion, rulesVersion)
 }
 
@@ -129,7 +128,7 @@ function reportAttack (attackData) {
   }
 
   if (limiter.isAllowed()) {
-    newTags[MANUAL_KEEP] = 'true'
+    keepTrace(rootSpan, SAMPLING_MECHANISM_APPSEC)
 
     standalone.sample(rootSpan)
   }
@@ -184,6 +183,8 @@ function finishRequest (req, res) {
   if (metricsQueue.size) {
     rootSpan.addTags(Object.fromEntries(metricsQueue))
 
+    keepTrace(rootSpan, SAMPLING_MECHANISM_APPSEC)
+
     standalone.sample(rootSpan)
 
     metricsQueue.clear()

packages/dd-trace/src/appsec/sdk/track_event.js

@@ -2,10 +2,11 @@
 
 const log = require('../../log')
 const { getRootSpan } = require('./utils')
-const { MANUAL_KEEP } = require('../../../../../ext/tags')
 const { setUserTags } = require('./set_user')
 const standalone = require('../standalone')
 const waf = require('../waf')
+const { SAMPLING_MECHANISM_APPSEC } = require('../../constants')
+const { keepTrace } = require('../../priority_sampler')
 
 function trackUserLoginSuccessEvent (tracer, user, metadata) {
   // TODO: better user check here and in _setUser() ?
@@ -55,9 +56,10 @@ function trackEvent (eventName, fields, sdkMethodName, rootSpan, mode) {
     return
   }
 
+  keepTrace(rootSpan, SAMPLING_MECHANISM_APPSEC)
+
   const tags = {
-    [`appsec.events.${eventName}.track`]: 'true',
-    [MANUAL_KEEP]: 'true'
+    [`appsec.events.${eventName}.track`]: 'true'
   }
 
   if (mode === 'sdk') {

packages/dd-trace/src/priority_sampler.js

@@ -108,6 +108,18 @@ class PrioritySampler {
     }
   }
 
+  setPriority (span, samplingPriority, mechanism = SAMPLING_MECHANISM_MANUAL) {
+    if (!span || !this.validate(samplingPriority)) return
+
+    const context = this._getContext(span)
+
+    context._sampling.priority = samplingPriority
+    context._sampling.mechanism = mechanism
+
+    const root = context._trace.started[0]
+    this._addDecisionMaker(root)
+  }
+
   _getContext (span) {
     return typeof span.context === 'function' ? span.context() : span
   }
@@ -201,6 +213,10 @@ class PrioritySampler {
       if (rule.match(span)) return rule
     }
   }
+
+  static keepTrace (span, mechanism) {
+    span?._prioritySampler?.setPriority(span, USER_KEEP, mechanism)
+  }
 }
 
 module.exports = PrioritySampler

packages/dd-trace/test/appsec/iast/vulnerability-reporter.spec.js

@@ -2,7 +2,8 @@ const { addVulnerability, sendVulnerabilities, clearCache, start, stop } =
   require('../../../src/appsec/iast/vulnerability-reporter')
 const VulnerabilityAnalyzer = require('../../../../dd-trace/src/appsec/iast/analyzers/vulnerability-analyzer')
 const appsecStandalone = require('../../../src/appsec/standalone')
-const { APPSEC_PROPAGATION_KEY } = require('../../../src/constants')
+const { APPSEC_PROPAGATION_KEY, SAMPLING_MECHANISM_APPSEC } = require('../../../src/constants')
+const { USER_KEEP } = require('../../../../../ext/priority')
 
 describe('vulnerability-reporter', () => {
   let vulnerabilityAnalyzer
@@ -82,9 +83,14 @@ describe('vulnerability-reporter', () => {
     describe('without rootSpan', () => {
       let fakeTracer
       let onTheFlySpan
+      let prioritySampler
 
       beforeEach(() => {
+        prioritySampler = {
+          setPriority: sinon.stub()
+        }
         onTheFlySpan = {
+          _prioritySampler: prioritySampler,
           finish: sinon.spy(),
           addTags: sinon.spy(),
           context () {
@@ -120,10 +126,11 @@ describe('vulnerability-reporter', () => {
           '_dd.iast.enabled': 1
         })
         expect(onTheFlySpan.addTags.secondCall).to.have.been.calledWithExactly({
-          'manual.keep': 'true',
           '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3410512655,' +
             '"evidence":{"value":"sha1"},"location":{"spanId":42,"path":"filename.js","line":73}}]}'
         })
+        expect(prioritySampler.setPriority)
+          .to.have.been.calledOnceWithExactly(onTheFlySpan, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
         expect(onTheFlySpan.finish).to.have.been.calledOnce
       })
 
@@ -140,10 +147,15 @@ describe('vulnerability-reporter', () => {
   describe('sendVulnerabilities', () => {
     let span
     let context
+    let prioritySampler
 
     beforeEach(() => {
       context = { _trace: { tags: {} } }
+      prioritySampler = {
+        setPriority: sinon.stub()
+      }
       span = {
+        _prioritySampler: prioritySampler,
         addTags: sinon.stub(),
         context: sinon.stub().returns(context)
       }
@@ -178,10 +190,10 @@ describe('vulnerability-reporter', () => {
         vulnerabilityAnalyzer._createVulnerability('INSECURE_HASHING', { value: 'sha1' }, 888))
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3254801297,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should send only valid vulnerabilities', () => {
@@ -191,10 +203,10 @@ describe('vulnerability-reporter', () => {
       iastContext.vulnerabilities.push({ invalid: 'vulnerability' })
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3254801297,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should send vulnerabilities with evidence, ranges and sources', () => {
@@ -239,7 +251,6 @@ describe('vulnerability-reporter', () => {
 
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[{"origin":"ORIGIN_TYPE_1","name":"PARAMETER_NAME_1","value":"joe"},' +
           '{"origin":"ORIGIN_TYPE_2","name":"PARAMETER_NAME_2","value":"[email protected]"}],' +
           '"vulnerabilities":[{"type":"SQL_INJECTION","hash":4676753086,' +
@@ -249,6 +260,7 @@ describe('vulnerability-reporter', () => {
           '[{"value":"SELECT id FROM u WHERE email = \'"},{"value":"[email protected]","source":1},{"value":"\';"}]},' +
           '"location":{"spanId":888,"path":"filename.js","line":99}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should send multiple vulnerabilities with same tainted source', () => {
@@ -293,7 +305,6 @@ describe('vulnerability-reporter', () => {
 
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[{"origin":"ORIGIN_TYPE_1","name":"PARAMETER_NAME_1","value":"joe"}],' +
           '"vulnerabilities":[{"type":"SQL_INJECTION","hash":4676753086,' +
           '"evidence":{"valueParts":[{"value":"SELECT * FROM u WHERE name = \'"},{"value":"joe","source":0},' +
@@ -302,6 +313,7 @@ describe('vulnerability-reporter', () => {
           '[{"value":"UPDATE u SET name=\'"},{"value":"joe","source":0},{"value":"\' WHERE id=1;"}]},' +
           '"location":{"spanId":888,"path":"filename.js","line":99}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should send once with multiple vulnerabilities', () => {
@@ -314,7 +326,6 @@ describe('vulnerability-reporter', () => {
         { path: '/path/to/file3.js', line: 3 }))
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[' +
           '{"type":"INSECURE_HASHING","hash":1697980169,"evidence":{"value":"sha1"},' +
             '"location":{"spanId":888,"path":"/path/to/file1.js","line":1}},' +
@@ -323,6 +334,7 @@ describe('vulnerability-reporter', () => {
           '{"type":"INSECURE_HASHING","hash":1755238473,"evidence":{"value":"md5"},' +
             '"location":{"spanId":-5,"path":"/path/to/file3.js","line":3}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should send once vulnerability with one vulnerability', () => {
@@ -332,10 +344,10 @@ describe('vulnerability-reporter', () => {
           { path: 'filename.js', line: 88 }))
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3410512691,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888,"path":"filename.js","line":88}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should not send duplicated vulnerabilities', () => {
@@ -348,10 +360,10 @@ describe('vulnerability-reporter', () => {
           { path: 'filename.js', line: 88 }))
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3410512691,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888,"path":"filename.js","line":88}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should not send duplicated vulnerabilities in multiple sends', () => {
@@ -365,10 +377,10 @@ describe('vulnerability-reporter', () => {
       sendVulnerabilities(iastContext.vulnerabilities, span)
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3410512691,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888,"path":"filename.js","line":88}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should not deduplicate vulnerabilities if not enabled', () => {
@@ -384,12 +396,12 @@ describe('vulnerability-reporter', () => {
         { value: 'sha1' }, 888, { path: 'filename.js', line: 88 }))
       sendVulnerabilities(iastContext.vulnerabilities, span)
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3410512691,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":888,"path":"filename.js","line":88}},' +
           '{"type":"INSECURE_HASHING","hash":3410512691,"evidence":{"value":"sha1"},"location":' +
           '{"spanId":888,"path":"filename.js","line":88}}]}'
       })
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
 
     it('should add _dd.p.appsec trace tag with standalone enabled', () => {
@@ -401,11 +413,12 @@ describe('vulnerability-reporter', () => {
       sendVulnerabilities(iastContext.vulnerabilities, span)
 
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3254801297,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":999}}]}'
       })
 
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
+
       expect(span.context()._trace.tags).to.have.property(APPSEC_PROPAGATION_KEY)
     })
 
@@ -418,11 +431,12 @@ describe('vulnerability-reporter', () => {
       sendVulnerabilities(iastContext.vulnerabilities, span)
 
       expect(span.addTags).to.have.been.calledOnceWithExactly({
-        'manual.keep': 'true',
         '_dd.iast.json': '{"sources":[],"vulnerabilities":[{"type":"INSECURE_HASHING","hash":3254801297,' +
           '"evidence":{"value":"sha1"},"location":{"spanId":999}}]}'
       })
 
+      expect(prioritySampler.setPriority).to.have.been.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
+
       expect(span.context()._trace.tags).to.not.have.property(APPSEC_PROPAGATION_KEY)
     })
   })
@@ -441,7 +455,8 @@ describe('vulnerability-reporter', () => {
       global.setInterval = sinon.spy(global.setInterval)
       global.clearInterval = sinon.spy(global.clearInterval)
       span = {
-        addTags: sinon.stub()
+        addTags: sinon.stub(),
+        keep: sinon.stub()
       }
     })