fix: various WebSocketRoute issues

- Do not allow two clients to route simultaneously.
- Cleanup after dispatcher disconnects.
- Ensure that routing works in context reuse scenarios.
- Test context reuse both in the launched browser with a single
  connection, and in a connected browser with different
  connections.

Author: dgozman

packages/playwright-core/src/client/network.ts

@@ -564,7 +564,8 @@ export class WebSocketRoute extends ChannelOwner<channels.WebSocketRouteChannel>
     if (this._connected)
       return;
     // Ensure that websocket is "open" and can send messages without an actual server connection.
-    await this._channel.ensureOpened();
+    // If this happens after the page has been closed, ignore the error.
+    await this._channel.ensureOpened().catch(() => {});
   }
 }
 

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

@@ -308,8 +308,8 @@ export abstract class BrowserContext extends SdkObject {
     return this.doSetHTTPCredentials(httpCredentials);
   }
 
-  hasBinding(name: string) {
-    return this._pageBindings.has(name);
+  getBindingClient(name: string): unknown | undefined {
+    return this._pageBindings.get(name)?.forClient;
   }
 
   async exposePlaywrightBindingIfNeeded() {
@@ -328,7 +328,7 @@ export abstract class BrowserContext extends SdkObject {
     return this._playwrightBindingExposed;
   }
 
-  async exposeBinding(progress: Progress, name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource): Promise<PageBinding> {
+  async exposeBinding(progress: Progress, name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource, forClient?: unknown): Promise<PageBinding> {
     if (this._pageBindings.has(name))
       throw new Error(`Function "${name}" has been already registered`);
     for (const page of this.pages()) {
@@ -337,6 +337,7 @@ export abstract class BrowserContext extends SdkObject {
     }
     await progress.race(this.exposePlaywrightBindingIfNeeded());
     const binding = new PageBinding(name, playwrightBinding, needsHandle);
+    binding.forClient = forClient;
     this._pageBindings.set(name, binding);
     progress.cleanupWhenAborted(() => this._pageBindings.delete(name));
     await progress.race(this.doAddInitScript(binding.initScript));
@@ -443,8 +444,8 @@ export abstract class BrowserContext extends SdkObject {
       this._options.httpCredentials = { username, password: password || '' };
   }
 
-  async addInitScript(progress: Progress | undefined, source: string, name?: string) {
-    const initScript = new InitScript(source, name);
+  async addInitScript(progress: Progress | undefined, source: string) {
+    const initScript = new InitScript(source);
     this.initScripts.push(initScript);
     progress?.cleanupWhenAborted(() => this.removeInitScripts([initScript]));
     const promise = this.doAddInitScript(initScript);

packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts

@@ -58,6 +58,7 @@ export class BrowserContextDispatcher extends Dispatcher<BrowserContext, channel
   private _clockPaused = false;
   private _requestInterceptor: RouteHandler;
   private _interceptionUrlMatchers: (string | RegExp)[] = [];
+  private _routeWebSocketInitScripts: InitScript[] = [];
 
   static from(parentScope: DispatcherScope, context: BrowserContext): BrowserContextDispatcher {
     const result = parentScope.connection.existingDispatcher<BrowserContextDispatcher>(context);
@@ -319,7 +320,7 @@ export class BrowserContextDispatcher extends Dispatcher<BrowserContext, channel
   async setWebSocketInterceptionPatterns(params: channels.PageSetWebSocketInterceptionPatternsParams, progress: Progress): Promise<void> {
     this._webSocketInterceptionPatterns = params.patterns;
     if (params.patterns.length)
-      await WebSocketRouteDispatcher.installIfNeeded(progress, this.connection, this._context);
+      this._routeWebSocketInitScripts.push(await WebSocketRouteDispatcher.install(progress, this.connection, this._context));
   }
 
   async storageState(params: channels.BrowserContextStorageStateParams, progress: Progress): Promise<channels.BrowserContextStorageStateResult> {
@@ -418,6 +419,8 @@ export class BrowserContextDispatcher extends Dispatcher<BrowserContext, channel
     this._bindings = [];
     this._context.removeInitScripts(this._initScritps).catch(() => {});
     this._initScritps = [];
+    this._routeWebSocketInitScripts.map(script => WebSocketRouteDispatcher.uninstall(this.connection, this._context, script).catch(() => {}));
+    this._routeWebSocketInitScripts = [];
     if (this._clockPaused)
       this._context.clock.resumeNoReply();
     this._clockPaused = false;

packages/playwright-core/src/server/dispatchers/pageDispatcher.ts

@@ -49,6 +49,7 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
   _webSocketInterceptionPatterns: channels.PageSetWebSocketInterceptionPatternsParams['patterns'] = [];
   private _bindings: PageBinding[] = [];
   private _initScripts: InitScript[] = [];
+  private _routeWebSocketInitScripts: InitScript[] = [];
   private _requestInterceptor: RouteHandler;
   private _interceptionUrlMatchers: (string | RegExp)[] = [];
   private _locatorHandlers = new Set<number>();
@@ -208,7 +209,7 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
   async setWebSocketInterceptionPatterns(params: channels.PageSetWebSocketInterceptionPatternsParams, progress: Progress): Promise<void> {
     this._webSocketInterceptionPatterns = params.patterns;
     if (params.patterns.length)
-      await WebSocketRouteDispatcher.installIfNeeded(progress, this.connection, this._page);
+      this._routeWebSocketInitScripts.push(await WebSocketRouteDispatcher.install(progress, this.connection, this._page));
   }
 
   async expectScreenshot(params: channels.PageExpectScreenshotParams, progress: Progress): Promise<channels.PageExpectScreenshotResult> {
@@ -367,6 +368,8 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
     this._bindings = [];
     this._page.removeInitScripts(this._initScripts).catch(() => {});
     this._initScripts = [];
+    this._routeWebSocketInitScripts.map(script => WebSocketRouteDispatcher.uninstall(this.connection, this._page, script).catch(() => {}));
+    this._routeWebSocketInitScripts = [];
     for (const uid of this._locatorHandlers)
       this._page.unregisterLocatorHandler(uid);
     this._locatorHandlers.clear();

packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts

@@ -29,6 +29,7 @@ import type { Frame } from '../frames';
 import type * as ws from '@injected/webSocketMock';
 import type * as channels from '@protocol/channels';
 import type { Progress } from '@protocol/progress';
+import type { InitScript, PageBinding } from '../page';
 
 export class WebSocketRouteDispatcher extends Dispatcher<SdkObject, channels.WebSocketRouteChannel, PageDispatcher | BrowserContextDispatcher> implements channels.WebSocketRouteChannel {
   _type_WebSocketRoute = true;
@@ -58,11 +59,14 @@ export class WebSocketRouteDispatcher extends Dispatcher<SdkObject, channels.Web
     (scope as any)._dispatchEvent('webSocketRoute', { webSocketRoute: this });
   }
 
-  static async installIfNeeded(progress: Progress, connection: DispatcherConnection, target: Page | BrowserContext) {
-    const kBindingName = '__pwWebSocketBinding';
+  static async install(progress: Progress, connection: DispatcherConnection, target: Page | BrowserContext): Promise<InitScript> {
     const context = target instanceof Page ? target.browserContext : target;
-    if (!context.hasBinding(kBindingName)) {
-      await context.exposeBinding(progress, kBindingName, false, (source, payload: ws.BindingPayload) => {
+    let data = context.getBindingClient(kBindingName) as BindingData | undefined;
+    if (data && data.connection !== connection)
+      throw new Error('Another client is already routing WebSockets');
+    if (!data) {
+      data = { counter: 0, connection, binding: null as any };
+      data.binding = await context.exposeBinding(progress, kBindingName, false, (source, payload: ws.BindingPayload) => {
         if (payload.type === 'onCreate') {
           const contextDispatcher = connection.existingDispatcher<BrowserContextDispatcher>(context);
           const pageDispatcher = contextDispatcher ? PageDispatcher.fromNullable(contextDispatcher, source.page) : undefined;
@@ -89,19 +93,27 @@ export class WebSocketRouteDispatcher extends Dispatcher<SdkObject, channels.Web
           dispatcher?._dispatchEvent('closePage', { code: payload.code, reason: payload.reason, wasClean: payload.wasClean });
         if (payload.type === 'onCloseServer')
           dispatcher?._dispatchEvent('closeServer', { code: payload.code, reason: payload.reason, wasClean: payload.wasClean });
-      });
+      }, data);
     }
+    ++data.counter;
+
+    return await target.addInitScript(progress, `
+      (() => {
+        const module = {};
+        ${rawWebSocketMockSource.source}
+        (module.exports.inject())(globalThis);
+      })();
+    `);
+  }
 
-    const kInitScriptName = 'webSocketMockSource';
-    if (!target.initScripts.find(s => s.name === kInitScriptName)) {
-      await target.addInitScript(progress, `
-        (() => {
-          const module = {};
-          ${rawWebSocketMockSource.source}
-          (module.exports.inject())(globalThis);
-        })();
-      `, kInitScriptName);
-    }
+  static async uninstall(connection: DispatcherConnection, target: Page | BrowserContext, initScript: InitScript) {
+    const context = target instanceof Page ? target.browserContext : target;
+    const data = context.getBindingClient(kBindingName) as BindingData | undefined;
+    if (!data || data.connection !== connection)
+      return;
+    if (--data.counter <= 0)
+      await context.removeExposedBindings([data.binding]);
+    await target.removeInitScripts([initScript]);
   }
 
   async connect(params: channels.WebSocketRouteConnectParams, progress: Progress) {
@@ -155,3 +167,6 @@ function matchesPattern(dispatcher: PageDispatcher | BrowserContextDispatcher, b
   }
   return false;
 }
+
+const kBindingName = '__pwWebSocketBinding';
+type BindingData = { counter: number, connection: DispatcherConnection, binding: PageBinding };

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

@@ -562,8 +562,8 @@ export class Page extends SdkObject {
     await this.delegate.bringToFront();
   }
 
-  async addInitScript(progress: Progress, source: string, name?: string) {
-    const initScript = new InitScript(source, name);
+  async addInitScript(progress: Progress, source: string) {
+    const initScript = new InitScript(source);
     this.initScripts.push(initScript);
     progress.cleanupWhenAborted(() => this.removeInitScripts([initScript]));
     await progress.race(this.delegate.addInitScript(initScript));
@@ -889,6 +889,7 @@ export class PageBinding {
   readonly initScript: InitScript;
   readonly needsHandle: boolean;
   readonly cleanupScript: string;
+  forClient?: unknown;
 
   constructor(name: string, playwrightFunction: frames.FunctionWithSource, needsHandle: boolean) {
     this.name = name;
@@ -924,13 +925,11 @@ export class PageBinding {
 
 export class InitScript {
   readonly source: string;
-  readonly name?: string;
 
-  constructor(source: string, name?: string) {
+  constructor(source: string) {
     this.source = `(() => {
       ${source}
     })();`;
-    this.name = name;
   }
 }