fix hapi plugin losing the scope when a payload is sent (#667)

packages/datadog-plugin-hapi/src/index.js

@@ -48,6 +48,32 @@ function createWrapDispatch (tracer, config) {
   }
 }
 
+function createWrapRebuild () {
+  return function wrapRebuild (rebuild) {
+    return function rebuildWithTrace (event) {
+      rebuild.apply(this, arguments)
+
+      this._cycle = this._cycle.map(wrapMiddleware)
+    }
+  }
+}
+
+function createWrapLifecycle () {
+  return function wrapLifecycle (lifecycle) {
+    return function lifecycleWithTrace () {
+      return lifecycle.apply(this, arguments).map(wrapMiddleware)
+    }
+  }
+}
+
+function wrapMiddleware (middleware) {
+  if (typeof middleware !== 'function') return middleware
+
+  return function (request, next) {
+    return web.reactivate(request.raw.req, () => middleware.apply(this, arguments))
+  }
+}
+
 module.exports = [
   {
     name: '@hapi/hapi',
@@ -60,6 +86,17 @@ module.exports = [
       this.unwrap(Request, 'generate')
     }
   },
+  {
+    name: '@hapi/hapi',
+    versions: ['>=17.9'],
+    file: 'lib/route.js',
+    patch (Route, tracer, config) {
+      this.wrap(Route.prototype, 'rebuild', createWrapRebuild(tracer, config))
+    },
+    unpatch (Route) {
+      this.unwrap(Route.prototype, 'rebuild')
+    }
+  },
   {
     name: 'hapi',
     versions: ['>=17.1'],
@@ -93,6 +130,28 @@ module.exports = [
       this.unwrap(Request.prototype, '_execute')
     }
   },
+  {
+    name: 'hapi',
+    versions: ['>=10.4'],
+    file: 'lib/route.js',
+    patch (Route, tracer, config) {
+      this.wrap(Route.prototype, 'rebuild', createWrapRebuild(tracer, config))
+    },
+    unpatch (Route) {
+      this.unwrap(Route.prototype, 'rebuild')
+    }
+  },
+  {
+    name: 'hapi',
+    versions: ['2 - 10.3'],
+    file: 'lib/route.js',
+    patch (Route, tracer, config) {
+      this.wrap(Route.prototype, 'lifecycle', createWrapLifecycle(tracer, config))
+    },
+    unpatch (Route) {
+      this.unwrap(Route.prototype, 'lifecycle')
+    }
+  },
   {
     name: '@hapi/hapi',
     versions: ['>=17.9'],

packages/datadog-plugin-hapi/test/index.spec.js

@@ -117,6 +117,27 @@ describe('Plugin', () => {
           .catch(done)
       })
 
+      it('should run the request handler in the request scope with a payload', done => {
+        server.route({
+          method: 'POST',
+          path: '/user/{id}',
+          handler: (request, h) => {
+            try {
+              expect(tracer.scope().active()).to.not.be.null
+              done()
+            } catch (e) {
+              done(e)
+            }
+
+            return handler(request, h)
+          }
+        })
+
+        axios
+          .post(`http://localhost:${port}/user/123`, {})
+          .catch(done)
+      })
+
       it('should run pre-handlers in the request scope', done => {
         server.route({
           method: 'GET',