feat: custom websocket support (#1696)

README.md

@@ -237,6 +237,28 @@ use the latest auth token, you must have some outside mechanism running that
 handles application-level authentication refreshing so that the websocket
 connection can simply grab the latest valid token or signed url.
 
+#### Customize Websockets with `createWebsocket` (Websocket Only)
+
+When you need to add a custom websocket subprotocol or header to open a connection
+through a proxy with custom authentication this callback allows you to create your own
+instance of a websocket which will be used in the mqtt client.
+
+```js
+  const createWebsocket = createWebsocket(url, websocketSubProtocols, options) => {
+    const subProtocols = [
+      websocketSubProtocols[0],
+      'myCustomSubprotocolOrOAuthToken',
+    ]
+    return new WebSocket(url, subProtocols)
+  }
+
+    const connection = await mqtt.connectAsync(<wss url>, {
+      ...,
+      createWebsocket: createWebsocket,
+    });
+```
+
+
 #### Enabling Reconnection with `reconnectPeriod` option
 
 To ensure that the mqtt client automatically tries to reconnect when the
@@ -429,6 +451,8 @@ The arguments are:
   - `transformWsUrl` : optional `(url, options, client) => url` function
     For ws/wss protocols only. Can be used to implement signing
     urls which upon reconnect can have become expired.
+  - `createWebsocket` : optional `url, websocketSubProtocols, options) => Websocket` function
+      For ws/wss protocols only. Can be used to implement a custom websocket subprotocol or implementation.
   - `resubscribe` : if connection is broken and reconnects,
     subscribed topics are automatically subscribed again (default `true`)
   - `messageIdProvider`: custom messageId provider. when `new UniqueMessageIdProvider()` is set, then non conflict messageId is provided.

src/lib/client.ts

@@ -218,6 +218,13 @@ export interface IClientOptions extends ISecureClientOptions {
 		client: MqttClient,
 	) => string
 
+	/** when defined this function will be called to create the Websocket instance, used to add custom protocols or websocket implementations */
+	createWebsocket?: (
+		url: string,
+		websocketSubProtocols: string[],
+		options: IClientOptions,
+	) => any
+
 	/** Custom message id provider */
 	messageIdProvider?: IMessageIdProvider
 

src/lib/connect/ws.ts

@@ -108,11 +108,16 @@ function createWebSocket(
 	debug(
 		`creating new Websocket for url: ${url} and protocol: ${websocketSubProtocol}`,
 	)
-	const socket = new WS(
-		url,
-		[websocketSubProtocol],
-		opts.wsOptions as ClientOptions,
-	)
+	let socket: WS
+	if (opts.createWebsocket) {
+		socket = opts.createWebsocket(url, [websocketSubProtocol], opts)
+	} else {
+		socket = new WS(
+			url,
+			[websocketSubProtocol],
+			opts.wsOptions as ClientOptions,
+		)
+	}
 	return socket
 }
 
@@ -123,7 +128,12 @@ function createBrowserWebSocket(client: MqttClient, opts: IClientOptions) {
 			: 'mqtt'
 
 	const url = buildUrl(opts, client)
-	const socket = new WebSocket(url, [websocketSubProtocol])
+	let socket: WebSocket
+	if (opts.createWebsocket) {
+		socket = opts.createWebsocket(url, [websocketSubProtocol], opts)
+	} else {
+		socket = new WebSocket(url, [websocketSubProtocol])
+	}
 	socket.binaryType = 'arraybuffer'
 	return socket
 }

test/websocket_client.ts

@@ -10,10 +10,14 @@ import { IClientOptions } from '../src/lib/client'
 
 const port = 9999
 const httpServer = http.createServer()
-
+let lastProcotols = new Set<string>()
 function attachWebsocketServer(httpServer2) {
 	const webSocketServer = new WebSocket.Server({
 		server: httpServer2,
+		handleProtocols: (protocols: Set<string>, request: any) => {
+			lastProcotols = protocols
+			return [...protocols][0]
+		},
 		perMessageDeflate: false,
 	})
 
@@ -132,6 +136,34 @@ describe('Websocket Client', () => {
 		})
 	})
 
+	it('should be able to create custom Websocket instance', function test(done) {
+		const baseUrl = 'ws://localhost:9999/mqtt'
+		let urlInCallback: string
+		const opts = makeOptions({
+			path: '/mqtt',
+			createWebsocket(
+				url: string,
+				websocketSubProtocols: string[],
+				options: IClientOptions,
+			) {
+				urlInCallback = url
+				assert.equal(url, baseUrl)
+				const subProtocols = [
+					websocketSubProtocols[0],
+					'myCustomSubprotocol',
+				]
+				return new WebSocket(url, subProtocols)
+			},
+		})
+		const client = mqtt.connect(opts)
+		client.on('connect', () => {
+			assert.equal((client.stream as any).url, urlInCallback)
+			assert.equal(baseUrl, urlInCallback)
+			assert.equal('myCustomSubprotocol', [...lastProcotols][1])
+			client.end(true, (err) => done(err))
+		})
+	})
+
 	it('should use mqttv3.1 as the protocol if using v3.1', function test(done) {
 		httpServer.once('client', (client) => {
 			assert.strictEqual(client.protocol, 'mqttv3.1')