feat: BiDi implementation of Puppeteer.connect for Firefox (#11451)

Co-authored-by: Maksim Sadym <sadym@google.com>
Co-authored-by: Alex Rudenko <OrKoN@users.noreply.github.com>

packages/puppeteer-core/src/bidi/BidiOverCdp.ts

@@ -80,15 +80,15 @@ export async function connectBidiOverCdp(
 class CdpConnectionAdapter {
   #cdp: CdpConnection;
   #adapters = new Map<CDPSession, CDPClientAdapter<CDPSession>>();
-  #browser: CDPClientAdapter<CdpConnection>;
+  #browserCdpConnection: CDPClientAdapter<CdpConnection>;
 
   constructor(cdp: CdpConnection) {
     this.#cdp = cdp;
-    this.#browser = new CDPClientAdapter(cdp);
+    this.#browserCdpConnection = new CDPClientAdapter(cdp);
   }
 
   browserClient(): CDPClientAdapter<CdpConnection> {
-    return this.#browser;
+    return this.#browserCdpConnection;
   }
 
   getCdpClient(id: string) {
@@ -97,7 +97,11 @@ class CdpConnectionAdapter {
       throw new Error(`Unknown CDP session with id ${id}`);
     }
     if (!this.#adapters.has(session)) {
-      const adapter = new CDPClientAdapter(session, id, this.#browser);
+      const adapter = new CDPClientAdapter(
+        session,
+        id,
+        this.#browserCdpConnection
+      );
       this.#adapters.set(session, adapter);
       return adapter;
     }
@@ -105,7 +109,7 @@ class CdpConnectionAdapter {
   }
 
   close() {
-    this.#browser.close();
+    this.#browserCdpConnection.close();
     for (const adapter of this.#adapters.values()) {
       adapter.close();
     }

packages/puppeteer-core/src/bidi/Browser.ts

@@ -253,7 +253,9 @@ export class BidiBrowser extends Browser {
     if (this.#connection.closed) {
       return;
     }
-    await this.#connection.send('browser.close', {});
+
+    // `browser.close` can close connection before the response is received.
+    await this.#connection.send('browser.close', {}).catch(debugError);
     await this.#closeCallback?.call(null);
     this.#connection.dispose();
   }

packages/puppeteer-core/src/bidi/BrowserConnector.ts

@@ -0,0 +1,133 @@
+/*
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * 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 type {BrowserCloseCallback} from '../api/Browser.js';
+import {Connection} from '../cdp/Connection.js';
+import type {ConnectionTransport} from '../common/ConnectionTransport.js';
+import type {
+  BrowserConnectOptions,
+  ConnectOptions,
+} from '../common/ConnectOptions.js';
+import {UnsupportedOperation} from '../common/Errors.js';
+import {debugError, DEFAULT_VIEWPORT} from '../common/util.js';
+
+import type {BidiBrowser} from './Browser.js';
+import type {BidiConnection} from './Connection.js';
+
+/**
+ * Users should never call this directly; it's called when calling `puppeteer.connect`
+ * with `protocol: 'webDriverBiDi'`. This method attaches Puppeteer to an existing browser
+ * instance. First it tries to connect to the browser using pure BiDi. If the protocol is
+ * not supported, connects to the browser using BiDi over CDP.
+ *
+ * @internal
+ */
+export async function _connectToBiDiBrowser(
+  connectionTransport: ConnectionTransport,
+  url: string,
+  options: BrowserConnectOptions & ConnectOptions
+): Promise<BidiBrowser> {
+  const {ignoreHTTPSErrors = false, defaultViewport = DEFAULT_VIEWPORT} =
+    options;
+
+  const {bidiConnection, closeCallback} = await getBiDiConnection(
+    connectionTransport,
+    url,
+    options
+  );
+  const BiDi = await import(/* webpackIgnore: true */ './bidi.js');
+  const bidiBrowser = await BiDi.BidiBrowser.create({
+    connection: bidiConnection,
+    closeCallback,
+    process: undefined,
+    defaultViewport: defaultViewport,
+    ignoreHTTPSErrors: ignoreHTTPSErrors,
+  });
+  return bidiBrowser;
+}
+
+/**
+ * Returns a BiDiConnection established to the endpoint specified by the options and a
+ * callback closing the browser. Callback depends on whether the connection is pure BiDi
+ * or BiDi over CDP.
+ * The method tries to connect to the browser using pure BiDi protocol, and falls back
+ * to BiDi over CDP.
+ */
+async function getBiDiConnection(
+  connectionTransport: ConnectionTransport,
+  url: string,
+  options: BrowserConnectOptions
+): Promise<{
+  bidiConnection: BidiConnection;
+  closeCallback: BrowserCloseCallback;
+}> {
+  const BiDi = await import(/* webpackIgnore: true */ './bidi.js');
+  const {ignoreHTTPSErrors = false, slowMo = 0, protocolTimeout} = options;
+
+  // Try pure BiDi first.
+  const pureBidiConnection = new BiDi.BidiConnection(
+    url,
+    connectionTransport,
+    slowMo,
+    protocolTimeout
+  );
+  try {
+    const result = await pureBidiConnection.send('session.status', {});
+    if ('type' in result && result.type === 'success') {
+      // The `browserWSEndpoint` points to an endpoint supporting pure WebDriver BiDi.
+      return {
+        bidiConnection: pureBidiConnection,
+        closeCallback: async () => {
+          await pureBidiConnection.send('browser.close', {}).catch(debugError);
+        },
+      };
+    }
+  } catch (e: any) {
+    if (!('name' in e && e.name === 'ProtocolError')) {
+      // Unexpected exception not related to BiDi / CDP. Rethrow.
+      throw e;
+    }
+  }
+  // Unbind the connection to avoid memory leaks.
+  pureBidiConnection.unbind();
+
+  // Fall back to CDP over BiDi reusing the WS connection.
+  const cdpConnection = new Connection(
+    url,
+    connectionTransport,
+    slowMo,
+    protocolTimeout
+  );
+
+  const version = await cdpConnection.send('Browser.getVersion');
+  if (version.product.toLowerCase().includes('firefox')) {
+    throw new UnsupportedOperation(
+      'Firefox is not supported in BiDi over CDP mode.'
+    );
+  }
+
+  // TODO: use other options too.
+  const bidiOverCdpConnection = await BiDi.connectBidiOverCdp(cdpConnection, {
+    acceptInsecureCerts: ignoreHTTPSErrors,
+  });
+  return {
+    bidiConnection: bidiOverCdpConnection,
+    closeCallback: async () => {
+      // In case of BiDi over CDP, we need to close browser via CDP.
+      await cdpConnection.send('Browser.close').catch(debugError);
+    },
+  };
+}

packages/puppeteer-core/src/bidi/Connection.ts

@@ -21,6 +21,7 @@ import type {ConnectionTransport} from '../common/ConnectionTransport.js';
 import {debug} from '../common/Debug.js';
 import {EventEmitter} from '../common/EventEmitter.js';
 import {debugError} from '../common/util.js';
+import {assert} from '../util/assert.js';
 
 import {type BrowsingContext, cdpSessions} from './BrowsingContext.js';
 
@@ -176,7 +177,7 @@ export class BidiConnection extends EventEmitter<BidiEvents> {
 
     this.#transport = transport;
     this.#transport.onmessage = this.onMessage.bind(this);
-    this.#transport.onclose = this.#onClose.bind(this);
+    this.#transport.onclose = this.unbind.bind(this);
   }
 
   get closed(): boolean {
@@ -191,6 +192,8 @@ export class BidiConnection extends EventEmitter<BidiEvents> {
     method: T,
     params: Commands[T]['params']
   ): Promise<{result: Commands[T]['returnType']}> {
+    assert(!this.#closed, 'Protocol error: Connection closed.');
+
     return this.#callbacks.create(method, this.#timeout, id => {
       const stringifiedMessage = JSON.stringify({
         id,
@@ -244,6 +247,15 @@ export class BidiConnection extends EventEmitter<BidiEvents> {
           return;
       }
     }
+    // Even if the response in not in BiDi protocol format but `id` is provided, reject
+    // the callback. This can happen if the endpoint supports CDP instead of BiDi.
+    if ('id' in object) {
+      this.#callbacks.reject(
+        (object as {id: number}).id,
+        `Protocol Error. Message is not in BiDi protocol format: '${message}'`,
+        object.message
+      );
+    }
     debugError(object);
   }
 
@@ -293,7 +305,12 @@ export class BidiConnection extends EventEmitter<BidiEvents> {
     this.#browsingContexts.delete(id);
   }
 
-  #onClose(): void {
+  /**
+   * Unbinds the connection, but keeps the transport open. Useful when the transport will
+   * be reused by other connection e.g. with different protocol.
+   * @internal
+   */
+  unbind(): void {
     if (this.#closed) {
       return;
     }
@@ -302,11 +319,15 @@ export class BidiConnection extends EventEmitter<BidiEvents> {
     this.#transport.onmessage = () => {};
     this.#transport.onclose = () => {};
 
+    this.#browsingContexts.clear();
     this.#callbacks.clear();
   }
 
+  /**
+   * Unbinds the connection and closes the transport.
+   */
   dispose(): void {
-    this.#onClose();
+    this.unbind();
     this.#transport.close();
   }
 }

packages/puppeteer-core/src/cdp/BrowserConnector.ts

@@ -14,31 +14,16 @@
  * limitations under the License.
  */
 
-import type {BidiBrowser} from '../bidi/Browser.js';
 import type {ConnectionTransport} from '../common/ConnectionTransport.js';
 import type {
   BrowserConnectOptions,
   ConnectOptions,
 } from '../common/ConnectOptions.js';
-import {UnsupportedOperation} from '../common/Errors.js';
+import {debugError, DEFAULT_VIEWPORT} from '../common/util.js';
-import {getFetch} from '../common/fetch.js';
-import {debugError} from '../common/util.js';
-import {isNode} from '../environment.js';
-import {assert} from '../util/assert.js';
-import {isErrorLike} from '../util/ErrorLike.js';
 
 import {CdpBrowser} from './Browser.js';
 import {Connection} from './Connection.js';
 
-const DEFAULT_VIEWPORT = Object.freeze({width: 800, height: 600});
-
-const getWebSocketTransportClass = async () => {
-  return isNode
-    ? (await import('../node/NodeWebSocketTransport.js')).NodeWebSocketTransport
-    : (await import('../common/BrowserWebSocketTransport.js'))
-        .BrowserWebSocketTransport;
-};
-
 /**
  * Users should never call this directly; it's called when calling
  * `puppeteer.connect` with `protocol: 'cdp'`.
@@ -46,6 +31,8 @@ const getWebSocketTransportClass = async () => {
  * @internal
  */
 export async function _connectToCdpBrowser(
+  connectionTransport: ConnectionTransport,
+  url: string,
   options: BrowserConnectOptions & ConnectOptions
 ): Promise<CdpBrowser> {
   const {
@@ -53,9 +40,16 @@ export async function _connectToCdpBrowser(
     defaultViewport = DEFAULT_VIEWPORT,
     targetFilter,
     _isPageTarget: isPageTarget,
+    slowMo = 0,
+    protocolTimeout,
   } = options;
 
-  const connection = await getCdpConnection(options);
+  const connection = new Connection(
+    url,
+    connectionTransport,
+    slowMo,
+    protocolTimeout
+  );
 
   const version = await connection.send('Browser.getVersion');
   const product = version.product.toLowerCase().includes('firefox')
@@ -80,112 +74,3 @@ export async function _connectToCdpBrowser(
   );
   return browser;
 }
-
-/**
- * Users should never call this directly; it's called when calling
- * `puppeteer.connect` with `protocol: 'webDriverBiDi'`.
- *
- * @internal
- */
-export async function _connectToBiDiOverCdpBrowser(
-  options: BrowserConnectOptions & ConnectOptions
-): Promise<BidiBrowser> {
-  const {ignoreHTTPSErrors = false, defaultViewport = DEFAULT_VIEWPORT} =
-    options;
-
-  const connection = await getCdpConnection(options);
-
-  const version = await connection.send('Browser.getVersion');
-  if (version.product.toLowerCase().includes('firefox')) {
-    throw new UnsupportedOperation(
-      'Firefox is not supported in BiDi over CDP mode.'
-    );
-  }
-
-  // TODO: use other options too.
-  const BiDi = await import(/* webpackIgnore: true */ '../bidi/bidi.js');
-  const bidiConnection = await BiDi.connectBidiOverCdp(connection, {
-    acceptInsecureCerts: ignoreHTTPSErrors,
-  });
-  const bidiBrowser = await BiDi.BidiBrowser.create({
-    connection: bidiConnection,
-    closeCallback: () => {
-      return connection.send('Browser.close').catch(debugError);
-    },
-    process: undefined,
-    defaultViewport: defaultViewport,
-    ignoreHTTPSErrors: ignoreHTTPSErrors,
-  });
-  return bidiBrowser;
-}
-
-async function getWSEndpoint(browserURL: string): Promise<string> {
-  const endpointURL = new URL('/json/version', browserURL);
-
-  const fetch = await getFetch();
-  try {
-    const result = await fetch(endpointURL.toString(), {
-      method: 'GET',
-    });
-    if (!result.ok) {
-      throw new Error(`HTTP ${result.statusText}`);
-    }
-    const data = await result.json();
-    return data.webSocketDebuggerUrl;
-  } catch (error) {
-    if (isErrorLike(error)) {
-      error.message =
-        `Failed to fetch browser webSocket URL from ${endpointURL}: ` +
-        error.message;
-    }
-    throw error;
-  }
-}
-
-/**
- * Returns a CDP connection for the given options.
- */
-async function getCdpConnection(
-  options: BrowserConnectOptions & ConnectOptions
-): Promise<Connection> {
-  const {
-    browserWSEndpoint,
-    browserURL,
-    transport,
-    headers = {},
-    slowMo = 0,
-    protocolTimeout,
-  } = options;
-
-  assert(
-    Number(!!browserWSEndpoint) + Number(!!browserURL) + Number(!!transport) ===
-      1,
-    'Exactly one of browserWSEndpoint, browserURL or transport must be passed to puppeteer.connect'
-  );
-
-  if (transport) {
-    return new Connection('', transport, slowMo, protocolTimeout);
-  } else if (browserWSEndpoint) {
-    const WebSocketClass = await getWebSocketTransportClass();
-    const connectionTransport: ConnectionTransport =
-      await WebSocketClass.create(browserWSEndpoint, headers);
-    return new Connection(
-      browserWSEndpoint,
-      connectionTransport,
-      slowMo,
-      protocolTimeout
-    );
-  } else if (browserURL) {
-    const connectionURL = await getWSEndpoint(browserURL);
-    const WebSocketClass = await getWebSocketTransportClass();
-    const connectionTransport: ConnectionTransport =
-      await WebSocketClass.create(connectionURL);
-    return new Connection(
-      connectionURL,
-      connectionTransport,
-      slowMo,
-      protocolTimeout
-    );
-  }
-  throw new Error('Invalid connection options');
-}

packages/puppeteer-core/src/common/BrowserConnector.ts

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * 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 type {Browser} from '../api/Browser.js';
+import {_connectToBiDiBrowser} from '../bidi/BrowserConnector.js';
+import {_connectToCdpBrowser} from '../cdp/BrowserConnector.js';
+import {isNode} from '../environment.js';
+import {assert} from '../util/assert.js';
+import {isErrorLike} from '../util/ErrorLike.js';
+
+import type {ConnectionTransport} from './ConnectionTransport.js';
+import type {ConnectOptions} from './ConnectOptions.js';
+import type {BrowserConnectOptions} from './ConnectOptions.js';
+import {getFetch} from './fetch.js';
+
+const getWebSocketTransportClass = async () => {
+  return isNode
+    ? (await import('../node/NodeWebSocketTransport.js')).NodeWebSocketTransport
+    : (await import('../common/BrowserWebSocketTransport.js'))
+        .BrowserWebSocketTransport;
+};
+
+/**
+ * Users should never call this directly; it's called when calling
+ * `puppeteer.connect`. This method attaches Puppeteer to an existing browser instance.
+ *
+ * @internal
+ */
+export async function _connectToBrowser(
+  options: ConnectOptions
+): Promise<Browser> {
+  const {connectionTransport, endpointUrl} =
+    await getConnectionTransport(options);
+
+  if (options.protocol === 'webDriverBiDi') {
+    const bidiBrowser = await _connectToBiDiBrowser(
+      connectionTransport,
+      endpointUrl,
+      options
+    );
+    return bidiBrowser;
+  } else {
+    const cdpBrowser = await _connectToCdpBrowser(
+      connectionTransport,
+      endpointUrl,
+      options
+    );
+    return cdpBrowser;
+  }
+}
+
+/**
+ * Establishes a websocket connection by given options and returns both transport and
+ * endpoint url the transport is connected to.
+ */
+async function getConnectionTransport(
+  options: BrowserConnectOptions & ConnectOptions
+): Promise<{connectionTransport: ConnectionTransport; endpointUrl: string}> {
+  const {browserWSEndpoint, browserURL, transport, headers = {}} = options;
+
+  assert(
+    Number(!!browserWSEndpoint) + Number(!!browserURL) + Number(!!transport) ===
+      1,
+    'Exactly one of browserWSEndpoint, browserURL or transport must be passed to puppeteer.connect'
+  );
+
+  if (transport) {
+    return {connectionTransport: transport, endpointUrl: ''};
+  } else if (browserWSEndpoint) {
+    const WebSocketClass = await getWebSocketTransportClass();
+    const connectionTransport: ConnectionTransport =
+      await WebSocketClass.create(browserWSEndpoint, headers);
+    return {
+      connectionTransport: connectionTransport,
+      endpointUrl: browserWSEndpoint,
+    };
+  } else if (browserURL) {
+    const connectionURL = await getWSEndpoint(browserURL);
+    const WebSocketClass = await getWebSocketTransportClass();
+    const connectionTransport: ConnectionTransport =
+      await WebSocketClass.create(connectionURL);
+    return {
+      connectionTransport: connectionTransport,
+      endpointUrl: connectionURL,
+    };
+  }
+  throw new Error('Invalid connection options');
+}
+
+async function getWSEndpoint(browserURL: string): Promise<string> {
+  const endpointURL = new URL('/json/version', browserURL);
+
+  const fetch = await getFetch();
+  try {
+    const result = await fetch(endpointURL.toString(), {
+      method: 'GET',
+    });
+    if (!result.ok) {
+      throw new Error(`HTTP ${result.statusText}`);
+    }
+    const data = await result.json();
+    return data.webSocketDebuggerUrl;
+  } catch (error) {
+    if (isErrorLike(error)) {
+      error.message =
+        `Failed to fetch browser webSocket URL from ${endpointURL}: ` +
+        error.message;
+    }
+    throw error;
+  }
+}

packages/puppeteer-core/src/common/Puppeteer.ts

@@ -15,11 +15,8 @@
  */
 
 import type {Browser} from '../api/Browser.js';
-import {
-  _connectToBiDiOverCdpBrowser,
-  _connectToCdpBrowser,
-} from '../cdp/BrowserConnector.js';
 
+import {_connectToBrowser} from './BrowserConnector.js';
 import type {ConnectOptions} from './ConnectOptions.js';
 import {
   type CustomQueryHandler,
@@ -131,10 +128,6 @@ export class Puppeteer {
    * @returns Promise which resolves to browser instance.
    */
   connect(options: ConnectOptions): Promise<Browser> {
-    if (options.protocol === 'webDriverBiDi') {
+    return _connectToBrowser(options);
-      return _connectToBiDiOverCdpBrowser(options);
-    } else {
-      return _connectToCdpBrowser(options);
-    }
   }
 }

packages/puppeteer-core/src/common/util.ts

@@ -46,6 +46,11 @@ import type {NetworkManagerEvents} from './NetworkManagerEvents.js';
  */
 export const debugError = debug('puppeteer:error');
 
+/**
+ * @internal
+ */
+export const DEFAULT_VIEWPORT = Object.freeze({width: 800, height: 600});
+
 /**
  * @internal
  */

packages/puppeteer-core/src/node/ProductLauncher.ts

@@ -31,7 +31,7 @@ import {CdpBrowser} from '../cdp/Browser.js';
 import {Connection} from '../cdp/Connection.js';
 import {TimeoutError} from '../common/Errors.js';
 import type {Product} from '../common/Product.js';
-import {debugError} from '../common/util.js';
+import {debugError, DEFAULT_VIEWPORT} from '../common/util.js';
 import type {Viewport} from '../common/Viewport.js';
 
 import type {
@@ -91,7 +91,7 @@ export abstract class ProductLauncher {
       handleSIGTERM = true,
       handleSIGHUP = true,
       ignoreHTTPSErrors = false,
-      defaultViewport = {width: 800, height: 600},
+      defaultViewport = DEFAULT_VIEWPORT,
       slowMo = 0,
       timeout = 30000,
       waitForInitialPage = true,
@@ -122,15 +122,15 @@ export abstract class ProductLauncher {
     });
 
     let browser: Browser;
-    let connection: Connection;
+    let cdpConnection: Connection;
     let closing = false;
 
-    const browserCloseCallback = async () => {
+    const browserCloseCallback: BrowserCloseCallback = async () => {
       if (closing) {
         return;
       }
       closing = true;
-      await this.closeBrowser(browserProcess, connection);
+      await this.closeBrowser(browserProcess, cdpConnection);
     };
 
     try {
@@ -148,13 +148,13 @@ export abstract class ProductLauncher {
         );
       } else {
         if (usePipe) {
-          connection = await this.createCdpPipeConnection(browserProcess, {
+          cdpConnection = await this.createCdpPipeConnection(browserProcess, {
             timeout,
             protocolTimeout,
             slowMo,
           });
         } else {
-          connection = await this.createCdpSocketConnection(browserProcess, {
+          cdpConnection = await this.createCdpSocketConnection(browserProcess, {
             timeout,
             protocolTimeout,
             slowMo,
@@ -163,7 +163,7 @@ export abstract class ProductLauncher {
         if (protocol === 'webDriverBiDi') {
           browser = await this.createBiDiOverCdpBrowser(
             browserProcess,
-            connection,
+            cdpConnection,
             browserCloseCallback,
             {
               timeout,
@@ -176,7 +176,7 @@ export abstract class ProductLauncher {
         } else {
           browser = await CdpBrowser._create(
             this.product,
-            connection,
+            cdpConnection,
             [],
             ignoreHTTPSErrors,
             defaultViewport,
@@ -234,12 +234,12 @@ export abstract class ProductLauncher {
    */
   protected async closeBrowser(
     browserProcess: ReturnType<typeof launch>,
-    connection?: Connection
+    cdpConnection?: Connection
   ): Promise<void> {
-    if (connection) {
+    if (cdpConnection) {
       // Attempt to close the browser gracefully
       try {
-        await connection.closeBrowser();
+        await cdpConnection.closeBrowser();
         await browserProcess.hasClosed();
       } catch (error) {
         debugError(error);

test/TestExpectations.json

@@ -2129,6 +2129,12 @@
     "parameters": ["chrome", "webDriverBiDi"],
     "expectations": ["PASS"]
   },
+  {
+    "testIdPattern": "[launcher.spec] Launcher specs Puppeteer Puppeteer.connect should be able to reconnect",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["firefox", "webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[launcher.spec] Launcher specs Puppeteer Puppeteer.connect should be able to reconnect to a disconnected browser",
     "platforms": ["darwin", "linux", "win32"],