feat(websocket): implement Web Sockets for Chromium & WebKit (#4234)

docs/api.md

@@ -23,6 +23,7 @@
 - [class: Response](#class-response)
 - [class: Selectors](#class-selectors)
 - [class: Route](#class-route)
+- [class: WebSocket](#class-websocket)
 - [class: TimeoutError](#class-timeouterror)
 - [class: Accessibility](#class-accessibility)
 - [class: Worker](#class-worker)
@@ -738,6 +739,7 @@ page.removeListener('request', logRequest);
 - [event: 'requestfailed'](#event-requestfailed)
 - [event: 'requestfinished'](#event-requestfinished)
 - [event: 'response'](#event-response)
+- [event: 'websocket'](#event-websocket)
 - [event: 'worker'](#event-worker)
 - [page.$(selector)](#pageselector)
 - [page.$$(selector)](#pageselector-1)
@@ -949,6 +951,11 @@ Emitted when a request finishes successfully after downloading the response body
 
 Emitted when [response] status and headers are received for a request. For a successful response, the sequence of events is `request`, `response` and `requestfinished`.
 
+#### event: 'websocket'
+- <[WebSocket]> websocket
+
+Emitted when <[WebSocket]> request is sent.
+
 #### event: 'worker'
 - <[Worker]>
 
@@ -4133,6 +4140,45 @@ await page.route('**/xhr_endpoint', route => route.fulfill({ path: 'mock_data.js
 - returns: <[Request]> A request to be routed.
 
 
+### class: WebSocket
+
+The [WebSocket] class represents websocket connections in the page.
+
+<!-- GEN:toc -->
+- [event: 'close'](#event-close-2)
+- [event: 'framereceived'](#event-framereceived)
+- [event: 'framesent'](#event-framesent)
+- [event: 'socketerror'](#event-socketerror)
+- [webSocket.url()](#websocketurl)
+<!-- GEN:stop -->
+
+#### event: 'close'
+
+Fired when the websocket closes.
+
+#### event: 'framereceived'
+- <[Object]> web socket frame data
+  - `payload` <[string]|[Buffer]> frame payload
+
+Fired when the websocket recieves a frame.
+
+#### event: 'framesent'
+- <[Object]> web socket frame data
+  - `payload` <[string]|[Buffer]> frame payload
+
+Fired when the websocket sends a frame.
+
+#### event: 'socketerror'
+- <[String]> the error message
+
+Fired when the websocket has an error.
+
+#### webSocket.url()
+- returns: <[string]>
+
+Contains the URL of the WebSocket.
+
+
 ### class: TimeoutError
 
 * extends: [Error]
@@ -4233,7 +4279,7 @@ for (const worker of page.workers())
 ```
 
 <!-- GEN:toc -->
-- [event: 'close'](#event-close-2)
+- [event: 'close'](#event-close-3)
 - [worker.evaluate(pageFunction[, arg])](#workerevaluatepagefunction-arg)
 - [worker.evaluateHandle(pageFunction[, arg])](#workerevaluatehandlepagefunction-arg)
 - [worker.url()](#workerurl)
@@ -4269,7 +4315,7 @@ If the function passed to the `worker.evaluateHandle` returns a [Promise], then
 ### class: BrowserServer
 
 <!-- GEN:toc -->
-- [event: 'close'](#event-close-3)
+- [event: 'close'](#event-close-4)
 - [browserServer.close()](#browserserverclose)
 - [browserServer.kill()](#browserserverkill)
 - [browserServer.process()](#browserserverprocess)

src/client/api.ts

@@ -29,7 +29,7 @@ export { TimeoutError } from '../utils/errors';
 export { Frame } from './frame';
 export { Keyboard, Mouse, Touchscreen } from './input';
 export { JSHandle } from './jsHandle';
-export { Request, Response, Route } from './network';
+export { Request, Response, Route, WebSocket } from './network';
 export { Page } from './page';
 export { Selectors } from './selectors';
 export { Video } from './video';

src/client/connection.ts

@@ -21,7 +21,7 @@ import { ChannelOwner } from './channelOwner';
 import { ElementHandle } from './elementHandle';
 import { Frame } from './frame';
 import { JSHandle } from './jsHandle';
-import { Request, Response, Route } from './network';
+import { Request, Response, Route, WebSocket } from './network';
 import { Page, BindingCall } from './page';
 import { Worker } from './worker';
 import { ConsoleMessage } from './consoleMessage';
@@ -226,6 +226,9 @@ export class Connection {
       case 'Selectors':
         result = new SelectorsOwner(parent, type, guid, initializer);
         break;
+      case 'WebSocket':
+        result = new WebSocket(parent, type, guid, initializer);
+        break;
       case 'Worker':
         result = new Worker(parent, type, guid, initializer);
         break;

src/client/events.ts

@@ -49,9 +49,17 @@ export const Events = {
     FrameNavigated: 'framenavigated',
     Load: 'load',
     Popup: 'popup',
+    WebSocket: 'websocket',
     Worker: 'worker',
   },
 
+  WebSocket: {
+    Close: 'close',
+    Error: 'socketerror',
+    FrameReceived: 'framereceived',
+    FrameSent: 'framesent',
+  },
+
   Worker: {
     Close: 'close',
   },

src/client/network.ts

@@ -23,6 +23,7 @@ import * as fs from 'fs';
 import * as mime from 'mime';
 import * as util from 'util';
 import { isString, headersObjectToArray, headersArrayToObject } from '../utils/utils';
+import { Events } from './events';
 
 export type NetworkCookie = {
   name: string,
@@ -312,6 +313,30 @@ export class Response extends ChannelOwner<channels.ResponseChannel, channels.Re
   }
 }
 
+export class WebSocket extends ChannelOwner<channels.WebSocketChannel, channels.WebSocketInitializer> {
+  static from(webSocket: channels.WebSocketChannel): WebSocket {
+    return (webSocket as any)._object;
+  }
+
+  constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.WebSocketInitializer) {
+    super(parent, type, guid, initializer);
+    this._channel.on('frameSent', (event: { opcode: number, data: string }) => {
+      const payload = event.opcode === 2 ? Buffer.from(event.data, 'base64') : event.data;
+      this.emit(Events.WebSocket.FrameSent, { payload });
+    });
+    this._channel.on('frameReceived', (event: { opcode: number, data: string }) => {
+      const payload = event.opcode === 2 ? Buffer.from(event.data, 'base64') : event.data;
+      this.emit(Events.WebSocket.FrameReceived, { payload });
+    });
+    this._channel.on('error', ({ error }) => this.emit(Events.WebSocket.Error, error));
+    this._channel.on('close', () => this.emit(Events.WebSocket.Close));
+  }
+
+  url(): string {
+    return this._initializer.url;
+  }
+}
+
 export function validateHeaders(headers: Headers) {
   for (const key of Object.keys(headers)) {
     const value = headers[key];

src/client/page.ts

@@ -31,7 +31,7 @@ import { Worker } from './worker';
 import { Frame, verifyLoadState, WaitForNavigationOptions } from './frame';
 import { Keyboard, Mouse, Touchscreen } from './input';
 import { assertMaxArguments, Func1, FuncOn, SmartHandle, serializeArgument, parseResult, JSHandle } from './jsHandle';
-import { Request, Response, Route, RouteHandler, validateHeaders } from './network';
+import { Request, Response, Route, RouteHandler, WebSocket, validateHeaders } from './network';
 import { FileChooser } from './fileChooser';
 import { Buffer } from 'buffer';
 import { ChromiumCoverage } from './chromiumCoverage';
@@ -130,6 +130,7 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
     this._channel.on('response', ({ response }) => this.emit(Events.Page.Response, Response.from(response)));
     this._channel.on('route', ({ route, request }) => this._onRoute(Route.from(route), Request.from(request)));
     this._channel.on('video', ({ relativePath }) => this.video()!._setRelativePath(relativePath));
+    this._channel.on('webSocket', ({ webSocket }) => this.emit(Events.Page.WebSocket, WebSocket.from(webSocket)));
     this._channel.on('worker', ({ worker }) => this._onWorker(Worker.from(worker)));
 
     if (this._browserContext._browserName === 'chromium') {

src/dispatchers/networkDispatchers.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { Request, Response, Route } from '../server/network';
+import { Request, Response, Route, WebSocket } from '../server/network';
 import * as channels from '../protocol/channels';
 import { Dispatcher, DispatcherScope, lookupNullableDispatcher, existingDispatcher } from './dispatcher';
 import { FrameDispatcher } from './frameDispatcher';
@@ -98,3 +98,15 @@ export class RouteDispatcher extends Dispatcher<Route, channels.RouteInitializer
     await this._object.abort(params.errorCode || 'failed');
   }
 }
+
+export class WebSocketDispatcher extends Dispatcher<WebSocket, channels.WebSocketInitializer> implements channels.WebSocketChannel {
+  constructor(scope: DispatcherScope, webSocket: WebSocket) {
+    super(scope, webSocket, 'WebSocket', {
+      url: webSocket.url(),
+    });
+    webSocket.on(WebSocket.Events.FrameSent, (event: { opcode: number, data: string }) => this._dispatchEvent('frameSent', event));
+    webSocket.on(WebSocket.Events.FrameReceived, (event: { opcode: number, data: string }) => this._dispatchEvent('frameReceived', event));
+    webSocket.on(WebSocket.Events.Error, (error: string) => this._dispatchEvent('error', { error }));
+    webSocket.on(WebSocket.Events.Close, () => this._dispatchEvent('close', {}));
+  }
+}

src/dispatchers/pageDispatcher.ts

@@ -25,7 +25,7 @@ import { ConsoleMessageDispatcher } from './consoleMessageDispatcher';
 import { DialogDispatcher } from './dialogDispatcher';
 import { DownloadDispatcher } from './downloadDispatcher';
 import { FrameDispatcher } from './frameDispatcher';
-import { RequestDispatcher, ResponseDispatcher, RouteDispatcher } from './networkDispatchers';
+import { RequestDispatcher, ResponseDispatcher, RouteDispatcher, WebSocketDispatcher } from './networkDispatchers';
 import { serializeResult, parseArgument, JSHandleDispatcher } from './jsHandleDispatcher';
 import { ElementHandleDispatcher, createHandle } from './elementHandlerDispatcher';
 import { FileChooser } from '../server/fileChooser';
@@ -72,6 +72,7 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageInitializer> i
     }));
     page.on(Page.Events.Response, response => this._dispatchEvent('response', { response: new ResponseDispatcher(this._scope, response) }));
     page.on(Page.Events.VideoStarted, (video: Video) => this._dispatchEvent('video', {  relativePath: video._relativePath }));
+    page.on(Page.Events.WebSocket, webSocket => this._dispatchEvent('webSocket', { webSocket: new WebSocketDispatcher(this._scope, webSocket) }));
     page.on(Page.Events.Worker, worker => this._dispatchEvent('worker', { worker: new WorkerDispatcher(this._scope, worker) }));
   }
 

src/protocol/channels.ts

@@ -694,6 +694,7 @@ export interface PageChannel extends Channel {
   on(event: 'response', callback: (params: PageResponseEvent) => void): this;
   on(event: 'route', callback: (params: PageRouteEvent) => void): this;
   on(event: 'video', callback: (params: PageVideoEvent) => void): this;
+  on(event: 'webSocket', callback: (params: PageWebSocketEvent) => void): this;
   on(event: 'worker', callback: (params: PageWorkerEvent) => void): this;
   setDefaultNavigationTimeoutNoReply(params: PageSetDefaultNavigationTimeoutNoReplyParams, metadata?: Metadata): Promise<PageSetDefaultNavigationTimeoutNoReplyResult>;
   setDefaultTimeoutNoReply(params: PageSetDefaultTimeoutNoReplyParams, metadata?: Metadata): Promise<PageSetDefaultTimeoutNoReplyResult>;
@@ -782,6 +783,9 @@ export type PageRouteEvent = {
 export type PageVideoEvent = {
   relativePath: string,
 };
+export type PageWebSocketEvent = {
+  webSocket: WebSocketChannel,
+};
 export type PageWorkerEvent = {
   worker: WorkerChannel,
 };
@@ -2185,6 +2189,31 @@ export type ResponseFinishedResult = {
   error?: string,
 };
 
+// ----------- WebSocket -----------
+export type WebSocketInitializer = {
+  url: string,
+};
+export interface WebSocketChannel extends Channel {
+  on(event: 'open', callback: (params: WebSocketOpenEvent) => void): this;
+  on(event: 'frameSent', callback: (params: WebSocketFrameSentEvent) => void): this;
+  on(event: 'frameReceived', callback: (params: WebSocketFrameReceivedEvent) => void): this;
+  on(event: 'error', callback: (params: WebSocketErrorEvent) => void): this;
+  on(event: 'close', callback: (params: WebSocketCloseEvent) => void): this;
+}
+export type WebSocketOpenEvent = {};
+export type WebSocketFrameSentEvent = {
+  opcode: number,
+  data: string,
+};
+export type WebSocketFrameReceivedEvent = {
+  opcode: number,
+  data: string,
+};
+export type WebSocketErrorEvent = {
+  error: string,
+};
+export type WebSocketCloseEvent = {};
+
 // ----------- ConsoleMessage -----------
 export type ConsoleMessageInitializer = {
   type: string,

src/protocol/protocol.yml

@@ -945,6 +945,10 @@ Page:
       parameters:
         relativePath: string
 
+    webSocket:
+      parameters:
+        webSocket: WebSocket
+
     worker:
       parameters:
         worker: Worker
@@ -1844,6 +1848,32 @@ Response:
         error: string?
 
 
+WebSocket:
+  type: interface
+
+  initializer:
+    url: string
+
+  events:
+    open:
+
+    frameSent:
+      parameters:
+        opcode: number
+        data: string
+
+    frameReceived:
+      parameters:
+        opcode: number
+        data: string
+
+    error:
+      parameters:
+        error: string
+
+    close:
+
+
 ConsoleMessage:
   type: interface
 

src/server/chromium/crNetworkManager.ts

@@ -55,6 +55,13 @@ export class CRNetworkManager {
       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)),
+      helper.addEventListener(session, 'Network.webSocketHandshakeResponseReceived', e => this._page._frameManager.onWebSocketResponse(e.requestId, e.response.status, e.response.statusText)),
+      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)),
     ];
   }