diff --git a/Cargo.lock b/Cargo.lock index 371c9ae8..c59c93a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -218,7 +218,7 @@ dependencies = [ "build_metadata", "bytes", "chrono", - "cookie_store 0.0.1", + "cookie_store 0.0.2", "core_types", "directories", "dirs", @@ -1112,13 +1112,15 @@ dependencies = [ [[package]] name = "cookie_store" -version = "0.0.1" +version = "0.0.2" dependencies = [ - "errors", + "cookie_store 0.21.1", + "log", "serde", - "serde_derive", "serde_json", + "url", "workspace-build-acceleration", + "wreq", ] [[package]] @@ -3790,7 +3792,7 @@ version = "0.0.1" dependencies = [ "anyhow", "cloudflare_errors", - "cookie_store 0.0.1", + "cookie_store 0.0.2", "dom_query 0.20.2", "errors", "filesys", diff --git a/crates/api_clients/midjourney_client/Cargo.toml b/crates/api_clients/midjourney_client/Cargo.toml index 1a417a70..5e91a069 100644 --- a/crates/api_clients/midjourney_client/Cargo.toml +++ b/crates/api_clients/midjourney_client/Cargo.toml @@ -17,7 +17,7 @@ workspace-build-acceleration.workspace = true # Internal cloudflare_errors.workspace = true -cookie_store.workspace = true +cookie_store = { workspace = true, features = ["wreq"] } errors.workspace = true # External diff --git a/crates/artcraftx/src/services/grok/windows/grok_login_window/grok_login_webview_extract_cookies.rs b/crates/artcraftx/src/services/grok/windows/grok_login_window/grok_login_webview_extract_cookies.rs index cfb9bf55..365531a2 100644 --- a/crates/artcraftx/src/services/grok/windows/grok_login_window/grok_login_webview_extract_cookies.rs +++ b/crates/artcraftx/src/services/grok/windows/grok_login_window/grok_login_webview_extract_cookies.rs @@ -1,4 +1,4 @@ -use cookie_store::cookie_store::CookieStore; +use cookie_store::cookie_store::{CapturedCookie, CookieStore}; use errors::AnyhowResult; use once_cell::sync::Lazy; use reqwest::Url; @@ -17,9 +17,14 @@ pub fn grok_login_webview_extract_cookies(webview: &WebviewWindow) -> AnyhowResu let mut cookie_store = CookieStore::empty(); let cookies = get_all_grok_cookies(webview)?; for cookie in cookies.iter() { - cookie_store.add_cookie_name_and_value( - cookie.name().to_string(), - cookie.value().to_string(), + cookie_store.insert_captured( + CapturedCookie { + name: cookie.name().to_string(), + value: cookie.value().to_string(), + maybe_domain: cookie.domain().map(str::to_string), + maybe_path: cookie.path().map(str::to_string), + }, + &ROOT_COOKIE_URL, ); } Ok(cookie_store) diff --git a/crates/artcraftx/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs b/crates/artcraftx/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs index 7d125a57..d367b515 100644 --- a/crates/artcraftx/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs +++ b/crates/artcraftx/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs @@ -1,4 +1,4 @@ -use cookie_store::cookie_store::CookieStore; +use cookie_store::cookie_store::{CapturedCookie, CookieStore}; use errors::AnyhowResult; use once_cell::sync::Lazy; use reqwest::Url; @@ -18,9 +18,14 @@ pub fn extract_midjourney_webview_cookies(webview: &WebviewWindow) -> AnyhowResu let mut cookie_store = CookieStore::empty(); let cookies = get_all_midjourney_cookies(webview)?; for cookie in cookies.iter() { - cookie_store.add_cookie_name_and_value( - cookie.name().to_string(), - cookie.value().to_string(), + cookie_store.insert_captured( + CapturedCookie { + name: cookie.name().to_string(), + value: cookie.value().to_string(), + maybe_domain: cookie.domain().map(str::to_string), + maybe_path: cookie.path().map(str::to_string), + }, + &WWW_COOKIE_URL, ); } Ok(cookie_store) diff --git a/crates/artcraftx/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs b/crates/artcraftx/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs index 1a4a0a82..003612da 100644 --- a/crates/artcraftx/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs +++ b/crates/artcraftx/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs @@ -1,4 +1,4 @@ -use cookie_store::cookie_store::CookieStore; +use cookie_store::cookie_store::{CapturedCookie, CookieStore}; use errors::AnyhowResult; use once_cell::sync::Lazy; use reqwest::Url; @@ -13,9 +13,14 @@ pub fn worldlabs_login_webview_extract_cookies(webview: &WebviewWindow) -> Anyho let mut cookie_store = CookieStore::empty(); let cookies = get_all_worldlabs_cookies(webview)?; for cookie in cookies.iter() { - cookie_store.add_cookie_name_and_value( - cookie.name().to_string(), - cookie.value().to_string(), + cookie_store.insert_captured( + CapturedCookie { + name: cookie.name().to_string(), + value: cookie.value().to_string(), + maybe_domain: cookie.domain().map(str::to_string), + maybe_path: cookie.path().map(str::to_string), + }, + &ROOT_COOKIE_URL, ); } Ok(cookie_store) diff --git a/crates/artcraftx/src/windows/login_window/login_window_thread.rs b/crates/artcraftx/src/windows/login_window/login_window_thread.rs index a846dcd3..0a12b2b0 100644 --- a/crates/artcraftx/src/windows/login_window/login_window_thread.rs +++ b/crates/artcraftx/src/windows/login_window/login_window_thread.rs @@ -119,7 +119,7 @@ fn check_login_window( session_cookie_names.iter().any(|name| cookie_store.has_cookie(name)) }; - let cookie_names = cookie_store.cookies.keys().cloned().collect::>().join(", "); + let cookie_names = cookie_store.cookie_names().join(", "); info!( "{} login check: host `{}`, path `{}`, cookies={} [{}], length={}, has_session={}", website, hostname, path, cookie_store.len(), cookie_names, cookie_length, has_captured_session, diff --git a/crates/artcraftx/src/windows/login_window/utils/extract_login_window_cookies.rs b/crates/artcraftx/src/windows/login_window/utils/extract_login_window_cookies.rs index 0c565a61..dfaedcc9 100644 --- a/crates/artcraftx/src/windows/login_window/utils/extract_login_window_cookies.rs +++ b/crates/artcraftx/src/windows/login_window/utils/extract_login_window_cookies.rs @@ -1,4 +1,4 @@ -use cookie_store::cookie_store::CookieStore; +use cookie_store::cookie_store::{CapturedCookie, CookieStore}; use errors::AnyhowResult; use reqwest::Url; use tauri::WebviewWindow; @@ -18,6 +18,10 @@ pub fn extract_login_window_cookies( webview: &WebviewWindow, cookie_urls: &[Url], ) -> AnyhowResult { + let Some(fallback_url) = cookie_urls.first() else { + return Ok(CookieStore::empty()); + }; + let hosts: Vec = cookie_urls .iter() .filter_map(|url| url.host_str().map(|host| host.to_ascii_lowercase())) @@ -26,9 +30,14 @@ pub fn extract_login_window_cookies( let mut cookie_store = CookieStore::empty(); for cookie in webview.cookies()?.iter() { if cookie_domain_matches(cookie.domain(), &hosts) { - cookie_store.add_cookie_name_and_value( - cookie.name().to_string(), - cookie.value().to_string(), + cookie_store.insert_captured( + CapturedCookie { + name: cookie.name().to_string(), + value: cookie.value().to_string(), + maybe_domain: cookie.domain().map(str::to_string), + maybe_path: cookie.path().map(str::to_string), + }, + fallback_url, ); } } diff --git a/crates/artcraftx/src/windows/login_window/utils/extract_user_info_from_cookies.rs b/crates/artcraftx/src/windows/login_window/utils/extract_user_info_from_cookies.rs index ad441907..32588b88 100644 --- a/crates/artcraftx/src/windows/login_window/utils/extract_user_info_from_cookies.rs +++ b/crates/artcraftx/src/windows/login_window/utils/extract_user_info_from_cookies.rs @@ -17,8 +17,8 @@ const USERNAME_CLAIM_KEYS: &[&str] = &["username", "preferred_username", "name", /// username/email so the saved credential can be told apart from other /// accounts. Returns `None` when nothing identifiable is found. pub fn extract_user_info_from_cookies(cookie_store: &CookieStore) -> Option { - for cookie in cookie_store.cookies.values() { - let Some(claims) = decode_jwt_payload(&cookie.value) else { + for (_name, value) in cookie_store.iter_name_values() { + let Some(claims) = decode_jwt_payload(value) else { continue; }; @@ -63,15 +63,17 @@ fn first_string_claim(claims: &Value, keys: &[&str]) -> Option { #[cfg(test)] mod tests { use super::*; + use reqwest::Url; // Payload: {"email":"echelon@gmail.com","username":"echelon","exp":1,"iat":1} const JWT_WITH_EMAIL: &str = "eyJhbGciOiJIUzI1NiJ9.eyJlbWFpbCI6ImVjaGVsb25AZ21haWwuY29tIiwidXNlcm5hbWUiOiJlY2hlbG9uIiwiZXhwIjoxLCJpYXQiOjF9.abc"; #[test] fn extracts_email_and_username_from_jwt_cookie() { + let url = site_url(); let mut store = CookieStore::empty(); - store.add_cookie_name_and_value("session".to_string(), JWT_WITH_EMAIL.to_string()); - store.add_cookie_name_and_value("analytics".to_string(), "not-a-jwt".to_string()); + store.insert_named("session", JWT_WITH_EMAIL, &url); + store.insert_named("analytics", "not-a-jwt", &url); let info = extract_user_info_from_cookies(&store).unwrap(); assert_eq!(info.email.as_deref(), Some("echelon@gmail.com")); @@ -81,7 +83,11 @@ mod tests { #[test] fn returns_none_without_identifiable_claims() { let mut store = CookieStore::empty(); - store.add_cookie_name_and_value("plain".to_string(), "opaque-token".to_string()); + store.insert_named("plain", "opaque-token", &site_url()); assert!(extract_user_info_from_cookies(&store).is_none()); } + + fn site_url() -> Url { + Url::parse("https://provider.example/").unwrap() + } } diff --git a/crates/lib/cookie_store/Cargo.toml b/crates/lib/cookie_store/Cargo.toml index 593e7a13..e92ff8fc 100644 --- a/crates/lib/cookie_store/Cargo.toml +++ b/crates/lib/cookie_store/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "cookie_store" edition = "2024" -version = "0.0.1" +version = "0.0.2" publish = false [package.metadata.cargo-udeps.ignore] @@ -11,17 +11,27 @@ normal = ["workspace-build-acceleration"] name = "cookie_store" path = "src/lib.rs" +[features] +default = [] +# Implements the wreq `CookieStore` trait on `SharedCookieStore` so a wreq +# client can read from and write Set-Cookie responses into a shared store. +wreq = ["dep:wreq"] + [dependencies] # Workspace hack (faster compile times) workspace-build-acceleration.workspace = true -# Internal -errors.workspace = true - # External +log.workspace = true serde.workspace = true -serde_derive.workspace = true serde_json.workspace = true +url.workspace = true +wreq = { workspace = true, optional = true } + +# NB: The RFC 6265 cookie store implementation this crate wraps, renamed to +# avoid clashing with this crate's own name. `serde_json` enables cookie +# serialization; we skip `public_suffix` to keep dependencies minimal. +rfc_cookie_store = { package = "cookie_store", version = "0.21.1", default-features = false, features = ["serde_json"] } [dev-dependencies] # None yet diff --git a/crates/lib/cookie_store/src/change_log.rs b/crates/lib/cookie_store/src/change_log.rs new file mode 100644 index 00000000..b760225c --- /dev/null +++ b/crates/lib/cookie_store/src/change_log.rs @@ -0,0 +1,91 @@ +use std::collections::VecDeque; +use std::time::SystemTime; + +/// Oldest records are dropped once the log exceeds this many entries. +pub const MAX_CHANGE_LOG_RECORDS: usize = 256; + +/// A bounded, in-memory history of cookie store mutations, for debugging +/// auth flows ("did the provider rotate the session cookie?", "when did we +/// lose it?"). Values are intentionally not recorded — only names, domains, +/// and what happened. +#[derive(Clone, Debug, Default)] +pub struct CookieChangeLog { + records: VecDeque, + next_sequence_number: u64, +} + +#[derive(Clone, Debug)] +pub struct CookieChangeRecord { + pub sequence_number: u64, + pub timestamp: SystemTime, + pub action: CookieChangeAction, + pub cookie_name: String, + pub maybe_domain: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CookieChangeAction { + /// A new cookie was stored. + Inserted, + /// An existing cookie's value was replaced. + Updated, + /// A server sent an expired Set-Cookie, deleting the stored cookie. + ExpiredByServer, + /// A cookie was refused (domain mismatch, malformed, etc.). + Rejected, + /// The whole store was cleared. + Cleared, +} + +impl CookieChangeLog { + pub fn record( + &mut self, + action: CookieChangeAction, + cookie_name: &str, + maybe_domain: Option, + ) { + self.records.push_back(CookieChangeRecord { + sequence_number: self.next_sequence_number, + timestamp: SystemTime::now(), + action, + cookie_name: cookie_name.to_owned(), + maybe_domain, + }); + self.next_sequence_number += 1; + while self.records.len() > MAX_CHANGE_LOG_RECORDS { + self.records.pop_front(); + } + } + + /// Records in chronological order, oldest first. + pub fn iter(&self) -> impl Iterator { + self.records.iter() + } + + pub fn len(&self) -> usize { + self.records.len() + } + + pub fn is_empty(&self) -> bool { + self.records.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn log_is_bounded_and_keeps_newest_records() { + let mut log = CookieChangeLog::default(); + for index in 0..(MAX_CHANGE_LOG_RECORDS + 10) { + log.record(CookieChangeAction::Inserted, &format!("cookie_{index}"), None); + } + + assert_eq!(log.len(), MAX_CHANGE_LOG_RECORDS); + let first = log.iter().next().unwrap(); + assert_eq!(first.sequence_number, 10); + let last = log.iter().last().unwrap(); + assert_eq!(last.cookie_name, format!("cookie_{}", MAX_CHANGE_LOG_RECORDS + 9)); + } +} diff --git a/crates/lib/cookie_store/src/cookie.rs b/crates/lib/cookie_store/src/cookie.rs deleted file mode 100644 index e292df40..00000000 --- a/crates/lib/cookie_store/src/cookie.rs +++ /dev/null @@ -1,34 +0,0 @@ - -#[derive(Clone)] -pub struct Cookie { - pub name: String, - pub value: String, - - // TODO: Later. - // pub domain: String, - // pub path: String, - // pub secure: bool, - // pub http_only: bool, - // pub expires: Option, -} - -impl Cookie { - pub fn new(name: String, value: String) -> Self { - Self { name, value } - } - - pub fn new_from_str(name: &str, value: &str) -> Self { - Self { - name: name.to_string(), - value: value.to_string(), - } - } - - pub fn name(&self) -> &str { - &self.name - } - - pub fn value(&self) -> &str { - &self.value - } -} diff --git a/crates/lib/cookie_store/src/cookie_store.rs b/crates/lib/cookie_store/src/cookie_store.rs index ca1c0757..3585476f 100644 --- a/crates/lib/cookie_store/src/cookie_store.rs +++ b/crates/lib/cookie_store/src/cookie_store.rs @@ -1,78 +1,376 @@ -use crate::cookie::Cookie; -use crate::serialized_cookie_store::{SerializableCookieStore, SerializedCookie}; -use std::collections::HashMap; +use crate::change_log::{CookieChangeAction, CookieChangeLog}; +use crate::serialized_cookie_store::SerializableCookieStore; +use log::warn; +use rfc_cookie_store::{CookieError, CookieStore as RfcCookieStore, RawCookie, StoreAction}; +use url::Url; -#[derive(Clone)] +/// A cookie jar with RFC 6265 storage semantics. +/// +/// Cookies are keyed by `(domain, path, name)` — not just name — and a repeat +/// `Set-Cookie` for the same key replaces the old value, while an expired +/// `Set-Cookie` deletes it (the standard server-side removal mechanism). +/// Every mutation is recorded in a bounded [`CookieChangeLog`] for debugging. +/// +/// This is a plain value type; wrap it in +/// [`crate::shared_cookie_store::SharedCookieStore`] to share it between +/// threads. +#[derive(Clone, Debug, Default)] pub struct CookieStore { - // TODO: This doesn't yet handle partial paths, subdomain rules, etc. - pub cookies: HashMap, + store: RfcCookieStore, + change_log: CookieChangeLog, +} + +/// A cookie captured from a browser or webview, with as much metadata as the +/// browser reported. +#[derive(Clone, Debug)] +pub struct CapturedCookie { + pub name: String, + pub value: String, + + /// Domain as reported by the browser. May carry a leading dot + /// (e.g. `.storyteller.ai`), which marks a parent-domain cookie. + pub maybe_domain: Option, + + pub maybe_path: Option, } impl CookieStore { pub fn empty() -> Self { + Self::default() + } + + pub (crate) fn from_rfc_store(store: RfcCookieStore) -> Self { Self { - cookies: HashMap::new(), + store, + change_log: CookieChangeLog::default(), } } - pub fn add_cookie(&mut self, cookie: Cookie) { - self.cookies.insert(cookie.name.clone(), cookie); + /// Insert a bare name/value cookie as a host-only cookie for `source_url`'s + /// host. Returns whether the cookie was stored (or expired an existing one). + pub fn insert_named(&mut self, name: &str, value: &str, source_url: &Url) -> bool { + let raw = RawCookie::new(name.to_owned(), value.to_owned()); + self.insert_raw(&raw, source_url) } - pub fn add_cookie_name_and_value(&mut self, name: String, value: String) { - self.add_cookie(Cookie { - name, - value, - }); + /// Insert a cookie captured from a browser or webview. + /// + /// When the capture carries a domain, the cookie is stored as a + /// domain-suffix cookie for it (so it matches subdomains, mirroring how the + /// browser scoped it). Without a domain it falls back to a host-only cookie + /// for `fallback_source_url`. + pub fn insert_captured(&mut self, captured: CapturedCookie, fallback_source_url: &Url) -> bool { + let maybe_domain = captured.maybe_domain + .as_deref() + .map(|domain| domain.trim_start_matches('.').to_ascii_lowercase()) + .filter(|domain| !domain.is_empty()); + + let Some(domain) = maybe_domain else { + return self.insert_named(&captured.name, &captured.value, fallback_source_url); + }; + + let source_url = match Url::parse(&format!("https://{domain}/")) { + Ok(url) => url, + Err(err) => { + warn!("Cannot build source URL for captured cookie {:?} (domain {:?}): {:?}", + captured.name, domain, err); + return false; + } + }; + + let mut raw = RawCookie::new(captured.name, captured.value); + raw.set_domain(domain); + raw.set_path(captured.maybe_path.unwrap_or_else(|| "/".to_owned())); + self.insert_raw(&raw, &source_url) } - pub fn get_cookie(&self, name: &str) -> Option<&Cookie> { - self.cookies.get(name) - } - - pub fn has_cookie(&self, name: &str) -> bool { - self.cookies.contains_key(name) + /// Apply one `Set-Cookie` header value received from `request_url`, + /// following the RFC 6265 storage model (replace / expire / reject). + pub fn apply_set_cookie_header(&mut self, set_cookie_value: &str, request_url: &Url) -> bool { + let cookie_name = set_cookie_value + .split('=') + .next() + .unwrap_or(set_cookie_value) + .trim() + .to_owned(); + let result = self.store.parse(set_cookie_value, request_url); + self.record_insert_result(result, &cookie_name, request_url) } - pub fn remove_cookie(&mut self, name: &str) { - self.cookies.remove(name); - } - - pub fn len(&self) -> usize { - self.cookies.len() + /// Apply several `Set-Cookie` header values received from `request_url`. + /// Returns how many were accepted. + pub fn apply_set_cookie_headers<'a>( + &mut self, + set_cookie_values: impl Iterator, + request_url: &Url, + ) -> usize { + set_cookie_values + .filter(|value| self.apply_set_cookie_header(value, request_url)) + .count() } pub fn clear_all(&mut self) { - self.cookies.clear(); + self.store.clear(); + self.change_log.record(CookieChangeAction::Cleared, "*", None); + } + + /// Whether an unexpired cookie with this name exists on any domain. + pub fn has_cookie(&self, name: &str) -> bool { + self.store.iter_unexpired().any(|cookie| cookie.name() == name) + } + + /// The value of the first unexpired cookie with this name, on any domain. + pub fn get_cookie_value(&self, name: &str) -> Option<&str> { + self.store + .iter_unexpired() + .find(|cookie| cookie.name() == name) + .map(|cookie| cookie.value()) + } + + /// `(name, value)` pairs of every unexpired cookie, across all domains. + pub fn iter_name_values(&self) -> impl Iterator { + self.store + .iter_unexpired() + .map(|cookie| (cookie.name(), cookie.value())) + } + + /// Names of every unexpired cookie, across all domains. + pub fn cookie_names(&self) -> Vec { + self.iter_name_values() + .map(|(name, _)| name.to_owned()) + .collect() + } + + pub fn len(&self) -> usize { + self.store.iter_unexpired().count() + } + + pub fn is_empty(&self) -> bool { + self.store.iter_unexpired().next().is_none() } /// NB: Just use this as a heuristic, and do not call it in a loop. pub fn calculate_approx_cookie_character_length(&self) -> usize { const COOKIE_SEP_LENGTH: usize = 2; // '=' and ';' - self.cookies - .values() - .map(|cookie| cookie.name.len() + cookie.value.len() + COOKIE_SEP_LENGTH) + self.iter_name_values() + .map(|(name, value)| name.len() + value.len() + COOKIE_SEP_LENGTH) .sum() } + /// Every unexpired cookie joined into one `Cookie:` header string, + /// regardless of domain or path. + /// + /// This is a blunt tool for scraping clients that captured cookies for one + /// known site and want to send them all back. It performs no request + /// matching; two same-named cookies on different domains both appear. + /// Prefer [`Self::cookie_header_for_url`] when the request URL is known. pub fn to_cookie_string(&self) -> String { - self - .cookies - .values() - .map(|cookie| format!("{}={}", cookie.name, cookie.value)) - .collect::>() - .join("; ") + self.iter_name_values() + .map(|(name, value)| format!("{}={}", name, value)) + .collect::>() + .join("; ") + } + + /// A `Cookie:` header value containing only the cookies that RFC 6265 + /// domain/path/secure matching would send to `request_url`, or `None` when + /// no cookies match. + pub fn cookie_header_for_url(&self, request_url: &Url) -> Option { + let header = self.store + .get_request_values(request_url) + .map(|(name, value)| format!("{}={}", name, value)) + .collect::>() + .join("; "); + if header.is_empty() { + None + } else { + Some(header) + } + } + + /// The recent mutation history, for debugging. + pub fn change_log(&self) -> &CookieChangeLog { + &self.change_log } pub fn to_serializable(&self) -> SerializableCookieStore { - let cookies = self - .cookies - .values() - .map(|cookie| SerializedCookie { - name: cookie.name.clone(), - value: cookie.value.clone(), - }) - .collect(); - SerializableCookieStore { cookies } + SerializableCookieStore::from_cookie_store(self) + } + + pub (crate) fn rfc_store(&self) -> &RfcCookieStore { + &self.store + } + + fn insert_raw(&mut self, raw: &RawCookie<'static>, source_url: &Url) -> bool { + let cookie_name = raw.name().to_owned(); + let result = self.store.insert_raw(raw, source_url); + self.record_insert_result(result, &cookie_name, source_url) + } + + fn record_insert_result( + &mut self, + result: Result, + cookie_name: &str, + request_url: &Url, + ) -> bool { + let maybe_domain = request_url.host_str().map(str::to_owned); + match result { + Ok(StoreAction::Inserted) => { + self.change_log.record(CookieChangeAction::Inserted, cookie_name, maybe_domain); + true + } + Ok(StoreAction::UpdatedExisting) => { + self.change_log.record(CookieChangeAction::Updated, cookie_name, maybe_domain); + true + } + Ok(StoreAction::ExpiredExisting) => { + self.change_log.record(CookieChangeAction::ExpiredByServer, cookie_name, maybe_domain); + true + } + Err(err) => { + warn!("Rejected cookie {:?} from {}: {:?}", cookie_name, request_url, err); + self.change_log.record(CookieChangeAction::Rejected, cookie_name, maybe_domain); + false + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::change_log::CookieChangeAction; + + const SITE_URL: &str = "https://www.example.com/"; + + mod insert_tests { + use super::*; + + #[test] + fn insert_named_stores_host_only_cookie() { + let mut store = CookieStore::empty(); + assert!(store.insert_named("session", "abc123", &site_url())); + + assert!(store.has_cookie("session")); + assert_eq!(store.get_cookie_value("session"), Some("abc123")); + assert_eq!(store.len(), 1); + } + + #[test] + fn insert_named_replaces_existing_value() { + let mut store = CookieStore::empty(); + store.insert_named("session", "old", &site_url()); + store.insert_named("session", "new", &site_url()); + + assert_eq!(store.len(), 1); + assert_eq!(store.get_cookie_value("session"), Some("new")); + } + + #[test] + fn insert_captured_with_parent_domain_matches_subdomains() { + let mut store = CookieStore::empty(); + let captured = CapturedCookie { + name: "auth".to_string(), + value: "token".to_string(), + maybe_domain: Some(".example.com".to_string()), + maybe_path: None, + }; + assert!(store.insert_captured(captured, &site_url())); + + let subdomain_url = Url::parse("https://api.example.com/v1/things").unwrap(); + assert_eq!(store.cookie_header_for_url(&subdomain_url), Some("auth=token".to_string())); + } + + #[test] + fn insert_captured_without_domain_falls_back_to_source_url() { + let mut store = CookieStore::empty(); + let captured = CapturedCookie { + name: "sid".to_string(), + value: "xyz".to_string(), + maybe_domain: None, + maybe_path: None, + }; + assert!(store.insert_captured(captured, &site_url())); + + assert_eq!(store.cookie_header_for_url(&site_url()), Some("sid=xyz".to_string())); + } + } + + mod set_cookie_header_tests { + use super::*; + + #[test] + fn apply_set_cookie_header_stores_cookie() { + let mut store = CookieStore::empty(); + assert!(store.apply_set_cookie_header("session=abc; Path=/; HttpOnly", &site_url())); + assert_eq!(store.get_cookie_value("session"), Some("abc")); + } + + #[test] + fn expired_set_cookie_deletes_existing_cookie() { + let mut store = CookieStore::empty(); + store.apply_set_cookie_header("session=abc; Path=/", &site_url()); + assert!(store.has_cookie("session")); + + let accepted = store.apply_set_cookie_header( + "session=gone; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT", + &site_url(), + ); + assert!(accepted); + assert!(!store.has_cookie("session")); + } + + #[test] + fn set_cookie_for_mismatched_domain_is_rejected() { + let mut store = CookieStore::empty(); + let accepted = store.apply_set_cookie_header("evil=1; Domain=other.com", &site_url()); + assert!(!accepted); + assert!(!store.has_cookie("evil")); + } + } + + mod header_rendering_tests { + use super::*; + + #[test] + fn to_cookie_string_joins_all_cookies() { + let mut store = CookieStore::empty(); + store.insert_named("a", "1", &site_url()); + store.insert_named("b", "2", &site_url()); + + let header = store.to_cookie_string(); + let mut parts = header.split("; ").collect::>(); + parts.sort(); + assert_eq!(parts, vec!["a=1", "b=2"]); + } + + #[test] + fn cookie_header_for_url_excludes_other_hosts() { + let mut store = CookieStore::empty(); + store.insert_named("mine", "1", &site_url()); + store.insert_named("theirs", "2", &Url::parse("https://other.example.org/").unwrap()); + + assert_eq!(store.cookie_header_for_url(&site_url()), Some("mine=1".to_string())); + } + } + + mod change_log_tests { + use super::*; + + #[test] + fn mutations_are_recorded() { + let mut store = CookieStore::empty(); + store.insert_named("session", "a", &site_url()); + store.insert_named("session", "b", &site_url()); + + let actions = store.change_log() + .iter() + .map(|record| record.action) + .collect::>(); + assert_eq!(actions, vec![CookieChangeAction::Inserted, CookieChangeAction::Updated]); + } + } + + fn site_url() -> Url { + Url::parse(SITE_URL).unwrap() } } diff --git a/crates/lib/cookie_store/src/lib.rs b/crates/lib/cookie_store/src/lib.rs index 754363b5..561085ad 100644 --- a/crates/lib/cookie_store/src/lib.rs +++ b/crates/lib/cookie_store/src/lib.rs @@ -1,7 +1,17 @@ //! cookie_store //! -//! Lightweight cookie jar +//! Cookie jar for scraping-style API clients. //! +//! Wraps the RFC 6265 `cookie_store` crate (imported as `rfc_cookie_store`) so +//! stored cookies follow real HTTP semantics: domain and path matching, +//! replacement on repeat Set-Cookie, and server-driven expiry. On top of that +//! this crate adds disk serialization, a bounded change log for debugging +//! auth flows, and a cheaply clonable thread-safe handle +//! ([`shared_cookie_store::SharedCookieStore`]). +//! +//! With the `wreq` feature enabled, `SharedCookieStore` implements wreq's +//! `CookieStore` trait, so a wreq client built with it as the cookie provider +//! transparently applies every `Set-Cookie` response to the shared store. // Never allow these #![forbid(private_bounds)] @@ -18,6 +28,10 @@ #![allow(dead_code)] #![allow(non_snake_case)] -pub mod cookie; +pub mod change_log; pub mod cookie_store; -pub mod serialized_cookie_store; \ No newline at end of file +pub mod serialized_cookie_store; +pub mod shared_cookie_store; + +#[cfg(feature = "wreq")] +pub mod wreq_cookie_provider; diff --git a/crates/lib/cookie_store/src/serialized_cookie_store.rs b/crates/lib/cookie_store/src/serialized_cookie_store.rs index 22956c11..e20e6986 100644 --- a/crates/lib/cookie_store/src/serialized_cookie_store.rs +++ b/crates/lib/cookie_store/src/serialized_cookie_store.rs @@ -1,19 +1,25 @@ -use crate::cookie::Cookie; use crate::cookie_store::CookieStore; +use rfc_cookie_store::Cookie as RfcCookie; +use rfc_cookie_store::CookieStore as RfcCookieStore; use serde::{Deserialize, Serialize}; +use std::convert::Infallible; use std::error::Error; use std::fmt::{Debug, Display, Formatter}; use std::path::Path; +/// Bumped when the on-disk shape changes. Version 1 was the legacy +/// name/value-only format; version 2 stores full RFC 6265 cookies +/// (domain, path, expiry, flags). Version 1 data does not deserialize — +/// callers already treat parse failures as "no saved cookies". +pub const SERIALIZED_COOKIE_STORE_VERSION: u32 = 2; + +/// The serializable snapshot of a [`CookieStore`]. Embed this in state +/// structs or write it to its own file with [`Self::write_to_file`]. #[derive(Serialize, Deserialize)] pub struct SerializableCookieStore { - pub (crate) cookies: Vec, -} - -#[derive(Serialize, Deserialize)] -pub (crate) struct SerializedCookie { - pub (crate) name: String, - pub (crate) value: String, + #[serde(default)] + pub version: u32, + cookies: Vec>, } #[derive(Debug)] @@ -34,33 +40,74 @@ impl Display for SerializableCookieStoreError { } impl SerializableCookieStore { + pub fn from_cookie_store(store: &CookieStore) -> Self { + let cookies = store.rfc_store() + .iter_unexpired() + .cloned() + .collect(); + Self { + version: SERIALIZED_COOKIE_STORE_VERSION, + cookies, + } + } + pub fn read_from_file>( file_path: P, ) -> Result { let content = std::fs::read_to_string(file_path) - .map_err(|err| SerializableCookieStoreError::IoError(err))?; + .map_err(SerializableCookieStoreError::IoError)?; serde_json::from_str(&content) - .map_err(|err| SerializableCookieStoreError::SerializationError(err)) + .map_err(SerializableCookieStoreError::SerializationError) } - + pub fn write_to_file>( &self, file_path: P, ) -> Result<(), SerializableCookieStoreError> { let serialized = serde_json::to_string(self) - .map_err(|err| SerializableCookieStoreError::SerializationError(err))?; + .map_err(SerializableCookieStoreError::SerializationError)?; std::fs::write(file_path, serialized) - .map_err(|err| SerializableCookieStoreError::IoError(err)) + .map_err(SerializableCookieStoreError::IoError) } - + pub fn to_cookie_store(&self) -> CookieStore { - let mut store = CookieStore::empty(); - for serialized_cookie in &self.cookies { - store.add_cookie(Cookie::new_from_str( - &serialized_cookie.name, - &serialized_cookie.value - )); - } - store + let cookies = self.cookies + .iter() + .cloned() + .map(Ok::, Infallible>); + let store = match RfcCookieStore::from_cookies(cookies, false) { + Ok(store) => store, + Err(infallible) => match infallible {}, + }; + CookieStore::from_rfc_store(store) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use url::Url; + + #[test] + fn round_trips_through_json() { + let site_url = Url::parse("https://www.example.com/").unwrap(); + let mut store = CookieStore::empty(); + store.insert_named("session", "abc123", &site_url); + store.apply_set_cookie_header("pref=dark; Domain=example.com; Path=/", &site_url); + + let serialized = serde_json::to_string(&store.to_serializable()).unwrap(); + let deserialized: SerializableCookieStore = serde_json::from_str(&serialized).unwrap(); + let restored = deserialized.to_cookie_store(); + + assert_eq!(restored.len(), 2); + assert_eq!(restored.get_cookie_value("session"), Some("abc123")); + assert_eq!(restored.get_cookie_value("pref"), Some("dark")); + } + + #[test] + fn legacy_version_1_data_fails_to_parse() { + let legacy_json = r#"{"cookies":[{"name":"session","value":"abc"}]}"#; + let result = serde_json::from_str::(legacy_json); + assert!(result.is_err()); } } diff --git a/crates/lib/cookie_store/src/shared_cookie_store.rs b/crates/lib/cookie_store/src/shared_cookie_store.rs new file mode 100644 index 00000000..4adbfd46 --- /dev/null +++ b/crates/lib/cookie_store/src/shared_cookie_store.rs @@ -0,0 +1,80 @@ +use crate::cookie_store::CookieStore; +use std::sync::{Arc, PoisonError, RwLock}; +use url::Url; + +/// A cheaply clonable, thread-safe handle to a [`CookieStore`]. +/// +/// Clones share the same underlying store, so a mutation made by one thread +/// (e.g. a Set-Cookie applied by an HTTP client) is visible to all holders. +/// With the `wreq` feature enabled this type implements wreq's `CookieStore` +/// trait and can be passed to `ClientBuilder::cookie_provider`. +#[derive(Clone, Debug, Default)] +pub struct SharedCookieStore { + inner: Arc>, +} + +impl SharedCookieStore { + pub fn empty() -> Self { + Self::default() + } + + pub fn new(store: CookieStore) -> Self { + Self { + inner: Arc::new(RwLock::new(store)), + } + } + + /// A point-in-time copy of the underlying store. + pub fn snapshot(&self) -> CookieStore { + self.with_read(|store| store.clone()) + } + + /// Swap the underlying store for all holders of this handle. + pub fn replace(&self, store: CookieStore) { + self.with_write(|current| *current = store); + } + + pub fn with_read(&self, reader: impl FnOnce(&CookieStore) -> R) -> R { + // NB: A panic mid-read cannot corrupt the store, so recover from poison. + let guard = self.inner.read().unwrap_or_else(PoisonError::into_inner); + reader(&guard) + } + + pub fn with_write(&self, writer: impl FnOnce(&mut CookieStore) -> R) -> R { + let mut guard = self.inner.write().unwrap_or_else(PoisonError::into_inner); + writer(&mut guard) + } + + pub fn has_cookie(&self, name: &str) -> bool { + self.with_read(|store| store.has_cookie(name)) + } + + pub fn to_cookie_string(&self) -> String { + self.with_read(|store| store.to_cookie_string()) + } + + pub fn cookie_header_for_url(&self, request_url: &Url) -> Option { + self.with_read(|store| store.cookie_header_for_url(request_url)) + } + + pub fn apply_set_cookie_header(&self, set_cookie_value: &str, request_url: &Url) -> bool { + self.with_write(|store| store.apply_set_cookie_header(set_cookie_value, request_url)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn clones_share_mutations() { + let site_url = Url::parse("https://www.example.com/").unwrap(); + let shared = SharedCookieStore::empty(); + let clone = shared.clone(); + + clone.apply_set_cookie_header("session=abc; Path=/", &site_url); + + assert!(shared.has_cookie("session")); + assert_eq!(shared.cookie_header_for_url(&site_url), Some("session=abc".to_string())); + } +} diff --git a/crates/lib/cookie_store/src/wreq_cookie_provider.rs b/crates/lib/cookie_store/src/wreq_cookie_provider.rs new file mode 100644 index 00000000..0f6c8ac4 --- /dev/null +++ b/crates/lib/cookie_store/src/wreq_cookie_provider.rs @@ -0,0 +1,60 @@ +//! Implements wreq's `CookieStore` trait on [`SharedCookieStore`], so a wreq +//! client can use our store as its cookie provider: +//! +//! ```ignore +//! let cookies = SharedCookieStore::empty(); +//! let client = wreq::Client::builder() +//! .cookie_provider(cookies.clone()) +//! .build()?; +//! // Every Set-Cookie the client receives now lands in `cookies`, +//! // and every request sends the matching stored cookies. +//! ``` + +use crate::shared_cookie_store::SharedCookieStore; +use log::warn; +use url::Url; +use wreq::Uri; +use wreq::cookie::{CookieStore as WreqCookieStore, Cookies}; +use wreq::header::HeaderValue; + +impl WreqCookieStore for SharedCookieStore { + fn set_cookies(&self, cookie_headers: &mut dyn Iterator, uri: &Uri) { + let Some(request_url) = uri_to_url(uri) else { + warn!("Cannot convert URI to URL for Set-Cookie storage: {}", uri); + return; + }; + self.with_write(|store| { + for header in cookie_headers { + match header.to_str() { + Ok(value) => { + store.apply_set_cookie_header(value, &request_url); + } + Err(err) => { + warn!("Set-Cookie header from {} is not valid UTF-8: {:?}", request_url, err); + } + } + } + }); + } + + fn cookies(&self, uri: &Uri) -> Cookies { + let Some(request_url) = uri_to_url(uri) else { + return Cookies::Empty; + }; + let maybe_header = self.cookie_header_for_url(&request_url); + match maybe_header { + Some(header) => match HeaderValue::from_str(&header) { + Ok(value) => Cookies::Compressed(value), + Err(err) => { + warn!("Stored cookies for {} do not form a valid header: {:?}", request_url, err); + Cookies::Empty + } + }, + None => Cookies::Empty, + } + } +} + +fn uri_to_url(uri: &Uri) -> Option { + Url::parse(&uri.to_string()).ok() +}