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,15 +97,19 @@ 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;
     }
     return this.#adapters.get(session)!;
   }
 
   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();
   }
 }