fix(ws): wrap event remove handler for message event removal (#7052)

* fix(ws): wrap event remove handler for message event removal

The remove handler has to understand what method was originally
wrapped and remove the wrapped method. Otherwise it would still
exist and be triggered.

Fixes: #7025

* fixup!

* chore: simplify code and add TODO comments

The instrumentation should be improved further for error handling
and preventing multiple message listeners creating individual spans.

Author: BridgeAR

packages/datadog-instrumentations/src/find-my-way.js

@@ -9,13 +9,13 @@ function wrapOn (on) {
   return function onWithTrace (method, path, opts) {
     const index = typeof opts === 'function' ? 2 : 3
     const handler = arguments[index]
-    const wrapper = shimmer.wrapFunction(handler, handler => function (req) {
-      routeChannel.publish({ req, route: path })
-
-      return handler.apply(this, arguments)
-    })
 
     if (typeof handler === 'function') {
+      const wrapper = shimmer.wrapFunction(handler, handler => function (req) {
+        routeChannel.publish({ req, route: path })
+
+        return handler.apply(this, arguments)
+      })
       arguments[index] = wrapper
     }
 
@@ -24,6 +24,7 @@ function wrapOn (on) {
 }
 
 addHook({ name: 'find-my-way', versions: ['>=1'] }, Router => {
+  // No need to wrap the off method as it would not contain the handler function.
   shimmer.wrap(Router.prototype, 'on', wrapOn)
 
   return Router

packages/datadog-instrumentations/src/ws.js

@@ -12,6 +12,9 @@ const producerCh = tracingChannel('ws:send')
 const receiverCh = tracingChannel('ws:receive')
 const closeCh = tracingChannel('ws:close')
 const emitCh = channel('tracing:ws:server:connect:emit')
+// TODO: Add a error channel / handle error events properly.
+
+const eventHandlerMap = new WeakMap()
 
 function wrapHandleUpgrade (handleUpgrade) {
   return function () {
@@ -67,24 +70,37 @@ function createWrapEmit (emit) {
 }
 
 function createWrappedHandler (handler) {
-  return function wrappedMessageHandler (data, binary) {
+  return shimmer.wrapFunction(handler, originalHandler => function (data, binary) {
     const byteLength = dataLength(data)
 
     const ctx = { data, binary, socket: this._sender?._socket, byteLength }
 
-    return receiverCh.traceSync(handler, ctx, this, data, binary)
-  }
+    return receiverCh.traceSync(originalHandler, ctx, this, data, binary)
+  })
 }
 
 function wrapListener (originalOn) {
   return function (eventName, handler) {
     if (eventName === 'message') {
-      return originalOn.call(this, eventName, createWrappedHandler(handler))
+      // Prevent multiple wrapping of the same handler in case the user adds the listener multiple times
+      const wrappedHandler = eventHandlerMap.get(handler) ?? createWrappedHandler(handler)
+      eventHandlerMap.set(handler, wrappedHandler)
+      return originalOn.call(this, eventName, wrappedHandler)
     }
     return originalOn.apply(this, arguments)
   }
 }
 
+function removeListener (originalOff) {
+  return function (eventName, handler) {
+    if (eventName === 'message') {
+      const wrappedHandler = eventHandlerMap.get(handler)
+      return originalOff.call(this, eventName, wrappedHandler)
+    }
+    return originalOff.apply(this, arguments)
+  }
+}
+
 function wrapClose (close) {
   return function (code, data) {
     // _closeFrameReceived is set to true when receiver receives a close frame from a peer
@@ -115,22 +131,24 @@ addHook({
   versions: ['>=8.0.0']
 }, ws => {
   shimmer.wrap(ws.prototype, 'send', wrapSend)
-  shimmer.wrap(ws.prototype, 'on', wrapListener)
   shimmer.wrap(ws.prototype, 'close', wrapClose)
+
+  // TODO: Do not wrap these methods. Instead, add a listener to the websocket instance when one is created.
+  // That way it avoids producing too many spans for the same websocket instance and less user code is impacted.
+  shimmer.wrap(ws.prototype, 'on', wrapListener)
+  shimmer.wrap(ws.prototype, 'addListener', wrapListener)
+  shimmer.wrap(ws.prototype, 'off', removeListener)
+  shimmer.wrap(ws.prototype, 'removeListener', removeListener)
+
   return ws
 })
 
-function detectType (data) {
-  if (typeof Blob !== 'undefined' && data instanceof Blob) return 'Blob'
-  if (typeof Buffer !== 'undefined' && Buffer.isBuffer(data)) return 'Buffer'
-  if (typeof data === 'string') return 'string'
-  return 'Unknown'
-}
-
 function dataLength (data) {
-  const type = detectType(data)
-  if (type === 'Blob') return data.size
-  if (type === 'Buffer') return data.length
-  if (type === 'string') return Buffer.byteLength(data)
-  return 0
+  if (typeof data === 'string') {
+    return Buffer.byteLength(data)
+  }
+  if (data instanceof Blob) {
+    return data.size
+  }
+  return data?.length ?? 0
 }

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

@@ -65,18 +65,34 @@ describe('Plugin', () => {
           agent.close({ ritmReset: false, wipe: true })
         })
 
-        it('should do automatic instrumentation', () => {
+        it('should do automatic instrumentation and remove broken handler', () => {
           wsServer.on('connection', (ws) => {
             connectionReceived = true
             ws.send('test message')
           })
 
-          client.on('message', (msg) => {
+          const brokenHandler = () => {
+            throw new Error('broken handler')
+          }
+
+          client.on('message', brokenHandler)
+
+          client.addListener('message', (msg) => {
             assert.strictEqual(msg.toString(), 'test message')
           })
 
-          return agent.assertSomeTraces(traces => {
-            assert.strictEqual(traces[0][0].name, 'web.request')
+          client.off('message', brokenHandler)
+
+          return agent.assertFirstTraceSpan({
+            name: 'websocket.send',
+            type: 'websocket',
+            resource: `websocket /${route}`,
+            service: 'some',
+            parent_id: 0n,
+            error: 0,
+            meta: {
+              'span.kind': 'producer',
+            }
           })
         })
 
@@ -103,7 +119,7 @@ describe('Plugin', () => {
           client.on('error', done)
         })
 
-        it('should instrument message sending', done => {
+        it('should instrument message sending and not double wrap the same handler', done => {
           wsServer.on('connection', ws => {
             connectionReceived = true
             ws.on('message', msg => {
@@ -116,32 +132,49 @@ describe('Plugin', () => {
             client.send('test message')
           })
 
-          client.on('message', (data) => {
+          const brokenHandler = () => {
+            throw new Error('broken handler')
+          }
+
+          client.on('message', brokenHandler)
+
+          const handler = (data) => {
             assert.strictEqual(data.toString(), 'test message')
             done()
-          })
+          }
+
+          client.addListener('message', handler)
+          client.on('message', handler)
+
+          const handlers = client.listeners('message')
+
+          assert.strictEqual(handlers[0].name, brokenHandler.name)
+          assert.strictEqual(handlers[1], handlers[2])
+
+          client.removeListener('message', brokenHandler)
+          client.removeListener('message', handler)
 
           client.on('error', done)
         })
 
-        it('should instrument message receiving', done => {
+        it('should instrument message receiving', () => {
           wsServer.on('connection', (ws) => {
             ws.on('message', (data) => {
               assert.strictEqual(data.toString(), 'test message from client')
             })
           })
-          agent.assertSomeTraces(traces => {
-            assert.strictEqual(traces[0][0].name, 'websocket.receive')
-            assert.strictEqual(traces[0][0].resource, `websocket /${route}`)
-          })
-            .then(done)
-            .catch(done)
 
           client.on('open', () => {
             client.send('test message from client')
           })
 
-          client.on('error', done)
+          return Promise.race([
+            once(client, 'error'),
+            agent.assertFirstTraceSpan({
+              name: 'websocket.receive',
+              resource: `websocket /${route}`
+            })
+          ])
         })
 
         it('should instrument connection close', () => {

packages/dd-trace/test/plugins/agent.js

@@ -5,10 +5,11 @@ const http = require('http')
 const path = require('path')
 const util = require('util')
 
-const msgpack = require('@msgpack/msgpack')
 const bodyParser = require('body-parser')
 const express = require('express')
+const msgpack = require('@msgpack/msgpack')
 const proxyquire = require('proxyquire')
+const semifies = require('semifies')
 
 const { assertObjectContains } = require('../../../../integration-tests/helpers')
 const { storage } = require('../../../datadog-core')
@@ -338,7 +339,7 @@ function runCallbackAgainstTraces (callback, options = {}, handlers) {
         enumerable: true
       })
       // Hack for the information to be fully visible.
-      error.message = util.inspect(error)
+      error.message = util.inspect(error, { depth: null })
       reject(error)
     }
   }, options.timeoutMs || 1000)
@@ -583,8 +584,11 @@ module.exports = {
         try {
           assertObjectContains(traces[0][0], callbackOrExpected)
         } catch (error) {
-          // eslint-disable-next-line no-console
-          console.error('Expected span %o did not match traces:\n%o', callbackOrExpected, traces)
+          // Enrich error with actual and expected traces for Node.js < 22.17.0
+          if (semifies(process.version, '<22.17.0')) {
+            error.actualTraces = util.inspect(traces, { depth: null })
+            error.expectedTraces = util.inspect(callbackOrExpected, { depth: null })
+          }
           throw error
         }
       } else {