fix: Added timeout for browserWS connection

package.json

@@ -93,6 +93,7 @@
     "webpack-cli": "^4.1.0"
   },
   "dependencies": {
+    "@types/ws": "^7.4.0",
     "atob": "^2.1.2",
     "axios": "^0.21.0",
     "boxen": "^4.2.0",

src/controllers/browser.ts

@@ -61,6 +61,7 @@ import { puppeteerConfig } from '../config/puppeteer.config';
 import StealthPlugin = require('puppeteer-extra-plugin-stealth');
 import { auth_InjectToken } from './auth';
 import { useragentOverride } from '../config/WAuserAgente';
+import { WebSocketTransport } from './websocket';
 
 export async function initWhatsapp(
   session: string,
@@ -136,28 +137,24 @@ export async function initBrowser(
 
   let browser = null;
   if (options.browserWS && options.browserWS != '') {
-    const checkConnection = new Promise(async (resolve) => {
-      const timeout = setTimeout(() => {
-        browser = 'connect';
-        resolve();
-      }, 10000);
-
-      puppeteer
+    const transport = await WebSocketTransport.create(
+      options.browserWS,
+      10000
+    ).catch(() => {
+      browser = 'connect';
+    });
+    if (transport) {
+      await puppeteer
         .connect({
-          browserWSEndpoint: options.browserWS,
+          transport: transport,
         })
         .then((e) => {
           browser = e;
         })
         .catch(() => {
           browser = 'connect';
-        })
-        .finally(() => {
-          clearTimeout(timeout);
-          resolve();
         });
-    });
-    await checkConnection;
+    }
   } else {
     await puppeteer
       .launch({

src/controllers/websocket.ts

@@ -0,0 +1,43 @@
+import { ConnectionTransport } from 'puppeteer';
+import * as WebSocket from 'ws';
+
+export class WebSocketTransport implements ConnectionTransport {
+  static create(url: string, timeout?: number): Promise<WebSocketTransport> {
+    return new Promise((resolve, reject) => {
+      const ws = new WebSocket(url, [], {
+        perMessageDeflate: false,
+        maxPayload: 256 * 1024 * 1024, // 256Mb
+        handshakeTimeout: timeout,
+      });
+
+      ws.addEventListener('open', () => resolve(new WebSocketTransport(ws)));
+      ws.addEventListener('error', reject);
+    });
+  }
+
+  private _ws: WebSocket;
+  onmessage?: (message: string) => void;
+  onclose?: () => void;
+
+  constructor(ws: WebSocket) {
+    this._ws = ws;
+    this._ws.addEventListener('message', (event) => {
+      if (this.onmessage) this.onmessage.call(null, event.data);
+    });
+    this._ws.addEventListener('close', () => {
+      if (this.onclose) this.onclose.call(null);
+    });
+    // Silently ignore all errors - we don't know what to do with them.
+    this._ws.addEventListener('error', () => {});
+    this.onmessage = null;
+    this.onclose = null;
+  }
+
+  send(message: string): void {
+    this._ws.send(message);
+  }
+
+  close(): void {
+    this._ws.close();
+  }
+}