fix AppSec SDK not triggering twice in a row (#5115)

DataDog · Jan 17, 2025 · 015a722 · 015a722
1 parent 3b8a6b9
commit 015a722
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 0 deletions.
diff --git a/packages/dd-trace/src/appsec/waf/waf_context_wrapper.js b/packages/dd-trace/src/appsec/waf/waf_context_wrapper.js
@@ -19,6 +19,7 @@ class WAFContextWrapper {
     this.rulesVersion = rulesVersion
     this.addressesToSkip = new Set()
     this.knownAddresses = knownAddresses
+    this.cachedUserIdActions = new Map()
   }
 
   run ({ persistent, ephemeral }, raspRule) {
@@ -27,6 +28,16 @@ class WAFContextWrapper {
       return
     }
 
+    // SPECIAL CASE FOR USER_ID
+    // TODO: make this universal
+    const userId = persistent?.[addresses.USER_ID] || ephemeral?.[addresses.USER_ID]
+    if (userId) {
+      const cachedAction = this.cachedUserIdActions.get(userId)
+      if (cachedAction) {
+        return cachedAction
+      }
+    }
+
     const payload = {}
     let payloadHasData = false
     const newAddressesToSkip = new Set(this.addressesToSkip)
@@ -79,6 +90,12 @@ class WAFContextWrapper {
 
       const blockTriggered = !!getBlockingAction(result.actions)
 
+      // SPECIAL CASE FOR USER_ID
+      // TODO: make this universal
+      if (userId && ruleTriggered && blockTriggered) {
+        this.setUserIdCache(userId, result)
+      }
+
       Reporter.reportMetrics({
         duration: result.totalRuntime / 1e3,
         durationExt: parseInt(end - start) / 1e3,
@@ -105,6 +122,26 @@ class WAFContextWrapper {
     }
   }
 
+  setUserIdCache (userId, result) {
+    // using old loops for speed
+    for (let i = 0; i < result.events.length; i++) {
+      const event = result.events[i]
+
+      for (let j = 0; j < event?.rule_matches?.length; j++) {
+        const match = event.rule_matches[j]
+
+        for (let k = 0; k < match?.parameters?.length; k++) {
+          const parameter = match.parameters[k]
+
+          if (parameter?.address === addresses.USER_ID) {
+            this.cachedUserIdActions.set(userId, result.actions)
+            return
+          }
+        }
+      }
+    }
+  }
+
   dispose () {
     this.ddwafContext.dispose()
   }

diff --git a/packages/dd-trace/test/appsec/sdk/user_blocking.spec.js b/packages/dd-trace/test/appsec/sdk/user_blocking.spec.js
@@ -227,6 +227,19 @@ describe('user_blocking', () => {
         }).then(done).catch(done)
         axios.get(`http://localhost:${port}/`)
       })
+
+      it('should return true action if userID was matched before with trackUserLoginSuccessEvent()', (done) => {
+        controller = (req, res) => {
+          tracer.appsec.trackUserLoginSuccessEvent({ id: 'blockedUser' })
+          const ret = tracer.appsec.isUserBlocked({ id: 'blockedUser' })
+          expect(ret).to.be.true
+          res.end()
+        }
+        agent.use(traces => {
+          expect(traces[0][0].meta).to.have.property('usr.id', 'blockedUser')
+        }).then(done).catch(done)
+        axios.get(`http://localhost:${port}/`)
+      })
     })
 
     describe('blockRequest', () => {