viewpoint-core 0.3.6

//! Storage state restoration functions.
//!
//! This module contains functions for restoring browser storage state
//! including localStorage and IndexedDB.

use std::sync::Arc;

use tracing::debug;
use viewpoint_cdp::CdpConnection;
use viewpoint_js::js;

use super::types::{IndexedDbDatabase, LocalStorageEntry};
use crate::error::ContextError;

/// Restore localStorage entries to a page.
///
/// # Errors
///
/// Returns an error if the storage state cannot be restored.
pub async fn restore_local_storage(
    connection: &Arc<CdpConnection>,
    session_id: &str,
    entries: &[LocalStorageEntry],
) -> Result<(), ContextError> {
    if entries.is_empty() {
        return Ok(());
    }

    let entries_json = serde_json::to_string(entries)
        .map_err(|e| ContextError::Internal(format!("Failed to serialize localStorage: {e}")))?;

    let js_code = js! {
        (function() {
            const entries = @{entries_json};
            for (const entry of entries) {
                localStorage.setItem(entry.name, entry.value);
            }
        })()
    };

    connection
        .send_command::<_, serde_json::Value>(
            "Runtime.evaluate",
            Some(viewpoint_cdp::protocol::runtime::EvaluateParams {
                expression: js_code,
                object_group: None,
                include_command_line_api: None,
                silent: Some(true),
                context_id: None,
                return_by_value: Some(true),
                await_promise: Some(false),
            }),
            Some(session_id),
        )
        .await?;

    debug!("Restored {} localStorage entries", entries.len());
    Ok(())
}

/// Restore IndexedDB databases to a page.
///
/// # Errors
///
/// Returns an error if the IndexedDB state cannot be restored.
pub async fn restore_indexed_db(
    connection: &Arc<CdpConnection>,
    session_id: &str,
    databases: &[IndexedDbDatabase],
) -> Result<(), ContextError> {
    if databases.is_empty() {
        return Ok(());
    }

    let db_json = serde_json::to_string(databases)
        .map_err(|e| ContextError::Internal(format!("Failed to serialize IndexedDB: {e}")))?;

    let js_code = js! {
        (async function() {
            const databases = @{db_json};

            for (const dbData of databases) {
                // Delete existing database to ensure clean state
                await new Promise((resolve, reject) => {
                    const request = indexedDB.deleteDatabase(dbData.name);
                    request.onerror = () => reject(request.error);
                    request.onsuccess = () => resolve();
                    request.onblocked = () => resolve(); // Proceed even if blocked
                });

                // Create database with schema
                const db = await new Promise((resolve, reject) => {
                    const request = indexedDB.open(dbData.name, dbData.version);
                    request.onerror = () => reject(request.error);
                    request.onupgradeneeded = (event) => {
                        const db = event.target.result;
                        for (const storeData of dbData.stores) {
                            const options = {};
                            if (storeData.keyPath) {
                                options.keyPath = storeData.keyPath.includes(",")
                                    ? storeData.keyPath.split(",")
                                    : storeData.keyPath;
                            }
                            if (storeData.autoIncrement) {
                                options.autoIncrement = true;
                            }

                            const store = db.createObjectStore(storeData.name, options);

                            // Create indexes
                            for (const indexData of (storeData.indexes || [])) {
                                store.createIndex(indexData.name,
                                    indexData.keyPath.includes(",")
                                        ? indexData.keyPath.split(",")
                                        : indexData.keyPath,
                                    { unique: indexData.unique, multiEntry: indexData.multiEntry }
                                );
                            }
                        }
                    };
                    request.onsuccess = () => resolve(request.result);
                });

                // Restore data
                for (const storeData of dbData.stores) {
                    if (storeData.entries.length > 0) {
                        const tx = db.transaction(storeData.name, "readwrite");
                        const store = tx.objectStore(storeData.name);

                        for (const entry of storeData.entries) {
                            if (storeData.keyPath) {
                                store.put(entry.value);
                            } else {
                                store.put(entry.value, entry.key);
                            }
                        }

                        await new Promise((resolve, reject) => {
                            tx.oncomplete = () => resolve();
                            tx.onerror = () => reject(tx.error);
                        });
                    }
                }

                db.close();
            }
        })()
    };

    connection
        .send_command::<_, serde_json::Value>(
            "Runtime.evaluate",
            Some(viewpoint_cdp::protocol::runtime::EvaluateParams {
                expression: js_code,
                object_group: None,
                include_command_line_api: None,
                silent: Some(true),
                context_id: None,
                return_by_value: Some(true),
                await_promise: Some(true),
            }),
            Some(session_id),
        )
        .await?;

    debug!("Restored {} IndexedDB databases", databases.len());
    Ok(())
}