Bug 1749074 - Cancel screenshots on page unload. r=sfoster

Differential Revision: https://phabricator.services.mozilla.com/D135522

browser/actors/ScreenshotsComponentChild.jsm

@@ -30,6 +30,19 @@ class ScreenshotsComponentChild extends JSWindowActorChild {
     return null;
   }
 
+  handleEvent(event) {
+    switch (event.type) {
+      case "keydown":
+        if (event.key === "Escape") {
+          this.requestCancelScreenshot();
+        }
+        break;
+      case "beforeunload":
+        this.requestCancelScreenshot();
+        break;
+    }
+  }
+
   /**
    * Send a request to cancel the screenshot to the parent process
    */
@@ -93,21 +106,12 @@ class ScreenshotsComponentChild extends JSWindowActorChild {
         this.document,
         this
       ));
-    this.document.addEventListener("keydown", this.handler);
+    this.document.addEventListener("keydown", this);
+    this.document.ownerGlobal.addEventListener("beforeunload", this);
     overlay.initialize();
     return true;
   }
 
-  /**
-   * Function to handle the escape key press to cancel screenshots overlay
-   * @param event The keydown event
-   */
-  handler = event => {
-    if (event.key === "Escape") {
-      this.requestCancelScreenshot();
-    }
-  };
-
   /**
    * Remove the screenshots overlay.
    *
@@ -115,7 +119,8 @@ class ScreenshotsComponentChild extends JSWindowActorChild {
    *   true when the overlay has been removed otherwise false
    */
   endScreenshotsOverlay() {
-    this.document.removeEventListener("keydown", this.handler);
+    this.document.removeEventListener("keydown", this);
+    this.document.ownerGlobal.removeEventListener("beforeunload", this);
     this._overlay?.tearDown();
     return true;
   }

browser/components/screenshots/ScreenshotsUtils.jsm

@@ -20,6 +20,14 @@ class ScreenshotsComponentParent extends JSWindowActorParent {
         ScreenshotsUtils.closePanel(browser);
     }
   }
+
+  didDestroy() {
+    // When restoring a crashed tab the browser is null
+    let browser = this.browsingContext.topFrameElement;
+    if (browser) {
+      ScreenshotsUtils.closePanel(browser, false);
+    }
+  }
 }
 
 var ScreenshotsUtils = {
@@ -123,16 +131,21 @@ var ScreenshotsUtils = {
   /**
    * Close the panel and call child actor to close the overlay
    * @param browser The current browser
+   * @param {bool} closeOverlay Whether or not to
+   * send a message to the child to close the overly.
+   * Defaults to true. Will be false when called from didDestroy.
    */
-  closePanel(browser) {
+  closePanel(browser, closeOverlay = true) {
     let buttonsPanel = browser.ownerDocument.querySelector(
       "#screenshotsPagePanel"
     );
     if (buttonsPanel && buttonsPanel.state !== "closed") {
       buttonsPanel.hidePopup();
     }
-    let actor = this.getActor(browser);
-    actor.sendQuery("Screenshots:HideOverlay");
+    if (closeOverlay) {
+      let actor = this.getActor(browser);
+      actor.sendQuery("Screenshots:HideOverlay");
+    }
   },
   /**
    * If the buttons panel exists and is open we will hide both the panel

browser/components/screenshots/tests/browser/browser.ini

@@ -7,9 +7,11 @@ prefs =
   extensions.screenshots.disabled=false
   screenshots.browser.component.enabled=true
 
+[browser_screenshots_page_unload.js]
 [browser_screenshots_test_escape.js]
 [browser_screenshots_test_full_page.js]
 skip-if = (!debug && os == 'win' && os_version == '6.1') # Bug 1746281
+[browser_screenshots_test_page_crash.js]
 [browser_screenshots_test_toggle_pref.js]
 [browser_screenshots_test_toolbar_button.js]
 [browser_screenshots_test_visible.js]

browser/components/screenshots/tests/browser/browser_screenshots_page_unload.js

@@ -0,0 +1,79 @@
+/* Any copyright is dedicated to the Public Domain.
+  http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+add_task(async function test() {
+  CustomizableUI.addWidgetToArea(
+    "screenshot-button",
+    CustomizableUI.AREA_NAVBAR
+  );
+  let screenshotBtn = document.getElementById("screenshot-button");
+  Assert.ok(screenshotBtn, "The screenshots button was added to the nav bar");
+
+  await BrowserTestUtils.withNewTab(
+    {
+      gBrowser,
+      url: TEST_PAGE,
+    },
+    async browser => {
+      await SpecialPowers.spawn(browser, [TEST_PAGE], url => {
+        let a = content.document.createElement("a");
+        a.id = "clickMe";
+        a.href = url;
+        a.textContent = "Click me to unload page";
+        content.document.querySelector("body").appendChild(a);
+      });
+
+      let helper = new ScreenshotsHelper(browser);
+
+      // click toolbar button so panel shows
+      helper.triggerUIFromToolbar();
+
+      let panel = gBrowser.selectedBrowser.ownerDocument.querySelector(
+        "#screenshotsPagePanel"
+      );
+      await BrowserTestUtils.waitForMutationCondition(
+        panel,
+        { attributes: true },
+        () => {
+          return BrowserTestUtils.is_visible(panel);
+        }
+      );
+      ok(BrowserTestUtils.is_visible(panel), "Panel buttons are visible");
+
+      await ContentTask.spawn(browser, null, async () => {
+        let screenshotsChild = content.windowGlobalChild.getActor(
+          "ScreenshotsComponent"
+        );
+        Assert.ok(screenshotsChild._overlay._initialized, "The overlay exists");
+      });
+
+      await SpecialPowers.spawn(browser, [], () => {
+        content.document.querySelector("#clickMe").click();
+      });
+
+      await BrowserTestUtils.waitForMutationCondition(
+        panel,
+        { attributes: true },
+        () => {
+          return BrowserTestUtils.is_hidden(panel);
+        }
+      );
+      ok(
+        BrowserTestUtils.is_hidden(panel),
+        "Panel buttons are hidden after page unload"
+      );
+
+      await ContentTask.spawn(browser, null, async () => {
+        let screenshotsChild = content.windowGlobalChild.getActor(
+          "ScreenshotsComponent"
+        );
+        Assert.ok(
+          !screenshotsChild._overlay._initialized,
+          "The overlay doesn't exist"
+        );
+      });
+    }
+  );
+});

browser/components/screenshots/tests/browser/browser_screenshots_test_page_crash.js

@@ -0,0 +1,77 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+add_task(async function test_fullpageScreenshot() {
+  CustomizableUI.addWidgetToArea(
+    "screenshot-button",
+    CustomizableUI.AREA_NAVBAR
+  );
+  let screenshotBtn = document.getElementById("screenshot-button");
+  Assert.ok(screenshotBtn, "The screenshots button was added to the nav bar");
+
+  await BrowserTestUtils.withNewTab(
+    {
+      gBrowser,
+      url: TEST_PAGE,
+    },
+    async browser => {
+      let helper = new ScreenshotsHelper(browser);
+
+      // click toolbar button so UI shows
+      helper.triggerUIFromToolbar();
+
+      let panel = gBrowser.selectedBrowser.ownerDocument.querySelector(
+        "#screenshotsPagePanel"
+      );
+      await BrowserTestUtils.waitForMutationCondition(
+        panel,
+        { attributes: true },
+        () => {
+          return BrowserTestUtils.is_visible(panel);
+        }
+      );
+      ok(BrowserTestUtils.is_visible(panel), "Panel buttons are visible");
+
+      await ContentTask.spawn(browser, null, async () => {
+        let screenshotsChild = content.windowGlobalChild.getActor(
+          "ScreenshotsComponent"
+        );
+        Assert.ok(screenshotsChild._overlay._initialized, "The overlay exists");
+      });
+
+      let waitForPanelHide = BrowserTestUtils.waitForMutationCondition(
+        panel,
+        { attributes: true },
+        () => {
+          return BrowserTestUtils.is_hidden(panel);
+        }
+      );
+
+      await BrowserTestUtils.crashFrame(browser);
+
+      await waitForPanelHide;
+      ok(
+        BrowserTestUtils.is_hidden(panel),
+        "Panel buttons are hidden after page crash"
+      );
+
+      await ContentTask.spawn(browser, null, async () => {
+        let screenshotsChild = content.windowGlobalChild.getActor(
+          "ScreenshotsComponent"
+        );
+        Assert.ok(!screenshotsChild._overlay, "The overlay doesn't exist");
+      });
+
+      let tab = gBrowser.getTabForBrowser(browser);
+
+      SessionStore.reviveCrashedTab(tab);
+
+      ok(
+        BrowserTestUtils.is_hidden(panel),
+        "Panel buttons are hidden after page crash"
+      );
+    }
+  );
+});