fix(appium): backport bidi close code validation (#22942)

packages/appium/lib/bidi-commands.ts

@@ -33,9 +33,27 @@ interface InitBiDiSocketResult {
   logSocketErr: LogSocketError;
 }
 
-const MIN_WS_CODE_VAL = 1000;
-const MAX_WS_CODE_VAL = 1015;
-const WS_FALLBACK_CODE = 1011; // server encountered an error while fulfilling request
+/**
+ * Close codes defined by RFC 6455, section 7.4.1 and the IANA registry.
+ * https://www.rfc-editor.org/rfc/rfc6455#section-7.4.1
+ */
+const WebSocketCloseCode = {
+  NORMAL_CLOSURE: 1000,
+  GOING_AWAY: 1001,
+  RESERVED: 1004,
+  NO_STATUS_RECEIVED: 1005,
+  ABNORMAL_CLOSURE: 1006,
+  INTERNAL_ERROR: 1011,
+  BAD_GATEWAY: 1014,
+} as const;
+// RFC 6455, section 7.4.2: https://www.rfc-editor.org/rfc/rfc6455#section-7.4.2
+const MIN_APPLICATION_CLOSE_CODE = 3000;
+const MAX_APPLICATION_CLOSE_CODE = 4999;
+const RESERVED_CLOSE_CODES: readonly number[] = [
+  WebSocketCloseCode.RESERVED,
+  WebSocketCloseCode.NO_STATUS_RECEIVED,
+  WebSocketCloseCode.ABNORMAL_CLOSURE,
+];
 const BIDI_EVENTS_MAP: WeakMap<AnyDriver, Record<string, number>> = new WeakMap();
 const MAX_LOGGED_DATA_LENGTH = 300;
 
@@ -338,15 +356,14 @@ function initBidiProxyHandlers(this: AnyDriver, proxyClient: WebSocket, ws: WebS
       `Upstream bidi socket closed connection (code ${code}, reason: '${reason}'). ` +
         `Closing proxy connection to client`,
     );
-    const intCode: number = typeof code === 'number' ? (code as number) : parseInt(code, 10);
-    if (Number.isNaN(intCode) || intCode < MIN_WS_CODE_VAL || intCode > MAX_WS_CODE_VAL) {
+    const closeCode = toSendableCloseCode(code);
+    if (closeCode !== code) {
       driverLog.warn(
         `Received code ${code} from upstream socket, but this is not a valid ` +
-          `websocket code. Rewriting to ${WS_FALLBACK_CODE} for ws compatibility`,
+          `websocket code. Rewriting to ${closeCode} for ws compatibility`,
       );
-      code = WS_FALLBACK_CODE;
     }
-    ws.close(code, reason);
+    ws.close(closeCode, reason);
   });
 
   proxyClient.on('error', (err) => {
@@ -413,7 +430,7 @@ function initBidiSocketHandlers(
     // If we're proxying, might as well close the upstream connection and clean it up
     if (proxyClient) {
       driverLog.debug('Also closing BiDi proxy socket connection');
-      proxyClient.close(code, reason);
+      proxyClient.close(toSendableCloseCode(code), reason);
     }
 
     const eventLogCounts = BIDI_EVENTS_MAP.get(bidiHandlerDriver);
@@ -548,4 +565,16 @@ async function assertIsOpen(ws: WebSocket, timeoutMs: number = 5000): Promise<We
   return ws;
 }
 
+/**
+ * Replace close codes that cannot be sent in a WebSocket close frame.
+ */
+function toSendableCloseCode(code: number): number {
+  const isStandard =
+    code >= WebSocketCloseCode.NORMAL_CLOSURE &&
+    code <= WebSocketCloseCode.BAD_GATEWAY &&
+    !RESERVED_CLOSE_CODES.includes(code);
+  const isApplication = code >= MIN_APPLICATION_CLOSE_CODE && code <= MAX_APPLICATION_CLOSE_CODE;
+  return isStandard || isApplication ? code : WebSocketCloseCode.INTERNAL_ERROR;
+}
+
 // #endregion

packages/appium/test/unit/appiumdriver.spec.ts

@@ -1,4 +1,5 @@
 import assert from 'node:assert/strict';
+import type {AddressInfo} from 'node:net';
 import {describe, it, beforeEach, afterEach, before} from 'node:test';
 
 import {BaseDriver} from '@appium/base-driver';
@@ -8,6 +9,7 @@ import type {Capabilities, Constraints, NSCapabilities, W3CCapabilities} from '@
 import {sleep} from 'asyncbox';
 import type {SinonMock, SinonSandbox, SinonStubbedMember} from 'sinon';
 import {createSandbox} from 'sinon';
+import WebSocket, {WebSocketServer} from 'ws';
 
 import type * as AppiumModule from '../../lib/appium';
 import {PLUGIN_TYPE, SESSION_DISCOVERY_FEATURE} from '../../lib/constants';
@@ -115,6 +117,54 @@ describe('AppiumDriver', function () {
 
       return [appium, mockFakeDriver];
     }
+    describe('proxied bidi socket', function () {
+      let appium: InstanceType<typeof AppiumModule.AppiumDriver>;
+      let mockFakeDriver: SinonMock;
+      let upstreamServer: WebSocketServer;
+      let appiumServer: WebSocketServer;
+      let client: WebSocket;
+      let upstreamSocket: WebSocket;
+      let sessionId: string;
+
+      async function listen(): Promise<WebSocketServer> {
+        const server = new WebSocketServer({port: 0, host: '127.0.0.1'});
+        await new Promise<void>((resolve) => server.once('listening', resolve));
+        return server;
+      }
+
+      beforeEach(async function () {
+        [appium, mockFakeDriver] = getDriverAndFakeDriver();
+        [sessionId] = (await appium.createSession(null as any, null as any, structuredClone(W3C_CAPS))).value!;
+        upstreamServer = await listen();
+        (fakeDriver as any)._bidiProxyUrl = `ws://127.0.0.1:${(upstreamServer.address() as AddressInfo).port}`;
+        const upstreamConnected = new Promise<WebSocket>((resolve) => upstreamServer.once('connection', resolve));
+        appiumServer = await listen();
+        appiumServer.on('connection', (ws, req) => appium.onBidiConnection(ws, req));
+        client = new WebSocket(`ws://127.0.0.1:${(appiumServer.address() as AddressInfo).port}/bidi/${sessionId}`);
+        await new Promise<void>((resolve) => client.once('open', resolve));
+        upstreamSocket = await upstreamConnected;
+      });
+
+      afterEach(async function () {
+        client.terminate();
+        await new Promise<void>((resolve) => appiumServer.close(() => resolve()));
+        await new Promise<void>((resolve) => upstreamServer.close(() => resolve()));
+        await appium.deleteSession(sessionId);
+        mockFakeDriver.restore();
+      });
+
+      it('should close the upstream socket when the client drops without a close frame', async function () {
+        const upstreamClosed = new Promise<number>((resolve) => upstreamSocket.once('close', resolve));
+        client.terminate();
+        assert.strictEqual(await upstreamClosed, 1011);
+      });
+
+      it('should close the client socket when the upstream drops without a close frame', async function () {
+        const clientClosed = new Promise<number>((resolve) => client.once('close', resolve));
+        upstreamSocket.terminate();
+        assert.strictEqual(await clientClosed, 1011);
+      });
+    });
     describe('configureGlobalFeatures', function () {
       let appium: InstanceType<typeof AppiumModule.AppiumDriver>;