fix(ws): handle socket availability gracefully (#6728)

This allows to receive the original error in case e.g., close is
called before the connection could be established.

Fixes: #6624

Author: BridgeAR

packages/datadog-instrumentations/src/ws.js

@@ -41,7 +41,7 @@ function wrapSend (send) {
 
     const [data, options, cb] = arguments
 
-    const ctx = { data, socket: this._sender._socket }
+    const ctx = { data, socket: this._sender?._socket }
 
     return typeof cb === 'function'
       ? producerCh.traceCallback(send, undefined, ctx, this, data, options, cb)
@@ -70,7 +70,7 @@ function createWrappedHandler (handler) {
   return function wrappedMessageHandler (data, binary) {
     const byteLength = dataLength(data)
 
-    const ctx = { data, binary, socket: this._sender._socket, byteLength }
+    const ctx = { data, binary, socket: this._sender?._socket, byteLength }
 
     return receiverCh.traceSync(handler, ctx, this, data, binary)
   }
@@ -93,7 +93,7 @@ function wrapClose (close) {
     // if both are true then the self is sending the close event
     const isPeerClose = this._closeFrameReceived === true && this._closeFrameSent === false
 
-    const ctx = { code, data, socket: this._sender._socket, isPeerClose }
+    const ctx = { code, data, socket: this._sender?._socket, isPeerClose }
 
     return closeCh.traceSync(close, ctx, this, ...arguments)
   }

packages/datadog-plugin-ws/src/close.js

@@ -17,7 +17,7 @@ class WSClosePlugin extends TracingPlugin {
     if (!traceWebsocketMessagesEnabled) return
 
     const { code, data, socket, isPeerClose } = ctx
-    if (!socket.spanContext) return
+    if (!socket?.spanContext) return
 
     const spanKind = isPeerClose ? 'consumer' : 'producer'
     const spanTags = socket.spanContext.spanTags

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

@@ -1,6 +1,11 @@
 'use strict'
 
+const assert = require('node:assert')
+const { once } = require('node:events')
+
 const { expect } = require('chai')
+const { describe, before, after, it } = require('mocha')
+
 const { withVersions } = require('../../dd-trace/test/setup/mocha')
 const agent = require('../../dd-trace/test/plugins/agent')
 
@@ -15,6 +20,32 @@ describe('Plugin', () => {
 
   describe('ws', () => {
     withVersions('ws', 'ws', '>=8.0.0', version => {
+      describe('regression tests', () => {
+        before(async () => {
+          await agent.load(['ws'], [{
+            service: 'some',
+            traceWebsocketMessagesEnabled: true
+          }])
+          WebSocket = require(`../../../versions/ws@${version}`).get()
+        })
+
+        it('should emit original error in case close is called before connection is established', async () => {
+          const socket = new WebSocket('wss://localhost:12345')
+
+          const errorPromise = once(socket, 'error')
+          socket.close()
+
+          const error = await errorPromise
+
+          // Some versions emit an array with an error, some directly emit the error
+          assert.strictEqual(error?.[0]?.message, 'WebSocket was closed before the connection was established')
+        })
+
+        after(async () => {
+          agent.close({ ritmReset: false, wipe: true })
+        })
+      })
+
       describe('when using WebSocket', () => {
         route = 'test'
         beforeEach(async () => {