From 897ee8fda10fa80093e58475f806a2fdf6b999eb Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 24 Apr 2025 18:48:38 -0400 Subject: [PATCH] cookie heuristic --- Cargo.lock | 2 + Cargo.toml | 1 + .../src/threads/sora_session_login_thread.rs | 5 ++ .../api_clients/openai_sora_client/Cargo.toml | 14 +-- .../src/utils/has_session_cookie.rs | 88 +++++++++++++++++-- 5 files changed, 96 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 10cbe3ac44..89ab40ea2e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7844,12 +7844,14 @@ dependencies = [ "anyhow", "base64 0.22.1", "chrono", + "cookie 0.18.1", "crockford", "errors", "filesys", "hex-literal", "idempotency", "log", + "once_cell", "rand 0.9.0", "reqwest 0.12.15", "serde", diff --git a/Cargo.toml b/Cargo.toml index 0a700a74f0..803c7a819a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -155,6 +155,7 @@ workspace-build-acceleration = { path = "crates/workspace-build-acceleration" } anyhow = "1.0.97" base64 = "0.22.1" chrono = { version = "0.4.40", features = ["serde"] } +cookie = "0.18.1" hex-literal = "1.0.0" log = "0.4.27" once_cell = "1.21.3" diff --git a/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs b/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs index fde729978f..391846714c 100644 --- a/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs +++ b/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs @@ -109,6 +109,11 @@ fn extract_cookies_to_file( sora_credential_manager: &SoraCredentialManager, ) -> AnyhowResult<()> { let new_cookies = get_all_sora_cookies_as_string(webview)?.trim().to_string(); + + if new_cookies.is_empty() { + return Ok(()); + } + let maybe_old_cookies = read_sora_cookies_from_disk(app_data_root); let mut should_write_cookies = maybe_old_cookies.is_none(); diff --git a/crates/lib/api_clients/openai_sora_client/Cargo.toml b/crates/lib/api_clients/openai_sora_client/Cargo.toml index 7a0025388d..f1ad20e2f4 100644 --- a/crates/lib/api_clients/openai_sora_client/Cargo.toml +++ b/crates/lib/api_clients/openai_sora_client/Cargo.toml @@ -19,24 +19,26 @@ workspace-build-acceleration = { workspace = true } # Internal crockford = { path = "../../data/crockford" } errors = { path = "../../errors" } -subprocess_common = { path = "../../subprocess_common" } idempotency = { path = "../../idempotency" } +subprocess_common = { path = "../../subprocess_common" } # External anyhow = { workspace = true } +base64 = { workspace = true } +chrono = { workspace = true} +cookie = { workspace = true } +hex-literal = { workspace = true } log = { workspace = true } +once_cell= { workspace = true } +rand = { workspace = true } reqwest = { workspace = true, features = ["multipart", "gzip"] } serde = { workspace = true } serde_derive = { workspace = true } serde_json = { workspace = true } sha3 = { workspace = true } -hex-literal = { workspace = true } -base64 = { workspace = true } +thiserror = { workspace = true } tokio = { version = "1.38.0", features = ["macros"] } toml = { workspace = true } -thiserror = { workspace = true } -chrono = { workspace = true} -rand = { workspace = true } [dev-dependencies] testing = { path = "../../testing" } diff --git a/crates/lib/api_clients/openai_sora_client/src/utils/has_session_cookie.rs b/crates/lib/api_clients/openai_sora_client/src/utils/has_session_cookie.rs index dd4d34aed8..f5a9c3ab2e 100644 --- a/crates/lib/api_clients/openai_sora_client/src/utils/has_session_cookie.rs +++ b/crates/lib/api_clients/openai_sora_client/src/utils/has_session_cookie.rs @@ -1,17 +1,92 @@ +use std::collections::HashSet; +use std::iter::Map; +use cookie::Cookie; +use once_cell::sync::Lazy; +use errors::AnyhowResult; pub const SESSION_COOKIE_NAME : &str = "__Secure-authjs.session-token"; -pub fn has_session_cookie(cookies: &str) -> bool { - // NB: This is not a valid check for cookies, just a cheap heuristic. - cookies.contains(SESSION_COOKIE_NAME) +/// Cookies Sora.com mints that are not the session cookie. We can ignore these. +/// This list is very likely to change over time. +const IRRELEVANT_COOKIE_NAMES: Lazy> = Lazy::new(|| { + HashSet::from([ + "cf_clearance", + "oai-did", + "_cfuvid", + "__cflb", + "__cf_bm", + "__Host-authjs.csrf-token", + "__Secure-authjs.callback-url", + "__Secure-authjs.pkce.code_verifier", + "__Secure-authjs.state", + ]) +}); + +/// Super naive heuristic for "logged out" total cookie length being ~1972 characters. +const SESSION_COOKIE_LENGTH_HEURISTIC : usize = 2000; + +//pub fn has_session_cookie(cookies: &str) -> bool { +// // NB: This is not a valid check for cookies, just a cheap heuristic. +// cookies.contains(SESSION_COOKIE_NAME) +//} + +#[derive(Copy,Clone,Eq,PartialEq,Debug)] +pub enum SessionCookiePresence { + Present, + MaybePresent, + Absent, +} + +/// OpenAI might change the name of the session cookie, and we don't want to break clients, +/// so we use some heuristics to determine if the cookie might be present. +pub fn has_session_cookie(cookies_header: &str) -> AnyhowResult { + let cookies_header = cookies_header.trim(); + if (cookies_header.is_empty()) { + return Ok(SessionCookiePresence::Absent); + } + + let cookies = Cookie::split_parse(cookies_header); + let mut unknown_cookie_count = 0; // NB: In the event they change their session cookie + + for cookie in cookies { + let cookie= cookie?; + let cookie_name = cookie.name(); + if cookie_name == SESSION_COOKIE_NAME { + return Ok(SessionCookiePresence::Present); + } else if !IRRELEVANT_COOKIE_NAMES.contains(cookie_name) { + unknown_cookie_count += 1; + } + } + + // This is a really stupid heuristic, but we want to be safe. + if cookies_header.len() > SESSION_COOKIE_LENGTH_HEURISTIC { + return Ok(SessionCookiePresence::MaybePresent); + } + + // Another stupid heuristic. + if unknown_cookie_count > 3 { + return Ok(SessionCookiePresence::MaybePresent); + } + + Ok(SessionCookiePresence::Absent) } #[cfg(test)] mod tests { - use crate::utils::has_session_cookie::has_session_cookie; + use crate::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence}; #[test] fn test() { + assert_eq!(has_session_cookie("").unwrap(), SessionCookiePresence::Absent); + assert_eq!(has_session_cookie("foo=bar").unwrap(), SessionCookiePresence::Absent); + + // Heuristic: unknown cookies. + assert_eq!(has_session_cookie("foo=bar;bin=baz;bat=ban;bash=barn").unwrap(), SessionCookiePresence::MaybePresent); + + // Heuristic: cookie length + let cookies = format!("cookie={}", "a".repeat(10000)); + assert_eq!(has_session_cookie(&cookies).unwrap(), SessionCookiePresence::MaybePresent); + let cookies = "__Host-authjs.csrf-token=a2783b5f89ee8dcb184f64723c3fdc91868f89840cb842fb20c3bd3\ d6b087d97%7C8e85a84ac3870dfab160d96c45fb8a0430350ae4b4755784c03c3b97a927eb06; __Secure-authjs.c\ allback-url=https%3A%2F%2Fsora.com%2F; __cflb=0H28vtyuPgbsDEUozcoHSWoAnvpSqAh9PUFn8MimKwT; __Se\ @@ -43,9 +118,6 @@ mod tests { c9-q1LsVirGizzhS_VTA5Ltnh4quwaIU3zkI_y3N9xV62ijUectvOPWONl42Qm9NUU4fQZynrUbTd9sSjYRNHfveIMDhbCj\ dDWjzHjMgCeUopSyrLLra1uC0NsJrBoq24QMQ00jHgKWOnqpdFTjPopZ-zDcN0E4FkYB8kKbrz4u4ZIn0fNNnE2bzlU8rbH"; - assert!(!has_session_cookie("")); - assert!(!has_session_cookie("foo")); - - assert!(has_session_cookie(cookies)); + assert_eq!(has_session_cookie(cookies).unwrap(), SessionCookiePresence::Present); } }