browser_commander/browser/webdriver/
storage.rs

1use super::ManagedWebDriver;
2use crate::browser::storage_state::StorageState;
3use anyhow::{anyhow, Result};
4use serde_json::json;
5
6impl ManagedWebDriver {
7    /// Restore portable state with native WebDriver and return to the current
8    /// URL. Local storage remains origin scoped; HttpOnly cookies use WebDriver.
9    pub async fn restore_state(&self, state: StorageState) -> Result<()> {
10        let previous = self.client.current_url().await?;
11        let restore=async {
12            for value in state.cookies {
13                let name=value["name"].as_str().ok_or_else(||anyhow!("cookie.name is required"))?;
14                let content=value["value"].as_str().ok_or_else(||anyhow!("cookie.value is required"))?;
15                let domain=value["domain"].as_str().ok_or_else(||anyhow!("cookie.domain is required"))?;
16                if let Some(bidi) = &self.bidi {
17                    let mut cookie=json!({"name":name,"value":{"type":"string","value":content},"domain":domain,
18                        "path":value["path"].as_str().unwrap_or("/"),"httpOnly":value["httpOnly"].as_bool().unwrap_or(false),
19                        "secure":value["secure"].as_bool().unwrap_or(false),
20                        "sameSite":value["sameSite"].as_str().unwrap_or("Lax").to_ascii_lowercase()});
21                    if let Some(expires)=value["expires"].as_f64().filter(|value|*value>0.0){cookie["expiry"]=json!(expires);}
22                    bidi.send("storage.setCookie",json!({"cookie":cookie})).await?;
23                    continue;
24                }
25                let scheme=if value["secure"].as_bool()==Some(true){"https"}else{"http"};
26                let origin=state.origins.iter().find(|origin|url::Url::parse(&origin.origin).is_ok_and(|url|url.host_str()==Some(domain.trim_start_matches('.'))))
27                    .map(|origin|origin.origin.clone()).unwrap_or_else(||format!("{scheme}://{}/",domain.trim_start_matches('.')));
28                self.client.goto(&origin).await?;
29                let mut cookie=cookie::Cookie::new(name.to_owned(),content.to_owned());
30                cookie.set_domain(domain.to_owned());
31                cookie.set_path(value["path"].as_str().unwrap_or("/").to_owned());
32                cookie.set_http_only(value["httpOnly"].as_bool().unwrap_or(false));
33                cookie.set_secure(value["secure"].as_bool().unwrap_or(false));
34                cookie.set_same_site(match value["sameSite"].as_str().unwrap_or("Lax") {"Strict"=>cookie::SameSite::Strict,"None"=>cookie::SameSite::None,_=>cookie::SameSite::Lax});
35                if let Some(expires)=value["expires"].as_f64().filter(|value| *value>0.0) {
36                    cookie.set_expires(time::OffsetDateTime::from_unix_timestamp(expires as i64)?);
37                }
38                self.client.add_cookie(cookie).await?;
39            }
40            for origin in state.origins {
41                self.client.goto(&origin.origin).await?;
42                self.client.execute("for(const item of arguments[0]) localStorage.setItem(item.name,item.value)",vec![serde_json::to_value(origin.local_storage)?]).await?;
43            }
44            Ok::<(),anyhow::Error>(())
45        }.await;
46        let returned = self.client.goto(previous.as_str()).await;
47        restore?;
48        returned?;
49        Ok(())
50    }
51
52    /// Export browser cookies through BiDi where available, otherwise the
53    /// current domain via W3C WebDriver, plus the active origin's localStorage.
54    pub async fn save_state(&self) -> Result<StorageState> {
55        let cookies = if let Some(bidi) = &self.bidi {
56            let response = bidi.send("storage.getCookies", json!({})).await?;
57            response["cookies"].as_array().ok_or_else(||anyhow!("Invalid BiDi cookies"))?.iter().map(|cookie|json!({
58                "name":cookie["name"],"value":cookie["value"]["value"],"domain":cookie["domain"],"path":cookie["path"],
59                "secure":cookie["secure"],"httpOnly":cookie["httpOnly"],"expires":cookie.get("expiry").cloned().unwrap_or(json!(-1)),
60                "sameSite":match cookie["sameSite"].as_str().unwrap_or("lax") {"strict"=>"Strict","none"=>"None",_=>"Lax"}
61            })).collect()
62        } else {
63            self.client.get_all_cookies().await?.into_iter().map(|cookie|json!({
64                "name":cookie.name(),"value":cookie.value(),"domain":cookie.domain().unwrap_or(""),"path":cookie.path().unwrap_or("/"),
65                "secure":cookie.secure().unwrap_or(false),"httpOnly":cookie.http_only().unwrap_or(false),
66                "expires":cookie.expires_datetime().map(|value|value.unix_timestamp()).unwrap_or(-1),
67                "sameSite":cookie.same_site().map(|value|format!("{value:?}")).unwrap_or_else(||"Lax".into())
68            })).collect()
69        };
70        let origin = self
71            .client
72            .execute("return location.origin", vec![])
73            .await?;
74        let origins = if origin
75            .as_str()
76            .is_some_and(|value| value.starts_with("http://") || value.starts_with("https://"))
77        {
78            let entries = self
79                .client
80                .execute(
81                    "return Object.entries(localStorage).map(([name,value])=>({name,value}))",
82                    vec![],
83                )
84                .await?;
85            serde_json::from_value(json!([{"origin":origin,"localStorage":entries}]))?
86        } else {
87            Vec::new()
88        };
89        Ok(StorageState { cookies, origins })
90    }
91}
browser_commander/browser/webdriver/storage.rs

browser_commander/browser/webdriver/
storage.rs
