From 93ba729918db6d656e4d082236d237216f330037 Mon Sep 17 00:00:00 2001 From: ishabi Date: Mon, 7 Oct 2024 09:24:31 +0200 Subject: [PATCH] use lru cache package --- .../src/appsec/api_security_sampler.js | 42 ++++++++++-- .../src/appsec/api_security_sampler_cache.js | 53 --------------- .../test/appsec/api_security_sampler.spec.js | 56 +++++++++++----- .../appsec/api_security_sampler_cache.spec.js | 65 ------------------- 4 files changed, 76 insertions(+), 140 deletions(-) delete mode 100644 packages/dd-trace/src/appsec/api_security_sampler_cache.js delete mode 100644 packages/dd-trace/test/appsec/api_security_sampler_cache.spec.js diff --git a/packages/dd-trace/src/appsec/api_security_sampler.js b/packages/dd-trace/src/appsec/api_security_sampler.js index 8c4b8fdc22..d53bd1d85a 100644 --- a/packages/dd-trace/src/appsec/api_security_sampler.js +++ b/packages/dd-trace/src/appsec/api_security_sampler.js @@ -1,16 +1,23 @@ 'use strict' -const ApiSecuritySamplerCache = require('./api_security_sampler_cache') +const crypto = require('node:crypto') +const LRUCache = require('lru-cache') const PrioritySampler = require('../priority_sampler') const web = require('../plugins/util/web') +const log = require('../log') + +const MAX_SIZE = 4096 +const DEFAULT_DELAY = 30 // 30s let enabled let sampledRequests -const prioritySampler = new PrioritySampler() +let prioritySampler function configure ({ apiSecurity }) { enabled = apiSecurity.enabled - sampledRequests = new ApiSecuritySamplerCache(apiSecurity.sampleDelay) + const ttl = parseSampleDelay(apiSecurity.sampleDelay) * 1000 + sampledRequests = new LRUCache({ max: MAX_SIZE, ttl }) + prioritySampler = new PrioritySampler() } function disable () { @@ -30,8 +37,8 @@ function sampleRequest (req, res) { return false } - const key = sampledRequests.computeKey(req, res) - const alreadySampled = sampledRequests.isSampled(key) + const key = computeKey(req, res) + const alreadySampled = sampledRequests.has(key) if (alreadySampled) return false @@ -40,8 +47,31 @@ function sampleRequest (req, res) { return true } +function isSampled (req, res) { + const key = computeKey(req, res) + return !!sampledRequests.has(key) +} + +function computeKey (req, res) { + const route = req.url + const method = req.method.toLowerCase() + const statusCode = res.statusCode + const str = route + statusCode + method + return crypto.createHash('md5').update(str).digest('hex') +} + +function parseSampleDelay (delay) { + if (typeof delay === 'number' && Number.isFinite(delay) && delay > 0) { + return delay + } else { + log.warn('Invalid delay value. Delay must be a positive number.') + return DEFAULT_DELAY + } +} + module.exports = { configure, disable, - sampleRequest + sampleRequest, + isSampled } diff --git a/packages/dd-trace/src/appsec/api_security_sampler_cache.js b/packages/dd-trace/src/appsec/api_security_sampler_cache.js deleted file mode 100644 index 4c328c4342..0000000000 --- a/packages/dd-trace/src/appsec/api_security_sampler_cache.js +++ /dev/null @@ -1,53 +0,0 @@ -'use strict' - -const crypto = require('node:crypto') -const log = require('../log') - -const MAX_SIZE = 4096 -const DEFAULT_DELAY = 30 // 30s - -class ApiSecuritySamplerCache extends Map { - constructor (delay) { - super() - this.delay = this._parseSampleDelay(delay) - } - - _parseSampleDelay (delay) { - if (typeof delay === 'number' && Number.isFinite(delay) && delay > 0) { - return delay - } else { - log.warn('Invalid delay value. Delay must be a positive number.') - return DEFAULT_DELAY - } - } - - computeKey (req, res) { - const route = req.url - const method = req.method.toLowerCase() - const statusCode = res.statusCode - const str = route + statusCode + method - return crypto.createHash('md5').update(str).digest('hex') - } - - isSampled (key) { - if (!super.has(key)) { - return false - } - const previous = super.get(key) - return Date.now() - previous < (this.delay * 1000) - } - - set (key) { - if (super.has(key)) { - super.delete(key) - } - - super.set(key, Date.now()) - if (super.size > MAX_SIZE) { - const oldestKey = super.keys().next().value - super.delete(oldestKey) - } - } -} - -module.exports = ApiSecuritySamplerCache 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 abcee98f3a..f2d0b0da7f 100644 --- a/packages/dd-trace/test/appsec/api_security_sampler.spec.js +++ b/packages/dd-trace/test/appsec/api_security_sampler.spec.js @@ -1,13 +1,17 @@ 'use strict' +const { performance } = require('node:perf_hooks') const proxyquire = require('proxyquire') describe('API Security Sampler', () => { - let apiSecuritySampler, webStub, clock, sampler, span + const req = { url: '/test', method: 'GET' } + const res = { statusCode: 200 } + let apiSecuritySampler, performanceNowStub, webStub, sampler, span beforeEach(() => { + performanceNowStub = sinon.stub(performance, 'now').returns(0) + webStub = { root: sinon.stub() } - clock = sinon.useFakeTimers(Date.now()) sampler = sinon.stub().returns({ isSampled: sinon.stub() @@ -18,15 +22,21 @@ describe('API Security Sampler', () => { '../priority_sampler': sampler }) - apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) + apiSecuritySampler.configure({ apiSecurity: { enabled: true } }) span = { context: sinon.stub().returns({}) } + + webStub.root.returns(span) + sampler().isSampled.returns(true) + + performanceNowStub.returns(performance.now() + 1) }) afterEach(() => { - clock.restore() + performanceNowStub.restore() + apiSecuritySampler.disable() }) it('should return false if not enabled', () => { @@ -40,31 +50,45 @@ describe('API Security Sampler', () => { }) 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 not sample before 30 seconds', () => { + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true + performanceNowStub.returns(performance.now() + 25000) + + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.false + expect(apiSecuritySampler.isSampled(req, res)).to.be.true + }) + + it('should sample after 30 seconds', () => { + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true + + performanceNowStub.returns(performance.now() + 35000) + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true }) 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 }) - it('should set enabled to false and clear the cache', () => { - const req = { url: '/test', method: 'GET' } - const res = { statusCode: 200 } + it('should remove oldest entry when max size is exceeded', () => { + const method = req.method + for (let i = 0; i < 4097; i++) { + expect(apiSecuritySampler.sampleRequest({ method, url: `/test${i}` }, res)).to.be.true + } - webStub.root.returns(span) - sampler().isSampled.returns(true) + expect(apiSecuritySampler.isSampled({ method, url: '/test0' }, res)).to.be.false + expect(apiSecuritySampler.isSampled({ method, url: '/test4096' }, res)).to.be.true + }) - apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } }) + it('should set enabled to false and clear the cache', () => { expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true apiSecuritySampler.disable() - expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false + expect(apiSecuritySampler.sampleRequest(req, res)).to.be.false }) }) diff --git a/packages/dd-trace/test/appsec/api_security_sampler_cache.spec.js b/packages/dd-trace/test/appsec/api_security_sampler_cache.spec.js deleted file mode 100644 index 868705cd45..0000000000 --- a/packages/dd-trace/test/appsec/api_security_sampler_cache.spec.js +++ /dev/null @@ -1,65 +0,0 @@ -'use strict' - -const ApiSecuritySamplerCache = require('../../src/appsec/api_security_sampler_cache') - -describe('ApiSecuritySamplerCache', () => { - let cache - let clock - - beforeEach(() => { - cache = new ApiSecuritySamplerCache() - clock = sinon.useFakeTimers(Date.now()) - }) - - afterEach(() => { - clock.restore() - }) - - describe('Standard behavior with default delay', () => { - it('should not sample when first seen', () => { - const req = { url: '/test', method: 'GET' } - const res = { status_code: 200 } - const key = cache.computeKey(req, res) - expect(cache.isSampled(key)).to.be.false - }) - - it('should sample within 30 seconds of first seen', () => { - const req = { url: '/test', method: 'GET' } - const res = { status_code: 200 } - const key = cache.computeKey(req, res) - cache.set(key) - clock.tick(29000) - expect(cache.isSampled(key)).to.be.true - }) - - it('should not sample after 30 seconds', () => { - const req = { url: '/test', method: 'GET' } - const res = { status_code: 200 } - const key = cache.computeKey(req, res) - cache.set(key) - clock.tick(31000) - expect(cache.isSampled(key)).to.be.false - }) - }) - - describe('Max size behavior', () => { - it('should remove oldest entry when max size is exceeded', () => { - const baseReq = { method: 'GET' } - const baseRes = { status_code: 200 } - - for (let i = 0; i < 4097; i++) { - const req = { ...baseReq, url: `test${i}` } - const key = cache.computeKey(req, baseRes) - cache.set(key) - } - - expect(cache.size).to.equal(4096) - - const firstKey = cache.computeKey({ ...baseReq, url: 'test0' }, baseRes) - expect(cache.isSampled(firstKey)).to.be.false - - const lastKey = cache.computeKey({ ...baseReq, url: 'test4096' }, baseRes) - expect(cache.isSampled(lastKey)).to.be.true - }) - }) -})