fix: improve robustness of various apis during navigation

src/screenshotter.ts

@@ -42,25 +42,10 @@ export class Screenshotter {
       let overridenViewport: types.Viewport | undefined;
       const viewport = this._page.viewport();
       let viewportSize: types.Size | undefined;
-      if (!viewport) {
+      if (!viewport)
-        viewportSize = await this._page.evaluate(() => ({
+        viewportSize = await this._getSize(false);
-          width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
-          height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight)
-        }));
-      }
       if (options.fullPage && !this._page._delegate.canScreenshotOutsideViewport()) {
-        const fullPageRect = await this._page.evaluate(() => ({
+        const fullPageRect = await this._getSize(true);
-          width: Math.max(
-              document.body.scrollWidth, document.documentElement.scrollWidth,
-              document.body.offsetWidth, document.documentElement.offsetWidth,
-              document.body.clientWidth, document.documentElement.clientWidth
-          ),
-          height: Math.max(
-              document.body.scrollHeight, document.documentElement.scrollHeight,
-              document.body.offsetHeight, document.documentElement.offsetHeight,
-              document.body.clientHeight, document.documentElement.clientHeight
-          )
-        }));
         overridenViewport = viewport ? { ...viewport, ...fullPageRect } : fullPageRect;
         await this._page.setViewport(overridenViewport);
       } else if (options.clip) {
@@ -129,6 +114,44 @@ export class Screenshotter {
       await platform.writeFileAsync(options.path, buffer);
     return buffer;
   }
+
+  private async _getSize(fullPage: boolean): Promise<{ width: number, height: number }> {
+    while (true) {
+      try {
+        const result = await this._page.evaluate((fullPage: boolean) => {
+          function calculate() {
+            return fullPage ? { width: Math.max(
+              document.body.scrollWidth, document.documentElement.scrollWidth,
+              document.body.offsetWidth, document.documentElement.offsetWidth,
+              document.body.clientWidth, document.documentElement.clientWidth
+            ), height: Math.max(
+              document.body.scrollHeight, document.documentElement.scrollHeight,
+              document.body.offsetHeight, document.documentElement.offsetHeight,
+              document.body.clientHeight, document.documentElement.clientHeight
+            )} : {
+              width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
+              height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight)
+            };
+          }
+          return new Promise<{ width: number, height: number }>(resolve => {
+            if (document.body && document.documentElement) {
+              resolve(calculate());
+            } else {
+              function listener() {
+                document.removeEventListener('DOMContentLoaded', listener);
+                resolve(calculate());
+              }
+              document.addEventListener('DOMContentLoaded', listener);
+            }
+          });
+        }, fullPage);
+        return result;
+      } catch (e) {
+        if (!(e instanceof Error) || !e.message.includes('context was destroyed'))
+          throw e;
+      }
+    }
+  }
 }
 
 const taskQueueSymbol = Symbol('TaskQueue');

src/webkit/wkAccessibility.ts

@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 import * as accessibility from '../accessibility';
-import { WKSession } from './wkConnection';
+import { Resender } from './wkConnection';
 import { Protocol } from './protocol';
 
-export async function getAccessibilityTree(session: WKSession) {
+export async function getAccessibilityTree(resender: Resender) {
+  return resender.sendWithRetries(async session => {
     const {axNode} = await session.send('Page.accessibilitySnapshot');
     return new WKAXNode(axNode);
+  });
 }
 
 class WKAXNode implements accessibility.AXNode {

src/webkit/wkConnection.ts

@@ -156,6 +156,11 @@ export class WKSession extends platform.EventEmitter {
   }
 }
 
+export interface Resender {
+  sendWithRetries<T>(action: (session: WKSession) => Promise<T>): Promise<T>;
+  sendToAllSessions(action: (session: WKSession) => Promise<any>): Promise<void>;
+}
+
 export function createProtocolError(error: Error, method: string, object: { error: { message: string; data: any; }; }): Error {
   let message = `Protocol error (${method}): ${object.error.message}`;
   if ('data' in object.error)
@@ -171,3 +176,7 @@ export function rewriteError(error: Error, message: string): Error {
 export function isSwappedOutError(e: Error) {
   return e.message.includes('Target was swapped out.');
 }
+
+export function isClosedError(e: Error) {
+  return e.message.includes('has been closed.');
+}

src/webkit/wkInput.ts

@@ -83,6 +83,8 @@ export class RawKeyboardImpl implements input.RawKeyboard {
   }
 
   async sendText(text: string): Promise<void> {
+    // TODO: it is impossible to guarantee the relative order of Page.insertText and other
+    // input commands.
     await this._session.send('Page.insertText', { text });
   }
 }

src/webkit/wkNetworkManager.ts

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { WKSession } from './wkConnection';
+import { WKSession, isSwappedOutError, isClosedError, Resender } from './wkConnection';
 import { Page } from '../page';
 import { helper, RegisteredListener, assert } from '../helper';
 import { Protocol } from './protocol';
@@ -27,14 +27,16 @@ import * as platform from '../platform';
 export class WKNetworkManager {
   private readonly _page: Page;
   private readonly _pageProxySession: WKSession;
+  private readonly _resender: Resender;
   private _session: WKSession;
   private readonly _requestIdToRequest = new Map<string, InterceptableRequest>();
   private _userCacheDisabled = false;
   private _sessionListeners: RegisteredListener[] = [];
 
-  constructor(page: Page, pageProxySession: WKSession) {
+  constructor(page: Page, pageProxySession: WKSession, resender: Resender) {
     this._page = page;
     this._pageProxySession = pageProxySession;
+    this._resender = resender;
   }
 
   async initializePageProxySession(credentials: types.Credentials | null) {
@@ -69,17 +71,11 @@ export class WKNetworkManager {
 
   async setCacheEnabled(enabled: boolean) {
     this._userCacheDisabled = !enabled;
-    await this._updateProtocolCacheDisabled();
+    await this._resender.sendToAllSessions(session => session.send('Network.setResourceCachingDisabled', { disabled: this._userCacheDisabled}));
   }
 
   async setRequestInterception(enabled: boolean): Promise<void> {
-    await this._session.send('Network.setInterceptionEnabled', { enabled, interceptRequests: enabled });
+    await this._resender.sendToAllSessions(session => session.send('Network.setInterceptionEnabled', { enabled, interceptRequests: enabled }));
-  }
-
-  async _updateProtocolCacheDisabled() {
-    await this._session.send('Network.setResourceCachingDisabled', {
-      disabled: this._userCacheDisabled
-    });
   }
 
   _onRequestWillBeSent(event: Protocol.Network.requestWillBeSentPayload) {
@@ -166,7 +162,7 @@ export class WKNetworkManager {
   }
 
   async setOfflineMode(value: boolean): Promise<void> {
-    await this._session.send('Network.setEmulateOfflineState', { offline: value });
+    await this._resender.sendToAllSessions(session => session.send('Network.setEmulateOfflineState', { offline: value }));
   }
 }
 
@@ -208,7 +204,12 @@ class InterceptableRequest implements network.RequestDelegate {
     const reason = errorReasons[errorCode];
     assert(reason, 'Unknown error code: ' + errorCode);
     await this._interceptedPromise;
+    try {
       await this._session.send('Network.interceptAsError', { requestId: this._requestId, reason });
+    } catch (e) {
+      if (!isSwappedOutError(e) && !isClosedError(e))
+        throw e;
+    }
   }
 
   async fulfill(response: { status: number; headers: network.Headers; contentType: string; body: (string | platform.BufferType); }) {

src/webkit/wkPage.ts

@@ -19,7 +19,7 @@ import * as frames from '../frames';
 import { debugError, helper, RegisteredListener } from '../helper';
 import * as dom from '../dom';
 import * as network from '../network';
-import { WKSession } from './wkConnection';
+import { WKSession, Resender } from './wkConnection';
 import { Events } from '../events';
 import { WKExecutionContext, EVALUATION_SCRIPT_URL } from './wkExecutionContext';
 import { WKNetworkManager } from './wkNetworkManager';
@@ -43,21 +43,24 @@ export class WKPage implements PageDelegate {
   _session: WKSession;
   readonly _page: Page;
   private readonly _pageProxySession: WKSession;
+  private readonly _resender: Resender;
   private readonly _networkManager: WKNetworkManager;
   private readonly _workers: WKWorkers;
   private readonly _contextIdToContext: Map<number, dom.FrameExecutionContext>;
   private _isolatedWorlds: Set<string>;
   private _sessionListeners: RegisteredListener[] = [];
   private readonly _bootstrapScripts: string[] = [];
+  private _defaultBackgroundColorOverride: { r: number; g: number; b: number; a: number; } | null = null;
 
-  constructor(browserContext: BrowserContext, pageProxySession: WKSession) {
+  constructor(browserContext: BrowserContext, pageProxySession: WKSession, resender: Resender) {
     this._pageProxySession = pageProxySession;
+    this._resender = resender;
     this.rawKeyboard = new RawKeyboardImpl(pageProxySession);
     this.rawMouse = new RawMouseImpl(pageProxySession);
     this._contextIdToContext = new Map();
     this._isolatedWorlds = new Set();
     this._page = new Page(this, browserContext);
-    this._networkManager = new WKNetworkManager(this._page, pageProxySession);
+    this._networkManager = new WKNetworkManager(this._page, pageProxySession, resender);
     this._workers = new WKWorkers(this._page);
   }
 
@@ -83,10 +86,6 @@ export class WKPage implements PageDelegate {
     this._networkManager.setSession(session);
     this._workers.setSession(session);
     this._isolatedWorlds = new Set();
-    // New bootstrap scripts may have been added during provisional load, push them
-    // again to be on the safe side.
-    if (this._bootstrapScripts.length)
-      this._setBootstrapScripts(session).catch(e => debugError(e));
   }
 
   // This method is called for provisional targets as well. The session passed as the parameter
@@ -109,11 +108,13 @@ export class WKPage implements PageDelegate {
     if (this._page._state.mediaType || this._page._state.colorScheme)
       promises.push(this._setEmulateMedia(session, this._page._state.mediaType, this._page._state.colorScheme));
     if (isProvisional)
-      promises.push(this._setBootstrapScripts(session));
+      promises.push(session.send('Page.setBootstrapScript', { source: this._bootstrapScripts.join(';') }));
     if (contextOptions.bypassCSP)
       promises.push(session.send('Page.setBypassCSP', { enabled: true }));
     if (this._page._state.extraHTTPHeaders !== null)
       promises.push(this._setExtraHTTPHeaders(session, this._page._state.extraHTTPHeaders));
+    if (this._defaultBackgroundColorOverride)
+      promises.push(session.send('Page.setDefaultBackgroundColorOverride', { color: this._defaultBackgroundColorOverride }));
     await Promise.all(promises).catch(e => {
       if (session.isDisposed())
         return;
@@ -222,7 +223,9 @@ export class WKPage implements PageDelegate {
   }
 
   async navigateFrame(frame: frames.Frame, url: string, referrer: string | undefined): Promise<frames.GotoResult> {
-    await this._session.send('Page.navigate', { url, frameId: frame._id, referrer });
+    await this._resender.sendWithRetries(async session => {
+      await session.send('Page.navigate', { url, frameId: frame._id, referrer });
+    });
     return {};  // We cannot get loaderId of cross-process navigation in advance.
   }
 
@@ -310,11 +313,11 @@ export class WKPage implements PageDelegate {
   }
 
   async setExtraHTTPHeaders(headers: network.Headers): Promise<void> {
-    await this._setExtraHTTPHeaders(this._session, headers);
+    await this._resender.sendToAllSessions(session => this._setExtraHTTPHeaders(session, headers));
   }
 
   async setEmulateMedia(mediaType: types.MediaType | null, colorScheme: types.ColorScheme | null): Promise<void> {
-    await this._setEmulateMedia(this._session, mediaType, colorScheme);
+    await this._resender.sendToAllSessions(session => this._setEmulateMedia(session, mediaType, colorScheme));
   }
 
   async setViewport(viewport: types.Viewport): Promise<void> {
@@ -326,12 +329,12 @@ export class WKPage implements PageDelegate {
     await this._pageProxySession.send('Emulation.setDeviceMetricsOverride', {width, height, fixedLayout, deviceScaleFactor: viewport.deviceScaleFactor || 1 });
   }
 
-  setCacheEnabled(enabled: boolean): Promise<void> {
+  async setCacheEnabled(enabled: boolean): Promise<void> {
-    return this._networkManager.setCacheEnabled(enabled);
+    await this._networkManager.setCacheEnabled(enabled);
   }
 
-  setRequestInterception(enabled: boolean): Promise<void> {
+  async setRequestInterception(enabled: boolean): Promise<void> {
-    return this._networkManager.setRequestInterception(enabled);
+    await this._networkManager.setRequestInterception(enabled);
   }
 
   async setOfflineMode(value: boolean) {
@@ -343,40 +346,45 @@ export class WKPage implements PageDelegate {
   }
 
   async reload(): Promise<void> {
-    await this._session.send('Page.reload');
+    await this._resender.sendWithRetries(async session => {
+      await session.send('Page.reload');
+    });
   }
 
   goBack(): Promise<boolean> {
-    return this._session.send('Page.goBack').then(() => true).catch(error => {
+    return this._resender.sendWithRetries(async session => {
+      return session.send('Page.goBack').then(() => true).catch(error => {
         if (error instanceof Error && error.message.includes(`Protocol error (Page.goBack): Failed to go`))
           return false;
         throw error;
       });
+    });
   }
 
   goForward(): Promise<boolean> {
-    return this._session.send('Page.goForward').then(() => true).catch(error => {
+    return this._resender.sendWithRetries(async session => {
+      return session.send('Page.goForward').then(() => true).catch(error => {
         if (error instanceof Error && error.message.includes(`Protocol error (Page.goForward): Failed to go`))
           return false;
         throw error;
       });
+    });
   }
 
   async exposeBinding(name: string, bindingFunction: string): Promise<void> {
     const script = `self.${name} = (param) => console.debug('${BINDING_CALL_MESSAGE}', {}, param); ${bindingFunction}`;
     this._bootstrapScripts.unshift(script);
-    await this._setBootstrapScripts(this._session);
+    await this._setBootstrapScripts();
     await Promise.all(this._page.frames().map(frame => frame.evaluate(script).catch(debugError)));
   }
 
   async evaluateOnNewDocument(script: string): Promise<void> {
     this._bootstrapScripts.push(script);
-    await this._setBootstrapScripts(this._session);
+    await this._setBootstrapScripts();
   }
 
-  private async _setBootstrapScripts(session: WKSession) {
+  private async _setBootstrapScripts() {
-    const source = this._bootstrapScripts.join(';');
+    await this._resender.sendToAllSessions(session => session.send('Page.setBootstrapScript', { source: this._bootstrapScripts.join(';') }));
-    await session.send('Page.setBootstrapScript', { source });
   }
 
   async closePage(runBeforeUnload: boolean): Promise<void> {
@@ -395,13 +403,15 @@ export class WKPage implements PageDelegate {
   }
 
   async setBackgroundColor(color?: { r: number; g: number; b: number; a: number; }): Promise<void> {
-    // TODO: line below crashes, sort it out.
+    this._defaultBackgroundColorOverride = color;
-    await this._session.send('Page.setDefaultBackgroundColorOverride', { color });
+    await this._resender.sendToAllSessions(session => session.send('Page.setDefaultBackgroundColorOverride', { color }));
   }
 
   async takeScreenshot(format: string, options: types.ScreenshotOptions, viewport: types.Viewport): Promise<platform.BufferType> {
     const rect = options.clip || { x: 0, y: 0, width: viewport.width, height: viewport.height };
-    const result = await this._session.send('Page.snapshotRect', { ...rect, coordinateSystem: options.fullPage ? 'Page' : 'Viewport' });
+    const result = await this._resender.sendWithRetries(session => {
+      return session.send('Page.snapshotRect', { ...rect, coordinateSystem: options.fullPage ? 'Page' : 'Viewport' });
+    });
     const prefix = 'data:image/png;base64,';
     let buffer = platform.Buffer.from(result.dataURL.substr(prefix.length), 'base64');
     if (format === 'jpeg')
@@ -491,7 +501,7 @@ export class WKPage implements PageDelegate {
   }
 
   async getAccessibilityTree() : Promise<accessibility.AXNode> {
-    return getAccessibilityTree(this._session);
+    return getAccessibilityTree(this._resender);
   }
 
   coverage(): Coverage | undefined {

src/webkit/wkPageProxy.ts

@@ -5,7 +5,7 @@
 import { BrowserContext } from '../browserContext';
 import { Page } from '../page';
 import { Protocol } from './protocol';
-import { WKSession } from './wkConnection';
+import { WKSession, Resender, isSwappedOutError } from './wkConnection';
 import { WKPage } from './wkPage';
 import { RegisteredListener, helper, assert, debugError } from '../helper';
 import { Events } from '../events';
@@ -15,7 +15,7 @@ import { Events } from '../events';
 // has undefined instead.
 const provisionalMessagesSymbol = Symbol('provisionalMessages');
 
-export class WKPageProxy {
+export class WKPageProxy implements Resender {
   private readonly _pageProxySession: WKSession;
   readonly _browserContext: BrowserContext;
   private _pagePromise: Promise<Page> | null = null;
@@ -24,6 +24,10 @@ export class WKPageProxy {
   private _firstTargetCallback: () => void;
   private readonly _sessions = new Map<string, WKSession>();
   private readonly _eventListeners: RegisteredListener[];
+  private readonly _disposedPromise: Promise<void>;
+  private _disposedCallback: () => void;
+  private _disposed = false;
+  private _waitForCommitCallbacks: (() => void)[] = [];
 
   constructor(pageProxySession: WKSession, browserContext: BrowserContext) {
     this._pageProxySession = pageProxySession;
@@ -35,6 +39,7 @@ export class WKPageProxy {
       helper.addEventListener(this._pageProxySession, 'Target.dispatchMessageFromTarget', this._onDispatchMessageFromTarget.bind(this)),
       helper.addEventListener(this._pageProxySession, 'Target.didCommitProvisionalTarget', this._onDidCommitProvisionalTarget.bind(this)),
     ];
+    this._disposedPromise = new Promise(f => this._disposedCallback = f);
 
     // Intercept provisional targets during cross-process navigation.
     this._pageProxySession.send('Target.setPauseOnStart', { pauseOnStart: true }).catch(e => {
@@ -45,6 +50,34 @@ export class WKPageProxy {
     });
   }
 
+  async sendWithRetries<T>(action: (session: WKSession) => Promise<T>): Promise<T> {
+    while (!this._disposed) {
+      try {
+        const result = await action(this._committedSession());
+        return result;
+      } catch (e) {
+        if (!isSwappedOutError(e))
+          throw e;
+      }
+      await Promise.race([
+        new Promise(f => this._waitForCommitCallbacks.push(f)),
+        this._disposedPromise,
+      ]);
+    }
+    return Promise.reject(new Error('The page has been closed.'));
+  }
+
+  async sendToAllSessions(action: (session: WKSession) => Promise<any>): Promise<void> {
+    const promises: Promise<any>[] = [];
+    for (const session of this._sessions.values()) {
+      promises.push(action(session).catch(e => {
+        if (!isSwappedOutError(e))
+          throw e;
+      }));
+    }
+    await Promise.all(promises);
+  }
+
   didClose() {
     if (this._wkPage)
       this._wkPage.didClose(false);
@@ -56,6 +89,8 @@ export class WKPageProxy {
     for (const session of this._sessions.values())
       session.dispose();
     this._sessions.clear();
+    this._disposed = true;
+    this._disposedCallback();
     if (this._wkPage)
       this._wkPage.didDisconnect();
   }
@@ -78,8 +113,7 @@ export class WKPageProxy {
     popupPageProxy.page().then(page => this._wkPage._page.emit(Events.Page.Popup, page));
   }
 
-  private async _initializeWKPage(): Promise<Page> {
+  private _committedSession(): WKSession {
-    await this._firstTargetPromise;
     let session: WKSession;
     for (const anySession of this._sessions.values()) {
       if (!(anySession as any)[provisionalMessagesSymbol]) {
@@ -88,7 +122,13 @@ export class WKPageProxy {
       }
     }
     assert(session, 'One non-provisional target session must exist');
-    this._wkPage = new WKPage(this._browserContext, this._pageProxySession);
+    return session;
+  }
+
+  private async _initializeWKPage(): Promise<Page> {
+    await this._firstTargetPromise;
+    const session = this._committedSession();
+    this._wkPage = new WKPage(this._browserContext, this._pageProxySession, this);
     this._wkPage.setSession(session);
     await Promise.all([
       this._wkPage._initializePageProxySession(),
@@ -155,5 +195,9 @@ export class WKPageProxy {
     for (const message of provisionalMessages)
       newSession.dispatchMessage(JSON.parse(message));
     this._wkPage.setSession(newSession);
+    const callbacks = this._waitForCommitCallbacks;
+    this._waitForCommitCallbacks = [];
+    for (const callback of callbacks)
+      callback();
   }
 }

test/assets/redirectloop1.html

@@ -0,0 +1,9 @@
+<script>
+  setTimeout(() => {
+    const iter = window.localStorage.iter || '';
+    window.localStorage.iter = iter + 'a';
+    if (iter.length === 10)
+      return;
+    window.location.href = window.location.href.replace('loop1', 'loop2');
+  }, 1);
+</script>>

test/assets/redirectloop2.html

@@ -0,0 +1,5 @@
+<script>
+  setTimeout(() => {
+    window.location.href = window.location.href.replace('loop2', 'loop1');
+  }, 1);
+</script>>

test/emulation.spec.js

@@ -167,6 +167,19 @@ module.exports.describe = function({testRunner, expect, playwright, FFOX, CHROME
       await page.emulateMedia({ colorScheme: 'bad' }).catch(e => error = e);
       expect(error.message).toBe('Unsupported color scheme: bad');
     });
+    it.skip(FFOX)('should work during navigation', async({page, server}) => {
+      await page.emulateMedia({ colorScheme: 'light' });
+      const navigated = page.goto(server.EMPTY_PAGE);
+      const schemes = ['dark', 'light'];
+      let scheme = 0;
+      for (let i = 0; i < 9; i++) {
+        page.emulateMedia({ colorScheme: schemes[scheme] });
+        scheme = 1 - scheme;
+        await new Promise(f => setTimeout(f, 1));
+      }
+      await navigated;
+      expect(await page.evaluate(() => matchMedia('(prefers-color-scheme: dark)').matches)).toBe(true);
+    });
   });
 
   describe.skip(FFOX || WEBKIT)('BrowserContext({timezoneId})', function() {

test/screenshot.spec.js

@@ -27,6 +27,14 @@ module.exports.describe = function({testRunner, expect, product, FFOX, CHROME, W
       const screenshot = await page.screenshot();
       expect(screenshot).toBeGolden('screenshot-sanity.png');
     });
+    it.skip(FFOX)('should work while navigating', async({page, server}) => {
+      await page.setViewport({width: 500, height: 500});
+      await page.goto(server.PREFIX + '/redirectloop1.html');
+      for (let i = 0; i < 10; i++) {
+        const screenshot = await page.screenshot({ fullPage: true });
+        expect(screenshot).toBeInstanceOf(Buffer);
+      }
+    });
     it('should clip rect', async({page, server}) => {
       await page.setViewport({width: 500, height: 500});
       await page.goto(server.PREFIX + '/grid.html');