chore: extract remote connection to a separate file (#13331)

This commit is contained in:
Yury Semikhatsky 2022-04-05 14:47:11 -07:00 committed by GitHub
parent c69c74caba
commit cdef8171c1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 156 additions and 130 deletions

View File

@ -0,0 +1,137 @@
/**
* Copyright (c) Microsoft Corporation.
*
* 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 WebSocket from 'ws';
import { DispatcherConnection, DispatcherScope, Root } from '../dispatchers/dispatcher';
import { PlaywrightDispatcher } from '../dispatchers/playwrightDispatcher';
import { Browser } from '../server/browser';
import { serverSideCallMetadata } from '../server/instrumentation';
import { createPlaywright, Playwright } from '../server/playwright';
import { gracefullyCloseAll } from '../utils/processLauncher';
import { registry } from '../utils/registry';
import { SocksProxy } from '../utils/socksProxy';
export class PlaywrightConnection {
private _ws: WebSocket;
private _onClose: () => void;
private _dispatcherConnection: DispatcherConnection;
private _cleanups: (() => Promise<void>)[] = [];
private _debugLog: (m: string) => void;
private _disconnected = false;
constructor(ws: WebSocket, enableSocksProxy: boolean, browserAlias: string | undefined, browser: Browser | undefined, log: (m: string) => void, onClose: () => void) {
this._ws = ws;
this._onClose = onClose;
this._debugLog = log;
this._dispatcherConnection = new DispatcherConnection();
this._dispatcherConnection.onmessage = message => {
if (ws.readyState !== ws.CLOSING)
ws.send(JSON.stringify(message));
};
ws.on('message', (message: string) => {
this._dispatcherConnection.dispatch(JSON.parse(Buffer.from(message).toString()));
});
ws.on('close', () => this._onDisconnect());
ws.on('error', error => this._onDisconnect(error));
new Root(this._dispatcherConnection, async scope => {
if (browser)
return await this._initPreLaunchedBrowserMode(scope, browser);
if (!browserAlias)
return await this._initPlaywrightConnectMode(scope, enableSocksProxy);
return await this._initLaunchBrowserMode(scope, enableSocksProxy, browserAlias);
});
}
private async _initPlaywrightConnectMode(scope: DispatcherScope, enableSocksProxy: boolean) {
this._debugLog(`engaged playwright.connect mode`);
const playwright = createPlaywright('javascript');
// Close all launched browsers on disconnect.
this._cleanups.push(() => gracefullyCloseAll());
const socksProxy = enableSocksProxy ? await this._enableSocksProxy(playwright) : undefined;
return new PlaywrightDispatcher(scope, playwright, socksProxy);
}
private async _initLaunchBrowserMode(scope: DispatcherScope, enableSocksProxy: boolean, browserAlias: string) {
this._debugLog(`engaged launch mode for "${browserAlias}"`);
const executable = registry.findExecutable(browserAlias);
if (!executable || !executable.browserName)
throw new Error(`Unsupported browser "${browserAlias}`);
const playwright = createPlaywright('javascript');
const socksProxy = enableSocksProxy ? await this._enableSocksProxy(playwright) : undefined;
const browser = await playwright[executable.browserName].launch(serverSideCallMetadata(), {
channel: executable.type === 'browser' ? undefined : executable.name,
});
// Close the browser on disconnect.
// TODO: it is technically possible to launch more browsers over protocol.
this._cleanups.push(() => browser.close());
browser.on(Browser.Events.Disconnected, () => {
// Underlying browser did close for some reason - force disconnect the client.
this.close({ code: 1001, reason: 'Browser closed' });
});
return new PlaywrightDispatcher(scope, playwright, socksProxy, browser);
}
private async _initPreLaunchedBrowserMode(scope: DispatcherScope, browser: Browser) {
this._debugLog(`engaged pre-launched mode`);
browser.on(Browser.Events.Disconnected, () => {
// Underlying browser did close for some reason - force disconnect the client.
this.close({ code: 1001, reason: 'Browser closed' });
});
const playwright = browser.options.rootSdkObject as Playwright;
const playwrightDispatcher = new PlaywrightDispatcher(scope, playwright, undefined, browser);
// In pre-launched mode, keep the browser and just cleanup new contexts.
// TODO: it is technically possible to launch more browsers over protocol.
this._cleanups.push(() => playwrightDispatcher.cleanup());
return playwrightDispatcher;
}
private async _enableSocksProxy(playwright: Playwright) {
const socksProxy = new SocksProxy();
playwright.options.socksProxyPort = await socksProxy.listen(0);
this._debugLog(`started socks proxy on port ${playwright.options.socksProxyPort}`);
this._cleanups.push(() => socksProxy.close());
return socksProxy;
}
private async _onDisconnect(error?: Error) {
this._disconnected = true;
this._debugLog(`disconnected. error: ${error}`);
// Avoid sending any more messages over closed socket.
this._dispatcherConnection.onmessage = () => {};
this._debugLog(`starting cleanup`);
for (const cleanup of this._cleanups)
await cleanup().catch(() => {});
this._onClose();
this._debugLog(`finished cleanup`);
}
async close(reason?: { code: number, reason: string }) {
if (this._disconnected)
return;
this._debugLog(`force closing connection: ${reason?.reason || ''} (${reason?.code || 0})`);
try {
this._ws.close(reason?.code, reason?.reason);
} catch (e) {
}
}
}

View File

@ -17,17 +17,19 @@
import debug from 'debug'; import debug from 'debug';
import * as http from 'http'; import * as http from 'http';
import WebSocket from 'ws'; import WebSocket from 'ws';
import { DispatcherConnection, DispatcherScope, Root } from '../dispatchers/dispatcher';
import { serverSideCallMetadata } from '../server/instrumentation';
import { createPlaywright, Playwright } from '../server/playwright';
import { Browser } from '../server/browser'; import { Browser } from '../server/browser';
import { gracefullyCloseAll } from '../utils/processLauncher'; import { PlaywrightConnection } from './playwrightConnection';
import { registry } from '../utils/registry';
import { PlaywrightDispatcher } from '../dispatchers/playwrightDispatcher';
import { SocksProxy } from '../utils/socksProxy';
const debugLog = debug('pw:server'); const debugLog = debug('pw:server');
let lastConnectionId = 0;
const kConnectionSymbol = Symbol('kConnection');
function newLogger() {
const id = ++lastConnectionId;
return (message: string) => debugLog(`[id=${id}] ${message}`);
}
export class PlaywrightServer { export class PlaywrightServer {
private _path: string; private _path: string;
private _maxClients: number; private _maxClients: number;
@ -76,8 +78,16 @@ export class PlaywrightServer {
ws.close(1013, 'Playwright Server is busy'); ws.close(1013, 'Playwright Server is busy');
return; return;
} }
const url = new URL('http://localhost' + (request.url || ''));
const browserHeader = request.headers['x-playwright-browser'];
const browserAlias = url.searchParams.get('browser') || (Array.isArray(browserHeader) ? browserHeader[0] : browserHeader);
const proxyHeader = request.headers['x-playwright-proxy'];
const proxyValue = url.searchParams.get('proxy') || (Array.isArray(proxyHeader) ? proxyHeader[0] : proxyHeader);
const enableSocksProxy = this._enableSocksProxy && proxyValue === '*';
this._clientsCount++; this._clientsCount++;
const connection = new Connection(ws, request, this._enableSocksProxy, this._browser, () => this._clientsCount--); const log = newLogger();
log(`serving connection: ${request.url}`);
const connection = new PlaywrightConnection(ws, enableSocksProxy, browserAlias, this._browser, log, () => this._clientsCount--);
(ws as any)[kConnectionSymbol] = connection; (ws as any)[kConnectionSymbol] = connection;
}); });
@ -92,7 +102,7 @@ export class PlaywrightServer {
const waitForClose = new Promise(f => server.close(f)); const waitForClose = new Promise(f => server.close(f));
// First disconnect all remaining clients. // First disconnect all remaining clients.
await Promise.all(Array.from(server.clients).map(async ws => { await Promise.all(Array.from(server.clients).map(async ws => {
const connection = (ws as any)[kConnectionSymbol] as Connection | undefined; const connection = (ws as any)[kConnectionSymbol] as PlaywrightConnection | undefined;
if (connection) if (connection)
await connection.close(); await connection.close();
try { try {
@ -107,124 +117,3 @@ export class PlaywrightServer {
debugLog('closed server'); debugLog('closed server');
} }
} }
let lastConnectionId = 0;
const kConnectionSymbol = Symbol('kConnection');
class Connection {
private _ws: WebSocket;
private _onClose: () => void;
private _dispatcherConnection: DispatcherConnection;
private _cleanups: (() => Promise<void>)[] = [];
private _id: number;
private _disconnected = false;
constructor(ws: WebSocket, request: http.IncomingMessage, enableSocksProxy: boolean, browser: Browser | undefined, onClose: () => void) {
this._ws = ws;
this._onClose = onClose;
this._id = ++lastConnectionId;
debugLog(`[id=${this._id}] serving connection: ${request.url}`);
this._dispatcherConnection = new DispatcherConnection();
this._dispatcherConnection.onmessage = message => {
if (ws.readyState !== ws.CLOSING)
ws.send(JSON.stringify(message));
};
ws.on('message', (message: string) => {
this._dispatcherConnection.dispatch(JSON.parse(Buffer.from(message).toString()));
});
ws.on('close', () => this._onDisconnect());
ws.on('error', error => this._onDisconnect(error));
new Root(this._dispatcherConnection, async scope => {
if (browser)
return await this._initPreLaunchedBrowserMode(scope, browser);
const url = new URL('http://localhost' + (request.url || ''));
const browserHeader = request.headers['x-playwright-browser'];
const browserAlias = url.searchParams.get('browser') || (Array.isArray(browserHeader) ? browserHeader[0] : browserHeader);
const proxyHeader = request.headers['x-playwright-proxy'];
const proxyValue = url.searchParams.get('proxy') || (Array.isArray(proxyHeader) ? proxyHeader[0] : proxyHeader);
if (!browserAlias)
return await this._initPlaywrightConnectMode(scope, enableSocksProxy && proxyValue === '*');
return await this._initLaunchBrowserMode(scope, enableSocksProxy && proxyValue === '*', browserAlias);
});
}
private async _initPlaywrightConnectMode(scope: DispatcherScope, enableSocksProxy: boolean) {
debugLog(`[id=${this._id}] engaged playwright.connect mode`);
const playwright = createPlaywright('javascript');
// Close all launched browsers on disconnect.
this._cleanups.push(() => gracefullyCloseAll());
const socksProxy = enableSocksProxy ? await this._enableSocksProxy(playwright) : undefined;
return new PlaywrightDispatcher(scope, playwright, socksProxy);
}
private async _initLaunchBrowserMode(scope: DispatcherScope, enableSocksProxy: boolean, browserAlias: string) {
debugLog(`[id=${this._id}] engaged launch mode for "${browserAlias}"`);
const executable = registry.findExecutable(browserAlias);
if (!executable || !executable.browserName)
throw new Error(`Unsupported browser "${browserAlias}`);
const playwright = createPlaywright('javascript');
const socksProxy = enableSocksProxy ? await this._enableSocksProxy(playwright) : undefined;
const browser = await playwright[executable.browserName].launch(serverSideCallMetadata(), {
channel: executable.type === 'browser' ? undefined : executable.name,
});
// Close the browser on disconnect.
// TODO: it is technically possible to launch more browsers over protocol.
this._cleanups.push(() => browser.close());
browser.on(Browser.Events.Disconnected, () => {
// Underlying browser did close for some reason - force disconnect the client.
this.close({ code: 1001, reason: 'Browser closed' });
});
return new PlaywrightDispatcher(scope, playwright, socksProxy, browser);
}
private async _initPreLaunchedBrowserMode(scope: DispatcherScope, browser: Browser) {
debugLog(`[id=${this._id}] engaged pre-launched mode`);
browser.on(Browser.Events.Disconnected, () => {
// Underlying browser did close for some reason - force disconnect the client.
this.close({ code: 1001, reason: 'Browser closed' });
});
const playwright = browser.options.rootSdkObject as Playwright;
const playwrightDispatcher = new PlaywrightDispatcher(scope, playwright, undefined, browser);
// In pre-launched mode, keep the browser and just cleanup new contexts.
// TODO: it is technically possible to launch more browsers over protocol.
this._cleanups.push(() => playwrightDispatcher.cleanup());
return playwrightDispatcher;
}
private async _enableSocksProxy(playwright: Playwright) {
const socksProxy = new SocksProxy();
playwright.options.socksProxyPort = await socksProxy.listen(0);
debugLog(`[id=${this._id}] started socks proxy on port ${playwright.options.socksProxyPort}`);
this._cleanups.push(() => socksProxy.close());
return socksProxy;
}
private async _onDisconnect(error?: Error) {
this._disconnected = true;
debugLog(`[id=${this._id}] disconnected. error: ${error}`);
// Avoid sending any more messages over closed socket.
this._dispatcherConnection.onmessage = () => {};
debugLog(`[id=${this._id}] starting cleanup`);
for (const cleanup of this._cleanups)
await cleanup().catch(() => {});
this._onClose();
debugLog(`[id=${this._id}] finished cleanup`);
}
async close(reason?: { code: number, reason: string }) {
if (this._disconnected)
return;
debugLog(`[id=${this._id}] force closing connection: ${reason?.reason || ''} (${reason?.code || 0})`);
try {
this._ws.close(reason?.code, reason?.reason);
} catch (e) {
}
}
}