fix(mitm): double-headers

ulixee · Apr 1, 2022 · a424c61 · a424c61
1 parent a4f79cd
commit a424c61
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 17 deletions.
diff --git a/core/lib/InjectedScripts.ts b/core/lib/InjectedScripts.ts
@@ -89,14 +89,12 @@ export default class InjectedScripts {
     return puppetPage.addNewDocumentScript(showInteractionScript, true);
   }
 
-  public static getIndexedDbStorageRestoreScript(restoreDBs: IIndexedDB[]): string {
+  public static getIndexedDbStorageRestoreScript(): string {
     return `(function restoreIndexedDB(dbs) {
 const exports = {}; // workaround for ts adding an exports variable
 ${stringifiedTypeSerializerClass};
 
 ${pageScripts.indexedDbRestore};
-restoreUserStorage(dbs);
-})(${JSON.stringify(restoreDBs)});`;
-
+})();`;
   }
 }
diff --git a/core/lib/UserProfile.ts b/core/lib/UserProfile.ts
@@ -73,12 +73,15 @@ export default class UserProfile {
 
   public static async installStorage(session: Session, page: IPuppetPage): Promise<void> {
     const { userProfile } = session;
-    const domStorage = userProfile.storage;
-    if (!domStorage) return;
-    const sessionId = session.id;
-    const origins = Object.keys(domStorage);
-
+    const domStorage: IDomStorage = {};
+    const origins: string[] = [];
+    for (const [origin, storage] of Object.entries(userProfile.storage)) {
+      const og = origin.toLowerCase();
+      origins.push(og);
+      domStorage[og] = storage;
+    }
     if (!origins.length) return;
+    const sessionId = session.id;
 
     const interceptorHandlers = session.mitmRequestSession.interceptorHandlers;
 
@@ -109,12 +112,19 @@ for (const [key,value] of ${JSON.stringify(localStorage)}) {
 }\n`;
             }
 
-            if (originStorage?.indexedDB) {
+            let readyClass = 'ready';
+            if (originStorage?.indexedDB?.length) {
+              readyClass ='';
               script += `\n\n 
-             ${InjectedScripts.getIndexedDbStorageRestoreScript(originStorage.indexedDB)}`;
+             ${InjectedScripts.getIndexedDbStorageRestoreScript()}
+             
+             const dbs = ${JSON.stringify(originStorage.indexedDB)};
+             restoreUserStorage(dbs).then(() => {
+               document.body.setAttribute('class', 'ready');
+             });`;
             }
 
-            res.end(`<html><body>
+            res.end(`<html><body class="${readyClass}">
 <h5>${url.origin}</h5>
 <script>
 ${script}
@@ -138,12 +148,24 @@ ${origins.map(x => `<iframe src="${x}"></iframe>`).join('\n')}
       for (const frame of page.frames) {
         if (frame === page.mainFrame) {
           // no loader is set, so need to have special handling
-          if (!page.mainFrame.activeLoader.lifecycle.load) {
-            await page.mainFrame.waitOn('frame-lifecycle', x => x.name === 'load');
+          if (!frame.activeLoader.lifecycle.load) {
+            await frame.waitOn('frame-lifecycle', x => x.name === 'load');
           }
           continue;
         }
         await frame.waitForLifecycleEvent('load');
+
+        await frame.evaluate(
+          `(async function() {
+            while (!document.querySelector("body.ready")) {
+              await new Promise(resolve => setTimeout(resolve, 20));
+            }
+          })()`,
+          true,
+          {
+            shouldAwaitExpression: true,
+          },
+        );
       }
     } finally {
       session.mitmRequestSession.interceptorHandlers = interceptorHandlers;

diff --git a/core/test/user-profile.test.ts b/core/test/user-profile.test.ts
@@ -232,7 +232,7 @@ document.querySelector('#session').innerHTML = [session1,session2,session3].join
       },
     };
     for (let i = 0; i < 100; i += 1) {
-      storage[`https://domain${i}.com`] = {
+      storage[`https://${(i === 1 ? 'D' : 'd')}omain${i}.com`] = {
         indexedDB: [],
         localStorage: [
           ['1', '2'],
@@ -324,7 +324,7 @@ localStorage.setItem('Test1', 'value1');
             localStorage: [['test', 'site1.org']],
             sessionStorage: [],
           },
-          'https://site2.org': {
+          'https://SITE2.org': {
             indexedDB: [],
             localStorage: [['test2', 'site2.org']],
             sessionStorage: [],

diff --git a/mitm/handlers/HttpRequestHandler.ts b/mitm/handlers/HttpRequestHandler.ts
@@ -222,9 +222,9 @@ export default class HttpRequestHandler extends BaseHttpHandler {
     const { serverToProxyResponse, proxyToClientResponse, requestSession, events } = context;
 
     proxyToClientResponse.statusCode = context.status;
+
     // write individually so we properly write header-lists
     for (const [key, value] of Object.entries(context.responseHeaders)) {
-      proxyToClientResponse.setHeader(key, value);
       try {
         proxyToClientResponse.setHeader(key, value);
       } catch (error) {

diff --git a/puppet-chrome/lib/Frame.ts b/puppet-chrome/lib/Frame.ts
@@ -503,6 +503,9 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
   private async createIsolatedWorld(): Promise<number> {
     try {
       if (!this.isAttached) return;
+      await new Promise(setImmediate);
+      if (this.isolatedContextId) return this.isolatedContextId;
+
       const isolatedWorld = await this.devtoolsSession.send(
         'Page.createIsolatedWorld',
         {