From 313a4b41ff52f4ff962b8493c5662e5d2002a9f7 Mon Sep 17 00:00:00 2001
From: rochdev <roch.devost@datadoghq.com>
Date: Fri, 25 May 2018 13:23:16 -0400
Subject: [PATCH] fix express res.end() handler not bound to tracer context

---
 src/plugins/express.js       | 21 ++++++++++-----------
 test/plugins/express.spec.js | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/plugins/express.js b/src/plugins/express.js
index 2ae0e7974af..dc8644ac081 100644
--- a/src/plugins/express.js
+++ b/src/plugins/express.js
@@ -24,10 +24,9 @@ function createWrapMethod (tracer, config) {
     }, span => {
       const originalEnd = res.end
 
-      res.end = function () {
-        res.end = originalEnd
-        const returned = res.end.apply(this, arguments)
-        const paths = tracer._context.get('express.paths')
+      res.end = tracer.bind(function () {
+        const returned = originalEnd.apply(this, arguments)
+        const paths = tracer.currentSpan().context()._express_paths
 
         if (paths) {
           span.setTag('resource.name', paths.join(''))
@@ -40,11 +39,11 @@ function createWrapMethod (tracer, config) {
         span.finish()
 
         return returned
-      }
+      })
 
       req._datadog_trace_patched = true
 
-      return next()
+      next()
     })
   }
 
@@ -65,15 +64,15 @@ function createWrapProcessParams (tracer, config) {
   return function wrapProcessParams (processParams) {
     return function processParamsWithTrace (layer, called, req, res, done) {
       const matchers = layer._datadog_matchers
-      let paths = context.get('express.paths') || []
+      const span = context.get('current')
+
+      if (matchers && span) {
+        const paths = span.context()._express_paths || []
 
-      if (matchers) {
         // Try to guess which path actually matched
         for (let i = 0; i < matchers.length; i++) {
           if (matchers[i].test(layer.path)) {
-            paths = paths.concat(matchers[i].path)
-
-            context.set('express.paths', paths)
+            span.context()._express_paths = paths.concat(matchers[i].path)
 
             break
           }
diff --git a/test/plugins/express.spec.js b/test/plugins/express.spec.js
index cc29651a0f7..9038c988e73 100644
--- a/test/plugins/express.spec.js
+++ b/test/plugins/express.spec.js
@@ -240,6 +240,32 @@ describe('Plugin', () => {
         })
       })
 
+      it('should bind the response to the current context', done => {
+        const app = express()
+
+        context.run(() => {
+          const send = context.bind(res => res.status(200).send())
+
+          app.get('/user', (req, res) => {
+            send(res)
+          })
+        })
+
+        getPort().then(port => {
+          agent
+            .use(traces => {
+              expect(traces[0][0]).to.have.property('resource', '/user')
+            })
+            .then(done)
+            .catch(done)
+
+          appListener = app.listen(port, 'localhost', () => {
+            axios.get(`http://localhost:${port}/user`)
+              .catch(done)
+          })
+        })
+      })
+
       it('should bind the next callback to the current context', done => {
         const app = express()