call sampleRequest in onResponseBody and incomingHttpEndTranslator

DataDog · Oct 9, 2024 · 448d9dd · 448d9dd
1 parent 4ebff24
commit 448d9dd
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 16 deletions.
diff --git a/packages/dd-trace/src/appsec/api_security_sampler.js b/packages/dd-trace/src/appsec/api_security_sampler.js
@@ -2,32 +2,29 @@
 
 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 { USER_KEEP, AUTO_KEEP, AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
+const { AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
 
 const MAX_SIZE = 4096
 const DEFAULT_DELAY = 30 // 30s
 
 let enabled
 let sampledRequests
-let prioritySampler
 
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
   const ttl = parseSampleDelay(apiSecurity.sampleDelay) * 1000
   sampledRequests = new LRUCache({ max: MAX_SIZE, ttl })
-  prioritySampler = new PrioritySampler()
 }
 
 function disable () {
   enabled = false
   sampledRequests?.clear()
 }
 
-function sampleRequest (req, res) {
-  if (!enabled) return false
+function sampleRequest (req, res, forceSample) {
+  if (!enabled || this.isSampled(req, res)) return false
 
   const rootSpan = web.root(req)
   if (!rootSpan) return false
@@ -38,17 +35,15 @@ function sampleRequest (req, res) {
     return false
   }
 
-  if (priority === AUTO_KEEP || priority === USER_KEEP) {
-    return sample(req, res)
+  if (!priority && !rootSpan._prioritySampler?.isSampled(rootSpan)) {
+    return false
   }
 
-  const isSampled = prioritySampler.isSampled(rootSpan)
-
-  if (!isSampled) {
-    return false
+  if (forceSample) {
+    sample(req, res)
   }
 
-  return sample(req, res)
+  return true
 }
 
 function sample (req, res) {
@@ -70,7 +65,7 @@ function isSampled (req, res) {
 function computeKey (req, res) {
   const route = req.route?.path || req.url
   const method = req.method.toLowerCase()
-  const statusCode = res.statusCode
+  const statusCode = res.statusCode === 304 ? 200 : res.statusCode
   const str = route + statusCode + method
   return crypto.createHash('md5').update(str).digest('hex')
 }

diff --git a/packages/dd-trace/src/appsec/index.js b/packages/dd-trace/src/appsec/index.js
@@ -132,7 +132,7 @@ function incomingHttpEndTranslator ({ req, res }) {
     persistent[addresses.HTTP_INCOMING_QUERY] = req.query
   }
 
-  if (apiSecuritySampler.sampleRequest(req, res)) {
+  if (apiSecuritySampler.sampleRequest(req, res, true)) {
     persistent[addresses.WAF_CONTEXT_PROCESSOR] = { 'extract-schema': true }
   }
 
@@ -202,7 +202,7 @@ function onRequestCookieParser ({ req, res, abortController, cookies }) {
 
 function onResponseBody ({ req, res, body }) {
   if (!body || typeof body !== 'object') return
-  if (!apiSecuritySampler.isSampled(req, res)) return
+  if (!apiSecuritySampler.sampleRequest(req, req.res)) return
 
   // we don't support blocking at this point, so no results needed
   waf.run({