From 82076392f79f6fb407f47309efe2de9810498b55 Mon Sep 17 00:00:00 2001 From: Simon Knott Date: Wed, 5 Feb 2025 12:20:37 +0100 Subject: [PATCH] extract --- .../src/server/browserContext.ts | 113 +--------------- .../src/server/injected/storageScript.ts | 121 ++++++++++++++++++ 2 files changed, 128 insertions(+), 106 deletions(-) create mode 100644 packages/playwright-core/src/server/injected/storageScript.ts diff --git a/packages/playwright-core/src/server/browserContext.ts b/packages/playwright-core/src/server/browserContext.ts index 3077c0c973..ca2a79b6d6 100644 --- a/packages/playwright-core/src/server/browserContext.ts +++ b/packages/playwright-core/src/server/browserContext.ts @@ -43,6 +43,7 @@ import type { Artifact } from './artifact'; import { Clock } from './clock'; import type { ClientCertificatesProxy } from './socksClientCertificatesInterceptor'; import { RecorderApp } from './recorder/recorderApp'; +import * as storageScript from './injected/storageScript'; export abstract class BrowserContext extends SdkObject { static Events = { @@ -513,76 +514,15 @@ export abstract class BrowserContext extends SdkObject { }; const originsToSave = new Set(this._origins); - async function _collectStorageScript() { - const idbResult = await Promise.all((await indexedDB.databases()).map(async dbInfo => { - if (!dbInfo.name) - throw new Error('Database name is empty'); - - function idbRequestToPromise(request: T) { - return new Promise((resolve, reject) => { - request.addEventListener('success', () => resolve(request.result)); - request.addEventListener('error', () => reject(request.error)); - }); - } - - const db = await idbRequestToPromise(indexedDB.open(dbInfo.name)); - const transaction = db.transaction(db.objectStoreNames, 'readonly'); - const stores = await Promise.all([...db.objectStoreNames].map(async storeName => { - const objectStore = transaction.objectStore(storeName); - - const keys = await idbRequestToPromise(objectStore.getAllKeys()); - const records = await Promise.all(keys.map(async key => { - return { - key: objectStore.keyPath === null ? key : undefined, - value: await idbRequestToPromise(objectStore.get(key)) - }; - })); - - const indexes = [...objectStore.indexNames].map(indexName => { - const index = objectStore.index(indexName); - return { - name: index.name, - keyPath: typeof index.keyPath === 'string' ? index.keyPath : undefined, - keyPathArray: Array.isArray(index.keyPath) ? index.keyPath : undefined, - multiEntry: index.multiEntry, - unique: index.unique, - }; - }); - - return { - name: storeName, - records: records, - indexes, - autoIncrement: objectStore.autoIncrement, - keyPath: typeof objectStore.keyPath === 'string' ? objectStore.keyPath : undefined, - keyPathArray: Array.isArray(objectStore.keyPath) ? objectStore.keyPath : undefined, - }; - })); - - return { - name: dbInfo.name, - version: dbInfo.version, - stores, - }; - })).catch(e => { - throw new Error('Unable to serialize IndexedDB: ' + e.message); - }); - - return { - localStorage: Object.keys(localStorage).map(name => ({ name, value: localStorage.getItem(name) })), - indexedDB: idbResult, - }; - } - // First try collecting storage stage from existing pages. for (const page of this.pages()) { const origin = page.mainFrame().origin(); if (!origin || !originsToSave.has(origin)) continue; try { - const storage = await page.mainFrame().nonStallingEvaluateInExistingContext(`(${_collectStorageScript.toString()})()`, 'utility'); + const storage: Awaited> = await page.mainFrame().nonStallingEvaluateInExistingContext(`(${storageScript.collect})()`, 'utility'); if (storage.localStorage.length || storage.indexedDB?.length) - result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB } as channels.OriginStorage); + result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB }); originsToSave.delete(origin); } catch { // When failed on the live page, we'll retry on the blank page below. @@ -600,9 +540,9 @@ export abstract class BrowserContext extends SdkObject { for (const origin of originsToSave) { const frame = page.mainFrame(); await frame.goto(internalMetadata, origin); - const storage = await frame.evaluateExpression(`(${_collectStorageScript.toString()})()`, { world: 'utility' }); + const storage: Awaited> = await frame.evaluateExpression(`(${storageScript.collect})()`, { world: 'utility' }); if (storage.localStorage.length || storage.indexedDB.length) - result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB } as channels.OriginStorage); + result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB }); } await page.close(internalMetadata); } @@ -665,47 +605,8 @@ export abstract class BrowserContext extends SdkObject { for (const originState of state.origins) { const frame = page.mainFrame(); await frame.goto(metadata, originState.origin); - - async function _restoreStorageState(originState: channels.OriginStorage) { - for (const { name, value } of (originState.localStorage || [])) - localStorage.setItem(name, value); - - function idbRequestToPromise(request: T) { - return new Promise((resolve, reject) => { - request.addEventListener('success', () => resolve(request.result)); - request.addEventListener('error', () => reject(request.error)); - }); - } - - await Promise.all((originState.indexedDB ?? []).map(async dbInfo => { - const openRequest = indexedDB.open(dbInfo.name, dbInfo.version); - openRequest.addEventListener('upgradeneeded', () => { - const db = openRequest.result; - for (const store of dbInfo.stores) { - const objectStore = db.createObjectStore(store.name, { autoIncrement: store.autoIncrement, keyPath: store.keyPathArray ?? store.keyPath }); - for (const index of store.indexes) - objectStore.createIndex(index.name, index.keyPathArray ?? index.keyPath!, { unique: index.unique, multiEntry: index.multiEntry }); - } - }); - - // after `upgradeneeded` finishes, `success` event is fired. - const db = await idbRequestToPromise(openRequest); - const transaction = db.transaction(db.objectStoreNames, 'readwrite'); - await Promise.all(dbInfo.stores.map(async store => { - const objectStore = transaction.objectStore(store.name); - await Promise.all(store.records.map(async record => { - await idbRequestToPromise( - objectStore.add( - record.value as any, // protocol says string, but this got deserialized above - objectStore.keyPath === null ? record.key : undefined - ) - ); - })); - })); - })); - } - - await frame.evaluateExpression(_restoreStorageState.toString(), { isFunction: true, world: 'utility' }, originState); + const args: Parameters = [originState]; + await frame.evaluateExpression(storageScript.restore.toString(), { isFunction: true, world: 'utility' }, args); } await page.close(internalMetadata); } diff --git a/packages/playwright-core/src/server/injected/storageScript.ts b/packages/playwright-core/src/server/injected/storageScript.ts new file mode 100644 index 0000000000..f19c57bec9 --- /dev/null +++ b/packages/playwright-core/src/server/injected/storageScript.ts @@ -0,0 +1,121 @@ +/** + * 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 type * as channels from '@protocol/channels'; + +export async function collect(): Promise> { + const idbResult = await Promise.all((await indexedDB.databases()).map(async dbInfo => { + if (!dbInfo.name) + throw new Error('Database name is empty'); + if (!dbInfo.version) + throw new Error('Database version is unset'); + + function idbRequestToPromise(request: T) { + return new Promise((resolve, reject) => { + request.addEventListener('success', () => resolve(request.result)); + request.addEventListener('error', () => reject(request.error)); + }); + } + + const db = await idbRequestToPromise(indexedDB.open(dbInfo.name)); + const transaction = db.transaction(db.objectStoreNames, 'readonly'); + const stores = await Promise.all([...db.objectStoreNames].map(async storeName => { + const objectStore = transaction.objectStore(storeName); + + const keys = await idbRequestToPromise(objectStore.getAllKeys()); + const records = await Promise.all(keys.map(async key => { + return { + key: objectStore.keyPath === null ? key : undefined, + value: await idbRequestToPromise(objectStore.get(key)) + }; + })); + + const indexes = [...objectStore.indexNames].map(indexName => { + const index = objectStore.index(indexName); + return { + name: index.name, + keyPath: typeof index.keyPath === 'string' ? index.keyPath : undefined, + keyPathArray: Array.isArray(index.keyPath) ? index.keyPath : undefined, + multiEntry: index.multiEntry, + unique: index.unique, + }; + }); + + return { + name: storeName, + records: records, + indexes, + autoIncrement: objectStore.autoIncrement, + keyPath: typeof objectStore.keyPath === 'string' ? objectStore.keyPath : undefined, + keyPathArray: Array.isArray(objectStore.keyPath) ? objectStore.keyPath : undefined, + }; + })); + + return { + name: dbInfo.name, + version: dbInfo.version, + stores, + }; + })).catch(e => { + throw new Error('Unable to serialize IndexedDB: ' + e.message); + }); + + return { + localStorage: Object.keys(localStorage).map(name => ({ name, value: localStorage.getItem(name)! })), + indexedDB: idbResult, + }; +} + +export async function restore(originState: channels.SetOriginStorage) { + for (const { name, value } of (originState.localStorage || [])) + localStorage.setItem(name, value); + + await Promise.all((originState.indexedDB ?? []).map(async dbInfo => { + const openRequest = indexedDB.open(dbInfo.name, dbInfo.version); + openRequest.addEventListener('upgradeneeded', () => { + const db = openRequest.result; + for (const store of dbInfo.stores) { + const objectStore = db.createObjectStore(store.name, { autoIncrement: store.autoIncrement, keyPath: store.keyPathArray ?? store.keyPath }); + for (const index of store.indexes) + objectStore.createIndex(index.name, index.keyPathArray ?? index.keyPath!, { unique: index.unique, multiEntry: index.multiEntry }); + } + }); + + function idbRequestToPromise(request: T) { + return new Promise((resolve, reject) => { + request.addEventListener('success', () => resolve(request.result)); + request.addEventListener('error', () => reject(request.error)); + }); + } + + // after `upgradeneeded` finishes, `success` event is fired. + const db = await idbRequestToPromise(openRequest); + const transaction = db.transaction(db.objectStoreNames, 'readwrite'); + await Promise.all(dbInfo.stores.map(async store => { + const objectStore = transaction.objectStore(store.name); + await Promise.all(store.records.map(async record => { + await idbRequestToPromise( + objectStore.add( + record.value as any, // protocol says string, but this got deserialized above + objectStore.keyPath === null ? record.key : undefined + ) + ); + })); + })); + })).catch(e => { + throw new Error('Unable to restore IndexedDB: ' + e.message); + }); +}