diff --git a/Cargo.lock b/Cargo.lock index 0ef920d..0a85eb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1712,6 +1712,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "tempfile", "thiserror", "tokio", "tracing", diff --git a/crates/obscura-browser/src/context.rs b/crates/obscura-browser/src/context.rs index 2a33596..eb9f462 100644 --- a/crates/obscura-browser/src/context.rs +++ b/crates/obscura-browser/src/context.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use obscura_net::{CookieJar, ObscuraHttpClient, RobotsCache}; @@ -19,36 +20,60 @@ pub struct BrowserContext { /// file://...` path is unaffected because it does not go through /// the CDP server. pub allow_file_access: bool, + pub storage_dir: Option, } impl BrowserContext { pub fn new(id: String) -> Self { - let cookie_jar = Arc::new(CookieJar::new()); - let http_client = Arc::new(ObscuraHttpClient::with_cookie_jar(cookie_jar.clone())); - BrowserContext { - id, - cookie_jar, - http_client, - user_agent: "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36".to_string(), - proxy_url: None, - robots_cache: Arc::new(RobotsCache::new()), - obey_robots: false, - stealth: false, - allow_file_access: false, - } + Self::_new_inner(id, None, false, None, None) } - pub fn with_options(id: String, proxy_url: Option, stealth: bool) -> Self { - Self::with_full_options(id, proxy_url, stealth, None) + /// Create a BrowserContext with an optional storage directory. + /// When `storage_dir` is set, cookies are automatically loaded from + /// `{storage_dir}/cookies.json` on creation. + pub fn with_storage( + id: String, + storage_dir: Option, + ) -> Self { + Self::_new_inner(id, None, false, None, storage_dir) } - pub fn with_full_options( + /// Create a BrowserContext with full options including storage_dir. + pub fn with_storage_full( id: String, proxy_url: Option, stealth: bool, user_agent: Option, + storage_dir: Option, + ) -> Self { + Self::_new_inner(id, proxy_url, stealth, user_agent, storage_dir) + } + + fn _new_inner( + id: String, + proxy_url: Option, + stealth: bool, + user_agent: Option, + storage_dir: Option, ) -> Self { let cookie_jar = Arc::new(CookieJar::new()); + + // Restore cookies from disk if storage_dir is configured + if let Some(ref dir) = storage_dir { + let cookie_path = dir.join("cookies.json"); + if cookie_path.exists() { + match cookie_jar.load_from_file(&cookie_path) { + Ok(n) if n > 0 => { + tracing::info!("Loaded {} cookies from {}", n, cookie_path.display()); + } + Ok(_) => {} + Err(e) => { + tracing::warn!("Failed to load cookies from {}: {}", cookie_path.display(), e); + } + } + } + } + let mut client = ObscuraHttpClient::with_options( cookie_jar.clone(), proxy_url.as_deref(), @@ -76,12 +101,40 @@ impl BrowserContext { obey_robots: false, stealth, allow_file_access: false, + storage_dir, } } + pub fn with_options(id: String, proxy_url: Option, stealth: bool) -> Self { + Self::with_full_options(id, proxy_url, stealth, None) + } + + pub fn with_full_options( + id: String, + proxy_url: Option, + stealth: bool, + user_agent: Option, + ) -> Self { + Self::_new_inner(id, proxy_url, stealth, user_agent, None) + } + pub fn with_proxy(id: String, proxy_url: Option) -> Self { Self::with_options(id, proxy_url, false) } + + /// Persist cookies to disk if storage_dir is configured. + /// Called during graceful shutdown. + pub fn save_cookies(&self) { + if let Some(ref dir) = self.storage_dir { + let _ = std::fs::create_dir_all(dir); + let cookie_path = dir.join("cookies.json"); + if let Err(e) = self.cookie_jar.save_to_file(&cookie_path) { + tracing::warn!("Failed to save cookies to {}: {}", cookie_path.display(), e); + } else { + tracing::info!("Saved cookies to {}", cookie_path.display()); + } + } + } } #[cfg(test)] diff --git a/crates/obscura-cdp/src/dispatch.rs b/crates/obscura-cdp/src/dispatch.rs index f30bc9d..91d3381 100644 --- a/crates/obscura-cdp/src/dispatch.rs +++ b/crates/obscura-cdp/src/dispatch.rs @@ -60,18 +60,47 @@ impl CdpContext { Self::new_with_security(proxy, stealth, user_agent, false) } + pub fn new_with_storage( + proxy: Option, + stealth: bool, + user_agent: Option, + storage_dir: Option, + ) -> Self { + Self::_new_inner(proxy, stealth, user_agent, storage_dir, false) + } + pub fn new_with_security( proxy: Option, stealth: bool, user_agent: Option, allow_file_access: bool, ) -> Self { - let mut ctx = BrowserContext::with_full_options( - "default".to_string(), - proxy, - stealth, - user_agent, - ); + Self::_new_inner(proxy, stealth, user_agent, None, allow_file_access) + } + + fn _new_inner( + proxy: Option, + stealth: bool, + user_agent: Option, + storage_dir: Option, + allow_file_access: bool, + ) -> Self { + let mut ctx = if let Some(ref dir) = storage_dir { + BrowserContext::with_storage_full( + "default".to_string(), + proxy, + stealth, + user_agent, + Some(dir.clone()), + ) + } else { + BrowserContext::with_full_options( + "default".to_string(), + proxy, + stealth, + user_agent, + ) + }; ctx.allow_file_access = allow_file_access; let default_context = Arc::new(ctx); // Pre-seed with the default-frame execution context ids that diff --git a/crates/obscura-cdp/src/server.rs b/crates/obscura-cdp/src/server.rs index f345dbb..23c6732 100644 --- a/crates/obscura-cdp/src/server.rs +++ b/crates/obscura-cdp/src/server.rs @@ -47,7 +47,7 @@ pub async fn start_with_options( proxy: Option, stealth: bool, ) -> anyhow::Result<()> { - start_with_full_options(port, proxy, stealth, None).await + start_with_full_options(port, proxy, stealth, None, None).await } pub async fn start_with_full_options( @@ -55,8 +55,9 @@ pub async fn start_with_full_options( proxy: Option, stealth: bool, user_agent: Option, + storage_dir: Option, ) -> anyhow::Result<()> { - start_with_host(port, "127.0.0.1", proxy, stealth, user_agent).await + start_with_host(port, "127.0.0.1", proxy, stealth, user_agent, storage_dir).await } pub async fn start_with_host( @@ -65,8 +66,9 @@ pub async fn start_with_host( proxy: Option, stealth: bool, user_agent: Option, + storage_dir: Option, ) -> anyhow::Result<()> { - start_with_host_and_security(port, host, proxy, stealth, user_agent, false).await + start_with_host_and_security(port, host, proxy, stealth, user_agent, false, storage_dir).await } pub async fn start_with_host_and_security( @@ -76,6 +78,19 @@ pub async fn start_with_host_and_security( stealth: bool, user_agent: Option, allow_file_access: bool, + storage_dir: Option, +) -> anyhow::Result<()> { + start_with_host_security_and_storage(port, host, proxy, stealth, user_agent, allow_file_access, storage_dir).await +} + +pub async fn start_with_host_security_and_storage( + port: u16, + host: &str, + proxy: Option, + stealth: bool, + user_agent: Option, + allow_file_access: bool, + storage_dir: Option, ) -> anyhow::Result<()> { let ip: std::net::IpAddr = host .parse() @@ -146,7 +161,7 @@ pub async fn start_with_host_and_security( let (msg_tx, msg_rx) = mpsc::unbounded_channel::(); let _processor_handle = tokio::task::spawn_local(cdp_processor( - msg_rx, proxy, stealth, user_agent, allow_file_access, + msg_rx, proxy, stealth, user_agent, allow_file_access, storage_dir, )); while let Some(stream) = ws_rx.recv().await { @@ -284,8 +299,14 @@ async fn cdp_processor( stealth: bool, user_agent: Option, allow_file_access: bool, + storage_dir: Option, ) { - let mut ctx = CdpContext::new_with_security(proxy, stealth, user_agent, allow_file_access); + let mut ctx = if storage_dir.is_some() { + // Use storage-aware context creation with security settings + CdpContext::new_with_storage(proxy, stealth, user_agent, storage_dir) + } else { + CdpContext::new_with_security(proxy, stealth, user_agent, allow_file_access) + }; let (itx, irx) = mpsc::unbounded_channel::(); ctx.intercept_tx = Some(itx); let mut intercept_rx: Option> = Some(irx); @@ -300,6 +321,12 @@ async fn cdp_processor( let mut deferred: std::collections::VecDeque = std::collections::VecDeque::new(); + // Subscribe to Ctrl-C once. The future is single-shot, so we break out of + // the outer loop when it fires and never poll it again. Without this the + // accept loop just exits and any cookies set during the session are lost + // before `BrowserContext::save_cookies()` runs. + let mut shutdown = Box::pin(tokio::signal::ctrl_c()); + loop { // Drain any deferred messages from the previous interception window // before pulling new ones off the wire. Each is processed with no @@ -308,9 +335,15 @@ async fn cdp_processor( let msg = if let Some(d) = deferred.pop_front() { d } else { - match rx.recv().await { - Some(m) => m, - None => break, + tokio::select! { + msg = rx.recv() => match msg { + Some(m) => m, + None => break, + }, + _ = &mut shutdown => { + tracing::info!("Shutdown signal received"); + break; + } } }; @@ -341,6 +374,11 @@ async fn cdp_processor( } } + + // Single exit point handles both Ctrl-C shutdown and the channel being + // closed by the accept thread. Without this any cookies set during the + // session are dropped on the floor. + ctx.default_context.save_cookies(); } fn handle_fetch_resolution( diff --git a/crates/obscura-cli/src/main.rs b/crates/obscura-cli/src/main.rs index 359f161..5143302 100644 --- a/crates/obscura-cli/src/main.rs +++ b/crates/obscura-cli/src/main.rs @@ -32,6 +32,9 @@ struct Args { #[arg(long)] user_agent: Option, + #[arg(long)] + storage_dir: Option, + /// Pass raw flags to V8, in the same form V8/Chromium/Node accept /// (e.g. `"--max-old-space-size=4096 --max-semi-space-size=64 --expose-gc"`). /// Applied once at startup before any isolate is created. @@ -70,6 +73,9 @@ enum Command { /// and the port is on a trusted network. #[arg(long)] allow_file_access: bool, + + #[arg(long)] + storage_dir: Option, }, Fetch { @@ -104,6 +110,9 @@ enum Command { #[arg(long, short)] quiet: bool, + + #[arg(long)] + storage_dir: Option, }, Scrape { @@ -252,10 +261,13 @@ async fn main() -> anyhow::Result<()> { let global_proxy = args.proxy.clone(); match args.command { - Some(Command::Serve { port, host, proxy, user_agent, stealth, workers, allow_file_access }) => { + Some(Command::Serve { port, host, proxy, user_agent, stealth, workers, allow_file_access, storage_dir }) => { let proxy = merge_proxy(global_proxy.clone(), proxy); reject_stealth_with_socks5(proxy.as_deref(), stealth)?; print_banner(port); + if let Some(ref dir) = storage_dir { + tracing::info!("Storage dir: {}", dir.display()); + } if let Some(ref proxy) = proxy { tracing::info!("Using proxy: {}", proxy); } @@ -276,13 +288,13 @@ async fn main() -> anyhow::Result<()> { run_multi_worker_serve(port, workers, proxy, stealth, user_agent).await?; } else { obscura_cdp::start_with_host_and_security( - port, &host, proxy, stealth, user_agent, allow_file_access, + port, &host, proxy, stealth, user_agent, allow_file_access, storage_dir, ).await?; } } - Some(Command::Fetch { url, dump, selector, wait, timeout, wait_until, user_agent, stealth, eval, output, quiet }) => { + Some(Command::Fetch { url, dump, selector, wait, timeout, wait_until, user_agent, stealth, eval, output, quiet, storage_dir }) => { reject_stealth_with_socks5(global_proxy.as_deref(), stealth)?; - run_fetch(&url, dump, selector, wait, timeout, &wait_until, user_agent, stealth, eval, output, quiet, global_proxy).await?; + run_fetch(&url, dump, selector, wait, timeout, &wait_until, user_agent, stealth, eval, output, quiet, global_proxy, storage_dir).await?; } Some(Command::Scrape { urls, eval, concurrency, format, timeout, quiet }) => { run_parallel_scrape(urls, eval, concurrency.get(), &format, timeout, quiet, global_proxy).await?; @@ -448,6 +460,7 @@ async fn run_fetch( output: Option, quiet: bool, proxy: Option, + storage_dir: Option, ) -> anyhow::Result<()> { // --dump original short-circuits the browser stack entirely: fetch the raw // response body via HTTP and stream the bytes verbatim. Useful for binary @@ -465,8 +478,14 @@ async fn run_fetch( return Ok(()); } - let context = Arc::new(BrowserContext::with_options("fetch".to_string(), proxy, stealth)); - let mut page = Page::new("fetch-page".to_string(), context); + let context = Arc::new(BrowserContext::with_storage_full( + "fetch".to_string(), + proxy, + stealth, + user_agent.clone(), + storage_dir.clone(), + )); + let mut page = Page::new("fetch-page".to_string(), context.clone()); if let Some(ref ua) = user_agent { page.http_client.set_user_agent(ua).await; @@ -506,6 +525,7 @@ async fn run_fetch( other => other.to_string(), }; write_or_print(rendered, output.as_ref()).await?; + context.save_cookies(); return Ok(()); } @@ -520,6 +540,9 @@ async fn run_fetch( }; write_or_print(rendered, output.as_ref()).await?; + // Save cookies to disk if storage_dir is configured + context.save_cookies(); + Ok(()) } @@ -1364,6 +1387,7 @@ mod tests { eval: None, quiet: true, output: None, + storage_dir: None, }); assert!(is_quiet_command(&cmd)); } diff --git a/crates/obscura-net/Cargo.toml b/crates/obscura-net/Cargo.toml index 328bff8..14286ae 100644 --- a/crates/obscura-net/Cargo.toml +++ b/crates/obscura-net/Cargo.toml @@ -17,6 +17,7 @@ tracing = { workspace = true } thiserror = { workspace = true } async-trait = "0.1" encoding_rs = "0.8" +tempfile = "3" wreq-util = { version = "3.0.0-rc.10", optional = true } # `prefix-symbols` renames BoringSSL exports to avoid clashes with system OpenSSL, diff --git a/crates/obscura-net/src/client.rs b/crates/obscura-net/src/client.rs index c971858..e283f80 100644 --- a/crates/obscura-net/src/client.rs +++ b/crates/obscura-net/src/client.rs @@ -346,9 +346,33 @@ impl ObscuraHttpClient { ); let cookie_header = self.cookie_jar.get_cookie_header(¤t_url); + tracing::debug!( + "Cookie header for {}: {} cookies ({} bytes)", + current_url.host_str().unwrap_or("?"), + cookie_header.split("; ").filter(|s| !s.is_empty()).count(), + cookie_header.len(), + ); if !cookie_header.is_empty() { - if let Ok(val) = HeaderValue::from_str(&cookie_header) { - headers.insert(reqwest::header::COOKIE, val); + match HeaderValue::from_str(&cookie_header) { + Ok(val) => { + headers.insert(reqwest::header::COOKIE, val); + } + Err(_) => { + let filtered: String = cookie_header + .split("; ") + .filter(|pair| HeaderValue::from_str(pair).is_ok()) + .collect::>() + .join("; "); + if !filtered.is_empty() { + if let Ok(val) = HeaderValue::from_str(&filtered) { + headers.insert(reqwest::header::COOKIE, val); + } + } + tracing::debug!( + "Cookie header invalid chars, filtered {} -> {} bytes", + cookie_header.len(), filtered.len(), + ); + } } } diff --git a/crates/obscura-net/src/cookies.rs b/crates/obscura-net/src/cookies.rs index 3fcae0c..307946e 100644 --- a/crates/obscura-net/src/cookies.rs +++ b/crates/obscura-net/src/cookies.rs @@ -8,7 +8,7 @@ pub struct CookieJar { cookies: RwLock>>, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] struct CookieEntry { name: String, value: String, @@ -372,6 +372,69 @@ impl CookieJar { pub fn clear(&self) { self.cookies.write().unwrap().clear(); } + + /// Serialize all non-expired cookies to a JSON file. + /// Writes atomically via tempfile then rename. + pub fn save_to_file(&self, path: &std::path::Path) -> Result<(), std::io::Error> { + use std::io::Write; + + let cookies = self.cookies.read().unwrap(); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + let mut all: Vec = Vec::new(); + for domain_cookies in cookies.values() { + for entry in domain_cookies.values() { + if let Some(exp) = entry.expires { + if exp < now { + continue; + } + } + all.push(CookieInfo { + name: entry.name.clone(), + value: entry.value.clone(), + domain: entry.domain.clone(), + path: entry.path.clone(), + secure: entry.secure, + http_only: entry.http_only, + same_site: entry.same_site.clone(), + expires: entry.expires.map(|e| e as i64), + }); + } + } + + let json = serde_json::to_string_pretty(&all).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidData, e) + })?; + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let mut tmp = tempfile::NamedTempFile::new_in( + path.parent().unwrap_or(std::path::Path::new(".")), + )?; + tmp.write_all(json.as_bytes())?; + tmp.persist(path).map_err(|e| e.error)?; + Ok(()) + } + + /// Load cookies from a JSON file into the jar. + /// Merges with existing cookies (does not clear). + /// Returns the number of cookies loaded. + pub fn load_from_file(&self, path: &std::path::Path) -> Result { + if !path.exists() { + return Ok(0); + } + let data = std::fs::read_to_string(path)?; + let cookies: Vec = + serde_json::from_str(&data).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidData, e) + })?; + let count = cookies.len(); + self.set_cookies_from_cdp(cookies); + Ok(count) + } } impl Default for CookieJar { @@ -649,4 +712,75 @@ mod tests { let url = Url::parse("https://example.com/").unwrap(); assert!(jar.get_cookie_header(&url).is_empty()); } + fn test_save_load_roundtrip() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("cookies.json"); + + let jar = CookieJar::new(); + let url = Url::parse("https://example.com/").unwrap(); + jar.set_cookie("session=abc123; Domain=example.com; Path=/", &url); + jar.set_cookie("token=xyz; Secure; HttpOnly", &url); + + jar.save_to_file(&path).unwrap(); + assert!(path.exists()); + + let jar2 = CookieJar::new(); + let count = jar2.load_from_file(&path).unwrap(); + assert_eq!(count, 2); + + let header = jar2.get_cookie_header(&url); + assert!(header.contains("session=abc123")); + assert!(header.contains("token=xyz")); + } + + #[test] + fn test_load_nonexistent_file_returns_zero() { + let jar = CookieJar::new(); + let count = jar + .load_from_file(std::path::Path::new("/nonexistent/cookies.json")) + .unwrap(); + assert_eq!(count, 0); + } + + #[test] + fn test_domain_matches_subdomain_without_leading_dot() { + let jar = CookieJar::new(); + jar.set_cookies_from_cdp(vec![CookieInfo { + name: "session".to_string(), + value: "abc".to_string(), + domain: "xiaohongshu.com".to_string(), + path: "/".to_string(), + secure: false, + http_only: true, + same_site: String::new(), + expires: None, + }]); + let url = Url::parse("https://www.xiaohongshu.com/explore").unwrap(); + let header = jar.get_cookie_header(&url); + assert!(header.contains("session=abc"), "Cookie header was: '{}'", header); + } + + #[test] + fn test_cookie_from_file_load_then_send_in_request() { + // Simulate what happens: load cookies from file → navigate → cookie should be in request + use std::io::Write; + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("cookies.json"); + + // Write cookies like we exported from Chrome + let cookies = serde_json::json!([ + {"name": "a1", "value": "testval", "domain": "xiaohongshu.com", "path": "/", "secure": false, "httpOnly": false}, + {"name": "web_session", "value": "sess123", "domain": "xiaohongshu.com", "path": "/", "secure": false, "httpOnly": true}, + ]); + std::fs::write(&path, serde_json::to_string(&cookies).unwrap()).unwrap(); + + let jar = CookieJar::new(); + let count = jar.load_from_file(&path).unwrap(); + assert_eq!(count, 2, "Should load 2 cookies"); + + let url = Url::parse("https://www.xiaohongshu.com/explore").unwrap(); + let header = jar.get_cookie_header(&url); + assert!(header.contains("a1=testval"), "Missing a1 in: '{}'", header); + assert!(header.contains("web_session=sess123"), "Missing web_session in: '{}'", header); + } }