From f28f07ef772082dfae2e00aa9d97872987272855 Mon Sep 17 00:00:00 2001
From: rochdev <roch.devost@datadoghq.com>
Date: Thu, 24 May 2018 10:30:48 -0400
Subject: [PATCH] fix express context being lost in some cases

---
 src/plugins/express.js       |  8 +++++
 test/plugins/express.spec.js | 61 ++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/src/plugins/express.js b/src/plugins/express.js
index 515be34a755..b7a51c3b609 100644
--- a/src/plugins/express.js
+++ b/src/plugins/express.js
@@ -42,6 +42,8 @@ function createWrapMethod (tracer, config) {
         return returned
       }
 
+      req._datadog_trace_patched = true
+
       return next()
     })
   }
@@ -84,6 +86,8 @@ function createWrapProcessParams (tracer, config) {
 }
 
 function createWrapRouterMethod (tracer) {
+  const context = tracer._context
+
   return function wrapRouterMethod (original) {
     return function methodWithTrace (fn) {
       const offset = this.stack.length
@@ -94,6 +98,10 @@ function createWrapRouterMethod (tracer) {
         const handle = layer.handle_request
 
         layer.handle_request = (req, res, next) => {
+          if (req._datadog_trace_patched) {
+            next = context.bind(next)
+          }
+
           return handle.call(layer, req, res, next)
         }
 
diff --git a/test/plugins/express.spec.js b/test/plugins/express.spec.js
index 198b6d3f6bd..7ab186f6df0 100644
--- a/test/plugins/express.spec.js
+++ b/test/plugins/express.spec.js
@@ -240,6 +240,67 @@ describe('Plugin', () => {
         })
       })
 
+      it('should bind the next callback to the current context', done => {
+        const app = express()
+
+        app.use((req, res, next) => {
+          context.run(() => {
+            context.set('foo', 'bar')
+            next()
+          })
+        })
+
+        app.get('/user', (req, res) => {
+          res.status(200).send(context.get('foo'))
+        })
+
+        getPort().then(port => {
+          appListener = app.listen(port, 'localhost', () => {
+            axios.get(`http://localhost:${port}/user`)
+              .then(res => {
+                expect(res.status).to.equal(200)
+                expect(res.data).to.be.empty
+                done()
+              })
+              .catch(done)
+          })
+        })
+      })
+
+      it('should only include paths for routes that matched', done => {
+        const app = express()
+        const router = express.Router()
+
+        router.use('/baz', (req, res, next) => next())
+        router.get('/user/:id', (req, res) => {
+          res.status(200).send()
+        })
+        router.use('/qux', (req, res, next) => next())
+
+        app.use('/foo', (req, res, next) => next())
+        app.use('/app', router)
+        app.use('/bar', (req, res, next) => next())
+
+        getPort().then(port => {
+          agent
+            .use(traces => {
+              expect(traces[0][0]).to.have.property('resource', '/app/user/:id')
+            })
+            .then(done)
+            .catch(done)
+
+          appListener = app.listen(port, 'localhost', () => {
+            axios.get(`http://localhost:${port}/app/user/123`)
+              .then(res => {
+                expect(res.status).to.equal(200)
+                expect(res.data).to.be.empty
+                done()
+              })
+              .catch(done)
+          })
+        })
+      })
+
       it('should extract its parent span from the headers', done => {
         const app = express()