fix: record sessionStorage when authenticating

harlan-zw · May 22, 2024 · 9fec05f · 9fec05f
1 parent 383f246
commit 9fec05f
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/packages/core/src/puppeteer/util.ts b/packages/core/src/puppeteer/util.ts
@@ -25,6 +25,17 @@ export async function setupPage(page: Page) {
       resolvedConfig.localStorage,
     )
   }
+  // set session storage
+  if (resolvedConfig.sessionStorage) {
+    await page.evaluateOnNewDocument(
+      (data) => {
+        sessionStorage.clear()
+        for (const key in data)
+          sessionStorage.setItem(key, data[key])
+      },
+      resolvedConfig.sessionStorage,
+    )
+  }
   if (resolvedConfig.cookies) {
     await page.setCookie(...resolvedConfig.cookies.map(cookie => ({ domain: resolvedConfig.site, ...cookie })))
       .catch(softErrorHandler('Failed to set cookies'))

diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
@@ -277,6 +277,12 @@ export interface ResolvedUserConfig {
    * @default {}
    */
   localStorage: Record<string, any>
+  /**
+   * Session storage to add to the browser context.
+   *
+   * @default {}
+   */
+  sessionStorage: Record<string, any>
   /**
    * Extra headers to provide for any HTTP requests.
    *

diff --git a/packages/core/src/unlighthouse.ts b/packages/core/src/unlighthouse.ts
@@ -145,6 +145,9 @@ export async function createUnlighthouse(userConfig: UserConfig, provider?: Prov
   if (resolvedConfig.hooks?.authenticate) {
     // do an authentication step
     await worker.cluster.execute({}, async (taskCtx) => {
+      logger.debug('Running authentication hook')
+      await taskCtx.page.setBypassCSP(true)
+
       await hooks.callHook('authenticate', taskCtx.page)
       // collect page authentication, either cookie or localStorage tokens
       const localStorageData = await taskCtx.page.evaluate(() => {
@@ -155,12 +158,28 @@ export async function createUnlighthouse(userConfig: UserConfig, provider?: Prov
             json[key] = localStorage.getItem(key)
         }
         return json
+      }).catch((e) => {
+        logger.warn('Failed to collect authentication localStorage.\n', e)
+        return {}
+      })
+      const sessionStorageData = await taskCtx.page.evaluate(() => {
+        const json: Record<string, any> = {}
+        for (let i = 0; i < sessionStorage.length; i++) {
+          const key = sessionStorage.key(i)
+          if (key)
+            json[key] = sessionStorage.getItem(key)
+        }
+        return json
+      }).catch((e) => {
+        logger.warn('Failed to collect authentication sessionStorage.\n', e)
+        return {}
       })
       const cookies = await taskCtx.page.cookies()
       // merge this into the config
       // @ts-expect-error untyped
       ctx.resolvedConfig.cookies = [...(ctx.resolvedConfig.cookies || []), ...cookies as any as ResolvedUserConfig['cookies']]
       ctx.resolvedConfig.localStorage = { ...ctx.resolvedConfig.localStorage, ...localStorageData }
+      ctx.resolvedConfig.sessionStorage = { ...ctx.resolvedConfig.sessionStorage, ...sessionStorageData }
     })
   }