chore: implement browserType.connect on the driver side (#8566)

src/client/browserType.ts

@@ -19,15 +19,13 @@ import { Browser } from './browser';
 import { BrowserContext, prepareBrowserContextParams } from './browserContext';
 import { ChannelOwner } from './channelOwner';
 import { LaunchOptions, LaunchServerOptions, ConnectOptions, LaunchPersistentContextOptions, BrowserContextOptions } from './types';
-import WebSocket from 'ws';
 import { Connection } from './connection';
 import { Events } from './events';
-import { TimeoutSettings } from '../utils/timeoutSettings';
 import { ChildProcess } from 'child_process';
 import { envObjectToArray } from './clientHelper';
-import { assert, headersObjectToArray, makeWaitForNextTask, getUserAgent } from '../utils/utils';
-import { kBrowserClosedError } from '../utils/errors';
+import { assert, headersObjectToArray, getUserAgent, ManualPromise } from '../utils/utils';
 import * as api from '../../types/types';
+import { kBrowserClosedError } from '../utils/errors';
 
 export interface BrowserServerLauncher {
   launchServer(options?: LaunchServerOptions): Promise<api.BrowserServer>;
@@ -42,7 +40,6 @@ export interface BrowserServer extends api.BrowserServer {
 }
 
 export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, channels.BrowserTypeInitializer> implements api.BrowserType {
-  private _timeoutSettings = new TimeoutSettings();
   _serverLauncher?: BrowserServerLauncher;
   _contexts = new Set<BrowserContext>();
 
@@ -126,48 +123,27 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
     assert(optionsOrWsEndpoint.wsEndpoint, 'options.wsEndpoint is required');
     return this._connect(optionsOrWsEndpoint.wsEndpoint, optionsOrWsEndpoint);
   }
+
   async _connect(wsEndpoint: string, params: Partial<ConnectOptions> = {}): Promise<Browser> {
     const logger = params.logger;
-    const paramsHeaders = Object.assign({'User-Agent': getUserAgent()}, params.headers);
-    return this._wrapApiCall(async () => {
-      const ws = new WebSocket(wsEndpoint, [], {
-        perMessageDeflate: false,
-        maxPayload: 256 * 1024 * 1024, // 256Mb,
-        handshakeTimeout: this._timeoutSettings.timeout(params),
-        headers: paramsHeaders,
-      });
-      const connection = new Connection(() => ws.close());
+    return await this._wrapApiCall(async (channel: channels.BrowserTypeChannel) => {
+      const timeoutPromise = new ManualPromise<Browser>();
+      const timer = params.timeout ? setTimeout(() => timeoutPromise.reject(new Error(`Timeout ${params.timeout}ms exceeded.`)), params.timeout) : undefined;
 
-      // The 'ws' module in node sometimes sends us multiple messages in a single task.
-      const waitForNextTask = params.slowMo
-        ? (cb: () => any) => setTimeout(cb, params.slowMo)
-        : makeWaitForNextTask();
-      connection.onmessage = message => {
-        // Connection should handle all outgoing message in disconnected().
-        if (ws.readyState !== WebSocket.OPEN)
-          return;
-        ws.send(JSON.stringify(message));
-      };
-      ws.addEventListener('message', event => {
-        waitForNextTask(() => {
-          try {
-            // Since we may slow down the messages, but disconnect
-            // synchronously, we might come here with a message
-            // after disconnect.
-            if (!connection.isDisconnected())
-              connection.dispatch(JSON.parse(event.data));
-          } catch (e) {
-            console.error(`Playwright: Connection dispatch error`);
-            console.error(e);
-            ws.close();
-          }
-        });
+      const { pipe } = await channel.connect({ wsEndpoint, headers: params.headers, timeout: params.timeout });
+      const connection = new Connection(() => pipe.close().catch(() => {}));
+      connection.onmessage = message => pipe.send({ message }).catch(() => { });
+      pipe.on('message', ({ message }) => {
+        try {
+          if (!connection!.isDisconnected())
+            connection!.dispatch(message);
+        } catch (e) {
+          console.error(`Playwright: Connection dispatch error`);
+          console.error(e);
+          pipe.close().catch(() => {});
+        }
       });
 
-      let timeoutCallback = (e: Error) => {};
-      const timeoutPromise = new Promise<Browser>((f, r) => timeoutCallback = r);
-      const timer = params.timeout ? setTimeout(() => timeoutCallback(new Error(`Timeout ${params.timeout}ms exceeded.`)), params.timeout) : undefined;
-
       const successPromise = new Promise<Browser>(async (fulfill, reject) => {
         if ((params as any).__testHookBeforeCreateBrowser) {
           try {
@@ -176,16 +152,17 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
             reject(e);
           }
         }
-        ws.addEventListener('open', async () => {
-          const prematureCloseListener = (event: { code: number, reason: string }) => {
-            reject(new Error(`WebSocket server disconnected (${event.code}) ${event.reason}`));
-          };
-          ws.addEventListener('close', prematureCloseListener);
-          const playwright = await connection.initializePlaywright();
+        const prematureCloseListener = (params: { error?: channels.SerializedError }) => {
+          reject(new Error(`WebSocket server disconnected ${params.error!.error?.message}`));
+        };
+        pipe.on('closed', prematureCloseListener);
+
+        pipe.on('opened', async () => {
+          const playwright = await connection!.initializePlaywright();
 
           if (!playwright._initializer.preLaunchedBrowser) {
             reject(new Error('Malformed endpoint. Did you use launchServer method?'));
-            ws.close();
+            pipe.close().catch(() => {});
             return;
           }
 
@@ -193,7 +170,7 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
           browser._logger = logger;
           browser._remoteType = 'owns-connection';
           browser._setBrowserType((playwright as any)[browser._name]);
-          const closeListener = () => {
+          const closeListener = (param: { error?: Error }) => {
             // Emulate all pages, contexts and the browser closing upon disconnect.
             for (const context of browser.contexts()) {
               for (const page of context.pages())
@@ -202,20 +179,18 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
             }
             browser._didClose();
             connection.didDisconnect(kBrowserClosedError);
+            if (param.error)
+              reject(new Error(param.error + '. Most likely ws endpoint is incorrect'));
           };
-          ws.removeEventListener('close', prematureCloseListener);
-          ws.addEventListener('close', closeListener);
+          pipe.off('closed', prematureCloseListener);
+          pipe.on('closed', closeListener);
           browser.on(Events.Browser.Disconnected, () => {
             playwright._cleanup();
-            ws.removeEventListener('close', closeListener);
-            ws.close();
+            pipe.off('closed', closeListener);
+            pipe.close().catch(() => {});
           });
           fulfill(browser);
         });
-        ws.addEventListener('error', event => {
-          ws.close();
-          reject(new Error(event.message + '. Most likely ws endpoint is incorrect'));
-        });
       });
 
       try {

src/client/connection.ts

@@ -38,6 +38,7 @@ import { Android, AndroidSocket, AndroidDevice } from './android';
 import { ParsedStackTrace } from '../utils/stackTrace';
 import { Artifact } from './artifact';
 import { EventEmitter } from 'events';
+import { JsonPipe } from './jsonPipe';
 
 class Root extends ChannelOwner<channels.RootChannel, {}> {
   constructor(connection: Connection) {
@@ -129,7 +130,7 @@ export class Connection extends EventEmitter {
     const object = this._objects.get(guid);
     if (!object)
       throw new Error(`Cannot find object to emit "${method}": ${guid}`);
-    object._channel.emit(method, this._replaceGuidsWithChannels(params));
+    object._channel.emit(method, object._type === 'JsonPipe' ? params : this._replaceGuidsWithChannels(params));
   }
 
   close() {
@@ -220,6 +221,9 @@ export class Connection extends EventEmitter {
       case 'JSHandle':
         result = new JSHandle(parent, type, guid, initializer);
         break;
+      case 'JsonPipe':
+        result = new JsonPipe(parent, type, guid, initializer);
+        break;
       case 'Page':
         result = new Page(parent, type, guid, initializer);
         break;

src/client/jsonPipe.ts

@@ -0,0 +1,32 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as channels from '../protocol/channels';
+import { ChannelOwner } from './channelOwner';
+
+export class JsonPipe extends ChannelOwner<channels.JsonPipeChannel, channels.JsonPipeInitializer> {
+  static from(jsonPipe: channels.JsonPipeChannel): JsonPipe {
+    return (jsonPipe as any)._object;
+  }
+
+  constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.JsonPipeInitializer) {
+    super(parent, type, guid, initializer);
+  }
+
+  channel() {
+    return this._channel;
+  }
+}

src/dispatchers/browserTypeDispatcher.ts

@@ -20,6 +20,9 @@ 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';
 
 export class BrowserTypeDispatcher extends Dispatcher<BrowserType, channels.BrowserTypeInitializer, channels.BrowserTypeEvents> implements channels.BrowserTypeChannel {
   constructor(scope: DispatcherScope, browserType: BrowserType) {
@@ -47,4 +50,33 @@ export class BrowserTypeDispatcher extends Dispatcher<BrowserType, channels.Brow
       defaultContext: browser._defaultContext ? new BrowserContextDispatcher(browserDispatcher._scope, browser._defaultContext) : undefined,
     };
   }
+
+  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 || 30000,
+      headers: paramsHeaders,
+    });
+    const pipe = new JsonPipeDispatcher(this._scope);
+    ws.on('open', () => pipe.wasOpened());
+    ws.on('close', () => pipe.wasClosed());
+    ws.on('error', error => pipe.wasClosed(error));
+    pipe.on('close', () => ws.close());
+    pipe.on('message', message => ws.send(JSON.stringify(message)));
+    ws.addEventListener('message', event => {
+      waitForNextTask(() => {
+        try {
+          pipe.dispatch(JSON.parse(event.data));
+        } catch (e) {
+          ws.close();
+        }
+      });
+    });
+    return { pipe };
+  }
 }

src/dispatchers/dispatcher.ts

@@ -48,7 +48,7 @@ export class Dispatcher<Type extends { guid: string }, Initializer, Events> exte
   private _parent: Dispatcher<any, any, {}> | undefined;
   // Only "isScope" channel owners have registered dispatchers inside.
   private _dispatchers = new Map<string, Dispatcher<any, any, {}>>();
-  private _disposed = false;
+  protected _disposed = false;
 
   readonly _guid: string;
   readonly _type: string;
@@ -91,7 +91,7 @@ export class Dispatcher<Type extends { guid: string }, Initializer, Events> exte
     this._connection.sendMessageToClient(this._guid, this._type, method as string, params, sdkObject);
   }
 
-  _dispose() {
+  protected _dispose() {
     assert(!this._disposed);
     this._disposed = true;
 

src/dispatchers/jsonPipeDispatcher.ts

@@ -0,0 +1,56 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the 'License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as channels from '../protocol/channels';
+import { Dispatcher, DispatcherScope } from './dispatcher';
+import { createGuid } from '../utils/utils';
+import { serializeError } from '../protocol/serializers';
+
+export class JsonPipeDispatcher extends Dispatcher<{ guid: string }, channels.JsonPipeInitializer, channels.JsonPipeEvents> implements channels.JsonPipeChannel {
+  constructor(scope: DispatcherScope) {
+    super(scope, { guid: 'jsonPipe@' + createGuid() }, 'JsonPipe', {});
+  }
+
+  async send(params: channels.JsonPipeSendParams): Promise<channels.JsonPipeSendResult> {
+    this.emit('message', params.message);
+  }
+
+  async close(): Promise<void> {
+    this.emit('close');
+    if (!this._disposed) {
+      this._dispatchEvent('closed', {});
+      this._dispose();
+    }
+  }
+
+  wasOpened(): void {
+    if (!this._disposed)
+      this._dispatchEvent('opened');
+  }
+
+  dispatch(message: Object) {
+    if (!this._disposed)
+      this._dispatchEvent('message', { message });
+  }
+
+  wasClosed(error?: Error): void {
+    if (!this._disposed) {
+      const params = error ? { error: serializeError(error) } : {};
+      this._dispatchEvent('closed', params);
+      this._dispose();
+    }
+  }
+}

src/dispatchers/pageDispatcher.ts

@@ -267,7 +267,7 @@ export class BindingCallDispatcher extends Dispatcher<{ guid: string }, channels
   private _promise: Promise<any>;
 
   constructor(scope: DispatcherScope, name: string, needsHandle: boolean, source: { context: BrowserContext, page: Page, frame: Frame }, args: any[]) {
-    super(scope, { guid: createGuid() }, 'BindingCall', {
+    super(scope, { guid: 'bindingCall@' + createGuid() }, 'BindingCall', {
       frame: lookupDispatcher<FrameDispatcher>(source.frame),
       name,
       args: needsHandle ? undefined : args.map(serializeResult),

src/dispatchers/streamDispatcher.ts

@@ -22,7 +22,7 @@ import { createGuid } from '../utils/utils';
 export class StreamDispatcher extends Dispatcher<{ guid: string, stream: stream.Readable }, channels.StreamInitializer, channels.StreamEvents> implements channels.StreamChannel {
   private _ended: boolean = false;
   constructor(scope: DispatcherScope, stream: stream.Readable) {
-    super(scope, { guid: createGuid(), stream }, 'Stream', {});
+    super(scope, { guid: 'stream@' + createGuid(), stream }, 'Stream', {});
     // In Node v12.9.0+ we can use readableEnded.
     stream.once('end', () => this._ended =  true);
     stream.once('error', () => this._ended =  true);

src/protocol/channels.ts

@@ -296,10 +296,25 @@ export type BrowserTypeInitializer = {
   name: string,
 };
 export interface BrowserTypeChannel extends Channel {
+  connect(params: BrowserTypeConnectParams, metadata?: Metadata): Promise<BrowserTypeConnectResult>;
   launch(params: BrowserTypeLaunchParams, metadata?: Metadata): Promise<BrowserTypeLaunchResult>;
   launchPersistentContext(params: BrowserTypeLaunchPersistentContextParams, metadata?: Metadata): Promise<BrowserTypeLaunchPersistentContextResult>;
   connectOverCDP(params: BrowserTypeConnectOverCDPParams, metadata?: Metadata): Promise<BrowserTypeConnectOverCDPResult>;
 }
+export type BrowserTypeConnectParams = {
+  wsEndpoint: string,
+  headers?: any,
+  slowMo?: number,
+  timeout?: number,
+};
+export type BrowserTypeConnectOptions = {
+  headers?: any,
+  slowMo?: number,
+  timeout?: number,
+};
+export type BrowserTypeConnectResult = {
+  pipe: JsonPipeChannel,
+};
 export type BrowserTypeLaunchParams = {
   channel?: string,
   executablePath?: string,
@@ -3506,6 +3521,37 @@ export type AndroidElementInfo = {
   selected: boolean,
 };
 
+// ----------- JsonPipe -----------
+export type JsonPipeInitializer = {};
+export interface JsonPipeChannel extends Channel {
+  on(event: 'opened', callback: (params: JsonPipeOpenedEvent) => void): this;
+  on(event: 'message', callback: (params: JsonPipeMessageEvent) => void): this;
+  on(event: 'closed', callback: (params: JsonPipeClosedEvent) => void): this;
+  send(params: JsonPipeSendParams, metadata?: Metadata): Promise<JsonPipeSendResult>;
+  close(params?: JsonPipeCloseParams, metadata?: Metadata): Promise<JsonPipeCloseResult>;
+}
+export type JsonPipeOpenedEvent = {};
+export type JsonPipeMessageEvent = {
+  message: any,
+};
+export type JsonPipeClosedEvent = {};
+export type JsonPipeSendParams = {
+  message: any,
+};
+export type JsonPipeSendOptions = {
+
+};
+export type JsonPipeSendResult = void;
+export type JsonPipeCloseParams = {};
+export type JsonPipeCloseOptions = {};
+export type JsonPipeCloseResult = void;
+
+export interface JsonPipeEvents {
+  'opened': JsonPipeOpenedEvent;
+  'message': JsonPipeMessageEvent;
+  'closed': JsonPipeClosedEvent;
+}
+
 export const commandsWithTracingSnapshots = new Set([
   'EventTarget.waitForEventInfo',
   'BrowserContext.waitForEventInfo',

src/protocol/protocol.yml

@@ -449,6 +449,15 @@ BrowserType:
 
   commands:
 
+    connect:
+      parameters:
+        wsEndpoint: string
+        headers: json?
+        slowMo: number?
+        timeout: number?
+      returns:
+        pipe: JsonPipe
+
     launch:
       parameters:
         $mixin: LaunchOptions
@@ -2829,3 +2838,25 @@ AndroidElementInfo:
     longClickable: boolean
     scrollable: boolean
     selected: boolean
+
+
+JsonPipe:
+  type: interface
+
+  commands:
+    send:
+      parameters:
+        message: json
+
+    close:
+
+  events:
+
+    opened:
+
+    message:
+      parameters:
+        message: json
+
+    closed:
+        error: SerializedError?

src/protocol/validator.ts

@@ -181,6 +181,12 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     source: tString,
     contentScript: tOptional(tBoolean),
   });
+  scheme.BrowserTypeConnectParams = tObject({
+    wsEndpoint: tString,
+    headers: tOptional(tAny),
+    slowMo: tOptional(tNumber),
+    timeout: tOptional(tNumber),
+  });
   scheme.BrowserTypeLaunchParams = tObject({
     channel: tOptional(tString),
     executablePath: tOptional(tString),
@@ -1341,6 +1347,10 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     scrollable: tBoolean,
     selected: tBoolean,
   });
+  scheme.JsonPipeSendParams = tObject({
+    message: tAny,
+  });
+  scheme.JsonPipeCloseParams = tOptional(tObject({}));
 
   return scheme;
 }

tests/browsertype-launch-server.spec.ts

@@ -57,7 +57,7 @@ it.describe('launch server', () => {
     const browserServer = await browserType.launchServer(browserOptions);
     const error = await browserType.connect({ wsEndpoint: browserServer.wsEndpoint() + '-foo' }).catch(e => e);
     await browserServer.close();
-    expect(error.message).toContain('Most likely ws endpoint is incorrect');
+    expect(error.message).toContain('WebSocket server disconnected');
   });
 
   it('should fire "close" event during kill', async ({browserType, browserOptions}) => {