From 8c083486a0525d58953ed47b6ebb4dfa54033361 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Mon, 11 May 2020 13:49:57 -0700
Subject: [PATCH] fix(launch): handle websocket connect exceptions (#2184)

---
 src/server/chromium.ts | 4 +++-
 src/server/firefox.ts  | 7 ++++---
 src/server/webkit.ts   | 4 +++-
 src/transport.ts       | 9 ++++++++-
 test/launcher.spec.js  | 9 +++++++++
 5 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/server/chromium.ts b/src/server/chromium.ts
index eb8c28fbd5..ac808b4b6f 100644
--- a/src/server/chromium.ts
+++ b/src/server/chromium.ts
@@ -130,7 +130,9 @@ export class Chromium extends AbstractBrowserType<CRBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<CRBrowser> {
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
       return CRBrowser.connect(transport, false, new RootLogger(options.logger), options);
     });
   }
diff --git a/src/server/firefox.ts b/src/server/firefox.ts
index c8a1ff6107..883dd61e67 100644
--- a/src/server/firefox.ts
+++ b/src/server/firefox.ts
@@ -148,9 +148,10 @@ export class Firefox extends AbstractBrowserType<FFBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<FFBrowser> {
-    const logger = new RootLogger(options.logger);
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
-      return FFBrowser.connect(transport, logger, false, options.slowMo);
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
+      return FFBrowser.connect(transport, new RootLogger(options.logger), false, options.slowMo);
     });
   }
 
diff --git a/src/server/webkit.ts b/src/server/webkit.ts
index fd504fedfa..5660d3277a 100644
--- a/src/server/webkit.ts
+++ b/src/server/webkit.ts
@@ -133,7 +133,9 @@ export class WebKit extends AbstractBrowserType<WKBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<WKBrowser> {
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
       return WKBrowser.connect(transport, new RootLogger(options.logger), options.slowMo);
     });
   }
diff --git a/src/transport.ts b/src/transport.ts
index e1349193a1..0f901fa2fd 100644
--- a/src/transport.ts
+++ b/src/transport.ts
@@ -129,7 +129,14 @@ export class WebSocketTransport implements ConnectionTransport {
   static connect<T>(url: string, onopen: (transport: ConnectionTransport) => Promise<T> | T): Promise<T> {
     const transport = new WebSocketTransport(url);
     return new Promise<T>((fulfill, reject) => {
-      transport._ws.addEventListener('open', async () => fulfill(await onopen(transport)));
+      transport._ws.addEventListener('open', async () => {
+        try {
+          fulfill(await onopen(transport));
+        } catch (e) {
+          try { transport._ws.close(); } catch (ignored) {}
+          reject(e);
+        }
+      });
       transport._ws.addEventListener('error', event => reject(new Error('WebSocket error: ' + event.message)));
     });
   }
diff --git a/test/launcher.spec.js b/test/launcher.spec.js
index 46203f8fce..db44884b31 100644
--- a/test/launcher.spec.js
+++ b/test/launcher.spec.js
@@ -290,6 +290,15 @@ describe('browserType.connect', function() {
     await browserServer._checkLeaks();
     await browserServer.close();
   });
+  it.slow()('should handle exceptions during connect', async({browserType, defaultBrowserOptions, server}) => {
+    const browserServer = await browserType.launchServer(defaultBrowserOptions);
+    const e = new Error('Dummy');
+    const __testHookBeforeCreateBrowser = () => { throw e };
+    const error = await browserType.connect({ wsEndpoint: browserServer.wsEndpoint(), __testHookBeforeCreateBrowser }).catch(e => e);
+    await browserServer._checkLeaks();
+    await browserServer.close();
+    expect(error).toBe(e);
+  });
 });
 
 describe('browserType.launchPersistentContext', function() {