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>
2026-03-28 13:41:53 -07:00
parent 747a3772e1
commit dc26ff7667
5 changed files with 414 additions and 42 deletions
@@ -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(&current_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)