diff --git a/Cargo.lock b/Cargo.lock index cd8421fe..46f74e43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1132,11 +1132,13 @@ dependencies = [ name = "cookie_store_wrapper" version = "0.0.2" dependencies = [ + "chrono", "cookie_store", "log", "serde", "serde_derive", "serde_json", + "time", "toml 1.0.7+spec-1.1.0", "url", "workspace-build-acceleration", diff --git a/crates/artcraftx/src/commands/credentials/add_web_credential_command.rs b/crates/artcraftx/src/commands/credentials/add_web_credential_command.rs index 97aa37e4..a0a0b53e 100644 --- a/crates/artcraftx/src/commands/credentials/add_web_credential_command.rs +++ b/crates/artcraftx/src/commands/credentials/add_web_credential_command.rs @@ -3,7 +3,9 @@ use crate::credentials::cookie_credential::CookieCredential; use crate::credentials::auth_credential::{AuthCredential, CredentialSecret}; use core_types::enums::generation_source::{CredentialKind, GenerationSource}; use crate::credentials::credential_user_info::CredentialUserInfo; +use crate::credentials::service_cookie_origin::cookie_origin_for_service; use crate::error::artcraftx_error::ArtcraftXError; +use cookie_store_wrapper::cookie_store::CookieStore; use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::data_dir::subdirectory::app_credentials_dir::AppCredentialsDir; use chrono::Utc; @@ -14,7 +16,7 @@ use tauri::State; /// The captured result of a web login, ready to persist. pub struct WebCredentialSave { pub service: GenerationSource, - pub cookie_header: String, + pub cookies: CookieStore, pub maybe_user_info: Option, } @@ -37,10 +39,10 @@ pub fn save_web_credential( }); let cookie = CookieCredential { - cookie_header: save.cookie_header, updated_at: Some(Utc::now()), failed_at: None, succeeded_at: None, + cookies: save.cookies, }; let credential = match existing { @@ -111,11 +113,21 @@ pub async fn add_web_credential_command( return Err("Cookie header must not be empty".to_string()); } + let Some(cookie_origin) = cookie_origin_for_service(service) else { + let message = format!("Service {} has no cookie origin", service); + error!("{}", message); + return Err(message); + }; + let cookies = CookieStore::from_cookie_header(&cookie_header, &cookie_origin); + if cookies.is_empty() { + return Err("Cookie header contained no name=value cookies".to_string()); + } + let credential = save_web_credential( app_data_root.credentials_dir(), WebCredentialSave { service, - cookie_header, + cookies, maybe_user_info: None, }, ).map_err(|err| { diff --git a/crates/artcraftx/src/commands/credentials/artcraft_login_command.rs b/crates/artcraftx/src/commands/credentials/artcraft_login_command.rs index a45ae3c3..d7eaf29c 100644 --- a/crates/artcraftx/src/commands/credentials/artcraft_login_command.rs +++ b/crates/artcraftx/src/commands/credentials/artcraft_login_command.rs @@ -10,6 +10,8 @@ use crate::commands::credentials::credential_payload::CredentialPayload; use crate::utils::services::artcraft_api_host::maybe_artcraft_api_host_for_service; use crate::credentials::cookie_credential::CookieCredential; use crate::credentials::auth_credential::{AuthCredential, CredentialSecret}; +use crate::credentials::service_cookie_origin::cookie_origin_for_service; +use cookie_store_wrapper::cookie_store::CookieStore; use crate::credentials::credential_user_info::CredentialUserInfo; use crate::events::basic_sendable_event_trait::BasicSendableEvent; use crate::events::functional_events::refresh_account_state_event::RefreshAccountStateEvent; @@ -128,11 +130,18 @@ fn save_session_credential( ) -> Result { let now = Utc::now(); + let Some(cookie_origin) = cookie_origin_for_service(service) else { + return Err(ArtcraftLoginCommandError { + error_type: ArtcraftLoginErrorType::BadRequest, + message: format!("Service {} has no cookie origin", service), + }); + }; let cookie = CookieCredential { - cookie_header: format!("session={}", signed_session), updated_at: Some(now), failed_at: None, succeeded_at: Some(now), + cookies: CookieStore::from_cookie_header( + &format!("session={}", signed_session), &cookie_origin), }; let user_info = if username_or_email.contains('@') { diff --git a/crates/artcraftx/src/commands/credentials/edit_web_credential_command.rs b/crates/artcraftx/src/commands/credentials/edit_web_credential_command.rs index efffd5b0..30b6f598 100644 --- a/crates/artcraftx/src/commands/credentials/edit_web_credential_command.rs +++ b/crates/artcraftx/src/commands/credentials/edit_web_credential_command.rs @@ -1,7 +1,9 @@ use crate::commands::credentials::credential_payload::CredentialPayload; use crate::credentials::cookie_credential::CookieCredential; use crate::credentials::auth_credential::CredentialSecret; +use crate::credentials::service_cookie_origin::cookie_origin_for_service; use crate::state::data_dir::app_data_root::AppDataRoot; +use cookie_store_wrapper::cookie_store::CookieStore; use chrono::Utc; use log::{error, info}; use serde_derive::Serialize; @@ -46,19 +48,27 @@ pub async fn edit_web_credential_command( error!("{}", message); return Err(message); }; - let existing_cookie_header = existing_cookie.cookie_header.clone(); + let existing_cookies = existing_cookie.cookies.clone(); let existing_failed_at = existing_cookie.failed_at; let existing_succeeded_at = existing_cookie.succeeded_at; if let Some(new_cookie_header) = cookie_header { let new_cookie_header = new_cookie_header.trim().to_string(); - if !new_cookie_header.is_empty() && new_cookie_header != existing_cookie_header { - credential.secret = CredentialSecret::Cookies(CookieCredential { - cookie_header: new_cookie_header, - updated_at: Some(Utc::now()), - failed_at: existing_failed_at, - succeeded_at: existing_succeeded_at, - }); + if !new_cookie_header.is_empty() { + let Some(cookie_origin) = cookie_origin_for_service(credential.service) else { + let message = format!("Service {} has no cookie origin", credential.service); + error!("{}", message); + return Err(message); + }; + let new_cookies = CookieStore::from_cookie_header(&new_cookie_header, &cookie_origin); + if new_cookies != existing_cookies { + credential.secret = CredentialSecret::Cookies(CookieCredential { + updated_at: Some(Utc::now()), + failed_at: existing_failed_at, + succeeded_at: existing_succeeded_at, + cookies: new_cookies, + }); + } } } diff --git a/crates/artcraftx/src/commands/generate/common/generation_credential.rs b/crates/artcraftx/src/commands/generate/common/generation_credential.rs index e2d53ee5..1d2c789b 100644 --- a/crates/artcraftx/src/commands/generate/common/generation_credential.rs +++ b/crates/artcraftx/src/commands/generate/common/generation_credential.rs @@ -37,7 +37,7 @@ pub fn storyteller_creds_from_credential( credential_not_usable(credential, "the account has no session cookies") })?; - StorytellerCredentialSet::parse_multi_cookie_header(&cookie.cookie_header) + StorytellerCredentialSet::parse_multi_cookie_header(&cookie.cookie_header()) .map_err(|err| { credential_not_usable( credential, diff --git a/crates/artcraftx/src/credentials/cookie_credential.rs b/crates/artcraftx/src/credentials/cookie_credential.rs index 3040d204..f13249c9 100644 --- a/crates/artcraftx/src/credentials/cookie_credential.rs +++ b/crates/artcraftx/src/credentials/cookie_credential.rs @@ -1,13 +1,10 @@ use chrono::{DateTime, Utc}; +use cookie_store_wrapper::cookie_store::CookieStore; use serde_derive::{Deserialize, Serialize}; /// A browser-cookie secret for a website integration. #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] pub struct CookieCredential { - /// Entire cookie header payload (all cookies), as sent in the - /// `Cookie:` HTTP request header. - pub cookie_header: String, - /// Last time the cookies were rewritten (e.g. refreshed by the app). #[serde(default, skip_serializing_if = "Option::is_none")] pub updated_at: Option>, @@ -19,15 +16,25 @@ pub struct CookieCredential { /// Last time a request with these cookies succeeded. #[serde(default, skip_serializing_if = "Option::is_none")] pub succeeded_at: Option>, + + /// The stored cookies with their RFC 6265 attributes. + /// NB: Serialized as an array of tables, so this field must stay LAST for + /// the TOML credential files to remain valid. + pub cookies: CookieStore, } impl CookieCredential { - pub fn new(cookie_header: impl Into) -> Self { + pub fn new(cookies: CookieStore) -> Self { Self { - cookie_header: cookie_header.into(), updated_at: None, failed_at: None, succeeded_at: None, + cookies, } } + + /// The stored cookies rendered as a `Cookie:` request-header string. + pub fn cookie_header(&self) -> String { + self.cookies.to_cookie_string() + } } diff --git a/crates/artcraftx/src/credentials/credential_toml.rs b/crates/artcraftx/src/credentials/credential_toml.rs index 65f28d67..93f855bf 100644 --- a/crates/artcraftx/src/credentials/credential_toml.rs +++ b/crates/artcraftx/src/credentials/credential_toml.rs @@ -20,11 +20,18 @@ use core_types::identifiers::credential_id::CredentialId; /// id = "credential_01j9dq3v5x8k2m7n4p6r9t0wxy" /// service = "higgsfield_cookies" /// -/// [cookie] -/// cookie_header = "session=abc123; other=value" -/// /// [user_info] /// username = "creator123" +/// +/// [[cookie.cookies]] +/// name = "session" +/// value = "abc123" +/// domain = "higgsfield.ai" +/// +/// [[cookie.cookies]] +/// name = "other" +/// value = "value" +/// domain = "higgsfield.ai" /// ``` #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] pub struct CredentialToml { @@ -136,12 +143,14 @@ mod tests { const HAND_WRITTEN_COOKIE_FILE: &str = r#" service = "artcraft_cookies" - [cookie] - cookie_header = "session=abc123" - [user_info] username = "creator123" email = "creator@example.com" + + [[cookie.cookies]] + name = "session" + value = "abc123" + domain = "storyteller.ai" "#; const HAND_WRITTEN_API_KEY_FILE: &str = r#" @@ -158,7 +167,7 @@ mod tests { fn parses_hand_written_cookie_file() { let file: CredentialToml = toml::from_str(HAND_WRITTEN_COOKIE_FILE).unwrap(); assert_eq!(file.service, GenerationSource::ArtcraftCookies); - assert_eq!(file.cookie.unwrap().cookie_header, "session=abc123"); + assert_eq!(file.cookie.unwrap().cookie_header(), "session=abc123"); let user_info = file.user_info.unwrap(); assert_eq!(user_info.username.as_deref(), Some("creator123")); assert_eq!(user_info.email.as_deref(), Some("creator@example.com")); @@ -211,10 +220,12 @@ mod tests { fn both_secrets_is_an_error() { let toml_text = r#" service = "fal_api" - [cookie] - cookie_header = "a=b" [api_key] api_key = "x" + [[cookie.cookies]] + name = "a" + value = "b" + domain = "example.com" "#; let file: CredentialToml = toml::from_str(toml_text).unwrap(); let result = file.into_auth_credential(PathBuf::from("test.toml")); @@ -225,8 +236,10 @@ mod tests { fn secret_kind_must_match_service_kind() { let toml_text = r#" service = "fal_api" - [cookie] - cookie_header = "a=b" + [[cookie.cookies]] + name = "a" + value = "b" + domain = "example.com" "#; let file: CredentialToml = toml::from_str(toml_text).unwrap(); let result = file.into_auth_credential(PathBuf::from("test.toml")); diff --git a/crates/artcraftx/src/credentials/mod.rs b/crates/artcraftx/src/credentials/mod.rs index 227c5426..7596b0f9 100644 --- a/crates/artcraftx/src/credentials/mod.rs +++ b/crates/artcraftx/src/credentials/mod.rs @@ -22,3 +22,4 @@ pub mod credential_toml; pub mod credential_user_info; pub mod find_service_credentials; pub mod login_website; +pub mod service_cookie_origin; diff --git a/crates/artcraftx/src/credentials/service_cookie_origin.rs b/crates/artcraftx/src/credentials/service_cookie_origin.rs new file mode 100644 index 00000000..f1903fb1 --- /dev/null +++ b/crates/artcraftx/src/credentials/service_cookie_origin.rs @@ -0,0 +1,57 @@ +use core_types::enums::generation_source::GenerationSource; +use reqwest::Url; + +/// The canonical site origin that a cookie-kind service's cookies belong to. +/// +/// Used to attribute a domain when cookies arrive as a bare `Cookie:` header +/// string (hand-entered credentials, session strings from our own backend) +/// rather than from a browser capture that knows each cookie's real domain. +/// Returns `None` for API-key services, which have no cookies. +pub fn cookie_origin_for_service(service: GenerationSource) -> Option { + let origin = match service { + GenerationSource::Artcraft + | GenerationSource::ArtcraftLocal + | GenerationSource::ArtcraftCookies => "https://storyteller.ai/", + GenerationSource::Grok + | GenerationSource::GrokCookies + | GenerationSource::XAiCookies => "https://grok.com/", + GenerationSource::Midjourney + | GenerationSource::MidjourneyCookies => "https://www.midjourney.com/", + GenerationSource::Sora + | GenerationSource::SoraCookies => "https://chatgpt.com/", + GenerationSource::WorldLabs + | GenerationSource::WorldLabsCookies => "https://worldlabs.ai/", + GenerationSource::HiggsfieldCookies => "https://higgsfield.ai/", + GenerationSource::MagnificCookies => "https://www.magnific.com/", + GenerationSource::OpenArtCookies => "https://openart.ai/", + GenerationSource::RunwayCookies => "https://app.runwayml.com/", + GenerationSource::ArtcraftApi + | GenerationSource::Fal + | GenerationSource::FalApi + | GenerationSource::OpenAiApi + | GenerationSource::ReplicateApi + | GenerationSource::XAiApi => return None, + }; + Some(Url::parse(origin).expect("static origin URL should parse")) +} + +#[cfg(test)] +mod tests { + use super::*; + use core_types::enums::generation_source::CredentialKind; + + #[test] + fn every_cookie_service_has_an_origin() { + for service in GenerationSource::all_variants() { + let origin = cookie_origin_for_service(service); + match service.kind() { + CredentialKind::Cookies => { + assert!(origin.is_some(), "cookie service {service} has no cookie origin"); + } + CredentialKind::ApiKey => { + assert!(origin.is_none(), "api-key service {service} should have no cookie origin"); + } + } + } + } +} diff --git a/crates/artcraftx/src/services/storyteller/state/storyteller_credential_manager.rs b/crates/artcraftx/src/services/storyteller/state/storyteller_credential_manager.rs index 75301bae..e7418d2c 100644 --- a/crates/artcraftx/src/services/storyteller/state/storyteller_credential_manager.rs +++ b/crates/artcraftx/src/services/storyteller/state/storyteller_credential_manager.rs @@ -83,7 +83,7 @@ impl StorytellerCredentialManager { let cookie = credential.cookies()?; - StorytellerCredentialSet::parse_multi_cookie_header(&cookie.cookie_header) + StorytellerCredentialSet::parse_multi_cookie_header(&cookie.cookie_header()) .ok() .flatten() .filter(|creds| !creds.is_empty()) diff --git a/crates/artcraftx/src/state/data_dir/subdirectory/app_credentials_dir.rs b/crates/artcraftx/src/state/data_dir/subdirectory/app_credentials_dir.rs index e67926c7..5eb607a7 100644 --- a/crates/artcraftx/src/state/data_dir/subdirectory/app_credentials_dir.rs +++ b/crates/artcraftx/src/state/data_dir/subdirectory/app_credentials_dir.rs @@ -194,7 +194,7 @@ mod tests { write_file( dir.path(), "artcraft_user1.toml", - "service = \"artcraft_cookies\"\n[cookie]\ncookie_header = \"a=b\"\n", + "service = \"artcraft_cookies\"\n[[cookie.cookies]]\nname = \"a\"\nvalue = \"b\"\ndomain = \"storyteller.ai\"\n", ); write_file( dir.path(), 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 0a12b2b0..b4ac9a9d 100644 --- a/crates/artcraftx/src/windows/login_window/login_window_thread.rs +++ b/crates/artcraftx/src/windows/login_window/login_window_thread.rs @@ -138,7 +138,7 @@ fn check_login_window( app_data_root.credentials_dir(), WebCredentialSave { service: site.credential_service(), - cookie_header: cookie_store.to_cookie_string(), + cookies: cookie_store, maybe_user_info, }, )?; diff --git a/crates/lib/cookie_store_wrapper/Cargo.toml b/crates/lib/cookie_store_wrapper/Cargo.toml index c5511ec6..72fcf11b 100644 --- a/crates/lib/cookie_store_wrapper/Cargo.toml +++ b/crates/lib/cookie_store_wrapper/Cargo.toml @@ -22,10 +22,15 @@ wreq = ["dep:wreq"] workspace-build-acceleration.workspace = true # External +chrono.workspace = true log.workspace = true serde.workspace = true serde_json.workspace = true url.workspace = true + +# NB: For converting cookie expiry timestamps; the `cookie` crate (via +# `cookie_store`) already depends on `time`, so this adds no new weight. +time = "0.3" wreq = { workspace = true, optional = true } # NB: The RFC 6265 cookie store implementation this crate wraps. `serde_json` diff --git a/crates/lib/cookie_store_wrapper/src/cookie_store.rs b/crates/lib/cookie_store_wrapper/src/cookie_store.rs index 5f9c7e79..fca28dfb 100644 --- a/crates/lib/cookie_store_wrapper/src/cookie_store.rs +++ b/crates/lib/cookie_store_wrapper/src/cookie_store.rs @@ -38,6 +38,16 @@ impl CookieStore { Self::default() } + /// Parse a raw `Cookie:` request-header string (`"a=1; b=2"`) into a new + /// store of host-only cookies for `source_url`'s host. This is the entry + /// point for hand-entered cookie headers, where per-cookie attributes are + /// not known. + pub fn from_cookie_header(cookie_header: &str, source_url: &Url) -> Self { + let mut store = Self::empty(); + store.insert_cookie_header(cookie_header, source_url); + store + } + pub (crate) fn from_rfc_store(store: CookieStoreRfc) -> Self { Self { store, @@ -83,6 +93,21 @@ impl CookieStore { self.insert_raw(&raw, &source_url) } + /// Insert every `name=value` pair of a raw `Cookie:` request-header string + /// as a host-only cookie for `source_url`'s host. Returns how many cookies + /// were stored. + pub fn insert_cookie_header(&mut self, cookie_header: &str, source_url: &Url) -> usize { + cookie_header + .split(';') + .filter_map(|pair| pair.trim().split_once('=')) + .filter(|(name, _)| !name.trim().is_empty()) + .map(|(name, value)| (name.trim().to_owned(), value.trim().to_owned())) + .collect::>() + .into_iter() + .filter(|(name, value)| self.insert_named(name, value, source_url)) + .count() + } + /// 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 { @@ -195,7 +220,7 @@ impl CookieStore { &self.store } - fn insert_raw(&mut self, raw: &RawCookie<'static>, source_url: &Url) -> bool { + pub (crate) 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) diff --git a/crates/lib/cookie_store_wrapper/src/serde_support.rs b/crates/lib/cookie_store_wrapper/src/serde_support.rs index 5656bf21..43de1bee 100644 --- a/crates/lib/cookie_store_wrapper/src/serde_support.rs +++ b/crates/lib/cookie_store_wrapper/src/serde_support.rs @@ -6,59 +6,89 @@ //! //! ```ignore //! #[derive(Serialize, Deserialize)] -//! struct MyProviderState { -//! #[serde(skip_serializing_if = "Option::is_none")] -//! user_cookies: Option, +//! struct CookieCredential { +//! updated_at: Option>, +//! cookies: CookieStore, // NB: keep array-of-table fields LAST for TOML //! } //! ``` //! -//! The wire shape is a versioned envelope: `{ version, cookies: [...] }`, -//! where each cookie carries its full RFC 6265 attributes (domain, path, -//! expiry, flags). Expired cookies are dropped on both save and load. The -//! change log is debug state and is not serialized. +//! The wire shape is a flat list of cookies, chosen to stay human-editable +//! when it lands in a TOML file someone opens in vim: +//! +//! ```toml +//! [[cookie.cookies]] +//! name = "session" +//! value = "abc123" +//! domain = "example.com" +//! secure = true +//! http_only = true +//! expires = "2027-01-01T00:00:00+00:00" +//! ``` +//! +//! Defaulted attributes (`path = "/"`, false flags, session-cookie expiry) +//! are omitted on write but accepted on read. Expired cookies are dropped on +//! both save and load. The change log is debug state and is not serialized. use crate::cookie_store::CookieStore; use crate::shared_cookie_store::SharedCookieStore; use ::cookie_store::Cookie as CookieRfc; -use ::cookie_store::CookieStore as CookieStoreRfc; +use ::cookie_store::{CookieDomain, CookieExpiration, RawCookie}; +use chrono::{DateTime, SecondsFormat, Utc}; +use log::warn; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use std::convert::Infallible; +use std::time::SystemTime; +use url::Url; -/// Bumped when the wire shape changes. Version 1 was the legacy -/// name/value-only format; version 2 stores full RFC 6265 cookies. -/// Version 1 data does not deserialize — callers already treat parse -/// failures as "no saved cookies". -pub const COOKIE_STORE_SERDE_VERSION: u32 = 2; +const DEFAULT_PATH: &str = "/"; -/// The on-the-wire shape of a serialized [`CookieStore`]. -#[derive(Serialize, Deserialize)] -struct CookieStoreEnvelope { - #[serde(default)] - version: u32, - cookies: Vec>, +/// The wire shape of one cookie. Public so callers can build tooling around +/// the stored format, but normal use is via `CookieStore`'s serde impls. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct SerializedCookie { + pub name: String, + pub value: String, + + /// The domain the cookie belongs to, without a leading dot. Subdomains + /// also match unless `host_only` is set. + pub domain: String, + + /// Restrict the cookie to exactly `domain` (no subdomains). + #[serde(default, skip_serializing_if = "is_false")] + pub host_only: bool, + + #[serde(default = "default_path", skip_serializing_if = "is_default_path")] + pub path: String, + + #[serde(default, skip_serializing_if = "is_false")] + pub secure: bool, + + #[serde(default, skip_serializing_if = "is_false")] + pub http_only: bool, + + /// RFC 3339 expiry time. Omitted for session cookies (which never expire + /// on disk). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expires: Option, } impl Serialize for CookieStore { fn serialize(&self, serializer: S) -> Result { - let envelope = CookieStoreEnvelope { - version: COOKIE_STORE_SERDE_VERSION, - cookies: self.rfc_store().iter_unexpired().cloned().collect(), - }; - envelope.serialize(serializer) + serializer.collect_seq( + self.rfc_store() + .iter_unexpired() + .filter_map(serialized_cookie_from_rfc) + ) } } impl<'de> Deserialize<'de> for CookieStore { fn deserialize>(deserializer: D) -> Result { - let envelope = CookieStoreEnvelope::deserialize(deserializer)?; - let cookies = envelope.cookies - .into_iter() - .map(Ok::, Infallible>); - let store = match CookieStoreRfc::from_cookies(cookies, false) { - Ok(store) => store, - Err(infallible) => match infallible {}, - }; - Ok(CookieStore::from_rfc_store(store)) + let cookies = Vec::::deserialize(deserializer)?; + let mut store = CookieStore::empty(); + for cookie in cookies { + insert_serialized_cookie(&mut store, cookie); + } + Ok(store) } } @@ -76,11 +106,107 @@ impl<'de> Deserialize<'de> for SharedCookieStore { } } +/// Compares the cookies (name, value, and attributes), ignoring the change +/// log and insertion order. +impl PartialEq for CookieStore { + fn eq(&self, other: &Self) -> bool { + sorted_serialized_cookies(self) == sorted_serialized_cookies(other) + } +} + +impl Eq for CookieStore {} + +fn serialized_cookie_from_rfc(cookie: &CookieRfc<'static>) -> Option { + let (domain, host_only) = match &cookie.domain { + CookieDomain::HostOnly(host) => (host.clone(), true), + CookieDomain::Suffix(suffix) => (suffix.clone(), false), + other => { + warn!("Cannot serialize cookie {:?} with domain {:?}; skipping.", cookie.name(), other); + return None; + } + }; + + let expires = match &cookie.expires { + CookieExpiration::AtUtc(at_utc) => { + let system_time = SystemTime::from(*at_utc); + Some(DateTime::::from(system_time).to_rfc3339_opts(SecondsFormat::Secs, false)) + } + CookieExpiration::SessionEnd => None, + }; + + Some(SerializedCookie { + name: cookie.name().to_owned(), + value: cookie.value().to_owned(), + domain, + host_only, + path: (*cookie.path).to_owned(), + secure: cookie.secure().unwrap_or(false), + http_only: cookie.http_only().unwrap_or(false), + expires, + }) +} + +fn insert_serialized_cookie(store: &mut CookieStore, cookie: SerializedCookie) { + let source_url = match Url::parse(&format!("https://{}/", cookie.domain)) { + Ok(url) => url, + Err(err) => { + warn!("Cannot restore cookie {:?} with domain {:?}: {:?}", cookie.name, cookie.domain, err); + return; + } + }; + + let mut raw = RawCookie::new(cookie.name, cookie.value); + raw.set_path(cookie.path); + raw.set_secure(cookie.secure); + raw.set_http_only(cookie.http_only); + if !cookie.host_only { + raw.set_domain(cookie.domain); + } + if let Some(expires) = &cookie.expires { + match DateTime::parse_from_rfc3339(expires) { + Ok(at) => { + let system_time = SystemTime::from(at.with_timezone(&Utc)); + raw.set_expires(time::OffsetDateTime::from(system_time)); + } + Err(err) => { + warn!("Cookie {:?} has unparseable expires {:?} ({:?}); treating as session cookie.", + raw.name(), expires, err); + } + } + } + + // NB: An expired cookie is rejected here (with a warning), which is the + // intended load behavior. + store.insert_raw(&raw, &source_url); +} + +fn sorted_serialized_cookies(store: &CookieStore) -> Vec { + let mut cookies = store.rfc_store() + .iter_unexpired() + .filter_map(serialized_cookie_from_rfc) + .collect::>(); + cookies.sort_by(|a, b| { + (&a.domain, &a.path, &a.name).cmp(&(&b.domain, &b.path, &b.name)) + }); + cookies +} + +fn default_path() -> String { + DEFAULT_PATH.to_owned() +} + +fn is_default_path(path: &str) -> bool { + path == DEFAULT_PATH +} + +fn is_false(value: &bool) -> bool { + !value +} + #[cfg(test)] mod tests { use super::*; use serde_derive::{Deserialize, Serialize}; - use url::Url; #[test] fn cookie_store_round_trips_through_json() { @@ -89,9 +215,11 @@ mod tests { let serialized = serde_json::to_string(&store).unwrap(); let restored: CookieStore = serde_json::from_str(&serialized).unwrap(); - assert_eq!(restored.len(), 2); + assert_eq!(restored.len(), 3); assert_eq!(restored.get_cookie_value("session"), Some("abc123")); assert_eq!(restored.get_cookie_value("pref"), Some("dark")); + assert_eq!(restored.get_cookie_value("keeper"), Some("forever")); + assert_eq!(restored, store); } #[test] @@ -99,6 +227,7 @@ mod tests { #[derive(Serialize, Deserialize)] struct CredentialFile { service: String, + // NB: Array-of-table fields go last so TOML stays valid. cookies: CookieStore, } @@ -111,8 +240,73 @@ mod tests { let restored: CredentialFile = toml::from_str(&serialized).unwrap(); assert_eq!(restored.service, "midjourney_cookies"); - assert_eq!(restored.cookies.get_cookie_value("session"), Some("abc123")); - assert_eq!(restored.cookies.get_cookie_value("pref"), Some("dark")); + assert_eq!(restored.cookies, file.cookies); + } + + #[test] + fn toml_output_is_human_editable() { + #[derive(Serialize)] + struct CredentialFile { + cookies: CookieStore, + } + + let site_url = Url::parse("https://www.example.com/").unwrap(); + let mut store = CookieStore::empty(); + store.apply_set_cookie_header("session=abc123; Secure; HttpOnly; Path=/", &site_url); + + let serialized = toml::to_string(&CredentialFile { cookies: store }).unwrap(); + + // The whole cookie reads as a flat, obvious table. + let expected = "\ +[[cookies]] +name = \"session\" +value = \"abc123\" +domain = \"www.example.com\" +host_only = true +secure = true +http_only = true +"; + assert_eq!(serialized, expected); + } + + #[test] + fn hand_written_toml_with_defaults_parses() { + #[derive(Deserialize)] + struct CredentialFile { + cookies: CookieStore, + } + + let hand_written = r#" + [[cookies]] + name = "session" + value = "abc123" + domain = "example.com" + "#; + + let file: CredentialFile = toml::from_str(hand_written).unwrap(); + assert_eq!(file.cookies.get_cookie_value("session"), Some("abc123")); + + // Domain cookies (the default) match subdomains. + let subdomain_url = Url::parse("https://api.example.com/").unwrap(); + assert_eq!( + file.cookies.cookie_header_for_url(&subdomain_url), + Some("session=abc123".to_string()), + ); + } + + #[test] + fn expires_round_trips_and_expired_cookies_are_dropped() { + let json = r#"[ + {"name": "fresh", "value": "1", "domain": "example.com", "expires": "2100-01-01T00:00:00+00:00"}, + {"name": "stale", "value": "2", "domain": "example.com", "expires": "2001-01-01T00:00:00+00:00"} + ]"#; + + let restored: CookieStore = serde_json::from_str(json).unwrap(); + assert!(restored.has_cookie("fresh")); + assert!(!restored.has_cookie("stale")); + + let reserialized = serde_json::to_string(&restored).unwrap(); + assert!(reserialized.contains("\"expires\":\"2100-01-01T00:00:00+00:00\"")); } #[test] @@ -137,6 +331,10 @@ mod tests { 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); + store.apply_set_cookie_header( + "keeper=forever; Path=/; Expires=Tue, 03 Aug 2100 00:38:37 GMT", + &site_url, + ); store } }