From 9a668aeab31fd0f927e5fcc9c8a679eaa64b21b7 Mon Sep 17 00:00:00 2001
From: Max Schmitt <max@schmitt.mx>
Date: Tue, 5 Nov 2024 11:13:33 +0100
Subject: [PATCH] fix(websocket): pass through underlying WebSocket protocol
 (#33446)

---
 .../src/server/injected/webSocketMock.ts          |  7 +++++++
 tests/library/route-web-socket.spec.ts            | 15 +++++++++++----
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/packages/playwright-core/src/server/injected/webSocketMock.ts b/packages/playwright-core/src/server/injected/webSocketMock.ts
index 7a4742bac3..6ca09ede93 100644
--- a/packages/playwright-core/src/server/injected/webSocketMock.ts
+++ b/packages/playwright-core/src/server/injected/webSocketMock.ts
@@ -331,6 +331,13 @@ export function inject(globalThis: GlobalThis) {
     _ensureOpened() {
       if (this.readyState !== WebSocketMock.CONNECTING)
         return;
+      this.extensions = this._ws?.extensions || '';
+      if (this._ws)
+        this.protocol = this._ws.protocol;
+      else if (Array.isArray(this._protocols))
+        this.protocol = this._protocols[0] || '';
+      else
+        this.protocol = this._protocols || '';
       this.readyState = WebSocketMock.OPEN;
       this.dispatchEvent(new Event('open', { cancelable: true }));
     }
diff --git a/tests/library/route-web-socket.spec.ts b/tests/library/route-web-socket.spec.ts
index 6af0cb4bc0..d3c3607c99 100644
--- a/tests/library/route-web-socket.spec.ts
+++ b/tests/library/route-web-socket.spec.ts
@@ -33,11 +33,11 @@ function withResolvers<T = void>() {
   return { promise, resolve };
 }
 
-async function setupWS(target: Page | Frame, port: number, binaryType: 'blob' | 'arraybuffer') {
+async function setupWS(target: Page | Frame, port: number, binaryType: 'blob' | 'arraybuffer', protocols?: string | string[]) {
   await target.goto('about:blank');
-  await target.evaluate(({ port, binaryType }) => {
+  await target.evaluate(({ port, binaryType, protocols }) => {
     window.log = [];
-    window.ws = new WebSocket('ws://localhost:' + port + '/ws');
+    window.ws = new WebSocket('ws://localhost:' + port + '/ws', protocols);
     window.ws.binaryType = binaryType;
     window.ws.addEventListener('open', () => window.log.push('open'));
     window.ws.addEventListener('close', event => window.log.push(`close code=${event.code} reason=${event.reason} wasClean=${event.wasClean}`));
@@ -53,7 +53,7 @@ async function setupWS(target: Page | Frame, port: number, binaryType: 'blob' |
       window.log.push(`message: data=${data} origin=${event.origin} lastEventId=${event.lastEventId}`);
     });
     window.wsOpened = new Promise(f => window.ws.addEventListener('open', () => f()));
-  }, { port, binaryType });
+  }, { port, binaryType, protocols });
 }
 
 for (const mock of ['no-mock', 'no-match', 'pass-through']) {
@@ -191,6 +191,13 @@ for (const mock of ['no-mock', 'no-match', 'pass-through']) {
       expect(closed.code).toBe(3002);
       expect(closed.reason.toString()).toBe('oops');
     });
+
+    test('should pass through the required protocol', async ({ page, server }) => {
+      await setupWS(page, server.PORT, 'blob', 'my-custom-protocol');
+      await page.evaluate(() => window.wsOpened);
+      const protocol = await page.evaluate(() => window.ws.protocol);
+      expect(protocol).toBe('my-custom-protocol');
+    });
   });
 }