chore(grid): add command line options (#24802)

packages/playwright-grid/src/cli.ts

@@ -23,14 +23,24 @@ program
 
 program
     .command('grid')
-    .action(function() {
+    .option('--port <port>', 'port to listen to, 3333 by default')
-      require('./grid/grid');
+    .option('--access-key <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() {
+    .option('--grid <url>', 'grid address', 'localhost:3333')
-      require('./node/node');
+    .option('--capacity <capacity>', 'node capacity', '1')
+    .option('--access-key <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);

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<string, Node>();
   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();
-})();

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();

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<string, string> = {};

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',
+      command: 'node ../../packages/playwright-grid/cli.js grid --port=3333 --access-key=secret',
-      url: 'http://localhost:3333',
+      stdout: 'pipe',
+      url: 'http://localhost:3333/secret',
       reuseExistingServer: !process.env.CI,
-      env: {
-        PLAYWRIGHT_GRID_PORT: '3333',
-      }
     }, {
-      command: 'node ../../packages/playwright-grid/cli.js node',
+      command: 'node ../../packages/playwright-grid/cli.js node --grid=ws://localhost:3333 --access-key=secret --capacity=2',
-      env: {
-        PLAYWRIGHT_GRID_ENDPOINT: 'ws://localhost:3333',
-        PLAYWRIGHT_GRID_NODE_CAPACITY: '2',
-      }
     },
     {
-      command: 'node ../../packages/playwright-grid/cli.js node',
+      command: 'node ../../packages/playwright-grid/cli.js node --grid=ws://localhost:3333 --access-key=secret --capacity=2',
-      env: {
-        PLAYWRIGHT_GRID_ENDPOINT: 'ws://localhost:3333',
-        PLAYWRIGHT_GRID_NODE_CAPACITY: '2',
-      }
     }
   ];
 }