fix: save_state captures cross-domain cookies and localStorage (#1064)
The Rust rewrite of save_state only captured cookies and localStorage for the current page's origin, silently dropping cross-domain data (e.g. SSO/CAS auth cookies). This was a regression from the JS version. Cookies: replace Network.getCookies with Network.getAllCookies to return cookies from all domains the browser has visited. localStorage: track visited origins in BrowserManager during navigation, then collect their localStorage via a temporary CDP target with Fetch interception (serves blank HTML to avoid real network requests). Co-authored-by: hyunjinee <leehj0110@kakao.com>
This commit is contained in:
jin.2
@@ -1960,6 +1960,7 @@ async fn handle_close(state: &mut DaemonState) -> Result<Value, String> {
|
||||
None,
|
||||
Some(session_name.as_str()),
|
||||
&state.session_id,
|
||||
mgr.visited_origins(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -2996,6 +2997,7 @@ async fn handle_state_save(cmd: &Value, state: &DaemonState) -> Result<Value, St
|
||||
path,
|
||||
state.session_name.as_deref(),
|
||||
&state.session_id,
|
||||
mgr.visited_origins(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
@@ -199,6 +199,8 @@ pub struct BrowserManager {
|
||||
default_timeout_ms: u64,
|
||||
/// Stored download path from launch options, re-applied to new contexts (e.g., recording)
|
||||
pub download_path: Option<String>,
|
||||
/// Origins visited during this session, used by save_state to collect cross-origin localStorage.
|
||||
visited_origins: HashSet<String>,
|
||||
}
|
||||
|
||||
const LIGHTPANDA_CDP_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
@@ -268,6 +270,7 @@ impl BrowserManager {
|
||||
active_page_index: 0,
|
||||
default_timeout_ms: 25_000,
|
||||
download_path: download_path.clone(),
|
||||
visited_origins: HashSet::new(),
|
||||
};
|
||||
manager.discover_and_attach_targets().await?;
|
||||
manager
|
||||
@@ -333,6 +336,7 @@ impl BrowserManager {
|
||||
active_page_index: 0,
|
||||
default_timeout_ms: 10_000,
|
||||
download_path: None, // CDP connections don't have a launch-time download path
|
||||
visited_origins: HashSet::new(),
|
||||
};
|
||||
|
||||
manager.discover_and_attach_targets().await?;
|
||||
@@ -494,6 +498,14 @@ impl BrowserManager {
|
||||
let page_url = self.get_url().await.unwrap_or_else(|_| url.to_string());
|
||||
let title = self.get_title().await.unwrap_or_default();
|
||||
|
||||
// Track visited origin for cross-origin localStorage collection in save_state
|
||||
if let Ok(parsed) = url::Url::parse(&page_url) {
|
||||
let origin = parsed.origin().ascii_serialization();
|
||||
if origin != "null" {
|
||||
self.visited_origins.insert(origin);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(page) = self.pages.get_mut(self.active_page_index) {
|
||||
page.url = page_url.clone();
|
||||
page.title = title.clone();
|
||||
@@ -1129,6 +1141,10 @@ impl BrowserManager {
|
||||
self.pages.clone()
|
||||
}
|
||||
|
||||
pub fn visited_origins(&self) -> &HashSet<String> {
|
||||
&self.visited_origins
|
||||
}
|
||||
|
||||
pub async fn set_download_behavior(&self, download_path: &str) -> Result<(), String> {
|
||||
let session_id = self.active_session_id()?;
|
||||
self.client
|
||||
@@ -1276,6 +1292,7 @@ async fn initialize_lightpanda_manager(
|
||||
active_page_index: 0,
|
||||
default_timeout_ms: 25_000,
|
||||
download_path: None,
|
||||
visited_origins: HashSet::new(),
|
||||
};
|
||||
|
||||
match discover_and_attach_lightpanda_targets(&mut manager, deadline).await {
|
||||
|
||||
@@ -24,6 +24,19 @@ pub struct Cookie {
|
||||
pub same_site: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn get_all_cookies(client: &CdpClient, session_id: &str) -> Result<Vec<Cookie>, String> {
|
||||
let result = client
|
||||
.send_command_no_params("Network.getAllCookies", Some(session_id))
|
||||
.await?;
|
||||
|
||||
let cookies: Vec<Cookie> = result
|
||||
.get("cookies")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(cookies)
|
||||
}
|
||||
|
||||
pub async fn get_cookies(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
|
||||
@@ -1599,6 +1599,142 @@ async fn e2e_state_management() {
|
||||
assert_success(&resp);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Cross-domain state save (issue #1060)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn e2e_save_state_cross_domain() {
|
||||
let mut state = DaemonState::new();
|
||||
|
||||
// Launch
|
||||
let resp = execute_command(
|
||||
&json!({ "id": "1", "action": "launch", "headless": true }),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
// Navigate to domain A and set cookie + localStorage
|
||||
let resp = execute_command(
|
||||
&json!({ "id": "2", "action": "navigate", "url": "https://httpbin.org/html" }),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
let resp = execute_command(
|
||||
&json!({
|
||||
"id": "3", "action": "cookies_set",
|
||||
"name": "domainA_cookie", "value": "from_httpbin"
|
||||
}),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
let resp = execute_command(
|
||||
&json!({
|
||||
"id": "4", "action": "storage_set",
|
||||
"type": "local", "key": "domainA_key", "value": "domainA_val"
|
||||
}),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
// Navigate to domain B and set cookie + localStorage
|
||||
let resp = execute_command(
|
||||
&json!({ "id": "5", "action": "navigate", "url": "https://example.com" }),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
let resp = execute_command(
|
||||
&json!({
|
||||
"id": "6", "action": "cookies_set",
|
||||
"name": "domainB_cookie", "value": "from_example"
|
||||
}),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
let resp = execute_command(
|
||||
&json!({
|
||||
"id": "7", "action": "storage_set",
|
||||
"type": "local", "key": "domainB_key", "value": "domainB_val"
|
||||
}),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
// Save state (currently on example.com)
|
||||
let tmp_state = std::env::temp_dir()
|
||||
.join("agent-browser-e2e-cross-domain-state.json")
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
let resp = execute_command(
|
||||
&json!({ "id": "8", "action": "state_save", "path": &tmp_state }),
|
||||
&mut state,
|
||||
)
|
||||
.await;
|
||||
assert_success(&resp);
|
||||
|
||||
// Read and verify saved state
|
||||
let saved = std::fs::read_to_string(&tmp_state).expect("State file should exist");
|
||||
let state_data: serde_json::Value = serde_json::from_str(&saved).unwrap();
|
||||
|
||||
// Verify BOTH domain cookies are present
|
||||
let cookies = state_data["cookies"].as_array().unwrap();
|
||||
let has_domain_a = cookies.iter().any(|c| c["name"] == "domainA_cookie");
|
||||
let has_domain_b = cookies.iter().any(|c| c["name"] == "domainB_cookie");
|
||||
assert!(
|
||||
has_domain_a,
|
||||
"Should include cross-domain cookie from httpbin.org: {:?}",
|
||||
cookies
|
||||
);
|
||||
assert!(
|
||||
has_domain_b,
|
||||
"Should include cookie from example.com: {:?}",
|
||||
cookies
|
||||
);
|
||||
|
||||
// Verify BOTH origins' localStorage are present
|
||||
let origins = state_data["origins"].as_array().unwrap();
|
||||
let has_origin_a = origins.iter().any(|o| {
|
||||
o["origin"].as_str().is_some_and(|s| s.contains("httpbin"))
|
||||
&& o["localStorage"]
|
||||
.as_array()
|
||||
.is_some_and(|ls| ls.iter().any(|e| e["name"] == "domainA_key"))
|
||||
});
|
||||
let has_origin_b = origins.iter().any(|o| {
|
||||
o["origin"].as_str().is_some_and(|s| s.contains("example"))
|
||||
&& o["localStorage"]
|
||||
.as_array()
|
||||
.is_some_and(|ls| ls.iter().any(|e| e["name"] == "domainB_key"))
|
||||
});
|
||||
assert!(
|
||||
has_origin_a,
|
||||
"Should include localStorage from httpbin.org origin: {:?}",
|
||||
origins
|
||||
);
|
||||
assert!(
|
||||
has_origin_b,
|
||||
"Should include localStorage from example.com origin: {:?}",
|
||||
origins
|
||||
);
|
||||
|
||||
// Clean up
|
||||
let _ = std::fs::remove_file(&tmp_state);
|
||||
|
||||
let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
|
||||
assert_success(&resp);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Domain filter
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
+246
-42
@@ -1,12 +1,17 @@
|
||||
use aes_gcm::{aead::Aead, aead::KeyInit, Aes256Gcm};
|
||||
use base64::Engine;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashSet;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
use super::cdp::types::EvaluateParams;
|
||||
use super::cdp::types::{
|
||||
AttachToTargetParams, AttachToTargetResult, CloseTargetParams, CreateTargetParams,
|
||||
CreateTargetResult, EvaluateParams,
|
||||
};
|
||||
use super::cookies::{self, Cookie};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
@@ -32,16 +37,223 @@ pub struct StorageEntry {
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
fn collect_frame_origins(tree: &Value, origins: &mut HashSet<String>) {
|
||||
if let Some(frame) = tree.get("frame") {
|
||||
if let Some(url_str) = frame.get("url").and_then(|v| v.as_str()) {
|
||||
if let Ok(parsed) = url::Url::parse(url_str) {
|
||||
let origin = parsed.origin().ascii_serialization();
|
||||
if origin != "null" && !origin.is_empty() {
|
||||
origins.insert(origin);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) {
|
||||
for child in children {
|
||||
collect_frame_origins(child, origins);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the JS-evaluated origin storage data into an OriginStorage struct.
|
||||
fn parse_origin_storage(data: &Value) -> Option<OriginStorage> {
|
||||
if !data.is_object() {
|
||||
return None;
|
||||
}
|
||||
let origin = data
|
||||
.get("origin")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
if origin.is_empty() || origin == "null" {
|
||||
return None;
|
||||
}
|
||||
let local_storage: Vec<StorageEntry> = data
|
||||
.get("localStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let session_storage: Vec<StorageEntry> = data
|
||||
.get("sessionStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
Some(OriginStorage {
|
||||
origin,
|
||||
local_storage,
|
||||
session_storage,
|
||||
})
|
||||
}
|
||||
|
||||
/// Evaluate the storage-collection JS snippet and parse the result.
|
||||
async fn eval_origin_storage(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
origin_js: &str,
|
||||
) -> Option<OriginStorage> {
|
||||
let result = client
|
||||
.send_command_typed::<_, super::cdp::types::EvaluateResult>(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: origin_js.to_string(),
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
let data = result.result.value.unwrap_or(Value::Null);
|
||||
parse_origin_storage(&data)
|
||||
}
|
||||
|
||||
/// Create a temporary CDP target, navigate it to each origin to collect localStorage,
|
||||
/// then close it. Uses Fetch interception to serve blank HTML instead of making real
|
||||
/// network requests.
|
||||
async fn collect_storage_via_temp_target(
|
||||
client: &CdpClient,
|
||||
origins: &[String],
|
||||
origin_js: &str,
|
||||
) -> Result<Vec<OriginStorage>, String> {
|
||||
let create_result: CreateTargetResult = client
|
||||
.send_command_typed(
|
||||
"Target.createTarget",
|
||||
&CreateTargetParams {
|
||||
url: "about:blank".to_string(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let target_id = create_result.target_id;
|
||||
|
||||
// Ensure the target is closed even if attach or later steps fail
|
||||
let result = collect_storage_in_target(client, &target_id, origins, origin_js).await;
|
||||
|
||||
let _ = client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Target.closeTarget",
|
||||
&CloseTargetParams { target_id },
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
async fn collect_storage_in_target(
|
||||
client: &CdpClient,
|
||||
target_id: &str,
|
||||
origins: &[String],
|
||||
origin_js: &str,
|
||||
) -> Result<Vec<OriginStorage>, String> {
|
||||
let attach_result: AttachToTargetResult = client
|
||||
.send_command_typed(
|
||||
"Target.attachToTarget",
|
||||
&AttachToTargetParams {
|
||||
target_id: target_id.to_string(),
|
||||
flatten: true,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let temp_session = &attach_result.session_id;
|
||||
|
||||
client
|
||||
.send_command_no_params("Page.enable", Some(temp_session))
|
||||
.await?;
|
||||
client
|
||||
.send_command_no_params("Runtime.enable", Some(temp_session))
|
||||
.await?;
|
||||
|
||||
// Blank HTML response body, pre-encoded to avoid repeated base64 work per request
|
||||
let blank_html_b64 = base64::engine::general_purpose::STANDARD.encode("<html></html>");
|
||||
|
||||
let _ = client
|
||||
.send_command(
|
||||
"Fetch.enable",
|
||||
Some(json!({ "patterns": [{ "urlPattern": "*" }] })),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut event_rx = client.subscribe();
|
||||
let mut results = Vec::new();
|
||||
|
||||
for target_origin in origins {
|
||||
let nav_url = format!("{}/", target_origin.trim_end_matches('/'));
|
||||
if client
|
||||
.send_command(
|
||||
"Page.navigate",
|
||||
Some(json!({ "url": nav_url })),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
// Fulfill intercepted requests with blank HTML until the page loads
|
||||
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5);
|
||||
let mut page_loaded = false;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
match tokio::time::timeout(tokio::time::Duration::from_secs(2), event_rx.recv()).await {
|
||||
Ok(Ok(evt)) if evt.session_id.as_deref() == Some(temp_session) => {
|
||||
if evt.method == "Fetch.requestPaused" {
|
||||
if let Some(request_id) =
|
||||
evt.params.get("requestId").and_then(|v| v.as_str())
|
||||
{
|
||||
let _ = client
|
||||
.send_command(
|
||||
"Fetch.fulfillRequest",
|
||||
Some(json!({
|
||||
"requestId": request_id,
|
||||
"responseCode": 200,
|
||||
"responseHeaders": [
|
||||
{ "name": "Content-Type", "value": "text/html" }
|
||||
],
|
||||
"body": &blank_html_b64
|
||||
})),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
} else if evt.method == "Page.loadEventFired" {
|
||||
page_loaded = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Ok(_)) => continue, // event for a different session
|
||||
Ok(Err(_)) => continue, // lagged or closed — retry within deadline
|
||||
Err(_) => break, // outer timeout elapsed
|
||||
}
|
||||
}
|
||||
|
||||
if !page_loaded {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(storage) = eval_origin_storage(client, temp_session, origin_js).await {
|
||||
if !storage.local_storage.is_empty() || !storage.session_storage.is_empty() {
|
||||
results.push(storage);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
pub async fn save_state(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
path: Option<&str>,
|
||||
session_name: Option<&str>,
|
||||
session_id_str: &str,
|
||||
visited_origins: &HashSet<String>,
|
||||
) -> Result<String, String> {
|
||||
let cookies = cookies::get_cookies(client, session_id, None).await?;
|
||||
let cookies = cookies::get_all_cookies(client, session_id).await?;
|
||||
|
||||
// Get current origin's storage
|
||||
let origin_js = r#"(() => {
|
||||
const result = { origin: location.origin, localStorage: [], sessionStorage: [] };
|
||||
try {
|
||||
@@ -59,46 +271,38 @@ pub async fn save_state(
|
||||
return result;
|
||||
})()"#;
|
||||
|
||||
let origin_result: super::cdp::types::EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: origin_js.to_string(),
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let origin_data = origin_result.result.value.unwrap_or(Value::Null);
|
||||
let origins = if origin_data.is_object() {
|
||||
let origin = origin_data
|
||||
.get("origin")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let local_storage: Vec<StorageEntry> = origin_data
|
||||
.get("localStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let session_storage: Vec<StorageEntry> = origin_data
|
||||
.get("sessionStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if !origin.is_empty() && origin != "null" {
|
||||
vec![OriginStorage {
|
||||
origin,
|
||||
local_storage,
|
||||
session_storage,
|
||||
}]
|
||||
} else {
|
||||
vec![]
|
||||
// Merge visited origins with current frame tree origins
|
||||
let mut all_origins = visited_origins.clone();
|
||||
if let Ok(tree_result) = client
|
||||
.send_command_no_params("Page.getFrameTree", Some(session_id))
|
||||
.await
|
||||
{
|
||||
if let Some(tree) = tree_result.get("frameTree") {
|
||||
collect_frame_origins(tree, &mut all_origins);
|
||||
}
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
}
|
||||
|
||||
// 1. Collect localStorage from the current page
|
||||
let mut origins = Vec::new();
|
||||
let mut current_origin = String::new();
|
||||
|
||||
if let Some(storage) = eval_origin_storage(client, session_id, origin_js).await {
|
||||
current_origin = storage.origin.clone();
|
||||
if !storage.local_storage.is_empty() || !storage.session_storage.is_empty() {
|
||||
origins.push(storage);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Collect localStorage from remaining origins via a disposable temp target
|
||||
all_origins.remove(¤t_origin);
|
||||
if !all_origins.is_empty() {
|
||||
let remaining: Vec<String> = all_origins.into_iter().collect();
|
||||
if let Ok(temp_origins) =
|
||||
collect_storage_via_temp_target(client, &remaining, origin_js).await
|
||||
{
|
||||
origins.extend(temp_origins);
|
||||
}
|
||||
}
|
||||
|
||||
let state = StorageState { cookies, origins };
|
||||
let json_str = serde_json::to_string_pretty(&state)
|
||||
|
||||
Reference in New Issue
Block a user