feat(websockets): implement support for ws on cr/wk (#542)

docs/api.md

@@ -129,6 +129,7 @@
   * [event: 'requestfailed'](#event-requestfailed)
   * [event: 'requestfinished'](#event-requestfinished)
   * [event: 'response'](#event-response)
+  * [event: 'websocket'](#event-websocket)
   * [event: 'workercreated'](#event-workercreated)
   * [event: 'workerdestroyed'](#event-workerdestroyed)
   * [page.$(selector)](#pageselector)
@@ -214,6 +215,17 @@
   * [response.statusText()](#responsestatustext)
   * [response.text()](#responsetext)
   * [response.url()](#responseurl)
+- [class: WebSocket](#class-websocket)
+  * [event: 'close'](#event-close-1)
+  * [event: 'error'](#event-error)
+  * [event: 'messageReceived'](#event-messagereceived)
+  * [event: 'messageSent'](#event-messagesent)
+  * [event: 'open'](#event-open)
+  * [webSocket.requestHeaders()](#websocketrequestheaders)
+  * [webSocket.responseHeaders()](#websocketresponseheaders)
+  * [webSocket.status()](#websocketstatus)
+  * [webSocket.statusText()](#websocketstatustext)
+  * [webSocket.url()](#websocketurl)
 - [class: TimeoutError](#class-timeouterror)
 - [class: Accessibility](#class-accessibility)
   * [accessibility.snapshot([options])](#accessibilitysnapshotoptions)
@@ -1824,6 +1836,11 @@ Emitted when a request finishes successfully.
 
 Emitted when a [response] is received.
 
+#### event: 'websocket'
+- <[WebSocket]>
+
+Emitted when a WebSocket request is made.
+
 #### event: 'workercreated'
 - <[Worker]>
 
@@ -3003,6 +3020,54 @@ Contains the status text of the response (e.g. usually an "OK" for a success).
 
 Contains the URL of the response.
 
+### class: WebSocket
+
+[WebSocket] class represents web sockets that are created by page.
+
+#### event: 'close'
+
+Emitted when web socket closes.
+
+#### event: 'error'
+<[string]>
+
+Emitted on error while establishing the connection, sending or receiving the web socket frame.
+
+#### event: 'messageReceived'
+- <[string]|[Buffer]>
+
+Emitted when web socket receives data.
+
+#### event: 'messageSent'
+- <[string]|[Buffer]>
+
+Emitted when web socket sends data.
+
+#### event: 'open'
+
+Emitted when web socket opens.
+
+#### webSocket.requestHeaders()
+- returns: <[Object]> An object with HTTP headers associated with the WebSocket upgrade request. All header names are lower-case.
+
+#### webSocket.responseHeaders()
+- returns: <[Object]> An object with HTTP headers associated with the WebSocket upgrade response. All header names are lower-case.
+
+#### webSocket.status()
+- returns: <[number]>
+
+Contains the status code of the web socket (e.g., 101 for a successful upgrade).
+
+#### webSocket.statusText()
+- returns: <[string]>
+
+Contains the status text of the web socket response (e.g. usually "Switching Protocols" for a successful upgrade).
+
+#### webSocket.url()
+- returns: <[string]>
+
+Contains the URL of the web socket.
+
 ### class: TimeoutError
 
 * extends: [Error]

package.json

@@ -10,7 +10,7 @@
   "playwright": {
     "chromium_revision": "733125",
     "firefox_revision": "1016",
-    "webkit_revision": "1102"
+    "webkit_revision": "1104"
   },
   "scripts": {
     "unit": "node test/test.js",

src/api.ts

@@ -24,7 +24,7 @@ export { TimeoutError } from './errors';
 export { Frame } from './frames';
 export { Keyboard, Mouse } from './input';
 export { JSHandle } from './javascript';
-export { Request, Response } from './network';
+export { Request, Response, WebSocket } from './network';
 export { Coverage, FileChooser, Page, Worker } from './page';
 
 export { CRBrowser as ChromiumBrowser } from './chromium/crBrowser';

src/chromium/crNetworkManager.ts

@@ -41,14 +41,24 @@ export class CRNetworkManager {
   constructor(client: CRSession, page: Page) {
     this._client = client;
     this._page = page;
+    this._eventListeners = this.instrumentNetworkEvents(client);
+  }
 
-    this._eventListeners = [
-      helper.addEventListener(client, 'Fetch.requestPaused', this._onRequestPaused.bind(this)),
-      helper.addEventListener(client, 'Fetch.authRequired', this._onAuthRequired.bind(this)),
-      helper.addEventListener(client, 'Network.requestWillBeSent', this._onRequestWillBeSent.bind(this)),
-      helper.addEventListener(client, 'Network.responseReceived', this._onResponseReceived.bind(this)),
-      helper.addEventListener(client, 'Network.loadingFinished', this._onLoadingFinished.bind(this)),
-      helper.addEventListener(client, 'Network.loadingFailed', this._onLoadingFailed.bind(this)),
+  instrumentNetworkEvents(session: CRSession): RegisteredListener[] {
+    return [
+      helper.addEventListener(session, 'Fetch.requestPaused', this._onRequestPaused.bind(this)),
+      helper.addEventListener(session, 'Fetch.authRequired', this._onAuthRequired.bind(this)),
+      helper.addEventListener(session, 'Network.requestWillBeSent', this._onRequestWillBeSent.bind(this)),
+      helper.addEventListener(session, 'Network.responseReceived', this._onResponseReceived.bind(this)),
+      helper.addEventListener(session, 'Network.loadingFinished', this._onLoadingFinished.bind(this)),
+      helper.addEventListener(session, 'Network.loadingFailed', this._onLoadingFailed.bind(this)),
+      helper.addEventListener(session, 'Network.webSocketCreated', e => this._page._frameManager.onWebSocketCreated(e.requestId, e.url)),
+      helper.addEventListener(session, 'Network.webSocketWillSendHandshakeRequest', e => this._page._frameManager.onWebSocketRequest(e.requestId, e.request.headers)),
+      helper.addEventListener(session, 'Network.webSocketHandshakeResponseReceived', e => this._page._frameManager.onWebSocketResponse(e.requestId, e.response.status, e.response.statusText, e.response.headers)),
+      helper.addEventListener(session, 'Network.webSocketFrameSent', e => e.response.payloadData && this._page._frameManager.onWebSocketFrameSent(e.requestId, e.response.opcode, e.response.payloadData)),
+      helper.addEventListener(session, 'Network.webSocketFrameReceived', e => e.response.payloadData && this._page._frameManager.webSocketFrameReceived(e.requestId, e.response.opcode, e.response.payloadData)),
+      helper.addEventListener(session, 'Network.webSocketClosed', e => this._page._frameManager.webSocketClosed(e.requestId)),
+      helper.addEventListener(session, 'Network.webSocketFrameError', e => this._page._frameManager.webSocketError(e.requestId, e.errorMessage)),
     ];
   }
 

src/chromium/crPage.ts

@@ -241,12 +241,7 @@ export class CRPage implements PageDelegate {
       this._page._addConsoleMessage(event.type, args, toConsoleMessageLocation(event.stackTrace));
     });
     session.on('Runtime.exceptionThrown', exception => this._page.emit(Events.Page.PageError, exceptionToError(exception.exceptionDetails)));
-    session.on('Fetch.requestPaused', event => this._networkManager._onRequestPaused(event));
-    session.on('Fetch.authRequired', event => this._networkManager._onAuthRequired(event));
-    session.on('Network.requestWillBeSent', event => this._networkManager._onRequestWillBeSent(event));
-    session.on('Network.responseReceived', event => this._networkManager._onResponseReceived(event));
-    session.on('Network.loadingFinished', event => this._networkManager._onLoadingFinished(event));
-    session.on('Network.loadingFailed', event => this._networkManager._onLoadingFailed(event));
+    this._networkManager.instrumentNetworkEvents(session);
   }
 
   _onDetachedFromTarget(event: Protocol.Target.detachedFromTargetPayload) {

src/events.ts

@@ -38,7 +38,16 @@ export const Events = {
     FrameNavigated: 'framenavigated',
     Load: 'load',
     Popup: 'popup',
+    WebSocket: 'websocket',
     WorkerCreated: 'workercreated',
     WorkerDestroyed: 'workerdestroyed',
+  },
+
+  WebSocket: {
+    Close: 'close',
+    Error: 'error',
+    MessageReceived: 'messageReceived',
+    MessageSent: 'messageSent',
+    Open: 'open',
   }
 };

src/frames.ts

@@ -58,6 +58,7 @@ export type WaitForOptions = types.TimeoutOptions & { waitFor?: types.Visibility
 export class FrameManager {
   private _page: Page;
   private _frames = new Map<string, Frame>();
+  private _webSockets = new Map<string, network.WebSocket>();
   private _mainFrame: Frame;
   readonly _lifecycleWatchers = new Set<LifecycleWatcher>();
 
@@ -117,6 +118,7 @@ export class FrameManager {
     frame._lastDocumentId = documentId;
     this.frameLifecycleEvent(frameId, 'clear');
     this.clearInflightRequests(frame);
+    this.clearWebSockets(frame);
     if (!initial) {
       for (const watcher of this._lifecycleWatchers)
         watcher._onCommittedNewDocumentNavigation(frame);
@@ -184,6 +186,13 @@ export class FrameManager {
       this._startNetworkIdleTimer(frame, 'networkidle2');
   }
 
+  clearWebSockets(frame: Frame) {
+    // TODO: attributet sockets to frames.
+    if (frame.parentFrame())
+      return;
+    this._webSockets.clear();
+  }
+
   requestStarted(request: network.Request) {
     this._inflightRequestStarted(request);
     const frame = request.frame();
@@ -223,6 +232,50 @@ export class FrameManager {
       this._page.emit(Events.Page.RequestFailed, request);
   }
 
+  onWebSocketCreated(requestId: string, url: string) {
+    const ws = new network.WebSocket(url);
+    this._webSockets.set(requestId, ws);
+  }
+
+  onWebSocketRequest(requestId: string, headers: network.Headers) {
+    const ws = this._webSockets.get(requestId);
+    if (ws) {
+      ws._requestSent(headers);
+      this._page.emit(Events.Page.WebSocket, ws);
+    }
+  }
+
+  onWebSocketResponse(requestId: string, status: number, statusText: string, headers: network.Headers) {
+    const ws = this._webSockets.get(requestId);
+    if (ws)
+      ws._responseReceived(status, statusText, headers);
+  }
+
+  onWebSocketFrameSent(requestId: string, opcode: number, data: string) {
+    const ws = this._webSockets.get(requestId);
+    if (ws)
+      ws._frameSent(opcode, data);
+  }
+
+  webSocketFrameReceived(requestId: string, opcode: number, data: string) {
+    const ws = this._webSockets.get(requestId);
+    if (ws)
+      ws._frameReceived(opcode, data);
+  }
+
+  webSocketClosed(requestId: string) {
+    const ws = this._webSockets.get(requestId);
+    if (ws)
+      ws._closed();
+    this._webSockets.delete(requestId);
+  }
+
+  webSocketError(requestId: string, errorMessage: string): void {
+    const ws = this._webSockets.get(requestId);
+    if (ws)
+      ws._error(errorMessage);
+  }
+
   provisionalLoadFailed(documentId: string, error: string) {
     for (const watcher of this._lifecycleWatchers)
       watcher._onProvisionalLoadFailed(documentId, error);

src/network.ts

@@ -17,6 +17,7 @@
 import * as frames from './frames';
 import { assert } from './helper';
 import * as platform from './platform';
+import { Events } from './events';
 
 export type NetworkCookie = {
   name: string,
@@ -314,6 +315,70 @@ export interface RequestDelegate {
   continue(overrides: { url?: string; method?: string; postData?: string; headers?: Headers; }): Promise<void>;
 }
 
+export class WebSocket extends platform.EventEmitter {
+  private _url: string;
+  _status: number | null = null;
+  _statusText: string | null = null;
+  _requestHeaders: Headers | null = null;
+  _responseHeaders: Headers | null = null;
+
+  constructor(url: string) {
+    super();
+    this._url = url;
+  }
+
+  url(): string {
+    return this._url;
+  }
+
+  status(): number | null {
+    return this._status;
+  }
+
+  statusText(): string  | null {
+    return this._statusText;
+  }
+
+  requestHeaders(): Headers | null {
+    return this._requestHeaders;
+  }
+
+  responseHeaders(): Headers | null {
+    return this._responseHeaders;
+  }
+
+  _requestSent(headers: Headers) {
+    this._requestHeaders = {};
+    for (const [name, value] of Object.entries(headers))
+      this._requestHeaders[name.toLowerCase()] = value;
+  }
+
+  _responseReceived(status: number, statusText: string, headers: Headers) {
+    this._status = status;
+    this._statusText = statusText;
+    this._responseHeaders = {};
+    for (const [name, value] of Object.entries(headers))
+      this._responseHeaders[name.toLowerCase()] = value;
+    this.emit(Events.WebSocket.Open);
+  }
+
+  _frameSent(opcode: number, data: string) {
+    this.emit(Events.WebSocket.MessageSent, opcode === 2 ? Buffer.from(data, 'base64') : data);
+  }
+
+  _frameReceived(opcode: number, data: string) {
+    this.emit(Events.WebSocket.MessageReceived, opcode === 2 ? Buffer.from(data, 'base64') : data);
+  }
+
+  _error(errorMessage: string) {
+    this.emit(Events.WebSocket.Error, errorMessage);
+  }
+
+  _closed() {
+    this.emit(Events.WebSocket.Close);
+  }
+}
+
 // List taken from https://www.iana.org/assignments/http-status-codes/http-status-codes.xhtml with extra 306 and 418 codes.
 export const STATUS_TEXTS: { [status: string]: string } = {
   '100': 'Continue',

src/webkit/wkNetworkManager.ts

@@ -51,6 +51,13 @@ export class WKNetworkManager {
       helper.addEventListener(session, 'Network.responseReceived', e => this._onResponseReceived(e)),
       helper.addEventListener(session, 'Network.loadingFinished', e => this._onLoadingFinished(e)),
       helper.addEventListener(session, 'Network.loadingFailed', e => this._onLoadingFailed(e)),
+      helper.addEventListener(session, 'Network.webSocketCreated', e => this._page._frameManager.onWebSocketCreated(e.requestId, e.url)),
+      helper.addEventListener(session, 'Network.webSocketWillSendHandshakeRequest', e => this._page._frameManager.onWebSocketRequest(e.requestId, e.request.headers)),
+      helper.addEventListener(session, 'Network.webSocketHandshakeResponseReceived', e => this._page._frameManager.onWebSocketResponse(e.requestId, e.response.status, e.response.statusText, e.response.headers)),
+      helper.addEventListener(session, 'Network.webSocketFrameSent', e => e.response.payloadData && this._page._frameManager.onWebSocketFrameSent(e.requestId, e.response.opcode, e.response.payloadData)),
+      helper.addEventListener(session, 'Network.webSocketFrameReceived', e => e.response.payloadData && this._page._frameManager.webSocketFrameReceived(e.requestId, e.response.opcode, e.response.payloadData)),
+      helper.addEventListener(session, 'Network.webSocketClosed', e => this._page._frameManager.webSocketClosed(e.requestId)),
+      helper.addEventListener(session, 'Network.webSocketFrameError', e => this._page._frameManager.webSocketError(e.requestId, e.errorMessage)),
     ];
   }
 

test/network.spec.js

@@ -343,5 +343,116 @@ module.exports.describe = function({testRunner, expect, FFOX, CHROMIUM, WEBKIT})
       expect(error.message).toBe('Expected value of header "foo" to be String, but "number" is found.');
     });
   });
-};
 
+  describe.skip(FFOX)('WebSocket', function() {
+    it('should work', async({page, server}) => {
+      const value = await page.evaluate((port) => {
+        let cb;
+        const result = new Promise(f => cb = f);
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('message', data => { ws.close(); cb(data.data); });
+        return result;
+      }, server.PORT);
+      expect(value).toBe('incoming');
+    });
+    it('should emit open/close events', async({page, server}) => {
+      let socketClosed;
+      const socketClosePromise = new Promise(f => socketClosed = f);
+      const log = [];
+      page.on('websocket', ws => {
+        ws.on('open', () => log.push(`open<${ws.url()}>`));
+        ws.on('close', () => { log.push('close'); socketClosed(); });
+      });
+      page.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => ws.close());
+      }, server.PORT);
+      await socketClosePromise;
+      expect(log.join(':')).toBe(`open<ws://localhost:${server.PORT}/ws>:close`);
+    });
+    it('should expose status', async({page, server}) => {
+      let callback;
+      const result = new Promise(f => callback = f);
+      page.on('websocket', ws => ws.on('open', () => callback(ws)));
+      page.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => ws.close());
+      }, server.PORT);
+      const ws = await result;
+      expect(ws.status()).toBe(101);
+      expect(ws.statusText()).toBe('Switching Protocols');
+    });
+    it('should emit error', async({page, server}) => {
+      let callback;
+      const result = new Promise(f => callback = f);
+      page.on('websocket', ws => ws.on('error', callback));
+      page.evaluate((port) => {
+        new WebSocket('ws://localhost:' + port + '/bogus-ws');
+      }, server.PORT);
+      const message = await result;
+      expect(message).toContain('Unexpected response code: 400');
+    });
+    it('should emit frame events', async({page, server}) => {
+      let socketClosed;
+      const socketClosePromise = new Promise(f => socketClosed = f);
+      const log = [];
+      page.on('websocket', ws => {
+        ws.on('open', () => log.push('open'));
+        ws.on('messageSent', d => log.push('sent<' + d + '>'));
+        ws.on('messageReceived', d => log.push('received<' + d + '>'));
+        ws.on('close', () => { log.push('close'); socketClosed(); });
+      });
+      page.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => ws.send('outgoing'));
+        ws.addEventListener('message', () => { ws.close(); });
+      }, server.PORT);
+      await socketClosePromise;
+      expect(log.join(':')).toBe('open:sent<outgoing>:received<incoming>:close');
+    });
+    it('should emit binary frame events', async({page, server}) => {
+      let doneCallback;
+      const donePromise = new Promise(f => doneCallback = f);
+      const sent = [];
+      page.on('websocket', ws => {
+        ws.on('close', doneCallback);
+        ws.on('messageSent', d => sent.push(d));
+      });
+      page.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => {
+          const binary = new Uint8Array(5);
+          for (let i = 0; i < 5; ++i)
+            binary[i] = i;
+          ws.send('text');
+          ws.send(binary);
+          ws.close();
+        });
+      }, server.PORT);
+      await donePromise;
+      expect(sent[0]).toBe('text');
+      for (let i = 0; i < 5; ++i)
+        expect(sent[1][i]).toBe(i);
+    });
+    it('should report headers', async({page, server}) => {
+      let socketClosed;
+      let requestHeaders;
+      let responseHeaders;
+      const socketClosePromise = new Promise(f => socketClosed = f);
+      page.on('websocket', ws => {
+        requestHeaders = ws.requestHeaders();
+        ws.on('open', () => {
+          responseHeaders = ws.responseHeaders();
+        });
+        ws.on('close', socketClosed);
+      });
+      page.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => ws.close());
+      }, server.PORT);
+      await socketClosePromise;
+      expect(requestHeaders['connection']).toBe('Upgrade');
+      expect(responseHeaders['upgrade']).toBe('websocket');
+    });
+  });
+};

test/workers.spec.js

@@ -113,5 +113,25 @@ module.exports.describe = function({testRunner, expect, FFOX, CHROMIUM, WEBKIT})
       expect(response.request()).toBe(request);
       expect(response.ok()).toBe(true);
     });
+    it.skip(FFOX)('should report web socket activity', async function({page, server}) {
+      const [worker] = await Promise.all([
+        page.waitForEvent('workercreated'),
+        page.goto(server.PREFIX + '/worker/worker.html'),
+      ]);
+      const log = [];
+      let socketClosed;
+      const socketClosePromise = new Promise(f => socketClosed = f);
+      page.on('websocket', ws => {
+        ws.on('open', () => log.push(`open<${ws.url()}>`));
+        ws.on('close', () => { log.push('close'); socketClosed(); });
+      });
+      worker.evaluate((port) => {
+        const ws = new WebSocket('ws://localhost:' + port + '/ws');
+        ws.addEventListener('open', () => ws.close());
+      }, server.PORT);
+
+      await socketClosePromise;
+      expect(log.join(':')).toBe(`open<ws://localhost:${server.PORT}/ws>:close`);
+    });
   });
 };

utils/testserver/index.js

@@ -63,7 +63,7 @@ class TestServer {
     else
       this._server = http.createServer(this._onRequest.bind(this));
     this._server.on('connection', socket => this._onSocket(socket));
-    this._wsServer = new WebSocketServer({server: this._server});
+    this._wsServer = new WebSocketServer({server: this._server, path: '/ws'});
     this._wsServer.on('connection', this._onWebSocketConnection.bind(this));
     this._server.listen(port);
     this._dirPath = dirPath;
@@ -264,8 +264,8 @@ class TestServer {
     });
   }
 
-  _onWebSocketConnection(connection) {
-    connection.send('opened');
+  _onWebSocketConnection(ws) {
+    ws.send('incoming');
   }
 }