From ca4869c931568d6dd957744b8e62a6a8fe794122 Mon Sep 17 00:00:00 2001 From: Pavel Date: Fri, 27 Mar 2020 15:42:10 -0700 Subject: [PATCH] feat(websocket): wrap firefox web socket too --- src/platform.ts | 2 +- src/server/chromium.ts | 7 +-- src/server/firefox.ts | 138 +++++++++++++++++++++++++++++++++++++++-- src/server/webkit.ts | 39 ++++-------- src/transport.ts | 17 +++++ 5 files changed, 164 insertions(+), 39 deletions(-) diff --git a/src/platform.ts b/src/platform.ts index b3c2451214bb68..8697bc102c15e3 100644 --- a/src/platform.ts +++ b/src/platform.ts @@ -312,7 +312,7 @@ export function guid(): string { // 'onmessage' handler must be installed synchronously when 'onopen' callback is invoked to // avoid missing incoming messages. -export async function connectToWebsocket(url: string, onopen: (transport: ConnectionTransport) => Promise): Promise { +export async function connectToWebsocket(url: string, onopen: (transport: ConnectionTransport) => Promise | T): Promise { const transport = new WebSocketTransport(url); return new Promise((fulfill, reject) => { transport._ws.addEventListener('open', async () => fulfill(await onopen(transport))); diff --git a/src/server/chromium.ts b/src/server/chromium.ts index 07d4a2d79136b2..b2569642ba3455 100644 --- a/src/server/chromium.ts +++ b/src/server/chromium.ts @@ -95,7 +95,6 @@ export class Chromium implements BrowserType { const chromeExecutable = executablePath || this._executablePath; if (!chromeExecutable) throw new Error(`No executable path is specified. Pass "executablePath" option directly.`); - let browserServer: BrowserServer | undefined = undefined; const { launchedProcess, gracefullyClose } = await launchProcess({ executablePath: chromeExecutable, args: chromeArguments, @@ -123,6 +122,7 @@ export class Chromium implements BrowserType { }); let transport: PipeTransport | undefined = undefined; + let browserServer: BrowserServer | undefined = undefined; transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream, () => browserServer!.close()); browserServer = new BrowserServer(launchedProcess, gracefullyClose, launchType === 'server' ? wrapTransportWithWebSocket(transport, port || 0) : null); return { browserServer, transport }; @@ -251,11 +251,6 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number socket.on('message', (message: string) => { const parsedMessage = JSON.parse(Buffer.from(message).toString()) as ProtocolRequest; - if (parsedMessage.method.startsWith('Backend')) { - // Add backend domain handler here. - return; - } - // If message has sessionId, pass through. if (parsedMessage.sessionId) { transport.send(parsedMessage); diff --git a/src/server/firefox.ts b/src/server/firefox.ts index b19a956ed23dc0..c5885be482e13e 100644 --- a/src/server/firefox.ts +++ b/src/server/firefox.ts @@ -18,6 +18,7 @@ import * as fs from 'fs'; import * as os from 'os'; import * as path from 'path'; +import * as ws from 'ws'; import { ConnectOptions, LaunchType } from '../browser'; import { BrowserContext } from '../browserContext'; import { TimeoutError } from '../errors'; @@ -29,6 +30,7 @@ import * as platform from '../platform'; import { BrowserServer } from './browserServer'; import { BrowserArgOptions, BrowserType, LaunchOptions } from './browserType'; import { launchProcess, waitForLine } from './processLauncher'; +import { ConnectionTransport, ProtocolRequest, SequenceNumberMixer } from '../transport'; const mkdtempAsync = platform.promisify(fs.mkdtemp); @@ -97,9 +99,9 @@ export class Firefox implements BrowserType { } if (!ignoreDefaultArgs) - firefoxArguments.push(...this._defaultArgs(options, launchType, userDataDir!, port || 0)); + firefoxArguments.push(...this._defaultArgs(options, launchType, userDataDir!, 0)); else if (Array.isArray(ignoreDefaultArgs)) - firefoxArguments.push(...this._defaultArgs(options, launchType, userDataDir!, port || 0).filter(arg => !ignoreDefaultArgs.includes(arg))); + firefoxArguments.push(...this._defaultArgs(options, launchType, userDataDir!, 0).filter(arg => !ignoreDefaultArgs.includes(arg))); else firefoxArguments.push(...args); @@ -107,7 +109,6 @@ export class Firefox implements BrowserType { if (!firefoxExecutable) throw new Error(`No executable path is specified. Pass "executablePath" option directly.`); - let browserServer: BrowserServer | undefined = undefined; const { launchedProcess, gracefullyClose } = await launchProcess({ executablePath: firefoxExecutable, args: firefoxArguments, @@ -140,7 +141,10 @@ export class Firefox implements BrowserType { const timeoutError = new TimeoutError(`Timed out after ${timeout} ms while trying to connect to Firefox!`); const match = await waitForLine(launchedProcess, launchedProcess.stdout, /^Juggler listening on (ws:\/\/.*)$/, timeout, timeoutError); - const browserWSEndpoint = match[1]; + const innerEndpoint = match[1]; + + let browserServer: BrowserServer | undefined = undefined; + const browserWSEndpoint = launchType === 'server' ? (await platform.connectToWebsocket(innerEndpoint, t => wrapTransportWithWebSocket(t, port || 0))) : innerEndpoint; browserServer = new BrowserServer(launchedProcess, gracefullyClose, browserWSEndpoint); return browserServer; } @@ -185,3 +189,129 @@ export class Firefox implements BrowserType { } } +function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): string { + const server = new ws.Server({ port }); + const guid = platform.guid(); + const idMixer = new SequenceNumberMixer<{id: number, socket: ws}>(); + const pendingBrowserContextCreations = new Set(); + const pendingBrowserContextDeletions = new Map(); + const browserContextIds = new Map(); + const sessionIds = new Map(); + const sockets = new Set(); + + transport.onmessage = message => { + if (typeof message.id === 'number') { + // Process command response. + const value = idMixer.take(message.id); + if (!value) + return; + const { id, socket } = value; + + if (socket.readyState === ws.CLOSING) { + if (pendingBrowserContextCreations.has(id)) { + transport.send({ + id: ++SequenceNumberMixer._lastSequenceNumber, + method: 'Browser.removeBrowserContext', + params: { browserContextId: message.result.browserContextId } + }); + } + return; + } + + if (pendingBrowserContextCreations.has(message.id)) { + // Browser.createBrowserContext response -> establish context attribution. + browserContextIds.set(message.result.browserContextId, socket); + pendingBrowserContextCreations.delete(message.id); + } + + const deletedContextId = pendingBrowserContextDeletions.get(message.id); + if (deletedContextId) { + // Browser.removeBrowserContext response -> remove context attribution. + browserContextIds.delete(deletedContextId); + pendingBrowserContextDeletions.delete(message.id); + } + + message.id = id; + socket.send(JSON.stringify(message)); + return; + } + + // Process notification response. + const { method, params, sessionId } = message; + if (sessionId) { + const socket = sessionIds.get(sessionId); + if (!socket || socket.readyState === ws.CLOSING) { + // Drop unattributed messages on the floor. + return; + } + socket.send(JSON.stringify(message)); + return; + } + if (method === 'Browser.attachedToTarget') { + const socket = browserContextIds.get(params.targetInfo.browserContextId); + if (!socket || socket.readyState === ws.CLOSING) { + // Drop unattributed messages on the floor. + return; + } + sessionIds.set(params.sessionId, socket); + socket.send(JSON.stringify(message)); + return; + } + if (method === 'Browser.detachedFromTarget') { + const socket = sessionIds.get(params.sessionId); + sessionIds.delete(params.sessionId); + if (socket && socket.readyState !== ws.CLOSING) + socket.send(JSON.stringify(message)); + return; + } + }; + + transport.onclose = () => { + for (const socket of sockets) { + socket.removeListener('close', (socket as any).__closeListener); + socket.close(undefined, 'Browser disconnected'); + } + server.close(); + transport.onmessage = undefined; + transport.onclose = undefined; + }; + + server.on('connection', (socket: ws, req) => { + if (req.url !== '/' + guid) { + socket.close(); + return; + } + sockets.add(socket); + + socket.on('message', (message: string) => { + const parsedMessage = JSON.parse(Buffer.from(message).toString()); + const { id, method, params, sessionId } = parsedMessage; + const seqNum = idMixer.generate({ id, socket }); + transport.send({ ...parsedMessage, id: seqNum }); + if (method === 'Browser.createBrowserContext') + pendingBrowserContextCreations.add(seqNum); + if (method === 'Browser.removeBrowserContext') + pendingBrowserContextDeletions.set(seqNum, params.browserContextId); + }); + + socket.on('close', (socket as any).__closeListener = () => { + for (const [browserContextId, s] of browserContextIds) { + if (s === socket) { + transport.send({ + id: ++SequenceNumberMixer._lastSequenceNumber, + method: 'Browser.removeBrowserContext', + params: { browserContextId } + }); + browserContextIds.delete(browserContextId); + } + } + sockets.delete(socket); + }); + }); + + const address = server.address(); + if (typeof address === 'string') + return address + '/' + guid; + return 'ws://127.0.0.1:' + address.port + '/' + guid; +} + diff --git a/src/server/webkit.ts b/src/server/webkit.ts index ac0ba1e4b36225..6021b744e3905c 100644 --- a/src/server/webkit.ts +++ b/src/server/webkit.ts @@ -25,7 +25,7 @@ import * as os from 'os'; import { helper } from '../helper'; import { kBrowserCloseMessageId } from '../webkit/wkConnection'; import { LaunchOptions, BrowserArgOptions, BrowserType } from './browserType'; -import { ConnectionTransport } from '../transport'; +import { ConnectionTransport, SequenceNumberMixer } from '../transport'; import * as ws from 'ws'; import { ConnectOptions, LaunchType } from '../browser'; import { BrowserServer } from './browserServer'; @@ -163,23 +163,6 @@ const mkdtempAsync = platform.promisify(fs.mkdtemp); const WEBKIT_PROFILE_PATH = path.join(os.tmpdir(), 'playwright_dev_profile-'); -class SequenceNumberMixer { - static _lastSequenceNumber = 1; - private _values = new Map(); - - generate(value: V): number { - const sequenceNumber = ++SequenceNumberMixer._lastSequenceNumber; - this._values.set(sequenceNumber, value); - return sequenceNumber; - } - - take(sequenceNumber: number): V | undefined { - const value = this._values.get(sequenceNumber); - this._values.delete(sequenceNumber); - return value; - } -} - function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): string { const server = new ws.Server({ port }); const guid = platform.guid(); @@ -265,6 +248,16 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number } }; + transport.onclose = () => { + for (const socket of sockets) { + socket.removeListener('close', (socket as any).__closeListener); + socket.close(undefined, 'Browser disconnected'); + } + server.close(); + transport.onmessage = undefined; + transport.onclose = undefined; + }; + server.on('connection', (socket: ws, req) => { if (req.url !== '/' + guid) { socket.close(); @@ -302,16 +295,6 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number }); }); - transport.onclose = () => { - for (const socket of sockets) { - socket.removeListener('close', (socket as any).__closeListener); - socket.close(undefined, 'Browser disconnected'); - } - server.close(); - transport.onmessage = undefined; - transport.onclose = undefined; - }; - const address = server.address(); if (typeof address === 'string') return address + '/' + guid; diff --git a/src/transport.ts b/src/transport.ts index 315ca41837ab67..90fff6349a1127 100644 --- a/src/transport.ts +++ b/src/transport.ts @@ -113,3 +113,20 @@ export class DeferWriteTransport implements ConnectionTransport { this._delegate.close(); } } + +export class SequenceNumberMixer { + static _lastSequenceNumber = 1; + private _values = new Map(); + + generate(value: V): number { + const sequenceNumber = ++SequenceNumberMixer._lastSequenceNumber; + this._values.set(sequenceNumber, value); + return sequenceNumber; + } + + take(sequenceNumber: number): V | undefined { + const value = this._values.get(sequenceNumber); + this._values.delete(sequenceNumber); + return value; + } +}