fix koa router double patched on every request (#873)

DataDog · Feb 25, 2020 · 5b3da77 · 5b3da77
1 parent b606766
commit 5b3da77
Showing 1 changed file with 20 additions and 45 deletions.
diff --git a/packages/datadog-plugin-koa/src/index.js b/packages/datadog-plugin-koa/src/index.js
@@ -67,35 +67,14 @@ function createWrapRegister (tracer, config) {
   }
 }
 
-function createWrapRoutes (tracer, config) {
-  return function wrapRoutes (routes) {
-    return function routesWithTrace () {
-      const dispatch = routes.apply(this, arguments)
-      const dispatchWithTrace = function (ctx, next) {
-        if (!ctx.router) {
-          let router
-
-          Object.defineProperty(ctx, 'router', {
-            set (value) {
-              router = value
-
-              for (const layer of router.stack) {
-                wrapStack(layer)
-              }
-            },
-
-            get () {
-              return router
-            }
-          })
-        }
-
-        return dispatch.apply(this, arguments)
-      }
+function createWrapRouterUse (tracer, config) {
+  return function wrapUse (use) {
+    return function useWithTrace () {
+      const router = use.apply(this, arguments)
 
-      dispatchWithTrace.router = dispatch.router
+      router.stack.forEach(wrapStack)
 
-      return dispatchWithTrace
+      return router
     }
   }
 }
@@ -104,16 +83,22 @@ function wrapStack (layer) {
   layer.stack = layer.stack.map(middleware => {
     if (typeof middleware !== 'function') return middleware
 
+    middleware = middleware._dd_original || middleware
+
     const wrappedMiddleware = wrapMiddleware(middleware)
 
-    return function (ctx, next) {
+    const handler = function (ctx, next) {
       if (!ctx || !web.active(ctx.req)) return middleware.apply(this, arguments)
 
       web.exitRoute(ctx.req)
       web.enterRoute(ctx.req, layer.path)
 
       return wrappedMiddleware.apply(this, arguments)
     }
+
+    handler._dd_original = middleware
+
+    return handler
   })
 }
 
@@ -164,34 +149,24 @@ module.exports = [
     name: '@koa/router',
     versions: ['>=8'],
     patch (Router, tracer, config) {
-      this.wrap(Router.prototype, 'routes', createWrapRoutes(tracer, config))
-      this.wrap(Router.prototype, 'middleware', createWrapRoutes(tracer, config))
-    },
-    unpatch (Router) {
-      this.unwrap(Router.prototype, 'routes')
-      this.unwrap(Router.prototype, 'middleware')
-    }
-  },
-  {
-    name: 'koa-router',
-    versions: ['>7'],
-    patch (Router, tracer, config) {
-      this.wrap(Router.prototype, 'routes', createWrapRoutes(tracer, config))
-      this.wrap(Router.prototype, 'middleware', createWrapRoutes(tracer, config))
+      this.wrap(Router.prototype, 'register', createWrapRegister(tracer, config))
+      this.wrap(Router.prototype, 'use', createWrapRouterUse(tracer, config))
     },
     unpatch (Router) {
-      this.unwrap(Router.prototype, 'routes')
-      this.unwrap(Router.prototype, 'middleware')
+      this.unwrap(Router.prototype, 'register')
+      this.unwrap(Router.prototype, 'use')
     }
   },
   {
     name: 'koa-router',
-    versions: ['7'],
+    versions: ['>=7'],
     patch (Router, tracer, config) {
       this.wrap(Router.prototype, 'register', createWrapRegister(tracer, config))
+      this.wrap(Router.prototype, 'use', createWrapRouterUse(tracer, config))
     },
     unpatch (Router) {
       this.unwrap(Router.prototype, 'register')
+      this.unwrap(Router.prototype, 'use')
     }
   }
 ]