browser(firefox): instrument websockets (#4287)

Author: aslushnikov

browser_patches/firefox/BUILD_NUMBER

@@ -1,2 +1,2 @@
-1197
-Changed: [email protected] Wed 28 Oct 2020 02:43:11 PM PDT
+1198
+Changed: [email protected] Thu 29 Oct 2020 04:23:02 PM PDT

browser_patches/firefox/juggler/content/FrameTree.js

@@ -24,6 +24,10 @@ class FrameTree {
     this._browsingContextGroup.__jugglerFrameTrees.add(this);
     this._scriptsToEvaluateOnNewDocument = new Map();
 
+    this._webSocketEventService = Cc[
+      "@mozilla.org/websocketevent/service;1"
+    ].getService(Ci.nsIWebSocketEventService);
+
     this._bindings = new Map();
     this._runtime = new Runtime(false /* isWorker */);
     this._workers = new Map();
@@ -200,7 +204,7 @@ class FrameTree {
 
     if (isStart) {
       // Starting a new navigation.
-      frame._pendingNavigationId = this._channelId(channel);
+      frame._pendingNavigationId = channelId(channel);
       frame._pendingNavigationURL = channel.URI.spec;
       this.emit(FrameTree.Events.NavigationStarted, frame);
     } else if (isTransferring || (isStop && frame._pendingNavigationId && !status)) {
@@ -241,14 +245,6 @@ class FrameTree {
     }
   }
 
-  _channelId(channel) {
-    if (channel instanceof Ci.nsIIdentChannel) {
-      const identChannel = channel.QueryInterface(Ci.nsIIdentChannel);
-      return String(identChannel.channelId);
-    }
-    return helper.generateId();
-  }
-
   _onDocShellCreated(docShell) {
     // Bug 1142752: sometimes, the docshell appears to be immediately
     // destroyed, bailout early to prevent random exceptions.
@@ -302,6 +298,11 @@ FrameTree.Events = {
   GlobalObjectCreated: 'globalobjectcreated',
   WorkerCreated: 'workercreated',
   WorkerDestroyed: 'workerdestroyed',
+  WebSocketCreated: 'websocketcreated',
+  WebSocketOpened: 'websocketopened',
+  WebSocketClosed: 'websocketclosed',
+  WebSocketFrameReceived: 'websocketframereceived',
+  WebSocketFrameSent: 'websocketframesent',
   NavigationStarted: 'navigationstarted',
   NavigationCommitted: 'navigationcommitted',
   NavigationAborted: 'navigationaborted',
@@ -332,6 +333,90 @@ class Frame {
 
     this._textInputProcessor = null;
     this._executionContext = null;
+
+    this._webSocketListenerInnerWindowId = 0;
+    // WebSocketListener calls frameReceived event before webSocketOpened.
+    // To avoid this, serialize event reporting.
+    this._webSocketInfos = new Map();
+
+    const dispatchWebSocketFrameReceived = (webSocketSerialID, frame) => this._frameTree.emit(FrameTree.Events.WebSocketFrameReceived, {
+      frameId: this._frameId,
+      wsid: webSocketSerialID + '',
+      opcode: frame.opCode,
+      data: frame.opCode !== 1 ? btoa(frame.payload) : frame.payload,
+    });
+    this._webSocketListener = {
+      QueryInterface: ChromeUtils.generateQI([Ci.nsIWebSocketEventListener, ]),
+
+      webSocketCreated: (webSocketSerialID, uri, protocols) => {
+        this._frameTree.emit(FrameTree.Events.WebSocketCreated, {
+          frameId: this._frameId,
+          wsid: webSocketSerialID + '',
+          requestURL: uri,
+        });
+        this._webSocketInfos.set(webSocketSerialID, {
+          opened: false,
+          pendingIncomingFrames: [],
+        });
+      },
+
+      webSocketOpened: (webSocketSerialID, effectiveURI, protocols, extensions, httpChannelId) => {
+        this._frameTree.emit(FrameTree.Events.WebSocketOpened, {
+          frameId: this._frameId,
+          requestId: httpChannelId + '',
+          wsid: webSocketSerialID + '',
+          effectiveURL: effectiveURI,
+        });
+        const info = this._webSocketInfos.get(webSocketSerialID);
+        info.opened = true;
+        for (const frame of info.pendingIncomingFrames)
+          dispatchWebSocketFrameReceived(webSocketSerialID, frame);
+      },
+
+      webSocketMessageAvailable: (webSocketSerialID, data, messageType) => {
+        // We don't use this event.
+      },
+
+      webSocketClosed: (webSocketSerialID, wasClean, code, reason) => {
+        this._webSocketInfos.delete(webSocketSerialID);
+        let error = '';
+        if (!wasClean) {
+          const keys = Object.keys(Ci.nsIWebSocketChannel);
+          for (const key of keys) {
+            if (Ci.nsIWebSocketChannel[key] === code)
+              error = key;
+          }
+        }
+        this._frameTree.emit(FrameTree.Events.WebSocketClosed, {
+          frameId: this._frameId,
+          wsid: webSocketSerialID + '',
+          error,
+        });
+      },
+
+      frameReceived: (webSocketSerialID, frame) => {
+        // Report only text and binary frames.
+        if (frame.opCode !== 1 && frame.opCode !== 2)
+          return;
+        const info = this._webSocketInfos.get(webSocketSerialID);
+        if (info.opened)
+          dispatchWebSocketFrameReceived(webSocketSerialID, frame);
+        else
+          info.pendingIncomingFrames.push(frame);
+      },
+
+      frameSent: (webSocketSerialID, frame) => {
+        // Report only text and binary frames.
+        if (frame.opCode !== 1 && frame.opCode !== 2)
+          return;
+        this._frameTree.emit(FrameTree.Events.WebSocketFrameSent, {
+          frameId: this._frameId,
+          wsid: webSocketSerialID + '',
+          opcode: frame.opCode,
+          data: frame.opCode !== 1 ? btoa(frame.payload) : frame.payload,
+        });
+      },
+    };
   }
 
   dispose() {
@@ -354,6 +439,12 @@ class Frame {
   }
 
   _onGlobalObjectCleared() {
+    const webSocketService = this._frameTree._webSocketEventService;
+    if (this._webSocketListenerInnerWindowId)
+      webSocketService.removeListener(this._webSocketListenerInnerWindowId, this._webSocketListener);
+    this._webSocketListenerInnerWindowId = this.domWindow().windowGlobalChild.innerWindowId;
+    webSocketService.addListener(this._webSocketListenerInnerWindowId, this._webSocketListener);
+
     if (this._executionContext)
       this._runtime.destroyExecutionContext(this._executionContext);
     this._executionContext = this._runtime.createExecutionContext(this.domWindow(), this.domWindow(), {
@@ -473,6 +564,15 @@ class Worker {
   }
 }
 
+function channelId(channel) {
+  if (channel instanceof Ci.nsIIdentChannel) {
+    const identChannel = channel.QueryInterface(Ci.nsIIdentChannel);
+    return String(identChannel.channelId);
+  }
+  return helper.generateId();
+}
+
+
 var EXPORTED_SYMBOLS = ['FrameTree'];
 this.FrameTree = FrameTree;
 

browser_patches/firefox/juggler/content/PageAgent.js

@@ -177,6 +177,11 @@ class PageAgent {
       helper.on(this._frameTree, 'pageready', () => this._browserPage.emit('pageReady', {})),
       helper.on(this._frameTree, 'workercreated', this._onWorkerCreated.bind(this)),
       helper.on(this._frameTree, 'workerdestroyed', this._onWorkerDestroyed.bind(this)),
+      helper.on(this._frameTree, 'websocketcreated', event => this._browserPage.emit('webSocketCreated', event)),
+      helper.on(this._frameTree, 'websocketopened', event => this._browserPage.emit('webSocketOpened', event)),
+      helper.on(this._frameTree, 'websocketframesent', event => this._browserPage.emit('webSocketFrameSent', event)),
+      helper.on(this._frameTree, 'websocketframereceived', event => this._browserPage.emit('webSocketFrameReceived', event)),
+      helper.on(this._frameTree, 'websocketclosed', event => this._browserPage.emit('webSocketClosed', event)),
       helper.addObserver(this._onWindowOpen.bind(this), 'webNavigation-createdNavigationTarget-from-js'),
       this._runtime.events.onErrorFromWorker((domWindow, message, stack) => {
         const frame = this._frameTree.frameForDocShell(domWindow.docShell);

browser_patches/firefox/juggler/protocol/PageHandler.js

@@ -110,6 +110,12 @@ class PageHandler {
         runtimeConsole: emitProtocolEvent('Runtime.console'),
         runtimeExecutionContextCreated: emitProtocolEvent('Runtime.executionContextCreated'),
         runtimeExecutionContextDestroyed: emitProtocolEvent('Runtime.executionContextDestroyed'),
+
+        webSocketCreated: emitProtocolEvent('Page.webSocketCreated'),
+        webSocketOpened: emitProtocolEvent('Page.webSocketOpened'),
+        webSocketClosed: emitProtocolEvent('Page.webSocketClosed'),
+        webSocketFrameReceived: emitProtocolEvent('Page.webSocketFrameReceived'),
+        webSocketFrameSent: emitProtocolEvent('Page.webSocketFrameSent'),
       }),
     ];
   }

browser_patches/firefox/juggler/protocol/Protocol.js

@@ -668,6 +668,34 @@ const Page = {
       screencastId: t.String,
       file: t.String,
     },
+    'webSocketCreated': {
+      frameId: t.String,
+      wsid: t.String,
+      requestURL: t.String,
+    },
+    'webSocketOpened': {
+      frameId: t.String,
+      requestId: t.String,
+      wsid: t.String,
+      effectiveURL: t.String,
+    },
+    'webSocketClosed': {
+      frameId: t.String,
+      wsid: t.String,
+      error: t.String,
+    },
+    'webSocketFrameSent': {
+      frameId: t.String,
+      wsid: t.String,
+      opcode: t.Number,
+      data: t.String,
+    },
+    'webSocketFrameReceived': {
+      frameId: t.String,
+      wsid: t.String,
+      opcode: t.Number,
+      data: t.String,
+    },
   },
 
   methods: {