fix(webkit): forward network messages from provisional to committed page (#527)

src/webkit/wkPage.ts

@@ -60,7 +60,7 @@ export class WKPage implements PageDelegate {
     this._session = undefined as any as WKSession;
   }
 
-  async _initializePageProxySession() {
+  private async _initializePageProxySession() {
     const promises : Promise<any>[] = [
       this._pageProxySession.send('Dialog.enable'),
       this._networkManager.initializePageProxySession(this._page._state.credentials)
@@ -87,20 +87,28 @@ export class WKPage implements PageDelegate {
       this._setBootstrapScripts(session).catch(e => debugError(e));
   }
 
+  async initialize() {
+    await Promise.all([
+      this._initializePageProxySession(),
+      this._initializeSession(this._session, ({frameTree}) => this._handleFrameTree(frameTree)),
+    ]);
+  }
+
   // This method is called for provisional targets as well. The session passed as the parameter
   // may be different from the current session and may be destroyed without becoming current.
-  async _initializeSession(session: WKSession, isProvisional: boolean) {
+  async _initializeSession(session: WKSession, resourceTreeHandler: (r: Protocol.Page.getResourceTreeReturnValue) => void) {
+    const isProvisional = this._session !== session;
     const promises : Promise<any>[] = [
       // Page agent must be enabled before Runtime.
       session.send('Page.enable'),
-      session.send('Page.getResourceTree').then(({frameTree}) => this._handleFrameTree(frameTree)),
+      session.send('Page.getResourceTree').then(resourceTreeHandler),
       // Resource tree should be received before first execution context.
       session.send('Runtime.enable'),
       session.send('Page.createIsolatedWorld', { name: UTILITY_WORLD_NAME, source: `//# sourceURL=${EVALUATION_SCRIPT_URL}` }),
       session.send('Console.enable'),
       session.send('Page.setInterceptFileChooserDialog', { enabled: true }),
       this._networkManager.initializeSession(session, this._page._state.interceptNetwork, this._page._state.offlineMode),
-      this._workers.initializeSession(session),
+      this._workers.initializeSession(session)
     ];
     const contextOptions = this._page.browserContext()._options;
     if (contextOptions.userAgent)

src/webkit/wkPageProxy.ts

@@ -22,17 +22,16 @@ import { WKSession } from './wkConnection';
 import { WKPage } from './wkPage';
 import { RegisteredListener, helper, assert, debugError } from '../helper';
 import { Events } from '../events';
+import { WKProvisionalPage } from './wkProvisionalPage';
 
-// We keep provisional messages on the session instace until provisional
-// target is committed. Non-provisional target (there should be just one)
-// has undefined instead.
-const provisionalMessagesSymbol = Symbol('provisionalMessages');
+const isPovisionalSymbol = Symbol('isPovisional');
 
 export class WKPageProxy {
   private readonly _pageProxySession: WKSession;
   readonly _browserContext: BrowserContext;
   private _pagePromise: Promise<Page> | null = null;
   private _wkPage: WKPage | null = null;
+  private _provisionalPage: WKProvisionalPage | null = null;
   private readonly _firstTargetPromise: Promise<void>;
   private _firstTargetCallback?: () => void;
   private readonly _sessions = new Map<string, WKSession>();
@@ -69,6 +68,10 @@ export class WKPageProxy {
     for (const session of this._sessions.values())
       session.dispose();
     this._sessions.clear();
+    if (this._provisionalPage) {
+      this._provisionalPage.dispose();
+      this._provisionalPage = null;
+    }
     if (this._wkPage)
       this._wkPage.didDisconnect();
   }
@@ -79,7 +82,7 @@ export class WKPageProxy {
 
   private _isProvisionalCrossProcessLoadInProgress() : boolean {
     for (const anySession of this._sessions.values()) {
-      if ((anySession as any)[provisionalMessagesSymbol])
+      if ((anySession as any)[isPovisionalSymbol])
         return true;
     }
     return false;
@@ -113,7 +116,7 @@ export class WKPageProxy {
     await this._firstTargetPromise;
     let session: WKSession | undefined;
     for (const anySession of this._sessions.values()) {
-      if (!(anySession as any)[provisionalMessagesSymbol]) {
+      if (!(anySession as any)[isPovisionalSymbol]) {
         session = anySession;
         break;
       }
@@ -121,10 +124,7 @@ export class WKPageProxy {
     assert(session, 'One non-provisional target session must exist');
     this._wkPage = new WKPage(this._browserContext, this._pageProxySession);
     this._wkPage.setSession(session!);
-    await Promise.all([
-      this._wkPage._initializePageProxySession(),
-      this._wkPage._initializeSession(session!, false),
-    ]);
+    await this._wkPage.initialize();
     return this._wkPage._page;
   }
 
@@ -144,9 +144,11 @@ export class WKPageProxy {
       this._firstTargetCallback = undefined;
     }
     if (targetInfo.isProvisional)
-      (session as any)[provisionalMessagesSymbol] = [];
-    if (targetInfo.isProvisional && this._wkPage)
-      this._wkPage._initializeSession(session, true);
+      (session as any)[isPovisionalSymbol] = true;
+    if (targetInfo.isProvisional && this._wkPage) {
+      assert(!this._provisionalPage);
+      this._provisionalPage = new WKProvisionalPage(session, this._wkPage);
+    }
     if (targetInfo.isPaused)
       this._pageProxySession.send('Target.resume', { targetId: targetInfo.targetId }).catch(debugError);
   }
@@ -157,6 +159,10 @@ export class WKPageProxy {
     if (session)
       session.dispose();
     this._sessions.delete(targetId);
+    if (this._provisionalPage && this._provisionalPage._session === session) {
+      this._provisionalPage.dispose();
+      this._provisionalPage = null;
+    }
     if (this._wkPage && this._wkPage._session === session && crashed)
       this._wkPage.didClose(crashed);
   }
@@ -165,11 +171,7 @@ export class WKPageProxy {
     const { targetId, message } = event;
     const session = this._sessions.get(targetId);
     assert(session, 'Unknown target: ' + targetId);
-    const provisionalMessages = (session as any)[provisionalMessagesSymbol];
-    if (provisionalMessages)
-      provisionalMessages.push(message);
-    else
-      session!.dispatchMessage(JSON.parse(message));
+    session!.dispatchMessage(JSON.parse(message));
   }
 
   private _onDidCommitProvisionalTarget(event: Protocol.Target.didCommitProvisionalTargetPayload) {
@@ -180,11 +182,13 @@ export class WKPageProxy {
     assert(oldSession, 'Unknown old target: ' + oldTargetId);
     // TODO: make some calls like screenshot catch swapped out error and retry.
     oldSession!.errorText = 'Target was swapped out.';
-    const provisionalMessages = (newSession as any)[provisionalMessagesSymbol];
-    assert(provisionalMessages, 'Committing target must be provisional');
-    (newSession as any)[provisionalMessagesSymbol] = undefined;
-    for (const message of provisionalMessages)
-      newSession!.dispatchMessage(JSON.parse(message));
-    this._wkPage!.setSession(newSession!);
+    (newSession as any)[isPovisionalSymbol] = undefined;
+    if (this._provisionalPage) {
+      this._provisionalPage.commit();
+      this._provisionalPage.dispose();
+      this._provisionalPage = null;
+    }
+    if (this._wkPage)
+      this._wkPage.setSession(newSession!);
   }
 }

src/webkit/wkProvisionalPage.ts

@@ -0,0 +1,63 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { WKSession } from './wkConnection';
+import { WKPage } from './wkPage';
+import { RegisteredListener, helper, assert } from '../helper';
+import { Protocol } from './protocol';
+
+export class WKProvisionalPage {
+  readonly _session: WKSession;
+  private readonly _wkPage: WKPage;
+  private _sessionListeners: RegisteredListener[] = [];
+  private _mainFrameId: string | null = null;
+
+  constructor(session: WKSession, page: WKPage) {
+    this._session = session;
+    this._wkPage = page;
+
+    this._sessionListeners = [
+      'Network.requestWillBeSent',
+      'Network.requestIntercepted',
+      'Network.responseReceived',
+      'Network.loadingFinished',
+      'Network.loadingFailed',
+    ].map(name => helper.addEventListener(this._session, name, args => this._onNetworkEvent(name, args)));
+
+    this._wkPage._initializeSession(session, ({frameTree}) => this._handleFrameTree(frameTree));
+  }
+
+  dispose() {
+    helper.removeEventListeners(this._sessionListeners);
+  }
+
+  commit() {
+    assert(this._mainFrameId);
+    this._wkPage._onFrameAttached(this._mainFrameId!, null);
+  }
+
+  private _onNetworkEvent(eventName: string, payload: any) {
+    // Pretend that the events happened in the same process.
+    if (payload.frameId)
+      payload.frameId = this._wkPage._page._frameManager.mainFrame()._id;
+    this._wkPage._session.emit(eventName, payload);
+  }
+
+  private _handleFrameTree(frameTree: Protocol.Page.FrameResourceTree) {
+    assert(!frameTree.frame.parentId);
+    this._mainFrameId = frameTree.frame.id;
+  }
+}