From ca8bb68f09fa94e3c8e3e3c6ff54da5f217d5ddf Mon Sep 17 00:00:00 2001 From: ishabi Date: Thu, 3 Oct 2024 16:39:29 +0200 Subject: [PATCH] use priority simpler to get span priority --- .../src/appsec/api_security_sampler.js | 16 ++-- .../test/appsec/api_security_sampler.spec.js | 87 +++++++++---------- packages/dd-trace/test/appsec/index.spec.js | 15 +++- 3 files changed, 59 insertions(+), 59 deletions(-) diff --git a/packages/dd-trace/src/appsec/api_security_sampler.js b/packages/dd-trace/src/appsec/api_security_sampler.js index 6baaed4f87..8c4b8fdc22 100644 --- a/packages/dd-trace/src/appsec/api_security_sampler.js +++ b/packages/dd-trace/src/appsec/api_security_sampler.js @@ -1,11 +1,12 @@ 'use strict' const ApiSecuritySamplerCache = require('./api_security_sampler_cache') +const PrioritySampler = require('../priority_sampler') const web = require('../plugins/util/web') -const { USER_KEEP, AUTO_KEEP } = require('../../../../ext/priority') let enabled let sampledRequests +const prioritySampler = new PrioritySampler() function configure ({ apiSecurity }) { enabled = apiSecurity.enabled @@ -23,27 +24,22 @@ function sampleRequest (req, res) { const rootSpan = web.root(req) if (!rootSpan) return false - const priority = getSpanPriority(rootSpan) + const isSampled = prioritySampler.isSampled(rootSpan) - if (priority !== AUTO_KEEP && priority !== USER_KEEP) { + if (!isSampled) { return false } const key = sampledRequests.computeKey(req, res) - const isSampled = sampledRequests.isSampled(key) + const alreadySampled = sampledRequests.isSampled(key) - if (isSampled) return false + if (alreadySampled) return false sampledRequests.set(key) return true } -function getSpanPriority (span) { - const spanContext = span.context?.() - return spanContext._sampling?.priority // default ?? -} - module.exports = { configure, disable, diff --git a/packages/dd-trace/test/appsec/api_security_sampler.spec.js b/packages/dd-trace/test/appsec/api_security_sampler.spec.js index 08b4355084..abcee98f3a 100644 --- a/packages/dd-trace/test/appsec/api_security_sampler.spec.js +++ b/packages/dd-trace/test/appsec/api_security_sampler.spec.js @@ -3,73 +3,68 @@ const proxyquire = require('proxyquire') describe('API Security Sampler', () => { - let apiSecuritySampler, webStub, clock + let apiSecuritySampler, webStub, clock, sampler, span beforeEach(() => { webStub = { root: sinon.stub() } clock = sinon.useFakeTimers(Date.now()) + sampler = sinon.stub().returns({ + isSampled: sinon.stub() + }) + apiSecuritySampler = proxyquire('../../src/appsec/api_security_sampler', { - '../plugins/util/web': webStub + '../plugins/util/web': webStub, + '../priority_sampler': sampler }) + + apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) + + span = { + context: sinon.stub().returns({}) + } }) afterEach(() => { clock.restore() }) - describe('sampleRequest', () => { - beforeEach(() => { - apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) - }) - - it('should return false if not enabled', () => { - apiSecuritySampler.disable() - expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false - }) - - it('should return false if no root span', () => { - webStub.root.returns(null) - expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false - }) + it('should return false if not enabled', () => { + apiSecuritySampler.disable() + expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false + }) - it('should return true and put request in cache if priority is AUTO_KEEP', () => { - const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) } - webStub.root.returns(rootSpan) - const req = { url: '/test', method: 'GET' } - const res = { statusCode: 200 } - expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true - }) + it('should return false if no root span', () => { + webStub.root.returns(null) + expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false + }) - it('should return true and put request in cache if priority is USER_KEEP', () => { - const rootSpan = { context: () => ({ _sampling: { priority: 1 } }) } - webStub.root.returns(rootSpan) - const req = { url: '/test', method: 'GET' } - const res = { statusCode: 200 } - expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true - }) + it('should return true and put request in cache if priority is AUTO_KEEP or USER_KEEP', () => { + webStub.root.returns(span) + sampler().isSampled.returns(true) + const req = { url: '/test', method: 'GET' } + const res = { statusCode: 200 } + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true + }) - it('should return false if priority is neither AUTO_KEEP nor USER_KEEP', () => { - const rootSpan = { context: () => ({ _sampling: { priority: 0 } }) } - webStub.root.returns(rootSpan) - expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false - }) + it('should return false if priority is neither AUTO_KEEP nor USER_KEEP', () => { + webStub.root.returns(span) + sampler().isSampled.returns(false) + expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false }) - describe('disable', () => { - it('should set enabled to false and clear the cache', () => { - const req = { url: '/test', method: 'GET' } - const res = { statusCode: 200 } + it('should set enabled to false and clear the cache', () => { + const req = { url: '/test', method: 'GET' } + const res = { statusCode: 200 } - const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) } - webStub.root.returns(rootSpan) + webStub.root.returns(span) + sampler().isSampled.returns(true) - apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) - expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true + apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true - apiSecuritySampler.disable() + apiSecuritySampler.disable() - expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false - }) + expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false }) }) diff --git a/packages/dd-trace/test/appsec/index.spec.js b/packages/dd-trace/test/appsec/index.spec.js index c83852750f..13f6b197f0 100644 --- a/packages/dd-trace/test/appsec/index.spec.js +++ b/packages/dd-trace/test/appsec/index.spec.js @@ -46,6 +46,7 @@ describe('AppSec Index', function () { let graphql let apiSecuritySampler let rasp + let sampler const RULES = { rules: [{ a: 1 }] } @@ -78,6 +79,10 @@ describe('AppSec Index', function () { root: sinon.stub() } + sampler = sinon.stub().returns({ + isSampled: sinon.stub() + }) + blocking = { setTemplates: sinon.stub() } @@ -103,7 +108,8 @@ describe('AppSec Index', function () { } apiSecuritySampler = proxyquire('../../src/appsec/api_security_sampler', { - '../plugins/util/web': web + '../plugins/util/web': web, + '../priority_sampler': sampler }) sinon.spy(apiSecuritySampler, 'sampleRequest') @@ -549,9 +555,12 @@ describe('AppSec Index', function () { statusCode: 201 } - const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) } + const span = { + context: sinon.stub().returns({}) + } - web.root.returns(rootSpan) + web.root.returns(span) + sampler().isSampled.returns(true) AppSec.incomingHttpEndTranslator({ req, res })