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.

Author: dgozman

packages/playwright-core/src/server/screenshotter.ts

@@ -222,11 +222,18 @@ export class Screenshotter {
     }));
   }
 
-  async _maskElements(progress: Progress, options: ScreenshotOptions) {
+  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 false;
+      return cleanup;
 
-    const framesToParsedSelectors: MultiMap<Frame, ParsedSelector> = new MultiMap();
     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)

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') ] });
+    });
   });
 });