feat(timetravel): cookies/storage for pagestate

core/dbs/SessionDb.ts

@@ -22,6 +22,7 @@ import TabsTable from '../models/TabsTable';
 import ResourceStatesTable from '../models/ResourceStatesTable';
 import SocketsTable from '../models/SocketsTable';
 import Core from '../index';
+import StorageChangesTable from '../models/StorageChangesTable';
 
 const { log } = Log(module);
 
@@ -52,6 +53,7 @@ export default class SessionDb {
   public readonly mouseEvents: MouseEventsTable;
   public readonly focusEvents: FocusEventsTable;
   public readonly scrollEvents: ScrollEventsTable;
+  public readonly storageChanges: StorageChangesTable;
   public readonly screenshots: ScreenshotsTable;
   public readonly devtoolsMessages: DevtoolsMessagesTable;
   public readonly tabs: TabsTable;
@@ -90,6 +92,7 @@ export default class SessionDb {
     this.scrollEvents = new ScrollEventsTable(this.db);
     this.sessionLogs = new SessionLogsTable(this.db);
     this.screenshots = new ScreenshotsTable(this.db);
+    this.storageChanges = new StorageChangesTable(this.db);
     this.devtoolsMessages = new DevtoolsMessagesTable(this.db);
 
     this.tables.push(
@@ -110,6 +113,7 @@ export default class SessionDb {
       this.sessionLogs,
       this.devtoolsMessages,
       this.screenshots,
+      this.storageChanges,
     );
 
     if (!readonly) {

core/injected-scripts/DomAssertions.ts

@@ -101,6 +101,7 @@ class DomAssertions {
     comparison: IAssertionAndResult['comparison'],
     result: T,
   ): boolean {
+    if (comparison === '!!') return !!value;
     if (comparison === '===') return value === result;
     if (comparison === '!==') return value !== result;
     if (comparison === '<=') return value <= result;

core/lib/PageStateListener.ts

@@ -95,19 +95,30 @@ export default class PageStateListener extends TypedEventEmitter<IPageStateEvent
       url: 0,
       resource: [],
       dom: 0,
+      storage: 0,
     };
     const { domAssertionsByFrameId, assertions, totalAssertions, minValidAssertions } =
       this.batchAssertionsById.get(batchId);
     for (const assertion of assertions) {
-      const [frameId, type, args, , result] = assertion;
-      if (type === 'url') {
+      const [frameId, assertType, args, comparison, result] = assertion;
+      if (assertType === 'url') {
         const frame = this.tab.frameEnvironmentsById.get(frameId) ?? this.tab.mainFrameEnvironment;
         const url = await frame.getUrl();
         if (url !== result) failCounts.url += 1;
       }
-      if (type === 'resource') {
-        const resource = this.tab.findResource(args[0]);
-        if (!resource) failCounts.resource.push(args[0]);
+      if (assertType === 'resource') {
+        const filter = args[0];
+        const resource = this.tab.findResource(filter);
+        if (comparison === '!!' && !resource) failCounts.resource.push(filter);
+      }
+      if (assertType === 'storage') {
+        const [filter, prop] = args;
+        const storage = this.tab.findStorageChange(filter);
+        if (comparison === '!!' && !storage) failCounts.storage += 1;
+
+        if (comparison === '===' && prop) {
+          if (!storage || storage[prop] !== result) failCounts.storage += 1;
+        }
       }
     }
 
@@ -117,7 +128,8 @@ export default class PageStateListener extends TypedEventEmitter<IPageStateEvent
       failCounts.dom += failedDomAssertions;
     }
 
-    const failedCount = failCounts.url + failCounts.resource.length + failCounts.dom;
+    const failedCount =
+      failCounts.url + failCounts.resource.length + failCounts.dom + failCounts.storage;
     const validAssertions = totalAssertions - failedCount;
     this.logger.stats('BatchAssert results', {
       batchId,

core/lib/Resources.ts

@@ -11,6 +11,9 @@ import MitmRequestContext from '@ulixee/hero-mitm/lib/MitmRequestContext';
 import { IPuppetResourceRequest } from '@ulixee/hero-interfaces/IPuppetNetworkEvents';
 import ResourcesTable from '../models/ResourcesTable';
 import Session from './Session';
+import { ICookie } from '@ulixee/hero-interfaces/ICookie';
+import { Cookie } from 'tough-cookie';
+import StorageChangesTable, { IStorageChangesEntry } from '../models/StorageChangesTable';
 
 const { log } = Log(module);
 
@@ -20,6 +23,7 @@ export default class Resources {
   } = {};
 
   public readonly resourcesById = new Map<number, IResourceMeta>();
+  public readonly cookiesByDomain = new Map<string, Record<string, ICookie>>();
 
   private readonly mitmErrorsByUrl = new Map<
     string,
@@ -31,12 +35,14 @@ export default class Resources {
 
   private readonly logger: IBoundLog;
   private readonly model: ResourcesTable;
+  private readonly cookiesModel: StorageChangesTable;
 
   constructor(
     private readonly session: Session,
     readonly browserRequestMatcher: BrowserRequestMatcher,
   ) {
     this.model = session.db.resources;
+    this.cookiesModel = session.db.storageChanges;
     this.logger = log.createChild(module, {
       sessionId: session.id,
     });
@@ -185,6 +191,9 @@ export default class Resources {
 
     if (isResponse) {
       this.resourcesById.set(resource.id, resource);
+
+      const responseEvent = resourceEvent as IRequestSessionResponseEvent;
+      this.recordCookies(tabId, responseEvent);
     }
     return resource;
   }
@@ -288,6 +297,66 @@ export default class Resources {
     return resource;
   }
 
+  private recordCookies(tabId: number, responseEvent: IRequestSessionResponseEvent): void {
+    const { response } = responseEvent;
+    if (!response?.headers) return;
+
+    let setCookie = response.headers['set-cookie'] ?? response.headers['Set-Cookie'];
+    if (!setCookie) return;
+
+    if (!Array.isArray(setCookie)) setCookie = [setCookie];
+    const defaultDomain = responseEvent.url.host;
+    for (const cookieHeader of setCookie) {
+      const cookie = Cookie.parse(cookieHeader, { loose: true });
+      let domain = cookie.domain || defaultDomain;
+      // restore stripped leading .
+      if (cookie.domain && cookieHeader.toLowerCase().includes(`domain=.${domain}`)) {
+        domain = `.${domain}`;
+      }
+      if (!this.cookiesByDomain.has(domain)) this.cookiesByDomain.set(domain, {});
+      const domainCookies = this.cookiesByDomain.get(domain);
+      let action: IStorageChangesEntry['action'] = 'add';
+      const existing = domainCookies[cookie.key];
+      if (existing) {
+        if (cookie.expires && cookie.expires < new Date()) {
+          action = 'remove';
+        } else {
+          action = 'update';
+        }
+      }
+
+      let finalCookie: ICookie;
+      if (action === 'remove') {
+        delete domainCookies[cookie.key];
+      } else {
+        finalCookie = {
+          name: cookie.key,
+          sameSite: cookie.sameSite as any,
+          url: responseEvent.url.href,
+          domain,
+          path: cookie.path,
+          httpOnly: cookie.httpOnly,
+          value: cookie.value,
+          secure: cookie.secure,
+          expires: cookie.expires instanceof Date ? cookie.expires.toISOString() : undefined,
+        };
+
+        if (areCookiesEqual(existing, finalCookie)) continue;
+
+        domainCookies[finalCookie.name] = finalCookie;
+      }
+      this.cookiesModel.insert(tabId, responseEvent.frameId, {
+        type: 'cookie' as any,
+        action,
+        securityOrigin: responseEvent.url.origin,
+        key: cookie.key,
+        value: finalCookie?.value,
+        meta: finalCookie,
+        timestamp: response.browserLoadedTime ?? response.timestamp,
+      });
+    }
+  }
+
   private matchesMitmError(
     tabId: number,
     url: string,
@@ -365,3 +434,15 @@ export default class Resources {
     return new Error('Resource failed to load, but the reason was not provided by devtools.');
   }
 }
+
+function areCookiesEqual(a: ICookie, b: ICookie): boolean {
+  if ((a && !b) || (b && !a)) return false;
+  if (a.name !== b.name) return false;
+  if (a.value !== b.value) return false;
+  if (a.expires !== b.expires) return false;
+  if (a.path !== b.path) return false;
+  if (a.secure !== b.secure) return false;
+  if (a.sameParty !== b.sameParty) return false;
+  if (a.httpOnly !== b.httpOnly) return false;
+  return true;
+}

core/lib/Tab.ts

@@ -41,6 +41,7 @@ import Resources from './Resources';
 import PageStateListener from './PageStateListener';
 import IScreenRecordingOptions from '@ulixee/hero-interfaces/IScreenRecordingOptions';
 import ScreenshotsTable from '../models/ScreenshotsTable';
+import { IStorageChangesEntry } from '../models/StorageChangesTable';
 
 const { log } = Log(module);
 
@@ -328,6 +329,12 @@ export default class Tab
     return null;
   }
 
+  public findStorageChange(
+    filter: Omit<IStorageChangesEntry, 'tabId' | 'timestamp' | 'value' | 'meta'>,
+  ): IStorageChangesEntry {
+    return this.session.db.storageChanges.findChange(this.id, filter);
+  }
+
   /////// DELEGATED FNS ////////////////////////////////////////////////////////////////////////////////////////////////
 
   public interact(...interactionGroups: IInteractionGroups): Promise<void> {
@@ -755,6 +762,8 @@ export default class Tab
     page.on('resource-failed', this.onResourceFailed.bind(this), true);
     page.on('navigation-response', this.onNavigationResourceResponse.bind(this), true);
 
+    page.on('dom-storage-updated', this.onStorageUpdated.bind(this), true);
+
     // websockets
     page.on('websocket-handshake', ev => {
       this.session.mitmRequestSession?.registerWebsocketHeaders(this.id, ev);
@@ -977,6 +986,10 @@ export default class Tab
     });
   }
 
+  private onStorageUpdated(event: IPuppetPageEvents['dom-storage-updated']): void {
+    this.session.db.storageChanges.insert(this.id, null, event);
+  }
+
   private getFrameForEventOrQueueForReady(
     type: keyof IPuppetPageEvents,
     event: IPuppetPageEvents[keyof IPuppetPageEvents] & { frameId: string },

core/lib/UserProfile.ts

@@ -1,5 +1,5 @@
 import IUserProfile from '@ulixee/hero-interfaces/IUserProfile';
-import IDomStorage from '@ulixee/hero-interfaces/IDomStorage';
+import IDomStorage, { IDomStorageForOrigin } from '@ulixee/hero-interfaces/IDomStorage';
 import Log from '@ulixee/commons/lib/Logger';
 import { IPuppetPage } from '@ulixee/hero-interfaces/IPuppetPage';
 import { assert } from '@ulixee/commons/lib/utils';
@@ -12,24 +12,46 @@ export default class UserProfile {
   public static async export(session: Session): Promise<IUserProfile> {
     const cookies = await session.browserContext.getCookies();
 
-    const storage: IDomStorage = {};
+    const exportedStorage: IDomStorage = { ...(session.options.userProfile?.storage ?? {}) };
     for (const tab of session.tabsById.values()) {
       const page = tab.puppetPage;
 
-      const dbs = await page.getIndexedDbDatabaseNames();
-      const frames = page.frames;
-      for (const { origin, frameId, databases } of dbs) {
-        const frame = frames.find(x => x.id === frameId);
-        storage[origin] = await frame?.evaluate(
-          `window.exportDomStorage(${JSON.stringify(databases)})`,
-          true,
-        );
+      for (const {
+        origin,
+        storageForOrigin,
+        databaseNames,
+        frame,
+      } of await page.domStorageTracker.getStorageByOrigin()) {
+        const originStorage = {
+          ...storageForOrigin,
+          indexedDB: storageForOrigin.indexedDB.map(x => ({
+            ...x,
+            data: { ...x.data },
+            objectStores: [...x.objectStores],
+          })),
+        };
+        exportedStorage[origin] = originStorage;
+
+        if (frame) {
+          const databases = JSON.stringify(databaseNames);
+          const liveData = await frame.evaluate<IDomStorageForOrigin>(
+            `window.exportDomStorage(${databases})`,
+            true,
+          );
+          originStorage.localStorage = liveData.localStorage;
+          originStorage.sessionStorage = liveData.sessionStorage;
+          for (const dbWithData of liveData.indexedDB) {
+            if (!dbWithData) continue;
+            const idx = originStorage.indexedDB.findIndex(x => x.name === dbWithData.name);
+            originStorage.indexedDB[idx] = dbWithData;
+          }
+        }
       }
     }
 
     return {
       cookies,
-      storage,
+      storage: exportedStorage,
       userAgentString: session.plugins.browserEmulator.userAgentString,
       deviceProfile: session.plugins.browserEmulator.deviceProfile,
     } as IUserProfile;
@@ -48,7 +70,11 @@ export default class UserProfile {
       return this;
     }
 
-    const parentLogId = log.info('UserProfile.install', { sessionId });
+    const parentLogId = log.info('UserProfile.install', {
+      sessionId,
+      cookies: cookies?.length,
+      storageDomains: origins?.length,
+    });
 
     let page: IPuppetPage;
     try {
@@ -59,6 +85,8 @@ export default class UserProfile {
       }
 
       if (hasStorage) {
+        session.browserContext.domStorage = {};
+
         // install scripts so we can restore storage
         await InjectedScripts.installDomStorageRestore(page);
 
@@ -73,11 +101,27 @@ export default class UserProfile {
             continue;
           }
 
-          await page.navigate(origin);
-          await page.mainFrame.evaluate(
-            `window.restoreUserStorage(${JSON.stringify(originStorage)})`,
-            true,
-          );
+          try {
+            await page.navigate(origin);
+            await page.mainFrame.evaluate(
+              `window.restoreUserStorage(${JSON.stringify(originStorage)})`,
+              true,
+            );
+
+            session.browserContext.domStorage[origin] = {
+              indexedDB: originStorage.indexedDB.map(x => ({
+                ...x,
+                data: { ...x.data },
+                objectStores: [...x.objectStores],
+              })),
+              sessionStorage: originStorage.sessionStorage.map(x => ({ ...x })),
+              localStorage: originStorage.localStorage.map(x => ({ ...x })),
+            };
+          } catch (error) {
+            throw new Error(
+              `Could not restore profile for origin ("${origin}") => ${error.message}`,
+            );
+          }
         }
       }
     } finally {