use lru cache package

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
 }

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
   })
 })