Automated Session Tracking (#5060)

Author: simon-id

.github/workflows/instrumentations.yml

@@ -30,6 +30,14 @@ jobs:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - uses: ./.github/actions/plugins/test
 
+  express-session:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: express-session
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: ./.github/actions/plugins/test
+
   multer:
     runs-on: ubuntu-latest
     env:

packages/datadog-instrumentations/src/express-session.js

@@ -0,0 +1,41 @@
+'use strict'
+
+const shimmer = require('../../datadog-shimmer')
+const { channel, addHook } = require('./helpers/instrument')
+
+const sessionMiddlewareFinishCh = channel('datadog:express-session:middleware:finish')
+
+function wrapSessionMiddleware (sessionMiddleware) {
+  return function wrappedSessionMiddleware (req, res, next) {
+    shimmer.wrap(arguments, 2, function wrapNext (next) {
+      return function wrappedNext () {
+        if (sessionMiddlewareFinishCh.hasSubscribers) {
+          const abortController = new AbortController()
+
+          sessionMiddlewareFinishCh.publish({ req, res, sessionId: req.sessionID, abortController })
+
+          if (abortController.signal.aborted) return
+        }
+
+        return next.apply(this, arguments)
+      }
+    })
+
+    return sessionMiddleware.apply(this, arguments)
+  }
+}
+
+function wrapSession (session) {
+  return function wrappedSession () {
+    const sessionMiddleware = session.apply(this, arguments)
+
+    return shimmer.wrapFunction(sessionMiddleware, wrapSessionMiddleware)
+  }
+}
+
+addHook({
+  name: 'express-session',
+  versions: ['>=1.5.0']
+}, session => {
+  return shimmer.wrapFunction(session, wrapSession)
+})

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -47,6 +47,7 @@ module.exports = {
   elasticsearch: () => require('../elasticsearch'),
   express: () => require('../express'),
   'express-mongo-sanitize': () => require('../express-mongo-sanitize'),
+  'express-session': () => require('../express-session'),
   fastify: () => require('../fastify'),
   'find-my-way': () => require('../find-my-way'),
   fs: () => require('../fs'),

packages/datadog-instrumentations/test/express-session.spec.js

@@ -0,0 +1,94 @@
+'use strict'
+
+const { assert } = require('chai')
+const dc = require('dc-polyfill')
+const axios = require('axios')
+const agent = require('../../dd-trace/test/plugins/agent')
+
+withVersions('express-session', 'express-session', version => {
+  describe('express-session instrumentation', () => {
+    const sessionMiddlewareCh = dc.channel('datadog:express-session:middleware:finish')
+    let port, server, subscriberStub, routeHandlerStub
+
+    before(() => {
+      return agent.load(['http'], { client: false })
+    })
+
+    before((done) => {
+      const express = require('../../../versions/express').get()
+      const expressSession = require(`../../../versions/express-session@${version}`).get()
+
+      const app = express()
+
+      app.use(expressSession({
+        secret: 'secret',
+        resave: false,
+        rolling: true,
+        saveUninitialized: true,
+        genid: () => 'sid_123'
+      }))
+
+      app.get('/', (req, res) => {
+        routeHandlerStub()
+
+        res.send('OK')
+      })
+
+      server = app.listen(0, () => {
+        port = server.address().port
+        done()
+      })
+    })
+
+    beforeEach(() => {
+      routeHandlerStub = sinon.stub()
+      subscriberStub = sinon.stub()
+
+      sessionMiddlewareCh.subscribe(subscriberStub)
+    })
+
+    afterEach(() => {
+      sessionMiddlewareCh.unsubscribe(subscriberStub)
+    })
+
+    after(() => {
+      server.close()
+      return agent.close({ ritmReset: false })
+    })
+
+    it('should not do anything when there are no subscribers', async () => {
+      sessionMiddlewareCh.unsubscribe(subscriberStub)
+
+      const res = await axios.get(`http://localhost:${port}/`)
+
+      assert.equal(res.data, 'OK')
+      sinon.assert.notCalled(subscriberStub)
+      sinon.assert.calledOnce(routeHandlerStub)
+    })
+
+    it('should call the subscriber when the middleware is called', async () => {
+      subscriberStub.callsFake(({ sessionId }) => {
+        assert.equal(sessionId, 'sid_123')
+      })
+
+      const res = await axios.get(`http://localhost:${port}/`)
+
+      assert.equal(res.data, 'OK')
+      sinon.assert.calledOnce(subscriberStub)
+      sinon.assert.calledOnce(routeHandlerStub)
+    })
+
+    it('should not call next when the subscriber calls abort()', async () => {
+      subscriberStub.callsFake(({ res, abortController }) => {
+        res.end('BLOCKED')
+        abortController.abort()
+      })
+
+      const res = await axios.get(`http://localhost:${port}/`)
+
+      assert.equal(res.data, 'BLOCKED')
+      sinon.assert.calledOnce(subscriberStub)
+      sinon.assert.notCalled(routeHandlerStub)
+    })
+  })
+})

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

@@ -22,6 +22,7 @@ module.exports = {
 
   USER_ID: 'usr.id',
   USER_LOGIN: 'usr.login',
+  USER_SESSION_ID: 'usr.session_id',
 
   WAF_CONTEXT_PROCESSOR: 'waf.context.processor',
 

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

@@ -15,6 +15,7 @@ module.exports = {
   incomingHttpRequestEnd: dc.channel('dd-trace:incomingHttpRequestEnd'),
   passportVerify: dc.channel('datadog:passport:verify:finish'),
   passportUser: dc.channel('datadog:passport:deserializeUser:finish'),
+  expressSession: dc.channel('datadog:express-session:middleware:finish'),
   queryParser: dc.channel('datadog:query:read:finish'),
   setCookieChannel: dc.channel('datadog:iast:set-cookie'),
   nextBodyParsed: dc.channel('apm:next:body-parsed'),

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

@@ -11,6 +11,7 @@ const {
   incomingHttpRequestEnd,
   passportVerify,
   passportUser,
+  expressSession,
   queryParser,
   nextBodyParsed,
   nextQueryParsed,
@@ -69,6 +70,7 @@ function enable (_config) {
     incomingHttpRequestEnd.subscribe(incomingHttpEndTranslator)
     passportVerify.subscribe(onPassportVerify) // possible optimization: only subscribe if collection mode is enabled
     passportUser.subscribe(onPassportDeserializeUser)
+    expressSession.subscribe(onExpressSession)
     queryParser.subscribe(onRequestQueryParsed)
     nextBodyParsed.subscribe(onRequestBodyParsed)
     nextQueryParsed.subscribe(onRequestQueryParsed)
@@ -213,6 +215,25 @@ function onPassportDeserializeUser ({ user, abortController }) {
   handleResults(results, store.req, store.req.res, rootSpan, abortController)
 }
 
+function onExpressSession ({ req, res, sessionId, abortController }) {
+  const rootSpan = web.root(req)
+  if (!rootSpan) {
+    log.warn('[ASM] No rootSpan found in onExpressSession')
+    return
+  }
+
+  const isSdkCalled = rootSpan.context()._tags['usr.session_id']
+  if (isSdkCalled) return
+
+  const results = waf.run({
+    persistent: {
+      [addresses.USER_SESSION_ID]: sessionId
+    }
+  }, req)
+
+  handleResults(results, req, res, rootSpan, abortController)
+}
+
 function onRequestQueryParsed ({ req, res, query, abortController }) {
   if (!query || typeof query !== 'object') return
 
@@ -327,6 +348,7 @@ function disable () {
   if (incomingHttpRequestEnd.hasSubscribers) incomingHttpRequestEnd.unsubscribe(incomingHttpEndTranslator)
   if (passportVerify.hasSubscribers) passportVerify.unsubscribe(onPassportVerify)
   if (passportUser.hasSubscribers) passportUser.unsubscribe(onPassportDeserializeUser)
+  if (expressSession.hasSubscribers) expressSession.unsubscribe(onExpressSession)
   if (queryParser.hasSubscribers) queryParser.unsubscribe(onRequestQueryParsed)
   if (nextBodyParsed.hasSubscribers) nextBodyParsed.unsubscribe(onRequestBodyParsed)
   if (nextQueryParsed.hasSubscribers) nextQueryParsed.unsubscribe(onRequestQueryParsed)

packages/dd-trace/src/appsec/sdk/set_user.js

@@ -26,11 +26,15 @@ function setUser (tracer, user) {
   setUserTags(user, rootSpan)
   rootSpan.setTag('_dd.appsec.user.collection_mode', 'sdk')
 
-  waf.run({
-    persistent: {
-      [addresses.USER_ID]: '' + user.id
-    }
-  })
+  const persistent = {
+    [addresses.USER_ID]: '' + user.id
+  }
+
+  if (user.session_id && typeof user.session_id === 'string') {
+    persistent[addresses.USER_SESSION_ID] = user.session_id
+  }
+
+  waf.run({ persistent })
 }
 
 module.exports = {

packages/dd-trace/src/remote_config/capabilities.js

@@ -24,6 +24,7 @@ module.exports = {
   APM_TRACING_SAMPLE_RULES: 1n << 29n,
   ASM_AUTO_USER_INSTRUM_MODE: 1n << 31n,
   ASM_ENDPOINT_FINGERPRINT: 1n << 32n,
+  ASM_SESSION_FINGERPRINT: 1n << 33n,
   ASM_NETWORK_FINGERPRINT: 1n << 34n,
   ASM_HEADER_FINGERPRINT: 1n << 35n,
   ASM_RASP_CMDI: 1n << 37n

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

@@ -93,6 +93,7 @@ function enableWafUpdate (appsecConfig) {
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_ENDPOINT_FINGERPRINT, true)
+    rc.updateCapabilities(RemoteConfigCapabilities.ASM_SESSION_FINGERPRINT, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_NETWORK_FINGERPRINT, true)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_HEADER_FINGERPRINT, true)
 
@@ -127,6 +128,7 @@ function disableWafUpdate () {
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_ENDPOINT_FINGERPRINT, false)
+    rc.updateCapabilities(RemoteConfigCapabilities.ASM_SESSION_FINGERPRINT, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_NETWORK_FINGERPRINT, false)
     rc.updateCapabilities(RemoteConfigCapabilities.ASM_HEADER_FINGERPRINT, false)