Skip to main content

browser_commander/browser/
storage_state.rs

1//! Playwright-compatible cookie and origin-scoped localStorage state.
2
3use std::path::{Path, PathBuf};
4
5use anyhow::{anyhow, Context};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9use crate::core::engine::EngineAdapter;
10
11/// A name and value from an origin's localStorage.
12#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
13pub struct StorageEntry {
14    /// Storage key.
15    pub name: String,
16    /// Storage value.
17    pub value: String,
18}
19
20/// localStorage captured from one origin.
21#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
22pub struct StorageOrigin {
23    /// Scheme, hostname and port.
24    pub origin: String,
25    /// localStorage values for that origin.
26    #[serde(rename = "localStorage")]
27    pub local_storage: Vec<StorageEntry>,
28}
29
30/// Portable cookies and localStorage in Playwright's storage-state format.
31#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
32pub struct StorageState {
33    /// Cookies use the Playwright/CDP cookie shape.
34    #[serde(default)]
35    pub cookies: Vec<Value>,
36    /// Origin-scoped localStorage.
37    #[serde(default)]
38    pub origins: Vec<StorageOrigin>,
39}
40
41/// A state object or a JSON file to read before browser startup.
42#[derive(Clone, Debug)]
43pub enum StorageStateInput {
44    /// Read the state from a JSON file.
45    Path(PathBuf),
46    /// Use this state directly.
47    Value(StorageState),
48}
49
50impl From<StorageState> for StorageStateInput {
51    fn from(value: StorageState) -> Self {
52        Self::Value(value)
53    }
54}
55
56impl From<PathBuf> for StorageStateInput {
57    fn from(value: PathBuf) -> Self {
58        Self::Path(value)
59    }
60}
61
62impl From<&Path> for StorageStateInput {
63    fn from(value: &Path) -> Self {
64        Self::Path(value.to_path_buf())
65    }
66}
67
68impl StorageStateInput {
69    /// Read and validate the state.
70    pub fn load(&self) -> anyhow::Result<StorageState> {
71        let state = match self {
72            Self::Path(path) => {
73                let contents = std::fs::read(path)
74                    .with_context(|| format!("could not read storage state {}", path.display()))?;
75                serde_json::from_slice(&contents)
76                    .with_context(|| format!("invalid storage state {}", path.display()))?
77            }
78            Self::Value(state) => state.clone(),
79        };
80        validate_storage_state(&state)?;
81        Ok(state)
82    }
83}
84
85fn validate_storage_state(state: &StorageState) -> anyhow::Result<()> {
86    if state.cookies.iter().any(|cookie| {
87        cookie.get("name").and_then(Value::as_str).is_none()
88            || cookie.get("value").and_then(Value::as_str).is_none()
89    }) {
90        return Err(anyhow!(
91            "storage state cookies must have string names and values"
92        ));
93    }
94    if state.origins.iter().any(|origin| origin.origin.is_empty()) {
95        return Err(anyhow!("storage state origins must have nonempty URLs"));
96    }
97    Ok(())
98}
99
100/// A page script that restores localStorage only at its matching origin.
101pub(crate) fn restore_script(state: &StorageState) -> anyhow::Result<String> {
102    let origins = serde_json::to_string(&state.origins)?;
103    Ok(format!(
104        "(() => {{ const origins = {origins}; \
105         const entry = origins.find(item => item.origin === globalThis.location.origin); \
106         if (!entry) return; for (const item of entry.localStorage) \
107         globalThis.localStorage.setItem(item.name, item.value); }})()"
108    ))
109}
110
111/// Export the current browser's cookies and current page's localStorage.
112pub async fn save_storage_state(
113    page: &dyn EngineAdapter,
114    file_path: Option<&Path>,
115) -> anyhow::Result<StorageState> {
116    let value = page.export_storage_state().await?;
117    let state: StorageState = serde_json::from_value(value)?;
118    validate_storage_state(&state)?;
119    if let Some(path) = file_path {
120        let mut encoded = serde_json::to_vec_pretty(&state)?;
121        encoded.push(b'\n');
122        std::fs::write(path, encoded)?;
123    }
124    Ok(state)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn loads_portable_state_from_object_and_file() {
133        let state = StorageState {
134            cookies: vec![
135                serde_json::json!({"name":"sid","value":"saved","domain":"example.test","path":"/"}),
136            ],
137            origins: vec![StorageOrigin {
138                origin: "https://example.test".to_string(),
139                local_storage: vec![StorageEntry {
140                    name: "theme".to_string(),
141                    value: "dark".to_string(),
142                }],
143            }],
144        };
145        let directory = std::env::temp_dir().join(format!(
146            "browser-commander-storage-state-{}",
147            std::process::id()
148        ));
149        std::fs::create_dir_all(&directory).unwrap();
150        let path = directory.join("state.json");
151        std::fs::write(&path, serde_json::to_vec(&state).unwrap()).unwrap();
152
153        assert_eq!(
154            StorageStateInput::Value(state.clone()).load().unwrap(),
155            state
156        );
157        assert_eq!(StorageStateInput::Path(path).load().unwrap(), state);
158        std::fs::remove_dir_all(directory).unwrap();
159    }
160}