From e09945c637d7ca8ee7ef9ce0e288584c6087d79e Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 29 Sep 2022 13:04:19 -0800
Subject: [PATCH] chore: handle unexpected ws response (#17714)

---
 packages/playwright-core/src/server/transport.ts | 16 ++++++++++++++--
 tests/library/browsertype-connect.spec.ts        |  7 +++++++
 tests/library/browsertype-launch-server.spec.ts  |  2 +-
 utils/testserver/index.ts                        |  5 +++++
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/packages/playwright-core/src/server/transport.ts b/packages/playwright-core/src/server/transport.ts
index 3a8ddca591..b3766acba7 100644
--- a/packages/playwright-core/src/server/transport.ts
+++ b/packages/playwright-core/src/server/transport.ts
@@ -17,6 +17,7 @@
 
 import { ws } from '../utilsBundle';
 import type { WebSocket } from '../utilsBundle';
+import type { ClientRequest, IncomingMessage } from 'http';
 import type { Progress } from './progress';
 import { makeWaitForNextTask } from '../utils';
 
@@ -62,15 +63,26 @@ export class WebSocketTransport implements ConnectionTransport {
         await transport.closeAndWait().catch(e => null);
     });
     await new Promise<WebSocketTransport>((fulfill, reject) => {
-      transport._ws.addEventListener('open', async () => {
+      transport._ws.on('open', async () => {
         progress.log(`<ws connected> ${url}`);
         fulfill(transport);
       });
-      transport._ws.addEventListener('error', event => {
+      transport._ws.on('error', event => {
         progress.log(`<ws connect error> ${url} ${event.message}`);
         reject(new Error('WebSocket error: ' + event.message));
         transport._ws.close();
       });
+      transport._ws.on('unexpected-response', (request: ClientRequest, response: IncomingMessage) => {
+        const chunks: Buffer[] = [];
+        const errorPrefix = `${url} ${response.statusCode} ${response.statusMessage}`;
+        response.on('data', chunk => chunks.push(chunk));
+        response.on('close', () => {
+          const error = chunks.length ? `${errorPrefix}\n${Buffer.concat(chunks)}` : errorPrefix;
+          progress.log(`<ws unexpected response> ${error}`);
+          reject(new Error('WebSocket error: ' + error));
+          transport._ws.close();
+        });
+      });
     });
     success = true;
     return transport;
diff --git a/tests/library/browsertype-connect.spec.ts b/tests/library/browsertype-connect.spec.ts
index 70f9320424..379713445d 100644
--- a/tests/library/browsertype-connect.spec.ts
+++ b/tests/library/browsertype-connect.spec.ts
@@ -62,6 +62,13 @@ test('should connect over wss', async ({ browserType, startRemoteServer, httpsSe
   }
 });
 
+test('should print HTTP error', async ({ browserType, server, mode }) => {
+  test.skip(mode !== 'default'); // Out of process transport does not allow us to set env vars dynamically.
+  const error = await browserType.connect(`ws://localhost:${server.PORT}/ws-401`).catch(e => e);
+  expect(error.message).toContain('401');
+  expect(error.message).toContain('Unauthorized body');
+});
+
 test('should be able to reconnect to a browser', async ({ browserType, startRemoteServer, server }) => {
   const remoteServer = await startRemoteServer();
   {
diff --git a/tests/library/browsertype-launch-server.spec.ts b/tests/library/browsertype-launch-server.spec.ts
index 3b1ecb8644..dcb69d614b 100644
--- a/tests/library/browsertype-launch-server.spec.ts
+++ b/tests/library/browsertype-launch-server.spec.ts
@@ -57,7 +57,7 @@ it.describe('launch server', () => {
     const browserServer = await browserType.launchServer();
     const error = await browserType.connect({ wsEndpoint: browserServer.wsEndpoint() + '-foo' }).catch(e => e);
     await browserServer.close();
-    expect(error.message).toContain('Unexpected server response: 400');
+    expect(error.message).toContain('400 Bad Request');
   });
 
   it('should fire "close" event during kill', async ({ browserType }) => {
diff --git a/utils/testserver/index.ts b/utils/testserver/index.ts
index 1f0d5bfa7a..5fac659ffc 100644
--- a/utils/testserver/index.ts
+++ b/utils/testserver/index.ts
@@ -75,6 +75,11 @@ export class TestServer {
     this._wsServer = new ws.WebSocketServer({ noServer: true });
     this._server.on('upgrade', async (request, socket, head) => {
       const pathname = url.parse(request.url!).path;
+      if (pathname === '/ws-401') {
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\nUnauthorized body');
+        socket.destroy();
+        return;
+      }
       if (pathname === '/ws-slow')
         await new Promise(f => setTimeout(f, 2000));
       if (!['/ws', '/ws-slow'].includes(pathname)) {