use route.path or url to generate the key

packages/datadog-instrumentations/src/express.js

@@ -28,7 +28,7 @@ function wrapResponseJson (json) {
         obj = arguments[1]
       }
 
-      responseJsonChannel.publish({ req: this.req, body: obj })
+      responseJsonChannel.publish({ req: this.req, res: this.res, body: obj })
     }
 
     return json.apply(this, arguments)

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

@@ -5,6 +5,7 @@ const LRUCache = require('lru-cache')
 const PrioritySampler = require('../priority_sampler')
 const web = require('../plugins/util/web')
 const log = require('../log')
+const { USER_KEEP, AUTO_KEEP, AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
 
 const MAX_SIZE = 4096
 const DEFAULT_DELAY = 30 // 30s
@@ -31,12 +32,26 @@ function sampleRequest (req, res) {
   const rootSpan = web.root(req)
   if (!rootSpan) return false
 
+  const priority = getSpanPriority(rootSpan)
+
+  if (priority === AUTO_REJECT || priority === USER_REJECT) {
+    return false
+  }
+
+  if (priority === AUTO_KEEP || priority === USER_KEEP) {
+    return sample(req, res)
+  }
+
   const isSampled = prioritySampler.isSampled(rootSpan)
 
   if (!isSampled) {
     return false
   }
 
+  return sample(req, res)
+}
+
+function sample (req, res) {
   const key = computeKey(req, res)
   const alreadySampled = sampledRequests.has(key)
 
@@ -53,7 +68,7 @@ function isSampled (req, res) {
 }
 
 function computeKey (req, res) {
-  const route = req.route.path
+  const route = req.route?.path || req.url
   const method = req.method.toLowerCase()
   const statusCode = res.statusCode
   const str = route + statusCode + method
@@ -69,6 +84,11 @@ function parseSampleDelay (delay) {
   }
 }
 
+function getSpanPriority (span) {
+  const spanContext = span.context?.()
+  return spanContext._sampling?.priority
+}
+
 module.exports = {
   configure,
   disable,

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

@@ -132,7 +132,7 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
-  if (req.route && typeof req.route.path === 'string' && apiSecuritySampler.sampleRequest(req, res)) {
+  if (apiSecuritySampler.sampleRequest(req, res)) {
     persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
   }
 
@@ -200,8 +200,9 @@ function onRequestCookieParser ({ req, res, abortController, cookies }) {
   handleResults(results, req, res, rootSpan, abortController)
 }
 
-function onResponseBody ({ req, body }) {
+function onResponseBody ({ req, res, body }) {
   if (!body || typeof body !== 'object') return
+  if (!apiSecuritySampler.isSampled(req, res)) return
 
   // we don't support blocking at this point, so no results needed
   waf.run({

packages/dd-trace/test/appsec/api_security_sampler.spec.js

@@ -2,6 +2,7 @@
 
 const { performance } = require('node:perf_hooks')
 const proxyquire = require('proxyquire')
+const { USER_KEEP, AUTO_KEEP, AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
 
 describe('API Security Sampler', () => {
   const req = { route: { path: '/test' }, method: 'GET' }
@@ -25,7 +26,9 @@ describe('API Security Sampler', () => {
     apiSecuritySampler.configure({ apiSecurity: { enabled: true } })
 
     span = {
-      context: sinon.stub().returns({})
+      context: sinon.stub().returns({
+        _sampling: { priority: AUTO_KEEP }
+      })
     }
 
     webStub.root.returns(span)
@@ -49,10 +52,37 @@ describe('API Security Sampler', () => {
     expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
   })
 
-  it('should return true and put request in cache if priority is AUTO_KEEP or USER_KEEP', () => {
+  it('should return false for AUTO_REJECT priority', () => {
+    span.context.returns({ _sampling: { priority: AUTO_REJECT } })
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.false
+  })
+
+  it('should return false for USER_REJECT priority', () => {
+    span.context.returns({ _sampling: { priority: USER_REJECT } })
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.false
+  })
+
+  it('should sample for AUTO_KEEP priority without checking prioritySampler', () => {
+    span.context.returns({ _sampling: { priority: AUTO_KEEP } })
+    sampler().isSampled.returns(false)
     expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
   })
 
+  it('should sample for USER_KEEP priority without checking prioritySampler', () => {
+    span.context.returns({ _sampling: { priority: USER_KEEP } })
+    sampler().isSampled.returns(false)
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
+  })
+
+  it('should use prioritySampler for undefined priority', () => {
+    span.context.returns({ _sampling: { priority: undefined } })
+    sampler().isSampled.returns(true)
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
+
+    sampler().isSampled.returns(false)
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.false
+  })
+
   it('should not sample before 30 seconds', () => {
     expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
     performanceNowStub.returns(performance.now() + 25000)
@@ -69,11 +99,6 @@ describe('API Security Sampler', () => {
     expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
   })
 
-  it('should return false if priority is neither AUTO_KEEP nor USER_KEEP', () => {
-    sampler().isSampled.returns(false)
-    expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
-  })
-
   it('should remove oldest entry when max size is exceeded', () => {
     const method = req.method
     for (let i = 0; i < 4097; i++) {
@@ -120,4 +145,38 @@ describe('API Security Sampler', () => {
     expect(apiSecuritySampler.isSampled(req, res200)).to.be.true
     expect(apiSecuritySampler.isSampled(req, res404)).to.be.true
   })
+
+  it('should use route.path when available', () => {
+    const reqWithRoute = {
+      route: { path: '/users/:id' },
+      url: '/users/123',
+      method: 'GET'
+    }
+
+    apiSecuritySampler.sampleRequest(reqWithRoute, res)
+    expect(apiSecuritySampler.isSampled(reqWithRoute, res)).to.be.true
+
+    const reqWithDifferentUrl = {
+      route: { path: '/users/:id' },
+      url: '/users/456',
+      method: 'GET'
+    }
+    expect(apiSecuritySampler.isSampled(reqWithDifferentUrl, res)).to.be.true
+  })
+
+  it('should fall back to url when route.path is not available', () => {
+    const reqWithUrl = {
+      url: '/users/123',
+      method: 'GET'
+    }
+
+    apiSecuritySampler.sampleRequest(reqWithUrl, res)
+    expect(apiSecuritySampler.isSampled(reqWithUrl, res)).to.be.true
+
+    const reqWithDifferentUrl = {
+      url: '/users/456',
+      method: 'GET'
+    }
+    expect(apiSecuritySampler.isSampled(reqWithDifferentUrl, res)).to.be.false
+  })
 })

packages/dd-trace/test/appsec/index.spec.js

@@ -112,6 +112,7 @@ describe('AppSec Index', function () {
       '../priority_sampler': sampler
     })
     sinon.spy(apiSecuritySampler, 'sampleRequest')
+    sinon.spy(apiSecuritySampler, 'isSampled')
 
     rasp = {
       enable: sinon.stub(),
@@ -588,18 +589,33 @@ describe('AppSec Index', function () {
       })
 
       it('should not do anything if body is not an object', () => {
-        responseBody.publish({ req: {}, body: 'string' })
-        responseBody.publish({ req: {}, body: null })
+        responseBody.publish({ req: {}, res: {}, body: 'string' })
+        responseBody.publish({ req: {}, res: {}, body: null })
 
+        expect(apiSecuritySampler.isSampled).to.not.been.called
         expect(waf.run).to.not.been.called
       })
 
-      it('should call to the waf if body is an object', () => {
+      it('should not call to the waf if it is not a sampled request', () => {
+        apiSecuritySampler.isSampled = apiSecuritySampler.isSampled.instantiateFake(() => false)
         const req = {}
+        const res = {}
+
+        responseBody.publish({ req, res, body: {} })
+
+        expect(apiSecuritySampler.isSampled).to.have.been.calledOnceWith(req)
+        expect(waf.run).to.not.been.called
+      })
+
+      it('should call to the waf if it is a sampled request', () => {
+        apiSecuritySampler.isSampled = apiSecuritySampler.isSampled.instantiateFake(() => true)
+        const req = {}
+        const res = {}
         const body = {}
 
-        responseBody.publish({ req, body })
+        responseBody.publish({ req, res, body })
 
+        expect(apiSecuritySampler.isSampled).to.have.been.calledOnceWith(req, res)
         expect(waf.run).to.been.calledOnceWith({
           persistent: {
             [addresses.HTTP_INCOMING_RESPONSE_BODY]: body