From 120de62798d8108bc626601033b97fc90a6dac78 Mon Sep 17 00:00:00 2001 From: Pavel Feldman Date: Fri, 4 Aug 2023 18:08:50 -0700 Subject: [PATCH] chore(grid): add command line options (#24802) --- packages/playwright-grid/src/cli.ts | 18 ++++++++++---- packages/playwright-grid/src/grid/grid.ts | 26 +++++++++++---------- packages/playwright-grid/src/node/node.ts | 17 +++++++------- packages/playwright-grid/src/node/worker.ts | 6 ++++- tests/library/playwright.config.ts | 25 +++++++------------- 5 files changed, 51 insertions(+), 41 deletions(-) diff --git a/packages/playwright-grid/src/cli.ts b/packages/playwright-grid/src/cli.ts index b478589c94..081257ab46 100644 --- a/packages/playwright-grid/src/cli.ts +++ b/packages/playwright-grid/src/cli.ts @@ -23,14 +23,24 @@ program program .command('grid') - .action(function() { - require('./grid/grid'); + .option('--port ', 'port to listen to, 3333 by default') + .option('--access-key ', 'access key to the grid') + .action(async opts => { + const port = opts.port || +(process.env.PLAYWRIGHT_GRID_PORT || '3333'); + const accessKey = opts.accessKey || process.env.PLAYWRIGHT_GRID_ACCESS_KEY; + const { Grid } = await import('./grid/grid.js'); + const grid = new Grid(port, accessKey); + grid.start(); }); program .command('node') - .action(function() { - require('./node/node'); + .option('--grid ', 'grid address', 'localhost:3333') + .option('--capacity ', 'node capacity', '1') + .option('--access-key ', 'access key to the grid', '') + .action(async opts => { + const { Node } = await import('./node/node.js'); + new Node(opts.grid, +opts.capacity, opts.accessKey); }); program.parse(process.argv); diff --git a/packages/playwright-grid/src/grid/grid.ts b/packages/playwright-grid/src/grid/grid.ts index 2188674bcf..12ce60e56d 100644 --- a/packages/playwright-grid/src/grid/grid.ts +++ b/packages/playwright-grid/src/grid/grid.ts @@ -24,8 +24,6 @@ import type { Capabilities } from '../common/capabilities'; import type http from 'http'; import type stream from 'stream'; -const PORT = +(process.env.PLAYWRIGHT_GRID_PORT || '3113'); - class WebSocketRequest { private _socketError: Error | undefined; @@ -218,18 +216,22 @@ class Node { } } -class ProxyServer { +export class Grid { private _server: HttpServer; private _wsServer: WebSocketServer; private _nodes = new Map(); private _log: debug.Debugger; private _clientRequests: ClientRequest[] = []; + private _port: number; + private _accessKey: string; - constructor() { + constructor(port: number, accessKey: string) { this._log = debug(`pw:grid:proxy`); + this._port = port; + this._accessKey = accessKey; this._server = new HttpServer(); - this._server.routePath('/', (request, response) => { + this._server.routePath('/' + this._accessKey, (request, response) => { response.statusCode = 200; response.setHeader('Content-Type', 'text/plain'); response.end(this._state()); @@ -241,9 +243,14 @@ class ProxyServer { ws.on('error', e => this._log(e)); }); this._server.server.on('upgrade', async (request, socket, head) => { + if (this._accessKey && request.headers['x-playwright-access-key'] !== this._accessKey) { + socket.destroy(); + return; + } + const url = new URL('http://internal' + request.url); const params = url.searchParams; - this._log(url.toString()); + this._log(url.pathname); if (url.pathname.startsWith('/registerNode')) { const nodeRequest = new WebSocketRequest(this._wsServer, request, socket, head); @@ -352,7 +359,7 @@ class ProxyServer { } async start() { - const url = await this._server.start(PORT); + const url = await this._server.start(this._port); // eslint-disable-next-line no-console console.log('Server is listening on: ' + url); } @@ -361,8 +368,3 @@ class ProxyServer { function createGuid(): string { return crypto.randomBytes(16).toString('hex'); } - -(async () => { - const proxy = new ProxyServer(); - await proxy.start(); -})(); diff --git a/packages/playwright-grid/src/node/node.ts b/packages/playwright-grid/src/node/node.ts index d2d8218cf3..67280205b7 100644 --- a/packages/playwright-grid/src/node/node.ts +++ b/packages/playwright-grid/src/node/node.ts @@ -21,18 +21,20 @@ import type { Capabilities } from '../common/capabilities'; const log = debug('pw:grid:node'); -const endpoint = process.env.PLAYWRIGHT_GRID_ENDPOINT || 'ws://localhost:3113'; -const capacity = parseInt(process.env.PLAYWRIGHT_GRID_NODE_CAPACITY || '1', 10); const caps: Capabilities = { platform: process.platform, }; -class Node { +export class Node { workerSeq = 0; - constructor() { + constructor(grid: string, capacity: number, accessKey: string) { log('node created'); - const ws = new WebSocket(endpoint + `/registerNode?capacity=${capacity}&caps=${JSON.stringify(caps)}`); + const ws = new WebSocket(grid + `/registerNode?capacity=${capacity}&caps=${JSON.stringify(caps)}`, { + headers: { + 'x-playwright-access-key': accessKey, + } + }); let nodeId = ''; ws.on('error', error => { log(error); @@ -54,7 +56,8 @@ class Node { ...process.env, PLAYWRIGHT_GRID_NODE_ID: nodeId, PLAYWRIGHT_GRID_WORKER_ID: workerId, - PLAYWRIGHT_GRID_ENDPOINT: endpoint, + PLAYWRIGHT_GRID_ENDPOINT: grid, + PLAYWRIGHT_GRID_ACCESS_KEY: accessKey, }, detached: true }); @@ -63,5 +66,3 @@ class Node { ws.on('close', () => process.exit(0)); } } - -new Node(); diff --git a/packages/playwright-grid/src/node/worker.ts b/packages/playwright-grid/src/node/worker.ts index 433066755a..189d2c59d4 100644 --- a/packages/playwright-grid/src/node/worker.ts +++ b/packages/playwright-grid/src/node/worker.ts @@ -29,7 +29,11 @@ class Worker { let browserName: 'chromium' | 'webkit' | 'firefox'; let launchOptions: any; - const ws = new WebSocket(process.env.PLAYWRIGHT_GRID_ENDPOINT + `/registerWorker?nodeId=${process.env.PLAYWRIGHT_GRID_NODE_ID}&workerId=${workerId}`); + const ws = new WebSocket(process.env.PLAYWRIGHT_GRID_ENDPOINT + `/registerWorker?nodeId=${process.env.PLAYWRIGHT_GRID_NODE_ID}&workerId=${workerId}`, { + headers: { + 'x-playwright-access-key': process.env.PLAYWRIGHT_GRID_ACCESS_KEY!, + } + }); dispatcherConnection.onmessage = message => ws.send(JSON.stringify(message)); ws.on('upgrade', response => { const headers: Record = {}; diff --git a/tests/library/playwright.config.ts b/tests/library/playwright.config.ts index 7840f9d659..ecba8a7e3c 100644 --- a/tests/library/playwright.config.ts +++ b/tests/library/playwright.config.ts @@ -78,30 +78,23 @@ if (mode === 'service2') { if (mode === 'service-grid') { connectOptions = { - wsEndpoint: 'ws://localhost:3333/', + wsEndpoint: 'ws://localhost:3333', timeout: 60 * 60 * 1000, + headers: { + 'x-playwright-access-key': 'secret' + } }; webServer = [ { - command: 'node ../../packages/playwright-grid/cli.js grid', - url: 'http://localhost:3333', + command: 'node ../../packages/playwright-grid/cli.js grid --port=3333 --access-key=secret', + stdout: 'pipe', + url: 'http://localhost:3333/secret', reuseExistingServer: !process.env.CI, - env: { - PLAYWRIGHT_GRID_PORT: '3333', - } }, { - command: 'node ../../packages/playwright-grid/cli.js node', - env: { - PLAYWRIGHT_GRID_ENDPOINT: 'ws://localhost:3333', - PLAYWRIGHT_GRID_NODE_CAPACITY: '2', - } + command: 'node ../../packages/playwright-grid/cli.js node --grid=ws://localhost:3333 --access-key=secret --capacity=2', }, { - command: 'node ../../packages/playwright-grid/cli.js node', - env: { - PLAYWRIGHT_GRID_ENDPOINT: 'ws://localhost:3333', - PLAYWRIGHT_GRID_NODE_CAPACITY: '2', - } + command: 'node ../../packages/playwright-grid/cli.js node --grid=ws://localhost:3333 --access-key=secret --capacity=2', } ]; }