From 88e0a3b72c8fd4264e847912668bff7f94131845 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Thu, 31 Mar 2022 15:45:39 -0700
Subject: [PATCH] cherry-pick(#13222): fix(screenshot): do not stall on
 hideHiglight attempt 2 (#13228)

It turns out that "non stalling evaluate" can stall in Chromium
in some weird conditions, like `document.open` after some weird
`iframe.src` value.

We now only hide highlight in those frames where we did install
highlight in the first place.
---
 .../src/server/screenshotter.ts               | 24 ++++++++++++-------
 tests/page/page-screenshot.spec.ts            | 13 ++++++++++
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/packages/playwright-core/src/server/screenshotter.ts b/packages/playwright-core/src/server/screenshotter.ts
index 4e82cfe38c..c544069445 100644
--- a/packages/playwright-core/src/server/screenshotter.ts
+++ b/packages/playwright-core/src/server/screenshotter.ts
@@ -222,11 +222,18 @@ export class Screenshotter {
     }));
   }
 
-  async _maskElements(progress: Progress, options: ScreenshotOptions) {
-    if (!options.mask || !options.mask.length)
-      return false;
-
+  async _maskElements(progress: Progress, options: ScreenshotOptions): Promise<() => Promise<void>> {
     const framesToParsedSelectors: MultiMap<Frame, ParsedSelector> = new MultiMap();
+
+    const cleanup = async () => {
+      await Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
+        await frame.hideHighlight();
+      }));
+    };
+
+    if (!options.mask || !options.mask.length)
+      return cleanup;
+
     await Promise.all((options.mask || []).map(async ({ frame, selector }) => {
       const pair = await frame.resolveFrameForSelectorNoWait(selector);
       if (pair)
@@ -237,8 +244,8 @@ export class Screenshotter {
     await Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
       await frame.maskSelectors(framesToParsedSelectors.get(frame));
     }));
-    progress.cleanupWhenAborted(() => this._page.hideHighlight());
-    return true;
+    progress.cleanupWhenAborted(cleanup);
+    return cleanup;
   }
 
   private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean | undefined, options: ScreenshotOptions): Promise<Buffer> {
@@ -252,14 +259,13 @@ export class Screenshotter {
     }
     progress.throwIfAborted(); // Avoid extra work.
 
-    const hasHighlight = await this._maskElements(progress, options);
+    const cleanupHighlight = await this._maskElements(progress, options);
     progress.throwIfAborted(); // Avoid extra work.
 
     const buffer = await this._page._delegate.takeScreenshot(progress, format, documentRect, viewportRect, options.quality, fitsViewport);
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
-    if (hasHighlight)
-      await this._page.hideHighlight();
+    await cleanupHighlight();
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
     if (shouldSetDefaultBackground)
diff --git a/tests/page/page-screenshot.spec.ts b/tests/page/page-screenshot.spec.ts
index 1247c28383..ff36b98cf0 100644
--- a/tests/page/page-screenshot.spec.ts
+++ b/tests/page/page-screenshot.spec.ts
@@ -444,6 +444,19 @@ it.describe('page screenshot', () => {
       await route.fulfill({ body: '' });
       await done;
     });
+
+    it('should work when subframe used document.open after a weird url', async ({ page, server }) => {
+      await page.goto(server.EMPTY_PAGE);
+      await page.evaluate(() => {
+        const iframe = document.createElement('iframe');
+        iframe.src = 'javascript:hi';
+        document.body.appendChild(iframe);
+        iframe.contentDocument.open();
+        iframe.contentDocument.write('Hello');
+        iframe.contentDocument.close();
+      });
+      await page.screenshot({ mask: [ page.locator('non-existent') ] });
+    });
   });
 });