From 044ebd7fd8b1e773fa70065717230bd2bd4fbfcb Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Thu, 23 Jan 2020 14:58:30 -0800
Subject: [PATCH] fix: delete contexts from the map on navigation (#602)

---
 src/firefox/ffPage.ts   |  8 ++++++++
 src/webkit/wkPage.ts    | 21 ++++++++++++++-------
 test/evaluation.spec.js | 14 ++++++++++++++
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/firefox/ffPage.ts b/src/firefox/ffPage.ts
index 8e9650d03c..3996d69716 100644
--- a/src/firefox/ffPage.ts
+++ b/src/firefox/ffPage.ts
@@ -52,6 +52,7 @@ export class FFPage implements PageDelegate {
     this._contextIdToContext = new Map();
     this._page = new Page(this, browserContext);
     this._networkManager = new FFNetworkManager(session, this._page);
+    this._page.on(Events.Page.FrameDetached, frame => this._removeContextsForFrame(frame));
     this._eventListeners = [
       helper.addEventListener(this._session, 'Page.eventFired', this._onEventFired.bind(this)),
       helper.addEventListener(this._session, 'Page.frameAttached', this._onFrameAttached.bind(this)),
@@ -122,6 +123,13 @@ export class FFPage implements PageDelegate {
     context.frame._contextDestroyed(context as dom.FrameExecutionContext);
   }
 
+  private _removeContextsForFrame(frame: frames.Frame) {
+    for (const [contextId, context] of this._contextIdToContext) {
+      if (context.frame === frame)
+        this._contextIdToContext.delete(contextId);
+    }
+  }
+
   _onNavigationStarted() {
   }
 
diff --git a/src/webkit/wkPage.ts b/src/webkit/wkPage.ts
index 9e057ea427..a47289d6fd 100644
--- a/src/webkit/wkPage.ts
+++ b/src/webkit/wkPage.ts
@@ -59,6 +59,7 @@ export class WKPage implements PageDelegate {
     this._page = new Page(this, browserContext);
     this._workers = new WKWorkers(this._page);
     this._session = undefined as any as WKSession;
+    this._page.on(Events.Page.FrameDetached, frame => this._removeContextsForFrame(frame, false));
   }
 
   private async _initializePageProxySession() {
@@ -248,13 +249,8 @@ export class WKPage implements PageDelegate {
 
   private _onFrameNavigated(framePayload: Protocol.Page.Frame, initial: boolean) {
     const frame = this._page._frameManager.frame(framePayload.id);
-    for (const [contextId, context] of this._contextIdToContext) {
-      if (context.frame === frame) {
-        (context._delegate as WKExecutionContext)._dispose();
-        this._contextIdToContext.delete(contextId);
-        frame._contextDestroyed(context);
-      }
-    }
+    assert(frame);
+    this._removeContextsForFrame(frame!, true);
     if (!framePayload.parentId)
       this._workers.clear();
     this._page._frameManager.frameCommittedNewDocumentNavigation(framePayload.id, framePayload.url, framePayload.name || '', framePayload.loaderId, initial);
@@ -268,6 +264,17 @@ export class WKPage implements PageDelegate {
     this._page._frameManager.frameDetached(frameId);
   }
 
+  private _removeContextsForFrame(frame: frames.Frame, notifyFrame: boolean) {
+    for (const [contextId, context] of this._contextIdToContext) {
+      if (context.frame === frame) {
+        (context._delegate as WKExecutionContext)._dispose();
+        this._contextIdToContext.delete(contextId);
+        if (notifyFrame)
+          frame._contextDestroyed(context);
+      }
+    }
+  }
+
   private _onExecutionContextCreated(contextPayload : Protocol.Runtime.ExecutionContextDescription) {
     if (this._contextIdToContext.has(contextPayload.id))
       return;
diff --git a/test/evaluation.spec.js b/test/evaluation.spec.js
index cba1938ddd..f092e33926 100644
--- a/test/evaluation.spec.js
+++ b/test/evaluation.spec.js
@@ -331,6 +331,20 @@ module.exports.describe = function({testRunner, expect, FFOX, CHROMIUM, WEBKIT})
       expect(await page.frames()[0].evaluate(() => document.body.textContent.trim())).toBe('');
       expect(await page.frames()[1].evaluate(() => document.body.textContent.trim())).toBe(`Hi, I'm frame`);
     });
+    it('should dispose context on navigation', async({page, server}) => {
+      await page.goto(server.PREFIX + '/frames/one-frame.html');
+      expect(page.frames().length).toBe(2);
+      expect(page._delegate._contextIdToContext.size).toBe(4);
+      await page.goto(server.EMPTY_PAGE);
+      expect(page._delegate._contextIdToContext.size).toBe(2);
+    });
+    it('should dispose context on cross-origin navigation', async({page, server}) => {
+      await page.goto(server.PREFIX + '/frames/one-frame.html');
+      expect(page.frames().length).toBe(2);
+      expect(page._delegate._contextIdToContext.size).toBe(4);
+      await page.goto(server.CROSS_PROCESS_PREFIX + '/empty.html');
+      expect(page._delegate._contextIdToContext.size).toBe(2);
+    });
     it('should execute after cross-site navigation', async({page, server}) => {
       await page.goto(server.EMPTY_PAGE);
       const mainFrame = page.mainFrame();