From c83ac444b457e106ebb6812f43a65beec65bd6a3 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Fri, 6 Nov 2020 12:30:16 -0800
Subject: [PATCH] api(websocket): do not send websocket frames without a
 listener (#4361)

Once user adds 'websocket' listener, we start sending the frames
for all websockets on the page.
---
 src/client/page.ts                | 22 ++++++++++++++--------
 src/dispatchers/pageDispatcher.ts |  4 ++++
 src/protocol/channels.ts          |  8 ++++++++
 src/protocol/protocol.yml         |  4 ++++
 src/protocol/validator.ts         |  3 +++
 src/server/frames.ts              |  4 ++++
 src/server/page.ts                |  5 +++++
 7 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/src/client/page.ts b/src/client/page.ts
index 75eb9051d7..0a1f1f118c 100644
--- a/src/client/page.ts
+++ b/src/client/page.ts
@@ -577,19 +577,19 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
   }
 
   on(event: string | symbol, listener: Listener): this {
-    if (event === Events.Page.FileChooser) {
-      if (!this.listenerCount(event))
-        this._channel.setFileChooserInterceptedNoReply({ intercepted: true });
-    }
+    if (event === Events.Page.FileChooser && !this.listenerCount(event))
+      this._channel.setFileChooserInterceptedNoReply({ intercepted: true });
+    if (event === Events.Page.WebSocket && !this.listenerCount(event))
+      this._channel.setWebSocketFramesReportingEnabledNoReply({ enabled: true });
     super.on(event, listener);
     return this;
   }
 
   addListener(event: string | symbol, listener: Listener): this {
-    if (event === Events.Page.FileChooser) {
-      if (!this.listenerCount(event))
-        this._channel.setFileChooserInterceptedNoReply({ intercepted: true });
-    }
+    if (event === Events.Page.FileChooser && !this.listenerCount(event))
+      this._channel.setFileChooserInterceptedNoReply({ intercepted: true });
+    if (event === Events.Page.WebSocket && !this.listenerCount(event))
+      this._channel.setWebSocketFramesReportingEnabledNoReply({ enabled: true });
     super.addListener(event, listener);
     return this;
   }
@@ -598,6 +598,9 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
     super.off(event, listener);
     if (event === Events.Page.FileChooser && !this.listenerCount(event))
       this._channel.setFileChooserInterceptedNoReply({ intercepted: false });
+    // Note: we do not stop reporting web socket frames, since
+    // user might not listen to 'websocket' anymore, but still have
+    // a functioning WebSocket object.
     return this;
   }
 
@@ -605,6 +608,9 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
     super.removeListener(event, listener);
     if (event === Events.Page.FileChooser && !this.listenerCount(event))
       this._channel.setFileChooserInterceptedNoReply({ intercepted: false });
+    // Note: we do not stop reporting web socket frames, since
+    // user might not listen to 'websocket' anymore, but still have
+    // a functioning WebSocket object.
     return this;
   }
 
diff --git a/src/dispatchers/pageDispatcher.ts b/src/dispatchers/pageDispatcher.ts
index 1bd863ba5c..bd9476e5d6 100644
--- a/src/dispatchers/pageDispatcher.ts
+++ b/src/dispatchers/pageDispatcher.ts
@@ -158,6 +158,10 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageInitializer> i
     await this._page._setFileChooserIntercepted(params.intercepted);
   }
 
+  async setWebSocketFramesReportingEnabledNoReply(params: channels.PageSetWebSocketFramesReportingEnabledNoReplyParams): Promise<void> {
+    this._page._setWebSocketFramesReportingEnabled(params.enabled);
+  }
+
   async keyboardDown(params: channels.PageKeyboardDownParams): Promise<void> {
     await this._page.keyboard.down(params.key);
   }
diff --git a/src/protocol/channels.ts b/src/protocol/channels.ts
index eac8166c44..f7aa66a625 100644
--- a/src/protocol/channels.ts
+++ b/src/protocol/channels.ts
@@ -727,6 +727,7 @@ export interface PageChannel extends Channel {
   setDefaultNavigationTimeoutNoReply(params: PageSetDefaultNavigationTimeoutNoReplyParams, metadata?: Metadata): Promise<PageSetDefaultNavigationTimeoutNoReplyResult>;
   setDefaultTimeoutNoReply(params: PageSetDefaultTimeoutNoReplyParams, metadata?: Metadata): Promise<PageSetDefaultTimeoutNoReplyResult>;
   setFileChooserInterceptedNoReply(params: PageSetFileChooserInterceptedNoReplyParams, metadata?: Metadata): Promise<PageSetFileChooserInterceptedNoReplyResult>;
+  setWebSocketFramesReportingEnabledNoReply(params: PageSetWebSocketFramesReportingEnabledNoReplyParams, metadata?: Metadata): Promise<PageSetWebSocketFramesReportingEnabledNoReplyResult>;
   addInitScript(params: PageAddInitScriptParams, metadata?: Metadata): Promise<PageAddInitScriptResult>;
   close(params: PageCloseParams, metadata?: Metadata): Promise<PageCloseResult>;
   emulateMedia(params: PageEmulateMediaParams, metadata?: Metadata): Promise<PageEmulateMediaResult>;
@@ -838,6 +839,13 @@ export type PageSetFileChooserInterceptedNoReplyOptions = {
 
 };
 export type PageSetFileChooserInterceptedNoReplyResult = void;
+export type PageSetWebSocketFramesReportingEnabledNoReplyParams = {
+  enabled: boolean,
+};
+export type PageSetWebSocketFramesReportingEnabledNoReplyOptions = {
+
+};
+export type PageSetWebSocketFramesReportingEnabledNoReplyResult = void;
 export type PageAddInitScriptParams = {
   source: string,
 };
diff --git a/src/protocol/protocol.yml b/src/protocol/protocol.yml
index 4d86793d35..4e357e4583 100644
--- a/src/protocol/protocol.yml
+++ b/src/protocol/protocol.yml
@@ -624,6 +624,10 @@ Page:
       parameters:
         intercepted: boolean
 
+    setWebSocketFramesReportingEnabledNoReply:
+      parameters:
+        enabled: boolean
+
     addInitScript:
       parameters:
         source: string
diff --git a/src/protocol/validator.ts b/src/protocol/validator.ts
index a6c46450ef..1d9a554a23 100644
--- a/src/protocol/validator.ts
+++ b/src/protocol/validator.ts
@@ -330,6 +330,9 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
   scheme.PageSetFileChooserInterceptedNoReplyParams = tObject({
     intercepted: tBoolean,
   });
+  scheme.PageSetWebSocketFramesReportingEnabledNoReplyParams = tObject({
+    enabled: tBoolean,
+  });
   scheme.PageAddInitScriptParams = tObject({
     source: tString,
   });
diff --git a/src/server/frames.ts b/src/server/frames.ts
index 69c430688f..e89ef8aad5 100644
--- a/src/server/frames.ts
+++ b/src/server/frames.ts
@@ -359,12 +359,16 @@ export class FrameManager {
   }
 
   onWebSocketFrameSent(requestId: string, opcode: number, data: string) {
+    if (!this._page._webSocketFramesReportingEnabled)
+      return;
     const ws = this._webSockets.get(requestId);
     if (ws)
       ws.frameSent(opcode, data);
   }
 
   webSocketFrameReceived(requestId: string, opcode: number, data: string) {
+    if (!this._page._webSocketFramesReportingEnabled)
+      return;
     const ws = this._webSockets.get(requestId);
     if (ws)
       ws.frameReceived(opcode, data);
diff --git a/src/server/page.ts b/src/server/page.ts
index 12b9c26bfd..dc0c0d0147 100644
--- a/src/server/page.ts
+++ b/src/server/page.ts
@@ -145,6 +145,7 @@ export class Page extends EventEmitter {
   _ownedContext: BrowserContext | undefined;
   readonly selectors: Selectors;
   _video: Video | null = null;
+  _webSocketFramesReportingEnabled = false;
 
   constructor(delegate: PageDelegate, browserContext: BrowserContext) {
     super();
@@ -419,6 +420,10 @@ export class Page extends EventEmitter {
     await this._delegate.setFileChooserIntercepted(enabled);
   }
 
+  _setWebSocketFramesReportingEnabled(enabled: boolean) {
+    this._webSocketFramesReportingEnabled = enabled;
+  }
+
   videoStarted(video: Video) {
     this._video = video;
     this.emit(Page.Events.VideoStarted, video);