Skip to content

Commit

Permalink
Replace manual.keep tag usage with an specific method to keep the tra…
Browse files Browse the repository at this point in the history
…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
  • Loading branch information
iunanua authored and rochdev committed Nov 6, 2024
1 parent ee6aaab commit fcf8346
Show file tree
Hide file tree
Showing 10 changed files with 192 additions and 65 deletions.
17 changes: 7 additions & 10 deletions integration-tests/standalone-asm.spec.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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')
}
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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)
})

Expand Down
6 changes: 4 additions & 2 deletions packages/dd-trace/src/appsec/iast/vulnerability-reporter.js
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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()
Expand Down
9 changes: 5 additions & 4 deletions packages/dd-trace/src/appsec/reporter.js
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -129,7 +128,7 @@ function reportAttack (attackData) {
}

if (limiter.isAllowed()) {
newTags[MANUAL_KEEP] = 'true'
keepTrace(rootSpan, SAMPLING_MECHANISM_APPSEC)

standalone.sample(rootSpan)
}
Expand Down Expand Up @@ -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()
Expand Down
8 changes: 5 additions & 3 deletions packages/dd-trace/src/appsec/sdk/track_event.js
Original file line number Diff line number Diff line change
Expand Up @@ -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() ?
Expand Down Expand Up @@ -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') {
Expand Down
16 changes: 16 additions & 0 deletions packages/dd-trace/src/priority_sampler.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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
43 changes: 29 additions & 14 deletions packages/dd-trace/test/appsec/iast/vulnerability-reporter.spec.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 () {
Expand Down Expand Up @@ -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
})

Expand All @@ -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)
}
Expand Down Expand Up @@ -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', () => {
Expand All @@ -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', () => {
Expand Down Expand Up @@ -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":"joe@mail.com"}],' +
'"vulnerabilities":[{"type":"SQL_INJECTION","hash":4676753086,' +
Expand All @@ -249,6 +260,7 @@ describe('vulnerability-reporter', () => {
'[{"value":"SELECT id FROM u WHERE email = \'"},{"value":"joe@mail.com","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', () => {
Expand Down Expand Up @@ -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},' +
Expand All @@ -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', () => {
Expand All @@ -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}},' +
Expand All @@ -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', () => {
Expand All @@ -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', () => {
Expand All @@ -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', () => {
Expand All @@ -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', () => {
Expand All @@ -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', () => {
Expand All @@ -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)
})

Expand All @@ -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)
})
})
Expand All @@ -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()
}
})

Expand Down
Loading

0 comments on commit fcf8346

Please sign in to comment.