chore: use transport for BrowserType.connect (#12196)

This gives us logging, ECONNRESET error handling and proper cleanup.

packages/playwright-core/src/dispatchers/browserTypeDispatcher.ts

@@ -20,12 +20,12 @@ import * as channels from '../protocol/channels';
 import { Dispatcher, DispatcherScope } from './dispatcher';
 import { BrowserContextDispatcher } from './browserContextDispatcher';
 import { CallMetadata } from '../server/instrumentation';
-import WebSocket from 'ws';
 import { JsonPipeDispatcher } from '../dispatchers/jsonPipeDispatcher';
-import { getUserAgent, makeWaitForNextTask } from '../utils/utils';
-import { ManualPromise } from '../utils/async';
+import { getUserAgent } from '../utils/utils';
 import * as socks from '../utils/socksProxy';
 import EventEmitter from 'events';
+import { ProgressController } from '../server/progress';
+import { WebSocketTransport } from '../server/transport';
 
 export class BrowserTypeDispatcher extends Dispatcher<BrowserType, channels.BrowserTypeChannel> implements channels.BrowserTypeChannel {
   _type_BrowserType = true;
@@ -55,54 +55,41 @@ export class BrowserTypeDispatcher extends Dispatcher<BrowserType, channels.Brow
     };
   }
 
-  async connect(params: channels.BrowserTypeConnectParams): Promise<channels.BrowserTypeConnectResult> {
-    const waitForNextTask = params.slowMo
-      ? (cb: () => any) => setTimeout(cb, params.slowMo)
-      : makeWaitForNextTask();
-    const paramsHeaders = Object.assign({ 'User-Agent': getUserAgent() }, params.headers || {});
-    const ws = new WebSocket(params.wsEndpoint, [], {
-      perMessageDeflate: false,
-      maxPayload: 256 * 1024 * 1024, // 256Mb,
-      handshakeTimeout: params.timeout,
-      headers: paramsHeaders,
-      followRedirects: true,
-    });
-    let socksInterceptor: SocksInterceptor | undefined;
-    const pipe = new JsonPipeDispatcher(this._scope);
-    const openPromise = new ManualPromise<{ pipe: JsonPipeDispatcher }>();
-    ws.on('open', () => openPromise.resolve({ pipe }));
-    ws.on('close', () => {
-      socksInterceptor?.cleanup();
-      pipe.wasClosed();
-    });
-    ws.on('error', error => {
-      socksInterceptor?.cleanup();
-      if (openPromise.isDone()) {
-        pipe.wasClosed(error);
-      } else {
-        pipe.dispose();
-        openPromise.reject(error);
-      }
-    });
-    pipe.on('close', () => {
-      socksInterceptor?.cleanup();
-      ws.close();
-    });
-    pipe.on('message', message => ws.send(JSON.stringify(message)));
-    ws.addEventListener('message', event => {
-      waitForNextTask(() => {
-        try {
-          const json = JSON.parse(event.data as string);
-          if (json.method === '__create__' && json.params.type === 'SocksSupport')
-            socksInterceptor = new SocksInterceptor(ws, params.socksProxyRedirectPortForTest, json.params.guid);
-          if (!socksInterceptor?.interceptMessage(json))
+  async connect(params: channels.BrowserTypeConnectParams, metadata: CallMetadata): Promise<channels.BrowserTypeConnectResult> {
+    const controller = new ProgressController(metadata, this._object);
+    controller.setLogName('browser');
+    return await controller.run(async progress => {
+      const paramsHeaders = Object.assign({ 'User-Agent': getUserAgent() }, params.headers || {});
+      const transport = await WebSocketTransport.connect(progress, params.wsEndpoint, paramsHeaders, true);
+      let socksInterceptor: SocksInterceptor | undefined;
+      const pipe = new JsonPipeDispatcher(this._scope);
+      transport.onmessage = json => {
+        if (json.method === '__create__' && json.params.type === 'SocksSupport')
+          socksInterceptor = new SocksInterceptor(transport, params.socksProxyRedirectPortForTest, json.params.guid);
+        if (socksInterceptor?.interceptMessage(json))
+          return;
+        const cb = () => {
+          try {
             pipe.dispatch(json);
-        } catch (e) {
-          ws.close();
-        }
+          } catch (e) {
+            transport.close();
+          }
+        };
+        if (params.slowMo)
+          setTimeout(cb, params.slowMo);
+        else
+          cb();
+      };
+      pipe.on('message', message => {
+        transport.send(message);
       });
-    });
-    return openPromise;
+      transport.onclose = () => {
+        socksInterceptor?.cleanup();
+        pipe.wasClosed();
+      };
+      pipe.on('close', () => transport.close());
+      return { pipe };
+    }, params.timeout || 0);
   }
 }
 
@@ -112,7 +99,7 @@ class SocksInterceptor {
   private _socksSupportObjectGuid: string;
   private _ids = new Set<number>();
 
-  constructor(ws: WebSocket, redirectPortForTest: number | undefined, socksSupportObjectGuid: string) {
+  constructor(transport: WebSocketTransport, redirectPortForTest: number | undefined, socksSupportObjectGuid: string) {
     this._handler = new socks.SocksProxyHandler(redirectPortForTest);
     this._socksSupportObjectGuid = socksSupportObjectGuid;
 
@@ -125,7 +112,7 @@ class SocksInterceptor {
           try {
             const id = --lastId;
             this._ids.add(id);
-            ws.send(JSON.stringify({ id, guid: socksSupportObjectGuid, method: prop, params, metadata: { stack: [], apiName: '', internal: true } }));
+            transport.send({ id, guid: socksSupportObjectGuid, method: prop, params, metadata: { stack: [], apiName: '', internal: true } } as any);
           } catch (e) {
           }
         };

packages/playwright-core/src/server/transport.ts

@@ -52,9 +52,9 @@ export class WebSocketTransport implements ConnectionTransport {
   onclose?: () => void;
   readonly wsEndpoint: string;
 
-  static async connect(progress: Progress, url: string, headers?: { [key: string]: string; }): Promise<WebSocketTransport> {
+  static async connect(progress: Progress, url: string, headers?: { [key: string]: string; }, followRedirects?: boolean): Promise<WebSocketTransport> {
     progress.log(`<ws connecting> ${url}`);
-    const transport = new WebSocketTransport(progress, url, headers);
+    const transport = new WebSocketTransport(progress, url, headers, followRedirects);
     let success = false;
     progress.cleanupWhenAborted(async () => {
       if (!success)
@@ -75,13 +75,14 @@ export class WebSocketTransport implements ConnectionTransport {
     return transport;
   }
 
-  constructor(progress: Progress, url: string, headers?: { [key: string]: string; }) {
+  constructor(progress: Progress, url: string, headers?: { [key: string]: string; }, followRedirects?: boolean) {
     this.wsEndpoint = url;
     this._ws = new WebSocket(url, [], {
       perMessageDeflate: false,
       maxPayload: 256 * 1024 * 1024, // 256Mb,
       handshakeTimeout: progress.timeUntilDeadline(),
-      headers
+      headers,
+      followRedirects,
     });
     this._progress = progress;
     // The 'ws' module in node sometimes sends us multiple messages in a single task.
@@ -102,12 +103,12 @@ export class WebSocketTransport implements ConnectionTransport {
     });
 
     this._ws.addEventListener('close', event => {
-      this._progress && this._progress.log(`<ws disconnected> ${url}`);
+      this._progress && this._progress.log(`<ws disconnected> ${url} code=${event.code} reason=${event.reason}`);
       if (this.onclose)
         this.onclose.call(null);
     });
     // Prevent Error: read ECONNRESET.
-    this._ws.addEventListener('error', () => {});
+    this._ws.addEventListener('error', error => this._progress && this._progress.log(`<ws error> ${error}`));
   }
 
   send(message: ProtocolRequest) {
@@ -120,6 +121,8 @@ export class WebSocketTransport implements ConnectionTransport {
   }
 
   async closeAndWait() {
+    if (this._ws.readyState === WebSocket.CLOSED)
+      return;
     const promise = new Promise(f => this._ws.once('close', f));
     this.close();
     await promise; // Make sure to await the actual disconnect.

tests/browsertype-connect.spec.ts

@@ -107,7 +107,7 @@ test('should timeout in socket while connecting', async ({ browserType, startRem
     wsEndpoint: `ws://localhost:${server.PORT}/ws-slow`,
     timeout: 1000,
   }).catch(e => e);
-  expect(e.message).toContain('browserType.connect: Opening handshake has timed out');
+  expect(e.message).toContain('browserType.connect: Timeout 1000ms exceeded');
 });
 
 test('should timeout in connect while connecting', async ({ browserType, startRemoteServer, server }) => {