fix(webkit): improve target swap handling (#175)

- Fix "page closed twice" race.
- Do not fire 'disconnected' on swapped out sessions.
- Use a different error for commands sent to swapped out targets.
  This allows callers to detect this situation and retry/throw/catch.
- Restore more state on swap: extra http headers, user agent, emulated media.

src/webkit/Browser.ts

@@ -184,17 +184,8 @@ export class Browser extends EventEmitter {
 
   async _onProvisionalTargetCommitted({oldTargetId, newTargetId}) {
     const oldTarget = this._targets.get(oldTargetId);
-    if (!oldTarget._pagePromise)
-      return;
-    const page = await oldTarget._pagePromise;
     const newTarget = this._targets.get(newTargetId);
-    const newSession = this._connection.session(newTargetId);
-    page._swapSessionOnNavigation(newSession);
-    newTarget._pagePromise = oldTarget._pagePromise;
-    newTarget._adoptPage(page);
-    // Old target should not be accessed by anyone. Reset page promise so that
-    // old target does not close the page on connection reset.
-    oldTarget._pagePromise = null;
+    newTarget._swappedIn(oldTarget, this._connection.session(newTargetId));
   }
 
   disconnect() {

src/webkit/Connection.ts

@@ -104,7 +104,6 @@ export class Connection extends EventEmitter {
     } else if (object.method === 'Target.targetDestroyed') {
       const session = this._sessions.get(object.params.targetId);
       if (session) {
-        // FIXME: this is a workaround for cross-origin navigation in WebKit.
         session._onClosed();
         this._sessions.delete(object.params.targetId);
       }
@@ -121,6 +120,7 @@ export class Connection extends EventEmitter {
       const oldSession = this._sessions.get(oldTargetId);
       if (!oldSession)
         throw new Error('Unknown old target: ' + oldTargetId);
+      oldSession._swappedOut = true;
     }
   }
 
@@ -158,6 +158,7 @@ export class TargetSession extends EventEmitter {
   private _callbacks = new Map<number, {resolve:(o: any) => void, reject: (e: Error) => void, error: Error, method: string}>();
   private _targetType: string;
   private _sessionId: string;
+  _swappedOut = false;
   on: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   addListener: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   off: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
@@ -194,7 +195,7 @@ export class TargetSession extends EventEmitter {
     }).catch(e => {
       // There is a possible race of the connection closure. We may have received
       // targetDestroyed notification before response for the command, in that
-      // case it's safe to swallow the exception.g
+      // case it's safe to swallow the exception.
       const callback = this._callbacks.get(innerId);
       assert(!callback, 'Callback was not rejected when target was destroyed.');
     });
@@ -218,11 +219,17 @@ export class TargetSession extends EventEmitter {
   }
 
   _onClosed() {
-    for (const callback of this._callbacks.values())
-      callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target closed.`));
+    for (const callback of this._callbacks.values()) {
+      // TODO: make some calls like screenshot catch swapped out error and retry.
+      if (this._swappedOut)
+        callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target was swapped out.`));
+      else
+        callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target closed.`));
+    }
     this._callbacks.clear();
     this._connection = null;
-    this.emit(TargetSessionEvents.Disconnected);
+    if (!this._swappedOut)
+      this.emit(TargetSessionEvents.Disconnected);
   }
 }
 
@@ -237,3 +244,7 @@ function rewriteError(error: Error, message: string): Error {
   error.message = message;
   return error;
 }
+
+export function isSwappedOutError(e: Error) {
+  return e.message.includes('Target was swapped out.');
+}

src/webkit/ExecutionContext.ts

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { TargetSession } from './Connection';
+import { TargetSession, isSwappedOutError } from './Connection';
 import { helper } from '../helper';
 import { valueFromRemoteObject, releaseObject } from './protocolHelper';
 import { Protocol } from './protocol';
@@ -25,7 +25,7 @@ export const EVALUATION_SCRIPT_URL = '__playwright_evaluation_script__';
 const SOURCE_URL_REGEX = /^[\040\t]*\/\/[@#] sourceURL=\s*(\S*?)\s*$/m;
 
 export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
-  private _globalObjectId?: string;
+  private _globalObjectId?: Promise<string>;
   _session: TargetSession;
   _contextId: number;
   private _contextDestroyedCallback: () => void;
@@ -45,8 +45,6 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
   }
 
   async evaluate(context: js.ExecutionContext, returnByValue: boolean, pageFunction: Function | string, ...args: any[]): Promise<any> {
-    const suffix = `//# sourceURL=${EVALUATION_SCRIPT_URL}`;
-
     if (helper.isString(pageFunction)) {
       const contextId = this._contextId;
       const expression: string = pageFunction as string;
@@ -58,18 +56,9 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
         emulateUserGesture: true
       }).then(response => {
         if (response.result.type === 'object' && response.result.className === 'Promise') {
-          const contextDiscarded = this._executionContextDestroyedPromise.then(() => ({
-            wasThrown: true,
-            result: {
-              description: 'Protocol error: Execution context was destroyed, most likely because of a navigation.'
-            } as Protocol.Runtime.RemoteObject
-          }));
           return Promise.race([
-            contextDiscarded,
-            this._session.send('Runtime.awaitPromise', {
-              promiseObjectId: response.result.objectId,
-              returnByValue: false
-            })
+            this._executionContextDestroyedPromise.then(() => contextDestroyedResult),
+            this._awaitPromise(response.result.objectId),
           ]);
         }
         return response;
@@ -78,32 +67,8 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
           throw new Error('Evaluation failed: ' + response.result.description);
         if (!returnByValue)
           return context._createHandle(response.result);
-        if (response.result.objectId) {
-          const serializeFunction = function() {
-            try {
-              return JSON.stringify(this);
-            } catch (e) {
-              if (e instanceof TypeError)
-                return void 0;
-              throw e;
-            }
-          };
-          return this._session.send('Runtime.callFunctionOn', {
-            // Serialize object using standard JSON implementation to correctly pass 'undefined'.
-            functionDeclaration: serializeFunction + '\n' + suffix + '\n',
-            objectId: response.result.objectId,
-            returnByValue
-          }).then(serializeResponse => {
-            if (serializeResponse.wasThrown)
-              throw new Error('Serialization failed: ' + serializeResponse.result.description);
-            // This is the case of too long property chain, not serializable to json string.
-            if (serializeResponse.result.type === 'undefined')
-              return undefined;
-            if (serializeResponse.result.type !== 'string')
-              throw new Error('Unexpected result of JSON.stringify: ' + JSON.stringify(serializeResponse, null, 2));
-            return JSON.parse(serializeResponse.result.value);
-          });
-        }
+        if (response.result.objectId)
+          return this._returnObjectByValue(response.result.objectId);
         return valueFromRemoteObject(response.result);
       }).catch(rewriteError);
     }
@@ -164,18 +129,9 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
     }
     return callFunctionOnPromise.then(response => {
       if (response.result.type === 'object' && response.result.className === 'Promise') {
-        const contextDiscarded = this._executionContextDestroyedPromise.then(() => ({
-          wasThrown: true,
-          result: {
-            description: 'Protocol error: Execution context was destroyed, most likely because of a navigation.'
-          } as Protocol.Runtime.RemoteObject
-        }));
         return Promise.race([
-          contextDiscarded,
-          this._session.send('Runtime.awaitPromise', {
-            promiseObjectId: response.result.objectId,
-            returnByValue: false
-          })
+          this._executionContextDestroyedPromise.then(() => contextDestroyedResult),
+          this._awaitPromise(response.result.objectId),
         ]);
       }
       return response;
@@ -184,32 +140,8 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
         throw new Error('Evaluation failed: ' + response.result.description);
       if (!returnByValue)
         return context._createHandle(response.result);
-      if (response.result.objectId) {
-        const serializeFunction = function() {
-          try {
-            return JSON.stringify(this);
-          } catch (e) {
-            if (e instanceof TypeError)
-              return void 0;
-            throw e;
-          }
-        };
-        return this._session.send('Runtime.callFunctionOn', {
-          // Serialize object using standard JSON implementation to correctly pass 'undefined'.
-          functionDeclaration: serializeFunction + '\n' + suffix + '\n',
-          objectId: response.result.objectId,
-          returnByValue
-        }).then(serializeResponse => {
-          if (serializeResponse.wasThrown)
-            throw new Error('Serialization failed: ' + serializeResponse.result.description);
-          // This is the case of too long property chain, not serializable to json string.
-          if (serializeResponse.result.type === 'undefined')
-            return undefined;
-          if (serializeResponse.result.type !== 'string')
-            throw new Error('Unexpected result of JSON.stringify: ' + JSON.stringify(serializeResponse, null, 2));
-          return JSON.parse(serializeResponse.result.value);
-        });
-      }
+      if (response.result.objectId)
+        return this._returnObjectByValue(response.result.objectId);
       return valueFromRemoteObject(response.result);
     }).catch(rewriteError);
 
@@ -263,14 +195,64 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
     }
   }
 
-  async _contextGlobalObjectId() {
+  private _contextGlobalObjectId() {
     if (!this._globalObjectId) {
-      const globalObject = await this._session.send('Runtime.evaluate', { expression: 'this', contextId: this._contextId });
-      this._globalObjectId = globalObject.result.objectId;
+      this._globalObjectId = this._session.send('Runtime.evaluate', {
+        expression: 'this',
+        contextId: this._contextId
+      }).catch(e => {
+        if (isSwappedOutError(e))
+          throw new Error('Execution context was destroyed, most likely because of a navigation.');
+        throw e;
+      }).then(response => {
+        return response.result.objectId;
+      });
     }
     return this._globalObjectId;
   }
 
+  private _awaitPromise(objectId: Protocol.Runtime.RemoteObjectId) {
+    return this._session.send('Runtime.awaitPromise', {
+      promiseObjectId: objectId,
+      returnByValue: false
+    }).catch(e => {
+      if (isSwappedOutError(e))
+        return contextDestroyedResult;
+      throw e;
+    });
+  }
+
+  private _returnObjectByValue(objectId: Protocol.Runtime.RemoteObjectId) {
+    const serializeFunction = function() {
+      try {
+        return JSON.stringify(this);
+      } catch (e) {
+        if (e instanceof TypeError)
+          return void 0;
+        throw e;
+      }
+    };
+    return this._session.send('Runtime.callFunctionOn', {
+      // Serialize object using standard JSON implementation to correctly pass 'undefined'.
+      functionDeclaration: serializeFunction + '\n' + suffix + '\n',
+      objectId: objectId,
+      returnByValue: true
+    }).catch(e => {
+      if (isSwappedOutError(e))
+        return contextDestroyedResult;
+      throw e;
+    }).then(serializeResponse => {
+      if (serializeResponse.wasThrown)
+        throw new Error('Serialization failed: ' + serializeResponse.result.description);
+      // This is the case of too long property chain, not serializable to json string.
+      if (serializeResponse.result.type === 'undefined')
+        return undefined;
+      if (serializeResponse.result.type !== 'string')
+        throw new Error('Unexpected result of JSON.stringify: ' + JSON.stringify(serializeResponse, null, 2));
+      return JSON.parse(serializeResponse.result.value);
+    });
+  }
+
   async getProperties(handle: js.JSHandle): Promise<Map<string, js.JSHandle>> {
     const response = await this._session.send('Runtime.getProperties', {
       objectId: toRemoteObject(handle).objectId,
@@ -328,9 +310,16 @@ export class ExecutionContextDelegate implements js.ExecutionContextDelegate {
     }
     return { value: arg };
   }
-
 }
 
+const suffix = `//# sourceURL=${EVALUATION_SCRIPT_URL}`;
+const contextDestroyedResult = {
+  wasThrown: true,
+  result: {
+    description: 'Protocol error: Execution context was destroyed, most likely because of a navigation.'
+  } as Protocol.Runtime.RemoteObject
+};
+
 function toRemoteObject(handle: js.JSHandle): Protocol.Runtime.RemoteObject {
   return handle._remoteObject as Protocol.Runtime.RemoteObject;
 }

src/webkit/FrameManager.ts

@@ -91,7 +91,7 @@ export class FrameManager extends EventEmitter implements frames.FrameDelegate {
     ];
   }
 
-  async _swapSessionOnNavigation(newSession) {
+  _swapSessionOnNavigation(newSession) {
     helper.removeEventListeners(this._sessionListeners);
     this.disconnectFromTarget();
     this._session = newSession;

src/webkit/NetworkManager.ts

@@ -63,6 +63,7 @@ export class NetworkManager extends EventEmitter {
 
   async initialize() {
     await this._sesssion.send('Network.enable');
+    await this._sesssion.send('Network.setExtraHTTPHeaders', { headers: this._extraHTTPHeaders });
   }
 
   async setExtraHTTPHeaders(extraHTTPHeaders: { [s: string]: string; }) {

src/webkit/Page.ts

@@ -20,7 +20,7 @@ import * as console from '../console';
 import * as dialog from '../dialog';
 import * as dom from '../dom';
 import * as frames from '../frames';
-import { assert, debugError, helper, RegisteredListener } from '../helper';
+import { assert, helper, RegisteredListener } from '../helper';
 import * as input from '../input';
 import { ClickOptions, mediaColorSchemes, mediaTypes, MultiClickOptions } from '../input';
 import * as js from '../javascript';
@@ -49,6 +49,7 @@ export class Page extends EventEmitter {
   private _frameManager: FrameManager;
   private _bootstrapScripts: string[] = [];
   _javascriptEnabled = true;
+  private _userAgent: string | null = null;
   private _viewport: types.Viewport | null = null;
   _screenshotter: Screenshotter;
   private _workers = new Map<string, Worker>();
@@ -57,14 +58,6 @@ export class Page extends EventEmitter {
   private _emulatedMediaType: string | undefined;
   private _fileChooserInterceptors = new Set<(chooser: FileChooser) => void>();
 
-  static async create(session: TargetSession, browserContext: BrowserContext, defaultViewport: types.Viewport | null): Promise<Page> {
-    const page = new Page(session, browserContext);
-    await page._initialize();
-    if (defaultViewport)
-      await page.setViewport(defaultViewport);
-    return page;
-  }
-
   constructor(session: TargetSession, browserContext: BrowserContext) {
     super();
     this._closedPromise = new Promise(f => this._closedCallback = f);
@@ -97,12 +90,16 @@ export class Page extends EventEmitter {
   }
 
   async _initialize() {
-    return Promise.all([
+    await Promise.all([
       this._frameManager.initialize(),
       this._session.send('Console.enable'),
       this._session.send('Dialog.enable'),
       this._session.send('Page.setInterceptFileChooserDialog', { enabled: true }),
     ]);
+    if (this._userAgent !== null)
+      await this._session.send('Page.overrideUserAgent', { value: this._userAgent });
+    if (this._emulatedMediaType !== undefined)
+      await this._session.send('Page.setEmulatedMedia', { media: this._emulatedMediaType || '' });
   }
 
   _setSession(newSession: TargetSession) {
@@ -130,8 +127,8 @@ export class Page extends EventEmitter {
 
   async _swapSessionOnNavigation(newSession: TargetSession) {
     this._setSession(newSession);
-    await this._frameManager._swapSessionOnNavigation(newSession);
-    await this._initialize().catch(e => debugError('failed to enable agents after swap: ' + e));
+    this._frameManager._swapSessionOnNavigation(newSession);
+    await this._initialize();
   }
 
   browser(): Browser {
@@ -230,6 +227,7 @@ export class Page extends EventEmitter {
   }
 
   async setUserAgent(userAgent: string) {
+    this._userAgent = userAgent;
     await this._session.send('Page.overrideUserAgent', { value: userAgent });
   }
 
@@ -326,9 +324,8 @@ export class Page extends EventEmitter {
     assert(!options.type || mediaTypes.has(options.type), 'Unsupported media type: ' + options.type);
     assert(!options.colorScheme || mediaColorSchemes.has(options.colorScheme), 'Unsupported color scheme: ' + options.colorScheme);
     assert(!options.colorScheme, 'Media feature emulation is not supported');
-    const media = typeof options.type === 'undefined' ? this._emulatedMediaType : options.type;
-    await this._session.send('Page.setEmulatedMedia', { media: media || '' });
-    this._emulatedMediaType = options.type;
+    this._emulatedMediaType = typeof options.type === 'undefined' ? this._emulatedMediaType : options.type;
+    await this._session.send('Page.setEmulatedMedia', { media: this._emulatedMediaType || '' });
   }
 
   async setViewport(viewport: types.Viewport) {

src/webkit/Target.ts

@@ -19,6 +19,7 @@ import { RegisteredListener } from '../helper';
 import { Browser, BrowserContext } from './Browser';
 import { Page } from './Page';
 import { Protocol } from './protocol';
+import { isSwappedOutError, TargetSession } from './Connection';
 
 const targetSymbol = Symbol('target');
 
@@ -26,7 +27,7 @@ export class Target {
   private _browserContext: BrowserContext;
   _targetId: string;
   private _type: 'page' | 'service-worker' | 'worker';
-  _pagePromise: Promise<Page> | null = null;
+  private _pagePromise: Promise<Page> | null = null;
   private _url: string;
   _initializedPromise: Promise<boolean>;
   _initializedCallback: (value?: unknown) => void;
@@ -52,16 +53,28 @@ export class Target {
       this._pagePromise.then(page => page._didClose());
   }
 
-  _adoptPage(page: Page) {
+  async _swappedIn(oldTarget: Target, session: TargetSession) {
+    this._pagePromise = oldTarget._pagePromise;
+    // Swapped out target should not be accessed by anyone. Reset page promise so that
+    // old target does not close the page on connection reset.
+    oldTarget._pagePromise = null;
+    if (!this._pagePromise)
+      return;
+    const page = await this._pagePromise;
     (page as any)[targetSymbol] = this;
+    page._swapSessionOnNavigation(session).catch(rethrowIfNotSwapped);
   }
 
   async page(): Promise<Page | null> {
     if (this._type === 'page' && !this._pagePromise) {
       const session = this.browser()._connection.session(this._targetId);
-      this._pagePromise = Page.create(session, this._browserContext, this.browser()._defaultViewport).then(page => {
-        this._adoptPage(page);
-        return page;
+      this._pagePromise = new Promise(async f => {
+        const page = new Page(session, this._browserContext);
+        await page._initialize().catch(rethrowIfNotSwapped);
+        if (this.browser()._defaultViewport)
+          await page.setViewport(this.browser()._defaultViewport);
+        (page as any)[targetSymbol] = this;
+        f(page);
       });
     }
     return this._pagePromise;
@@ -83,3 +96,8 @@ export class Target {
     return this._browserContext;
   }
 }
+
+function rethrowIfNotSwapped(e: Error) {
+  if (!isSwappedOutError(e))
+    throw e;
+}