Make multitab work

packages/extension/src/contentScript/index.ts

@@ -62,6 +62,11 @@ class WebpageContentScript {
    */
   tabId: number | null = null;
 
+  /**
+   * Main frame id.
+   */
+  frameId: string | null = null;
+
   /**
    * TabId of the current Tab
    */
@@ -141,6 +146,7 @@ class WebpageContentScript {
 
       if (message?.payload?.type === 'SERVICEWORKER::WEBPAGE::TABID_STORAGE') {
         this.tabId = message.payload.tabId;
+        this.frameId = message.payload.frameId;
       }
 
       if (message?.payload?.type === 'DEVTOOL::WEBPAGE::GET_JS_COOKIES') {
@@ -164,11 +170,15 @@ class WebpageContentScript {
    * @param tabId The tabID whose cookies have to be fetched.
    */
   async getAndProcessJSCookies(tabId: string) {
+    if (!this.frameId) {
+      return;
+    }
     //@ts-ignore
     const jsCookies = await cookieStore.getAll();
     await processAndStoreDocumentCookies({
       tabUrl: window.location.href,
       tabId,
+      frameId: this.frameId,
       documentCookies: jsCookies,
     });
   }

packages/extension/src/serviceWorker/index.ts

@@ -36,7 +36,6 @@ import getQueryParams from '../utils/getQueryParams';
 import reloadCurrentTab from '../utils/reloadCurrentTab';
 import createCookieFromAuditsIssue from '../utils/createCookieFromAuditsIssue';
 import attachCDP from './attachCDP';
-import { getCurrentTabId } from '../utils/getCurrentTabId';
 
 let cookieDB: CookieDatabase | null = null;
 let syncCookieStore: SynchnorousCookieStore | undefined;
@@ -79,6 +78,7 @@ const ALLOWED_EVENTS = [
   'Audits.issueAdded',
   'Network.requestWillBeSent',
   'Page.frameAttached',
+  'Page.frameNavigated',
 ];
 
 const tabIdToFrames: {
@@ -98,12 +98,13 @@ chrome.tabs.onCreated.addListener(async (tab) => {
   unParsedResponseHeaders[tab.id.toString()] = {};
   requestIdToCDPURLMapping[tab.id.toString()] = {};
   auditsIssueForTab[tab.id.toString()] = {};
-  tabIdToFrames[tab.id.toString()] = new Set();
 
   if (!syncCookieStore) {
     syncCookieStore = new SynchnorousCookieStore();
   }
 
+  const targets = await chrome.debugger.getTargets();
+
   if (tabMode && tabMode !== 'unlimited') {
     const doesTabExist = tabToRead;
     if (
@@ -116,14 +117,18 @@ chrome.tabs.onCreated.addListener(async (tab) => {
     tabToRead = tab.id.toString();
     syncCookieStore?.addTabData(tab.id);
 
-    const targets = await chrome.debugger.getTargets();
     const currentTab = targets.filter(
       ({ tabId }) => tabId && tab.id && tabId === tab.id
     );
 
-    tabIdToFrames[tab.id.toString()].add(currentTab[0].id);
+    syncCookieStore.updateFrameIdSet(tab.id, currentTab[0].id);
   } else {
     syncCookieStore?.addTabData(tab.id);
+    const currentTab = targets.filter(
+      ({ tabId }) => tabId && tab.id && tabId === tab.id
+    );
+
+    syncCookieStore.updateFrameIdSet(tab.id, currentTab[0].id);
   }
 });
 
@@ -132,6 +137,12 @@ chrome.tabs.onCreated.addListener(async (tab) => {
  * @see https://developer.chrome.com/docs/extensions/reference/api/tabs#event-onRemoved
  */
 chrome.tabs.onRemoved.addListener((tabId) => {
+  delete unParsedRequestHeaders[tabId];
+  delete unParsedResponseHeaders[tabId];
+  delete requestIdToCDPURLMapping[tabId];
+  delete auditsIssueForTab[tabId];
+  delete tabIdToFrames[tabId];
+
   syncCookieStore?.removeTabData(tabId);
 });
 
@@ -198,6 +209,13 @@ chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
 
   if (changeInfo.status === 'loading' && tab.url) {
     syncCookieStore?.removeCookieData(tabId);
+
+    syncCookieStore?.updateFrameIdSet(
+      tabId,
+      (await chrome.debugger.getTargets()).filter(
+        (target) => target.tabId && target.tabId === tabId
+      )[0].id
+    );
   }
 
   try {
@@ -206,6 +224,9 @@ chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
       payload: {
         type: 'SERVICEWORKER::WEBPAGE::TABID_STORAGE',
         tabId,
+        frameId: (
+          await chrome.debugger.getTargets()
+        ).filter((target) => target.tabId && target.tabId === tabId)[0].id,
       },
     });
   } catch (error) {
@@ -222,15 +243,6 @@ chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
   } catch (error) {
     //Fail silently
   }
-  if (!tab.url) {
-    return;
-  }
-
-  syncCookieStore?.updateUrl(tabId, tab.url);
-
-  if (changeInfo.status === 'loading' && tab.url) {
-    syncCookieStore?.removeCookieData(tabId);
-  }
 });
 
 /**
@@ -331,9 +343,10 @@ chrome.debugger.onEvent.addListener(async (source, method, params) => {
     return;
   }
 
-  let tabId = await getCurrentTabId();
+  let tabId = '';
   try {
     const targets = await chrome.debugger.getTargets();
+
     await Promise.all(
       targets.map(async ({ id }) => {
         await chrome.debugger.attach({ targetId: id }, '1.3');
@@ -348,6 +361,38 @@ chrome.debugger.onEvent.addListener(async (source, method, params) => {
 
   if (source?.tabId) {
     tabId = source?.tabId?.toString();
+  } else if (source.targetId && method !== 'Page.frameAttached') {
+    const tab = Object.keys(tabIdToFrames).filter(
+      (key) => source.targetId && tabIdToFrames[key].has(source.targetId)
+    );
+    tabId = tab[0];
+  }
+
+  if (method === 'Page.frameAttached' && params) {
+    const frameData = params as Protocol.Page.FrameAttachedEvent;
+    if (source.tabId) {
+      syncCookieStore?.updateFrameIdSet(source.tabId, frameData.frameId);
+    } else if (source.targetId) {
+      Object.values(syncCookieStore?.tabs ?? {}).forEach(({ frameIdSet }) => {
+        if (source.targetId && frameIdSet.has(source.targetId)) {
+          frameIdSet.add(frameData.frameId);
+        }
+      });
+    }
+  }
+
+  if (method === 'Page.frameNavigated' && params) {
+    const {
+      frame: { parentId = '', id },
+    } = params as Protocol.Page.FrameNavigatedEvent;
+
+    if (parentId) {
+      Object.values(syncCookieStore?.tabs ?? {}).forEach(({ frameIdSet }) => {
+        if (frameIdSet.has(parentId)) {
+          frameIdSet.add(id);
+        }
+      });
+    }
   }
 
   if (!tabId) {
@@ -364,17 +409,6 @@ chrome.debugger.onEvent.addListener(async (source, method, params) => {
     cookieDB = await fetchDictionary();
   }
 
-  if (method === 'Page.frameAttached' && params) {
-    const frameData = params as Protocol.Page.FrameAttachedEvent;
-    if (frameData?.parentFrameId) {
-      Object.values(tabIdToFrames).forEach((setOfFrames) => {
-        if (setOfFrames.has(frameData.parentFrameId)) {
-          setOfFrames.add(frameData.frameId);
-        }
-      });
-    }
-  }
-
   if (method === 'Network.requestWillBeSent' && params) {
     const request = params as Protocol.Network.RequestWillBeSentEvent;
     // To get domain from the request URL if not given in the cookie line.
@@ -770,6 +804,19 @@ chrome.runtime.onMessage.addListener(async (request) => {
       })
     );
   }
+
+  if (request?.type === 'SERVICE_WORKER_DEVTOOLS_GET_FRAME_IDS') {
+    const requestedTabId = request?.payload?.tabId;
+    chrome.runtime.sendMessage({
+      type: 'ServiceWorker::DevTools::NEW_FRAMES_DATA',
+      payload: {
+        tabId: requestedTabId,
+        frames: Array.from(
+          syncCookieStore?.tabs[requestedTabId].frameIdSet ?? new Set()
+        ),
+      },
+    });
+  }
 });
 
 /**

packages/extension/src/store/synchnorousCookieStore.ts

@@ -48,9 +48,7 @@ class SynchnorousCookieStore {
       devToolsOpenState: boolean;
       popupOpenState: boolean;
       newUpdates: number;
-      frameIdSet: {
-        [frameUrl: string]: string[];
-      };
+      frameIdSet: Set<string>;
     };
   } = {};
 
@@ -190,7 +188,7 @@ class SynchnorousCookieStore {
    * @param {number} tabId Tab id.
    * @returns {string | null} The url of the tab if exists else null.
    */
-  getFrameIDSet(tabId: number): { [frameUrl: string]: string[] } | null {
+  getFrameIDSet(tabId: number): Set<string> | null {
     if (!this.tabs[tabId]) {
       return null;
     }
@@ -201,16 +199,13 @@ class SynchnorousCookieStore {
   /**
    * Update FrameId set for a given url for a given tab.
    * @param {number} tabId The url whose url needs to be update.
-   * @param {string} frameUrl The URL where frameId needs to be added.
    * @param {string} frameIdToAdd The new frameId to be added.
    */
-  updateFrameIdSet(tabId: number, frameUrl: string, frameIdToAdd: string) {
+  updateFrameIdSet(tabId: number, frameIdToAdd: string) {
     if (!this.tabs[tabId]) {
       return;
     } else {
-      this.tabs[tabId].frameIdSet[frameUrl] = [
-        ...new Set([...this.tabs[tabId].frameIdSet[frameUrl], frameIdToAdd]),
-      ];
+      this.tabs[tabId].frameIdSet.add(frameIdToAdd);
     }
   }
 
@@ -336,7 +331,7 @@ class SynchnorousCookieStore {
       devToolsOpenState: false,
       popupOpenState: false,
       newUpdates: 0,
-      frameIdSet: {},
+      frameIdSet: new Set(),
     };
   }
 

packages/extension/src/utils/processAndStoreDocumentCookies.ts

@@ -32,12 +32,14 @@ import { fetchDictionary } from './fetchCookieDictionary';
 interface ProcessAndStoreDucmentCookies {
   tabUrl: string;
   tabId: string;
+  frameId: string;
   documentCookies: ParsedCookie[];
 }
 
 const processAndStoreDocumentCookies = async ({
   tabUrl,
   tabId,
+  frameId,
   documentCookies,
 }: ProcessAndStoreDucmentCookies) => {
   try {
@@ -80,7 +82,7 @@ const processAndStoreDocumentCookies = async ({
           url: tabUrl,
           headerType: 'javascript', // @todo Update headerType name.
           isFirstParty: isFirstPartyCookie || null,
-          frameIdList: [0],
+          frameIdList: [frameId],
         };
       }
     );

packages/extension/src/view/devtools/stateProviders/syncCookieStore/index.tsx

@@ -199,6 +199,66 @@ export const Provider = ({ children }: PropsWithChildren) => {
     []
   );
 
+  useEffect(() => {
+    if (!tabFrames) {
+      return;
+    }
+
+    chrome.runtime.sendMessage({
+      type: 'SERVICE_WORKER_DEVTOOLS_GET_FRAME_IDS',
+      payload: {
+        tabId: chrome.devtools.inspectedWindow.tabId,
+      },
+    });
+  }, [tabFrames]);
+
+  /**
+   * Add more tabframes to the existing ones.
+   */
+  const addMoreTabIds = useCallback(
+    async (tabIds: string[]) => {
+      if (!tabFrames) {
+        return;
+      }
+
+      const modifiedTabFrames = tabFrames;
+
+      await Promise.all(
+        tabIds.map(async (id) => {
+          //@ts-ignore
+          try {
+            const frames = await chrome.debugger.sendCommand(
+              { targetId: id },
+              'Page.getFrameTree'
+            );
+            if (
+              //@ts-ignore
+              frames?.frameTree?.frame?.securityOrigin &&
+              //@ts-ignore
+              frames.frameTree.frame.id
+            ) {
+              modifiedTabFrames[
+                //@ts-ignore
+                frames?.frameTree?.frame?.securityOrigin
+              ].frameIds = [
+                ...new Set([
+                  //@ts-ignore
+                  ...modifiedTabFrames[frames.frameTree.frame.securityOrigin]
+                    .frameIds,
+                  //@ts-ignore
+                  frames.frameTree.frame.id,
+                ]),
+              ];
+            }
+          } catch (error) {
+            // fail silently
+          }
+        })
+      );
+      setTabFrames(modifiedTabFrames);
+    },
+    [tabFrames]
+  );
   /**
    * Stores object with frame URLs as keys and boolean values indicating if the frame contains cookies.
    */
@@ -329,6 +389,7 @@ export const Provider = ({ children }: PropsWithChildren) => {
         cookieData?: TabCookies;
         tabToRead?: string;
         tabMode?: string;
+        frames?: string[];
       };
     }) => {
       if (message.type === 'ServiceWorker::SET_TAB_TO_READ') {
@@ -378,6 +439,10 @@ export const Provider = ({ children }: PropsWithChildren) => {
         }
       }
 
+      if (message.type === 'ServiceWorker::DevTools::NEW_FRAMES_DATA') {
+        await addMoreTabIds(message.payload.frames ?? []);
+      }
+
       if (message.type === 'ServiceWorker::DevTools::TABS_RELOADED') {
         setSettingsChanged(false);
       }
@@ -388,7 +453,7 @@ export const Provider = ({ children }: PropsWithChildren) => {
     return () => {
       chrome.runtime.onMessage.removeListener(listener);
     };
-  }, [getAllFramesForCurrentTab, tabId, setSettingsChanged]);
+  }, [getAllFramesForCurrentTab, tabId, setSettingsChanged, addMoreTabIds]);
 
   const tabUpdateListener = useCallback(
     async (_tabId: number, changeInfo: chrome.tabs.TabChangeInfo) => {