From 62622fc1252540eaafc97ea4f9b429b29f60f08b Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 30 Oct 2025 22:26:40 -0400 Subject: [PATCH] grok credential holder --- Cargo.lock | 1 + Cargo.toml | 1 + .../src/credentials/grok_user_data.rs | 22 +++ .../grok_client/src/credentials/mod.rs | 3 +- .../src/datatypes/api/user_email.rs | 6 + .../grok_client/src/datatypes/api/user_id.rs | 7 + crates/desktop/artcraft/Cargo.toml | 1 + .../grok/state/grok_credential_holder.rs | 30 ++++ .../grok/state/grok_credential_manager.rs | 139 +++++++++--------- .../grok/state/grok_serializable_state.rs | 49 ++---- .../artcraft/src/services/grok/state/mod.rs | 3 +- 11 files changed, 153 insertions(+), 109 deletions(-) create mode 100644 crates/api_clients/grok_client/src/credentials/grok_user_data.rs create mode 100644 crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs diff --git a/Cargo.lock b/Cargo.lock index 567f9015dc..0e9594bcd6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -662,6 +662,7 @@ dependencies = [ "errors", "filesys", "futures 0.3.31", + "grok_client", "idempotency", "image", "images", diff --git a/Cargo.toml b/Cargo.toml index d5f0852888..4922216a02 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -198,6 +198,7 @@ errors = { path = "crates/lib/errors" } fal_client = { path = "crates/api_clients/fal_client" } file_to_base64_url = { path = "crates/lib/files/file_to_base64_url" } filesys = { path = "crates/lib/files/filesys" } +grok_client = { path = "crates/api_clients/grok_client" } idempotency = { path = "crates/lib/idempotency" } images = { path = "crates/lib/files/images" } jwt_light = { path = "crates/lib/jwt_light" } diff --git a/crates/api_clients/grok_client/src/credentials/grok_user_data.rs b/crates/api_clients/grok_client/src/credentials/grok_user_data.rs new file mode 100644 index 0000000000..8b9f28b563 --- /dev/null +++ b/crates/api_clients/grok_client/src/credentials/grok_user_data.rs @@ -0,0 +1,22 @@ +use crate::datatypes::api::user_email::UserEmail; +use crate::datatypes::api::user_id::UserId; + +#[derive(Clone)] +pub struct GrokUserData { + /// From index.html + /// Not strictly needed to sign requests, but typically needed to generate URLs. + pub user_id: UserId, + + /// From index.html + /// Not needed, but returned alongside other details. + pub user_email: Option, +} + +impl GrokUserData { + pub fn from_id_and_email(id: String, email: Option) -> Self { + Self { + user_id: UserId(id), + user_email: email.map(|email| UserEmail(email)), + } + } +} diff --git a/crates/api_clients/grok_client/src/credentials/mod.rs b/crates/api_clients/grok_client/src/credentials/mod.rs index a7eb1c1fa3..4d7ea21059 100644 --- a/crates/api_clients/grok_client/src/credentials/mod.rs +++ b/crates/api_clients/grok_client/src/credentials/mod.rs @@ -1,3 +1,4 @@ +pub mod grok_client_secrets; pub mod grok_cookies; pub mod grok_full_credentials; -pub mod grok_client_secrets; +pub mod grok_user_data; diff --git a/crates/api_clients/grok_client/src/datatypes/api/user_email.rs b/crates/api_clients/grok_client/src/datatypes/api/user_email.rs index 8524151d88..a6ba535a04 100644 --- a/crates/api_clients/grok_client/src/datatypes/api/user_email.rs +++ b/crates/api_clients/grok_client/src/datatypes/api/user_email.rs @@ -1,3 +1,9 @@ /// Type for User Emails #[derive(Clone, Debug)] pub struct UserEmail(pub String); + +impl UserEmail { + pub fn to_string(&self) -> String { + self.0.clone() + } +} diff --git a/crates/api_clients/grok_client/src/datatypes/api/user_id.rs b/crates/api_clients/grok_client/src/datatypes/api/user_id.rs index 61581f17bb..c7e522fddc 100644 --- a/crates/api_clients/grok_client/src/datatypes/api/user_id.rs +++ b/crates/api_clients/grok_client/src/datatypes/api/user_id.rs @@ -3,3 +3,10 @@ /// User IDs are UUIDs. #[derive(Clone, Debug)] pub struct UserId(pub String); + +impl UserId { + pub fn to_string(&self) -> String { + self.0.clone() + } +} + diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index 81025533cc..053bd3a9be 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -32,6 +32,7 @@ cookie_store.workspace = true enums.workspace = true errors.workspace = true filesys.workspace = true +grok_client.workspace = true idempotency.workspace = true images.workspace = true memory_store.workspace = true diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs new file mode 100644 index 0000000000..110d6dc798 --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs @@ -0,0 +1,30 @@ +use cookie_store::cookie_store::CookieStore; +use grok_client::credentials::grok_client_secrets::GrokClientSecrets; +use grok_client::credentials::grok_user_data::GrokUserData; + +#[derive(Clone)] +pub struct GrokCredentialHolder { + /// Directly off the Tauri browser session. + /// Read once, write to disk + /// The Grok client consumes a string-only form (rather than this cookie jar) + pub browser_cookies: Option, + + /// Email, user id, etc. + /// These can be persisted. + pub grok_user_data: Option, + + /// These have a different lifecycle than the other pieces and are + /// subject to change when the website changes. + /// We won't persist these. + pub grok_client_secrets: Option, +} + +impl GrokCredentialHolder { + pub fn empty() -> Self { + Self { + browser_cookies: None, + grok_user_data: None, + grok_client_secrets: None, + } + } +} diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs index 66c8e06bb3..5388f6add7 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs @@ -1,7 +1,10 @@ use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::grok::state::grok_credential_holder::GrokCredentialHolder; use crate::services::grok::state::grok_serializable_state::{GrokSerializableState, SERIALIZABLE_GROK_STATE_VERSION}; use cookie_store::cookie_store::CookieStore; use errors::AnyhowResult; +use grok_client::credentials::grok_client_secrets::GrokClientSecrets; +use grok_client::credentials::grok_user_data::GrokUserData; use log::{info, warn}; use std::fs::read_to_string; use std::sync::{Arc, RwLock}; @@ -9,7 +12,8 @@ use std::sync::{Arc, RwLock}; #[derive(Clone)] pub struct GrokCredentialManager { // TODO: Put last write/last read timestamps on these. - cookies: Arc>>, + credential_data: Arc>, + //user_info: Arc>>, app_data_root: AppDataRoot, } @@ -17,94 +21,97 @@ pub struct GrokCredentialManager { impl GrokCredentialManager { pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self { Self { - cookies: Arc::new(RwLock::new(None)), + credential_data: Arc::new(RwLock::new(GrokCredentialHolder::empty())), app_data_root: app_data_root.clone(), } } pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { - let mut cookies; - //let mut user_info; + let mut credential_data; let result = GrokSerializableState::read_from_disk(app_data_root); match result { Err(err) => { warn!("Failed to read grok state from disk: {:?}", err); - cookies = Arc::new(RwLock::new(None)); - //user_info = Arc::new(RwLock::new(None)); + credential_data = Arc::new(RwLock::new(GrokCredentialHolder::empty())); }, Ok(None) => { - cookies = Arc::new(RwLock::new(None)); - //user_info = Arc::new(RwLock::new(None)); + credential_data = Arc::new(RwLock::new(GrokCredentialHolder::empty())); } Ok(Some(state)) => { let maybe_cookies = state.user_cookies .map(|cookies| cookies.to_cookie_store()); - //let maybe_user_info = state.user_info - // .map(|info| info.to_user_info()); - cookies = Arc::new(RwLock::new(maybe_cookies)); - //user_info = Arc::new(RwLock::new(maybe_user_info)); + + let mut user_data = None; + if let Some(user_id) = state.user_id { + user_data = Some(GrokUserData::from_id_and_email(user_id, state.user_email)); + } + + credential_data = Arc::new(RwLock::new(GrokCredentialHolder { + browser_cookies: maybe_cookies, + grok_user_data: user_data, + grok_client_secrets: None, // NB: We don't want to keep this on disk. It goes stale. + })); } }; Self { - cookies, + credential_data, app_data_root: app_data_root.clone(), } } pub fn maybe_copy_cookie_store(&self) -> anyhow::Result> { - match self.cookies.read() { + match self.credential_data.read() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - Ok(store) => Ok(store.clone()), + Ok(holder) => Ok(holder.browser_cookies.clone()), } } - //pub fn maybe_copy_user_info(&self) -> anyhow::Result> { - // match self.user_info.read() { - // Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), - // Ok(info) => Ok(info.clone()), - // } - //} + pub fn maybe_copy_client_secrets(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.grok_client_secrets.clone()), + } + } pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> { - match self.cookies.write() { + match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - Ok(mut current_store) => { - *current_store = Some(store); + Ok(mut holder) => { + holder.browser_cookies = Some(store); Ok(()) } } } - //pub fn replace_user_info(&self, user_info: MidjourneyUserInfo) -> anyhow::Result<()> { - // match self.user_info.write() { - // Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - // Ok(mut current_info) => { - // *current_info = Some(user_info); - // Ok(()) - // } - // } - //} + pub fn replace_client_secrets(&self, secrets: GrokClientSecrets) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + holder.grok_client_secrets = Some(secrets); + Ok(()) + } + } + } // NB: This is just a heuristic. We'll add better checks later. pub fn session_appears_active(&self) -> anyhow::Result { - let maybe_cookies = match self.cookies.read() { + let holder = match self.credential_data.read() { Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), Ok(store) => store.clone(), }; - //let maybe_user_info = match self.user_info.read() { - // Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), - // Ok(info) => info.clone(), - //}; - - if maybe_cookies.is_none() { + if holder.browser_cookies.is_none() { return Ok(false); } - let cookies = maybe_cookies.unwrap(); - //let user_info = maybe_user_info.unwrap(); + let cookies = match holder.browser_cookies.as_ref() { + None => return Ok(false), + Some(cookies) => cookies, + }; + + let has_user_data = holder.grok_user_data.is_some(); // NB: We consider the session active if we have auth cookies and a user id. //let has_user_id = user_info.user_id.is_some(); @@ -119,51 +126,36 @@ impl GrokCredentialManager { // TODO: Heuristic should count. // TODO: Consolidate with "login window thread" logic. // TODO: Check timestamp of last valid request. - Ok(maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie) + Ok(has_user_data || (maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie)) } pub fn clear_credentials(&self) -> anyhow::Result<()> { - match self.cookies.write() { + match self.credential_data.write() { Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - Ok(mut store) => *store = None, + Ok(mut store) => *store = GrokCredentialHolder::empty(), } - //match self.user_info.write() { - // Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - // Ok(mut info) => *info = None, - //} Ok(()) } pub fn persist_to_disk(&self) -> anyhow::Result<()> { - let maybe_cookies = match self.cookies.read() { + let creds = match self.credential_data.read() { Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), Ok(store) => store.clone(), }; - //let maybe_user_info = match self.user_info.read() { - // Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), - // Ok(info) => info.clone(), - //}; - - let maybe_cookies_header= maybe_cookies - .as_ref() - .map(|cookies| cookies.to_cookie_string()); - - let maybe_cookies= maybe_cookies - .map(|cookies| cookies.to_serializable()); - - //let maybe_user_info= maybe_user_info - // .map(|info| info.to_serializable()); - - if maybe_cookies.is_none() { - info!("Nothing to write to disk, skipping."); - return Ok(()); - - } - let state = GrokSerializableState { version: SERIALIZABLE_GROK_STATE_VERSION, - user_cookies: maybe_cookies, + user_cookies: creds.browser_cookies + .as_ref() + .map(|cookies| cookies.to_serializable()), + user_id: creds.grok_user_data + .as_ref() + .map(|data| data.user_id.to_string()), + user_email: creds.grok_user_data + .as_ref() + .map(|data| data.user_email.as_ref()) + .flatten() + .map(|email| email.to_string()), }; let path = self.app_data_root.credentials_dir().get_grok_state_path(); @@ -172,7 +164,8 @@ impl GrokCredentialManager { std::fs::write(path, serialized)?; // TODO: This is just for building the client and testing. - if let Some(cookies_header) = maybe_cookies_header { + if let Some(cookies) = creds.browser_cookies.as_ref() { + let cookies_header = cookies.to_cookie_string(); let path = self.app_data_root.credentials_dir().get_grok_cookies_path(); std::fs::write(path, cookies_header)?; } diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs b/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs index 0121fb9c4e..cb1aed6817 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs @@ -1,55 +1,36 @@ -use std::fs::read_to_string; +use crate::core::state::data_dir::app_data_root::AppDataRoot; use cookie_store::serialized_cookie_store::SerializableCookieStore; use serde_derive::{Deserialize, Serialize}; -use crate::core::state::data_dir::app_data_root::AppDataRoot; +use std::fs::read_to_string; pub (super) const SERIALIZABLE_GROK_STATE_VERSION: u32 = 1; +/// We only need to serialize Grok cookies. +/// Most of the other state (verification_token, etc.) can be pulled at runtime. +/// We will cache the `user_id` and `user_email`. #[derive(Serialize, Deserialize)] pub (super) struct GrokSerializableState { pub (super) version: u32, #[serde(skip_serializing_if = "Option::is_none")] pub (super) user_cookies: Option, - - //#[serde(skip_serializing_if = "Option::is_none")] - //pub (super) user_info: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) user_id: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) user_email: Option, } -//#[derive(Serialize, Deserialize)] -//pub (super) struct SerializableMidjourneyUserInfo { -// #[serde(skip_serializing_if = "Option::is_none")] -// pub (super) user_id: Option, -// -// #[serde(skip_serializing_if = "Option::is_none")] -// pub (super) email: Option, -// -// #[serde(skip_serializing_if = "Option::is_none")] -// pub (super) websocket_token: Option, -//} - -//impl SerializableMidjourneyUserInfo { -// pub fn to_user_info(&self) -> MidjourneyUserInfo { -// MidjourneyUserInfo { -// user_id: self.user_id -// .as_deref() -// .map(|id| MidjourneyUserId::from_str(id)), -// email: self.email.clone(), -// websocket_token: self.websocket_token.clone(), -// } -// } -//} - - impl GrokSerializableState { pub fn read_from_disk(root: &AppDataRoot) -> anyhow::Result> { - let midjourney_state_path= root.credentials_dir().get_grok_state_path(); + let state_path = root.credentials_dir().get_grok_state_path(); - if !midjourney_state_path.exists() { + if !state_path.exists() { return Ok(None); } - let contents = read_to_string(&midjourney_state_path)? + let contents = read_to_string(&state_path)? .trim() .to_string(); @@ -57,4 +38,4 @@ impl GrokSerializableState { Ok(Some(state)) } -} \ No newline at end of file +} diff --git a/crates/desktop/artcraft/src/services/grok/state/mod.rs b/crates/desktop/artcraft/src/services/grok/state/mod.rs index dc62d868a1..9ea2b969bd 100644 --- a/crates/desktop/artcraft/src/services/grok/state/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/state/mod.rs @@ -1,2 +1,3 @@ pub mod grok_credential_manager; -pub mod grok_serializable_state; \ No newline at end of file +pub mod grok_serializable_state; +mod grok_credential_holder; \ No newline at end of file