mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-20 17:02:53 -04:00
artcraft-native: generate with Midjourney or fal (provider choice) (#2115)
Images can run on Midjourney (own account, via a sign-in window) or fal (own API key), chosen from the model selector's Select Provider submenu. Jobs finish into the library. The Midjourney credential manager and the router prompt helper move into shared crates used by both desktop apps.
This commit is contained in:
@@ -6,3 +6,4 @@ pub mod image_generation_draft_context;
|
||||
pub mod image_generation_draft_or_request;
|
||||
pub mod image_generation_request;
|
||||
pub mod providers;
|
||||
pub mod router_image_request_to_artcraft_prompt;
|
||||
|
||||
+12
-9
@@ -1,21 +1,24 @@
|
||||
//! The ArtCraft prompt record for an image request sent to another provider (Midjourney, fal),
|
||||
//! created first so the results can be attributed to it.
|
||||
|
||||
use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
|
||||
use artcraft_router::api::router_aspect_ratio::RouterAspectRatio;
|
||||
use artcraft_router::api::router_image_model::RouterImageModel;
|
||||
use artcraft_router::api::router_resolution::RouterResolution;
|
||||
use artcraft_router::api::router_provider::RouterProvider;
|
||||
use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
|
||||
use enums::common::generation::common_aspect_ratio::CommonAspectRatio as EnumsAspectRatio;
|
||||
use enums::common::generation::common_generation_mode::CommonGenerationMode;
|
||||
use enums::common::generation::common_model_type::CommonModelType;
|
||||
use enums::common::generation::common_resolution::CommonResolution as EnumsResolution;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use uuid_utils::uuid::generate_random_uuid;
|
||||
|
||||
use crate::api::router_aspect_ratio::RouterAspectRatio;
|
||||
use crate::api::router_image_model::RouterImageModel;
|
||||
use crate::api::router_provider::RouterProvider;
|
||||
use crate::api::router_resolution::RouterResolution;
|
||||
use crate::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
|
||||
|
||||
pub fn router_image_request_to_artcraft_prompt(
|
||||
request: &GenerateImageRequestBuilder,
|
||||
) -> CreatePromptRequest {
|
||||
CreatePromptRequest {
|
||||
uuid_idempotency_token: generate_random_uuid(),
|
||||
uuid_idempotency_token: uuid::Uuid::new_v4().to_string(),
|
||||
positive_prompt: request.prompt.clone(),
|
||||
negative_prompt: None,
|
||||
model_type: image_model_to_common_model_type(request.model),
|
||||
@@ -120,7 +123,7 @@ fn router_resolution_to_enums(res: RouterResolution) -> EnumsResolution {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
|
||||
fn base_builder() -> GenerateImageRequestBuilder {
|
||||
GenerateImageRequestBuilder {
|
||||
@@ -165,7 +168,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn edit_mode_when_images_present() {
|
||||
use artcraft_router::api::image_list_ref::ImageListRef;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
let builder = GenerateImageRequestBuilder {
|
||||
image_inputs: Some(ImageListRef::Urls(vec!["https://example.com/img.jpg".to_string()])),
|
||||
..base_builder()
|
||||
@@ -14,14 +14,19 @@ browser_emulation.workspace = true
|
||||
cloudflare_errors.workspace = true
|
||||
cookie_store_wrapper = { workspace = true, features = ["wreq"] }
|
||||
errors.workspace = true
|
||||
uuid_utils.workspace = true
|
||||
|
||||
# External
|
||||
anyhow.workspace = true
|
||||
base64.workspace = true
|
||||
chrono.workspace = true
|
||||
futures.workspace = true
|
||||
log.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
toml.workspace = true
|
||||
url.workspace = true
|
||||
wreq-util.workspace = true
|
||||
wreq.workspace = true
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
//! Reading Midjourney's Firebase auth token (the `__Host-Midjourney.AuthUserTokenV3_i` cookie).
|
||||
//! Its JWT payload is decoded WITHOUT verifying the signature: it's only used to learn the user
|
||||
//! id, the email and when the token expires, never to trust anything.
|
||||
|
||||
use base64::Engine;
|
||||
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::credentials::cookie_names::AUTH_COOKIE_NAME_I;
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
|
||||
/// The claim holding the Midjourney user id (the `singleplayer_{...}` channel id is built from
|
||||
/// it).
|
||||
const MIDJOURNEY_ID_CLAIM: &str = "midjourney_id";
|
||||
/// Refresh this long before the token expires (seconds).
|
||||
const REFRESH_MARGIN_SECS: i64 = 120;
|
||||
|
||||
/// The (unverified) payload of a `header.payload.signature` JWT.
|
||||
pub fn jwt_claims(token: &str) -> Option<Value> {
|
||||
let mut parts = token.split('.');
|
||||
let (_header, payload, _signature) = (parts.next()?, parts.next()?, parts.next()?);
|
||||
if parts.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
serde_json::from_slice(&BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?).ok()
|
||||
}
|
||||
|
||||
/// Whether the token is missing, unreadable, expired or about to expire.
|
||||
pub fn token_needs_refresh(maybe_token: Option<&str>, now: i64) -> bool {
|
||||
maybe_token.and_then(jwt_claims).and_then(|claims| claims.get("exp")?.as_i64()).is_none_or(|expires| expires <= now + REFRESH_MARGIN_SECS)
|
||||
}
|
||||
|
||||
/// The Midjourney user id from the auth token's claims.
|
||||
pub fn user_id_from_token(token: &str) -> Option<MidjourneyUserId> {
|
||||
let claims = jwt_claims(token)?;
|
||||
claims.get(MIDJOURNEY_ID_CLAIM).and_then(Value::as_str).filter(|id| !id.is_empty()).map(MidjourneyUserId::from_str)
|
||||
}
|
||||
|
||||
/// The Midjourney user id from a raw `Cookie:` header (`"a=1; b=2"`), without hitting the
|
||||
/// Cloudflare-gated index page.
|
||||
pub fn extract_midjourney_user_id_from_cookie_header(cookie_header: &str) -> Option<MidjourneyUserId> {
|
||||
let token = cookie_header.split(';').filter_map(|pair| pair.trim().split_once('=')).find(|(name, _)| *name == AUTH_COOKIE_NAME_I).map(|(_, value)| value)?;
|
||||
user_id_from_token(token)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Payload: {"midjourney_id":"26c9d38d-1f71-49c7-a356-b29dac58b54c","name":"u1"}
|
||||
const JWT: &str = "eyJhbGciOiJIUzI1NiJ9.eyJtaWRqb3VybmV5X2lkIjoiMjZjOWQzOGQtMWY3MS00OWM3LWEzNTYtYjI5ZGFjNThiNTRjIiwibmFtZSI6InUxIn0.sig";
|
||||
|
||||
#[test]
|
||||
fn extracts_user_id_from_auth_cookie() {
|
||||
let header = format!("_fbp=abc; {}={}; other=x", AUTH_COOKIE_NAME_I, JWT);
|
||||
let user_id = extract_midjourney_user_id_from_cookie_header(&header).unwrap();
|
||||
assert_eq!(user_id.as_str(), "26c9d38d-1f71-49c7-a356-b29dac58b54c");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_without_auth_cookie() {
|
||||
assert!(extract_midjourney_user_id_from_cookie_header("_fbp=abc; other=x").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refreshes_before_expiry_but_does_not_refresh_a_fresh_token() {
|
||||
assert!(!token_needs_refresh(Some(&token(5000)), 1000));
|
||||
assert!(token_needs_refresh(Some(&token(1120)), 1000));
|
||||
assert!(token_needs_refresh(Some(&token(999)), 1000));
|
||||
assert!(token_needs_refresh(Some("invalid"), 1000));
|
||||
assert!(token_needs_refresh(None, 1000));
|
||||
}
|
||||
|
||||
fn token(exp: i64) -> String {
|
||||
format!("header.{}.signature", BASE64_URL_SAFE_NO_PAD.encode(serde_json::json!({ "exp": exp }).to_string()))
|
||||
}
|
||||
}
|
||||
+7
-3
@@ -1,13 +1,17 @@
|
||||
//! The saved Midjourney login (`midjourney_cookies.toml`): the browser cookies captured at
|
||||
//! sign-in, the account they belong to, and when they were last refreshed or failed.
|
||||
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
|
||||
use browser_emulation::browser_profile::BrowserProfile;
|
||||
use chrono::{DateTime, Utc};
|
||||
use cookie_store_wrapper::cookie_store::CookieStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
use tempfile::NamedTempFile;
|
||||
use uuid_utils::uuid::generate_random_uuid;
|
||||
|
||||
use crate::services::midjourney::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
|
||||
use crate::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
|
||||
|
||||
/// Midjourney's slice of ArtCraftX's credential format. Other accounts retain
|
||||
/// their existing storage. Websocket tokens are deliberately never persisted.
|
||||
@@ -0,0 +1,287 @@
|
||||
//! The Midjourney login shared by the desktop apps: kept in one file, refreshed before it
|
||||
//! expires, and turned into a [`MidjourneySession`] for each submit or poll.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, bail};
|
||||
use browser_emulation::browser_profile::BrowserProfile;
|
||||
use chrono::Utc;
|
||||
use cookie_store_wrapper::cookie_store::CookieStore;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::Mutex;
|
||||
use url::Url;
|
||||
|
||||
use crate::credentials::auth_token::{extract_midjourney_user_id_from_cookie_header, jwt_claims, token_needs_refresh};
|
||||
use crate::credentials::cookie_names::{AUTH_COOKIE_NAME_I, AUTH_COOKIE_NAME_R};
|
||||
use crate::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
|
||||
use crate::credentials::credential_file::{MidjourneyAccountInfo, MidjourneyCredential};
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use crate::endpoints::refresh_token::{DEFAULT_FIREBASE_API_KEY, RefreshTokenArgs, RefreshTokenRequest, refresh_token};
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use crate::recipes::get_user_info::{GetUserInfoArgs, get_user_info};
|
||||
|
||||
/// Where the auth cookies are scoped.
|
||||
pub const MIDJOURNEY_ORIGIN: &str = "https://www.midjourney.com/";
|
||||
/// Token refresh and account lookups give up after this long.
|
||||
const NETWORK_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MidjourneyCredentialManager {
|
||||
credential: Arc<RwLock<Option<MidjourneyCredential>>>,
|
||||
// Serialize refresh, login, and logout so an old refresh cannot resurrect a
|
||||
// signed-out session or overwrite credentials captured by a new login.
|
||||
mutation: Arc<Mutex<()>>,
|
||||
credential_path: PathBuf,
|
||||
/// Older formats, removed on logout so they can't be migrated back.
|
||||
legacy_paths: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
/// What a submit or poll needs: the cookies to send, whose account it is, and the browser to
|
||||
/// present (the one the cookies were captured with).
|
||||
pub struct MidjourneySession {
|
||||
pub credential_id: String,
|
||||
pub cookie_header: String,
|
||||
pub user_id: MidjourneyUserId,
|
||||
pub browser: BrowserProfile,
|
||||
}
|
||||
|
||||
impl MidjourneyCredentialManager {
|
||||
/// A manager that saves to `credential_path`, starting from `credential`.
|
||||
pub fn new(credential_path: PathBuf, legacy_paths: Vec<PathBuf>, credential: Option<MidjourneyCredential>) -> Self {
|
||||
Self { credential: Arc::new(RwLock::new(credential)), mutation: Arc::new(Mutex::new(())), credential_path, legacy_paths }
|
||||
}
|
||||
|
||||
/// A manager for the login saved at `credential_path`, if there is one.
|
||||
pub fn load(credential_path: PathBuf, legacy_paths: Vec<PathBuf>) -> anyhow::Result<Self> {
|
||||
let credential = if credential_path.exists() {
|
||||
let credential = MidjourneyCredential::load(&credential_path)?;
|
||||
credential.save(&credential_path)?; // Persist a stable ID for hand-written files.
|
||||
Some(credential)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(Self::new(credential_path, legacy_paths, credential))
|
||||
}
|
||||
|
||||
pub fn credential_path(&self) -> &PathBuf {
|
||||
&self.credential_path
|
||||
}
|
||||
|
||||
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
|
||||
Ok(self.snapshot()?.map(|credential| credential.cookie.cookies))
|
||||
}
|
||||
|
||||
pub fn maybe_account_info(&self) -> anyhow::Result<Option<MidjourneyAccountInfo>> {
|
||||
Ok(self.snapshot()?.and_then(|credential| credential.maybe_user_info))
|
||||
}
|
||||
|
||||
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
|
||||
Ok(self.snapshot()?.is_some_and(|credential| credential.cookie.maybe_relogin_required_since.is_none() && (cookie_store_has_auth_cookies(&credential.cookie.cookies) || !token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_I), Utc::now().timestamp()))))
|
||||
}
|
||||
|
||||
pub async fn save_login(&self, cookies: CookieStore) -> anyhow::Result<()> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
let mut credential = MidjourneyCredential::new(cookies);
|
||||
populate_account_info(&mut credential);
|
||||
self.store(credential)
|
||||
}
|
||||
|
||||
pub async fn clear_credentials(&self) -> anyhow::Result<()> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
// Remove every format, so a later startup cannot migrate a logged-out account.
|
||||
for path in self.legacy_paths.iter().chain([&self.credential_path]) {
|
||||
match std::fs::remove_file(path) {
|
||||
Ok(()) => {},
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {},
|
||||
Err(err) => return Err(err.into()),
|
||||
}
|
||||
}
|
||||
*self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve identity and renew expiring Firebase credentials before sending
|
||||
/// generation or polling requests. The refresh token can outlive the ID JWT.
|
||||
pub async fn session(&self) -> anyhow::Result<Option<MidjourneySession>> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
let Some(mut credential) = self.snapshot()? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if credential.cookie.maybe_relogin_required_since.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
if token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_I), Utc::now().timestamp()) {
|
||||
let Some(refresh) = credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_R).map(str::to_owned) else {
|
||||
self.mark_relogin(&mut credential)?;
|
||||
return Ok(None);
|
||||
};
|
||||
let refreshed = tokio::time::timeout(NETWORK_TIMEOUT, refresh_token(RefreshTokenArgs { request: RefreshTokenRequest { refresh_token: &refresh, api_key: DEFAULT_FIREBASE_API_KEY }, browser: Some(credential.browser()) })).await?;
|
||||
match refreshed {
|
||||
Ok(response) => {
|
||||
let origin = Url::parse(MIDJOURNEY_ORIGIN)?;
|
||||
// Preserve the host-only Secure scope required by __Host- cookies.
|
||||
for (name, value) in [(AUTH_COOKIE_NAME_I, response.id_token), (AUTH_COOKIE_NAME_R, response.refresh_token)] {
|
||||
credential.cookie.cookies.apply_set_cookie_header(&format!("{name}={value}; Path=/; Secure; HttpOnly"), &origin);
|
||||
}
|
||||
credential.cookie.maybe_updated_at = Some(Utc::now());
|
||||
credential.cookie.maybe_succeeded_at = Some(Utc::now());
|
||||
populate_account_info(&mut credential);
|
||||
self.store(credential.clone())?;
|
||||
},
|
||||
Err(err) if is_revoked_refresh(&err) => {
|
||||
self.mark_relogin(&mut credential)?;
|
||||
return Ok(None);
|
||||
},
|
||||
Err(err) => {
|
||||
credential.cookie.maybe_failed_at = Some(Utc::now());
|
||||
self.store(credential)?;
|
||||
return Err(err.into());
|
||||
},
|
||||
}
|
||||
}
|
||||
populate_account_info(&mut credential);
|
||||
let cookie_header = credential.cookie.cookies.cookie_header_for_url(&Url::parse(MIDJOURNEY_ORIGIN)?).unwrap_or_default();
|
||||
let browser = credential.browser();
|
||||
let mut maybe_user_id = extract_midjourney_user_id_from_cookie_header(&cookie_header).or_else(|| credential.maybe_user_info.as_ref()?.maybe_user_id.as_deref().map(MidjourneyUserId::from_str));
|
||||
if maybe_user_id.is_none() {
|
||||
let info = tokio::time::timeout(NETWORK_TIMEOUT, get_user_info(GetUserInfoArgs { cookie_header: &cookie_header, hostname: None, browser: Some(browser.clone()) })).await??;
|
||||
maybe_user_id = info.user_id;
|
||||
credential.maybe_user_info = Some(MidjourneyAccountInfo { maybe_user_id: maybe_user_id.as_ref().map(|id| id.as_str().to_string()), maybe_email: info.email, maybe_username: None });
|
||||
self.store(credential.clone())?;
|
||||
}
|
||||
let Some(user_id) = maybe_user_id else {
|
||||
bail!("Could not resolve the Midjourney account; please log in again")
|
||||
};
|
||||
Ok(Some(MidjourneySession { credential_id: credential.id, cookie_header, user_id, browser }))
|
||||
}
|
||||
|
||||
fn snapshot(&self) -> anyhow::Result<Option<MidjourneyCredential>> {
|
||||
Ok(self.credential.read().map_err(|_| anyhow!("Midjourney credential lock poisoned"))?.clone())
|
||||
}
|
||||
|
||||
fn store(&self, credential: MidjourneyCredential) -> anyhow::Result<()> {
|
||||
credential.save(&self.credential_path)?;
|
||||
*self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = Some(credential);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn mark_relogin(&self, credential: &mut MidjourneyCredential) -> anyhow::Result<()> {
|
||||
let now = Utc::now();
|
||||
credential.cookie.maybe_failed_at = Some(now);
|
||||
credential.cookie.maybe_relogin_required_since.get_or_insert(now);
|
||||
self.store(credential.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Fills the account's id and email from the auth token's claims.
|
||||
pub fn populate_account_info(credential: &mut MidjourneyCredential) {
|
||||
let Some(claims) = credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_I).and_then(jwt_claims) else {
|
||||
return;
|
||||
};
|
||||
let info = credential.maybe_user_info.get_or_insert_with(MidjourneyAccountInfo::default);
|
||||
if let Some(id) = claims.get("midjourney_id").and_then(Value::as_str).filter(|id| !id.is_empty()) {
|
||||
info.maybe_user_id = Some(id.to_string());
|
||||
}
|
||||
if let Some(email) = claims.get("email").and_then(Value::as_str) {
|
||||
info.maybe_email = Some(email.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn is_revoked_refresh(error: &MidjourneyError) -> bool {
|
||||
let MidjourneyError::Api(MidjourneyApiError::InvalidRequest(body) | MidjourneyApiError::Unauthorized(body)) = error else {
|
||||
return false;
|
||||
};
|
||||
["TOKEN_EXPIRED", "INVALID_REFRESH_TOKEN", "USER_DISABLED", "USER_NOT_FOUND"].iter().any(|code| body.contains(code))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine;
|
||||
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
|
||||
const TEST_USER_ID: &str = "00000000-0000-4000-8000-000000000001";
|
||||
|
||||
mod persistence {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_session_without_refresh_requires_login_and_new_login_clears_it() {
|
||||
let (_dir, manager) = manager();
|
||||
manager.save_login(cookies(10)).await.unwrap();
|
||||
assert!(manager.session().await.unwrap().is_none());
|
||||
assert!(!manager.session_appears_active().unwrap());
|
||||
let dead = MidjourneyCredential::load(manager.credential_path()).unwrap();
|
||||
assert!(dead.cookie.maybe_relogin_required_since.is_some());
|
||||
manager.save_login(cookies(Utc::now().timestamp() + 3600)).await.unwrap();
|
||||
assert!(manager.session().await.unwrap().is_some());
|
||||
assert!(manager.session_appears_active().unwrap());
|
||||
assert!(manager.snapshot().unwrap().unwrap().cookie.maybe_relogin_required_since.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn loads_a_saved_login_and_logout_removes_every_format() {
|
||||
let (dir, manager) = manager();
|
||||
manager.save_login(cookies(Utc::now().timestamp() + 3600)).await.unwrap();
|
||||
let legacy = dir.path().join("midjourney_state.json");
|
||||
std::fs::write(&legacy, "{}").unwrap();
|
||||
let reloaded = MidjourneyCredentialManager::load(manager.credential_path().clone(), vec![legacy.clone()]).unwrap();
|
||||
let session = reloaded.session().await.unwrap().unwrap();
|
||||
assert_eq!(session.user_id.as_str(), TEST_USER_ID);
|
||||
reloaded.clear_credentials().await.unwrap();
|
||||
assert!(!legacy.exists());
|
||||
assert!(!manager.credential_path().exists());
|
||||
assert!(reloaded.session().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cookie_attributes_and_user_agent_survive_toml_round_trip() {
|
||||
let (_dir, manager) = manager();
|
||||
let origin = Url::parse(MIDJOURNEY_ORIGIN).unwrap();
|
||||
let mut cookies = CookieStore::empty();
|
||||
cookies.apply_set_cookie_header("session=host; Path=/; Secure; HttpOnly; SameSite=Lax", &origin);
|
||||
cookies.apply_set_cookie_header("session=domain; Domain=midjourney.com; Path=/api; Secure", &origin);
|
||||
let credential = MidjourneyCredential::new(cookies);
|
||||
let path = manager.credential_path();
|
||||
credential.save(path).unwrap();
|
||||
let loaded = MidjourneyCredential::load(path).unwrap();
|
||||
assert_eq!(loaded.cookie.cookies, credential.cookie.cookies);
|
||||
assert_eq!(loaded.browser(), credential.browser());
|
||||
assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&origin).as_deref(), Some("session=host"));
|
||||
assert!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("http://www.midjourney.com/").unwrap()).is_none());
|
||||
assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("https://midjourney.com/api/test").unwrap()).as_deref(), Some("session=domain"));
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
assert_eq!(std::fs::metadata(path).unwrap().permissions().mode() & 0o777, 0o600);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn temporary_refresh_failures_do_not_force_logout() {
|
||||
assert!(!is_revoked_refresh(&MidjourneyApiError::NetworkError("offline".into()).into()));
|
||||
assert!(!is_revoked_refresh(&MidjourneyApiError::InvalidRequest("API_KEY_INVALID".into()).into()));
|
||||
assert!(is_revoked_refresh(&MidjourneyApiError::InvalidRequest("INVALID_REFRESH_TOKEN".into()).into()));
|
||||
}
|
||||
|
||||
fn manager() -> (TempDir, MidjourneyCredentialManager) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = MidjourneyCredentialManager::new(dir.path().join("midjourney_cookies.toml"), Vec::new(), None);
|
||||
(dir, manager)
|
||||
}
|
||||
|
||||
fn cookies(exp: i64) -> CookieStore {
|
||||
CookieStore::from_cookie_header(&format!("{AUTH_COOKIE_NAME_I}={}", token(exp)), &Url::parse(MIDJOURNEY_ORIGIN).unwrap())
|
||||
}
|
||||
|
||||
fn token(exp: i64) -> String {
|
||||
let claims = serde_json::json!({"exp": exp, "midjourney_id": TEST_USER_ID});
|
||||
format!("header.{}.signature", BASE64_URL_SAFE_NO_PAD.encode(claims.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,7 @@
|
||||
pub mod auth_token;
|
||||
pub mod cookie_names;
|
||||
pub mod cookie_store_has_auth_cookies;
|
||||
pub mod credential_file;
|
||||
pub mod credential_manager;
|
||||
pub mod jwt_decode;
|
||||
pub mod midjourney_user_id;
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ pub const MIDJOURNEY_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X
|
||||
/// traffic (submit, imagine, index page, image downloads, websocket upgrade).
|
||||
///
|
||||
/// Safari fingerprint + the exact [`MIDJOURNEY_USER_AGENT`], to stay consistent
|
||||
/// with the macOS WKWebView login window that captured the cookies.
|
||||
/// with the login webview that captured the cookies.
|
||||
pub fn midjourney_browser_profile() -> BrowserProfile {
|
||||
BrowserProfile::safari_macos_with_user_agent(MIDJOURNEY_USER_AGENT)
|
||||
}
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod get_image_url;
|
||||
pub mod image_downloader_client;
|
||||
pub mod image_downloader_client;
|
||||
pub mod midjourney_browser_profile;
|
||||
@@ -1,2 +1 @@
|
||||
pub mod router_image_request_to_artcraft_prompt;
|
||||
pub mod router_video_request_to_artcraft_prompt;
|
||||
|
||||
+1
-1
@@ -19,7 +19,7 @@ use crate::core::api_adapters::models::image::tauri_image_model_to_generation_mo
|
||||
use crate::core::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model;
|
||||
use crate::core::commands::enqueue::generate_error::GenerateError;
|
||||
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
|
||||
use crate::core::commands::generate::common::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
|
||||
use artcraft_router::generate::generate_image::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
|
||||
use crate::core::commands::generate::generate_image::providers::artcraft_router::utils::convert_enums_to_router::{convert_aspect_ratio, convert_quality, convert_resolution};
|
||||
use crate::core::commands::generate::generate_image::providers::artcraft_router::utils::map_media_files_to_urls::map_media_file_tokens_to_cdn_urls;
|
||||
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@ use crate::core::api_adapters::models::image::tauri_image_model_to_generation_mo
|
||||
use crate::core::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model;
|
||||
use crate::core::commands::enqueue::generate_error::{BadInputReason, GenerateError};
|
||||
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
|
||||
use crate::core::commands::generate::common::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
|
||||
use artcraft_router::generate::generate_image::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
|
||||
use crate::core::commands::generate::generate_image::providers::artcraft_router::utils::convert_enums_to_router::convert_aspect_ratio;
|
||||
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
|
||||
use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel;
|
||||
|
||||
+82
-297
@@ -1,337 +1,126 @@
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::services::midjourney::state::midjourney_credential::{MidjourneyAccountInfo, MidjourneyCredential};
|
||||
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
|
||||
use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyState;
|
||||
use crate::services::midjourney::utils::extract_midjourney_user_id_from_cookies::extract_midjourney_user_id_from_cookie_header;
|
||||
use anyhow::{anyhow, bail};
|
||||
use base64::prelude::{Engine, BASE64_URL_SAFE_NO_PAD};
|
||||
use browser_emulation::browser_profile::BrowserProfile;
|
||||
use chrono::Utc;
|
||||
//! The app's Midjourney login: the shared manager (`midjourney_client`, also used by the native
|
||||
//! app) kept in the app's credentials directory, plus the one-time migration from the legacy
|
||||
//! `midjourney_state.json`.
|
||||
|
||||
use std::ops::Deref;
|
||||
|
||||
use cookie_store_wrapper::cookie_store::CookieStore;
|
||||
use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
|
||||
use midjourney_client::credentials::credential_file::{MidjourneyAccountInfo, MidjourneyCredential};
|
||||
use midjourney_client::credentials::credential_manager::{MIDJOURNEY_ORIGIN, MidjourneyCredentialManager as SharedCredentialManager, populate_account_info};
|
||||
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use midjourney_client::endpoints::refresh_token::{refresh_token, RefreshTokenArgs, RefreshTokenRequest, DEFAULT_FIREBASE_API_KEY};
|
||||
use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
|
||||
use midjourney_client::error::midjourney_error::MidjourneyError;
|
||||
use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoArgs};
|
||||
use serde_json::Value;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Mutex;
|
||||
use url::Url;
|
||||
|
||||
const AUTH_ID: &str = "__Host-Midjourney.AuthUserTokenV3_i";
|
||||
const AUTH_REFRESH: &str = "__Host-Midjourney.AuthUserTokenV3_r";
|
||||
const MIDJOURNEY_ORIGIN: &str = "https://www.midjourney.com/";
|
||||
pub use midjourney_client::credentials::credential_manager::MidjourneySession;
|
||||
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
|
||||
use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyState;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MidjourneyCredentialManager {
|
||||
credential: Arc<RwLock<Option<MidjourneyCredential>>>,
|
||||
// Serialize refresh, login, and logout so an old refresh cannot resurrect a
|
||||
// signed-out session or overwrite credentials captured by a new login.
|
||||
mutation: Arc<Mutex<()>>,
|
||||
app_data_root: AppDataRoot,
|
||||
}
|
||||
pub struct MidjourneyCredentialManager(SharedCredentialManager);
|
||||
|
||||
pub struct MidjourneySession {
|
||||
pub credential_id: String,
|
||||
pub cookie_header: String,
|
||||
pub user_id: MidjourneyUserId,
|
||||
pub browser: BrowserProfile,
|
||||
impl Deref for MidjourneyCredentialManager {
|
||||
type Target = SharedCredentialManager;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl MidjourneyCredentialManager {
|
||||
pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self {
|
||||
Self { credential: Arc::new(RwLock::new(None)), mutation: Arc::new(Mutex::new(())), app_data_root: app_data_root.clone() }
|
||||
let dir = app_data_root.credentials_dir();
|
||||
Self(SharedCredentialManager::new(dir.get_midjourney_credential_path(), vec![dir.get_midjourney_state_path()], None))
|
||||
}
|
||||
|
||||
pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
|
||||
let manager = Self::initialize_empty(app_data_root);
|
||||
match load_or_migrate(app_data_root) {
|
||||
Ok(credential) => {
|
||||
if let Ok(mut current) = manager.credential.write() {
|
||||
*current = credential;
|
||||
}
|
||||
Ok(manager) => manager,
|
||||
Err(err) => {
|
||||
log::warn!("Could not load Midjourney credentials: {}", err);
|
||||
Self::initialize_empty(app_data_root)
|
||||
},
|
||||
Err(err) => log::warn!("Could not load Midjourney credentials: {}", err),
|
||||
}
|
||||
manager
|
||||
}
|
||||
|
||||
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
|
||||
Ok(self.snapshot()?.map(|credential| credential.cookie.cookies))
|
||||
}
|
||||
|
||||
pub fn maybe_copy_user_info(&self) -> anyhow::Result<Option<MidjourneyUserInfo>> {
|
||||
Ok(self.snapshot()?.and_then(|credential| credential.maybe_user_info).map(|info| MidjourneyUserInfo { user_id: info.maybe_user_id.map(MidjourneyUserId::from_string), email: info.maybe_email, websocket_token: None }))
|
||||
}
|
||||
|
||||
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
|
||||
Ok(self.snapshot()?.is_some_and(|credential| credential.cookie.maybe_relogin_required_since.is_none() && (cookie_store_has_auth_cookies(&credential.cookie.cookies) || !token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_ID), Utc::now().timestamp()))))
|
||||
}
|
||||
|
||||
pub async fn save_login(&self, cookies: CookieStore) -> anyhow::Result<()> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
let mut credential = MidjourneyCredential::new(cookies);
|
||||
populate_account_info(&mut credential);
|
||||
self.store(credential)
|
||||
}
|
||||
|
||||
pub async fn clear_credentials(&self) -> anyhow::Result<()> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
// Remove both formats, so a later startup cannot migrate a logged-out account.
|
||||
for path in [self.app_data_root.credentials_dir().get_midjourney_state_path(), self.credential_path()] {
|
||||
match std::fs::remove_file(path) {
|
||||
Ok(()) => {},
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {},
|
||||
Err(err) => return Err(err.into()),
|
||||
}
|
||||
}
|
||||
*self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve identity and renew expiring Firebase credentials before sending
|
||||
/// generation or polling requests. The refresh token can outlive the ID JWT.
|
||||
pub async fn session(&self) -> anyhow::Result<Option<MidjourneySession>> {
|
||||
let _guard = self.mutation.lock().await;
|
||||
let Some(mut credential) = self.snapshot()? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if credential.cookie.maybe_relogin_required_since.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
if token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_ID), Utc::now().timestamp()) {
|
||||
let Some(refresh) = credential.cookie.cookies.get_cookie_value(AUTH_REFRESH).map(str::to_owned) else {
|
||||
self.mark_relogin(&mut credential)?;
|
||||
return Ok(None);
|
||||
};
|
||||
let refreshed = tokio::time::timeout(Duration::from_secs(30), refresh_token(RefreshTokenArgs { request: RefreshTokenRequest { refresh_token: &refresh, api_key: DEFAULT_FIREBASE_API_KEY }, browser: Some(credential.browser()) })).await?;
|
||||
match refreshed {
|
||||
Ok(response) => {
|
||||
let origin = Url::parse(MIDJOURNEY_ORIGIN)?;
|
||||
// Preserve the host-only Secure scope required by __Host- cookies.
|
||||
for (name, value) in [(AUTH_ID, response.id_token), (AUTH_REFRESH, response.refresh_token)] {
|
||||
credential.cookie.cookies.apply_set_cookie_header(&format!("{name}={value}; Path=/; Secure; HttpOnly"), &origin);
|
||||
}
|
||||
credential.cookie.maybe_updated_at = Some(Utc::now());
|
||||
credential.cookie.maybe_succeeded_at = Some(Utc::now());
|
||||
populate_account_info(&mut credential);
|
||||
self.store(credential.clone())?;
|
||||
},
|
||||
Err(err) if is_revoked_refresh(&err) => {
|
||||
self.mark_relogin(&mut credential)?;
|
||||
return Ok(None);
|
||||
},
|
||||
Err(err) => {
|
||||
credential.cookie.maybe_failed_at = Some(Utc::now());
|
||||
self.store(credential)?;
|
||||
return Err(err.into());
|
||||
},
|
||||
}
|
||||
}
|
||||
populate_account_info(&mut credential);
|
||||
let cookie_header = credential.cookie.cookies.cookie_header_for_url(&Url::parse(MIDJOURNEY_ORIGIN)?).unwrap_or_default();
|
||||
let browser = credential.browser();
|
||||
let mut maybe_user_id = extract_midjourney_user_id_from_cookie_header(&cookie_header).or_else(|| credential.maybe_user_info.as_ref()?.maybe_user_id.as_deref().map(MidjourneyUserId::from_str));
|
||||
if maybe_user_id.is_none() {
|
||||
let info = tokio::time::timeout(Duration::from_secs(30), get_user_info(GetUserInfoArgs { cookie_header: &cookie_header, hostname: None, browser: Some(browser.clone()) })).await??;
|
||||
maybe_user_id = info.user_id;
|
||||
credential.maybe_user_info = Some(MidjourneyAccountInfo { maybe_user_id: maybe_user_id.as_ref().map(|id| id.as_str().to_string()), maybe_email: info.email, maybe_username: None });
|
||||
self.store(credential.clone())?;
|
||||
}
|
||||
let Some(user_id) = maybe_user_id else {
|
||||
bail!("Could not resolve the Midjourney account; please log in again")
|
||||
};
|
||||
Ok(Some(MidjourneySession { credential_id: credential.id, cookie_header, user_id, browser }))
|
||||
}
|
||||
|
||||
fn snapshot(&self) -> anyhow::Result<Option<MidjourneyCredential>> {
|
||||
Ok(self.credential.read().map_err(|_| anyhow!("Midjourney credential lock poisoned"))?.clone())
|
||||
}
|
||||
|
||||
fn store(&self, credential: MidjourneyCredential) -> anyhow::Result<()> {
|
||||
credential.save(&self.credential_path())?;
|
||||
*self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = Some(credential);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn mark_relogin(&self, credential: &mut MidjourneyCredential) -> anyhow::Result<()> {
|
||||
let now = Utc::now();
|
||||
credential.cookie.maybe_failed_at = Some(now);
|
||||
credential.cookie.maybe_relogin_required_since.get_or_insert(now);
|
||||
self.store(credential.clone())
|
||||
}
|
||||
|
||||
fn credential_path(&self) -> std::path::PathBuf {
|
||||
self.app_data_root.credentials_dir().get_midjourney_credential_path()
|
||||
Ok(self.maybe_account_info()?.map(|info| MidjourneyUserInfo { user_id: info.maybe_user_id.map(MidjourneyUserId::from_string), email: info.maybe_email, websocket_token: None }))
|
||||
}
|
||||
}
|
||||
|
||||
fn load_or_migrate(root: &AppDataRoot) -> anyhow::Result<Option<MidjourneyCredential>> {
|
||||
let path = root.credentials_dir().get_midjourney_credential_path();
|
||||
if path.exists() {
|
||||
let credential = MidjourneyCredential::load(&path)?;
|
||||
credential.save(&path)?; // Persist a stable ID for hand-written files.
|
||||
return Ok(Some(credential));
|
||||
fn load_or_migrate(root: &AppDataRoot) -> anyhow::Result<MidjourneyCredentialManager> {
|
||||
let dir = root.credentials_dir();
|
||||
let (path, legacy_path) = (dir.get_midjourney_credential_path(), dir.get_midjourney_state_path());
|
||||
if path.exists() || !legacy_path.exists() {
|
||||
return Ok(MidjourneyCredentialManager(SharedCredentialManager::load(path, vec![legacy_path])?));
|
||||
}
|
||||
let legacy_path = root.credentials_dir().get_midjourney_state_path();
|
||||
if !legacy_path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let legacy: SerializableMidjourneyState = serde_json::from_str(&std::fs::read_to_string(legacy_path)?)?;
|
||||
let legacy: SerializableMidjourneyState = serde_json::from_str(&std::fs::read_to_string(&legacy_path)?)?;
|
||||
let Some(cookies) = legacy.user_cookies else {
|
||||
return Ok(None);
|
||||
return Ok(MidjourneyCredentialManager::initialize_empty(root));
|
||||
};
|
||||
let cookies = CookieStore::from_cookie_header(&cookies.to_cookie_store().to_cookie_string(), &Url::parse(MIDJOURNEY_ORIGIN)?);
|
||||
let mut credential = MidjourneyCredential::new(cookies);
|
||||
credential.maybe_user_info = legacy.user_info.map(|info| MidjourneyAccountInfo { maybe_user_id: info.user_id, maybe_email: info.email, maybe_username: None });
|
||||
populate_account_info(&mut credential);
|
||||
credential.save(&path)?;
|
||||
Ok(Some(credential))
|
||||
}
|
||||
|
||||
fn populate_account_info(credential: &mut MidjourneyCredential) {
|
||||
let Some(claims) = credential.cookie.cookies.get_cookie_value(AUTH_ID).and_then(jwt_claims) else {
|
||||
return;
|
||||
};
|
||||
let info = credential.maybe_user_info.get_or_insert_with(MidjourneyAccountInfo::default);
|
||||
if let Some(id) = claims.get("midjourney_id").and_then(Value::as_str).filter(|id| !id.is_empty()) {
|
||||
info.maybe_user_id = Some(id.to_string());
|
||||
}
|
||||
if let Some(email) = claims.get("email").and_then(Value::as_str) {
|
||||
info.maybe_email = Some(email.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn token_needs_refresh(maybe_token: Option<&str>, now: i64) -> bool {
|
||||
maybe_token.and_then(jwt_claims).and_then(|claims| claims.get("exp")?.as_i64()).is_none_or(|expires| expires <= now + 120)
|
||||
}
|
||||
|
||||
fn jwt_claims(token: &str) -> Option<Value> {
|
||||
let payload = token.split('.').nth(1)?;
|
||||
serde_json::from_slice(&BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?).ok()
|
||||
}
|
||||
|
||||
fn is_revoked_refresh(error: &MidjourneyError) -> bool {
|
||||
let MidjourneyError::Api(MidjourneyApiError::InvalidRequest(body) | MidjourneyApiError::Unauthorized(body)) = error else {
|
||||
return false;
|
||||
};
|
||||
["TOKEN_EXPIRED", "INVALID_REFRESH_TOKEN", "USER_DISABLED", "USER_NOT_FOUND"].iter().any(|code| body.contains(code))
|
||||
Ok(MidjourneyCredentialManager(SharedCredentialManager::new(path, vec![legacy_path], Some(credential))))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine};
|
||||
use chrono::Utc;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
|
||||
const TEST_USER_ID: &str = "00000000-0000-4000-8000-000000000001";
|
||||
const AUTH_ID: &str = "__Host-Midjourney.AuthUserTokenV3_i";
|
||||
const AUTH_REFRESH: &str = "__Host-Midjourney.AuthUserTokenV3_r";
|
||||
|
||||
mod persistence {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn migrates_legacy_cookies_and_logout_does_not_touch_artcraft() {
|
||||
let (_dir, root) = root();
|
||||
let legacy = root.credentials_dir().get_midjourney_state_path();
|
||||
let artcraft = root.credentials_dir().get_storyteller_session_cookie_file_path();
|
||||
std::fs::write(&artcraft, "existing-artcraft-session").unwrap();
|
||||
std::fs::write(
|
||||
&legacy,
|
||||
serde_json::json!({
|
||||
"version": 1,
|
||||
"user_cookies": {"cookies": [
|
||||
{"name": AUTH_ID, "value": token(Utc::now().timestamp() + 3600)},
|
||||
{"name": AUTH_REFRESH, "value": "test-refresh"}
|
||||
]},
|
||||
"user_info": {"user_id": TEST_USER_ID, "email": "test@example.invalid", "websocket_token": "old-socket-token"}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
let manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
let session = manager.session().await.unwrap().unwrap();
|
||||
assert_eq!(session.user_id.as_str(), TEST_USER_ID);
|
||||
let contents = std::fs::read_to_string(root.credentials_dir().get_midjourney_credential_path()).unwrap();
|
||||
assert!(contents.contains("midjourney_cookies"));
|
||||
assert!(!contents.contains("old-socket-token"));
|
||||
assert!(legacy.exists(), "migration preserves the old file until logout");
|
||||
let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
assert_eq!(reloaded.session().await.unwrap().unwrap().credential_id, session.credential_id);
|
||||
manager.clear_credentials().await.unwrap();
|
||||
assert!(!legacy.exists());
|
||||
assert!(!root.credentials_dir().get_midjourney_credential_path().exists());
|
||||
assert_eq!(std::fs::read_to_string(artcraft).unwrap(), "existing-artcraft-session");
|
||||
let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
assert!(reloaded.session().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_session_without_refresh_requires_login_and_new_login_clears_it() {
|
||||
let (_dir, root) = root();
|
||||
let manager = MidjourneyCredentialManager::initialize_empty(&root);
|
||||
manager.save_login(cookies(10)).await.unwrap();
|
||||
assert!(manager.session().await.unwrap().is_none());
|
||||
assert!(!manager.session_appears_active().unwrap());
|
||||
let dead = MidjourneyCredential::load(&root.credentials_dir().get_midjourney_credential_path()).unwrap();
|
||||
assert!(dead.cookie.maybe_relogin_required_since.is_some());
|
||||
manager.save_login(cookies(Utc::now().timestamp() + 3600)).await.unwrap();
|
||||
assert!(manager.session().await.unwrap().is_some());
|
||||
assert!(manager.session_appears_active().unwrap());
|
||||
assert!(manager.snapshot().unwrap().unwrap().cookie.maybe_relogin_required_since.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cookie_attributes_and_user_agent_survive_toml_round_trip() {
|
||||
let (_dir, root) = root();
|
||||
let origin = Url::parse(MIDJOURNEY_ORIGIN).unwrap();
|
||||
let mut cookies = CookieStore::empty();
|
||||
cookies.apply_set_cookie_header("session=host; Path=/; Secure; HttpOnly; SameSite=Lax", &origin);
|
||||
cookies.apply_set_cookie_header("session=domain; Domain=midjourney.com; Path=/api; Secure", &origin);
|
||||
let credential = MidjourneyCredential::new(cookies);
|
||||
let path = root.credentials_dir().get_midjourney_credential_path();
|
||||
credential.save(&path).unwrap();
|
||||
let loaded = MidjourneyCredential::load(&path).unwrap();
|
||||
assert_eq!(loaded.cookie.cookies, credential.cookie.cookies);
|
||||
assert_eq!(loaded.browser(), credential.browser());
|
||||
assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&origin).as_deref(), Some("session=host"));
|
||||
assert!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("http://www.midjourney.com/").unwrap()).is_none());
|
||||
assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("https://midjourney.com/api/test").unwrap()).as_deref(), Some("session=domain"));
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
assert_eq!(std::fs::metadata(path).unwrap().permissions().mode() & 0o777, 0o600);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupt_modern_credential_does_not_revive_legacy_state() {
|
||||
let (_dir, root) = root();
|
||||
std::fs::write(root.credentials_dir().get_midjourney_credential_path(), "not valid toml").unwrap();
|
||||
std::fs::write(root.credentials_dir().get_midjourney_state_path(), "{}").unwrap();
|
||||
assert!(load_or_migrate(&root).is_err());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn migrates_legacy_cookies_and_logout_does_not_touch_artcraft() {
|
||||
let (_dir, root) = root();
|
||||
let legacy = root.credentials_dir().get_midjourney_state_path();
|
||||
let artcraft = root.credentials_dir().get_storyteller_session_cookie_file_path();
|
||||
std::fs::write(&artcraft, "existing-artcraft-session").unwrap();
|
||||
std::fs::write(
|
||||
&legacy,
|
||||
serde_json::json!({
|
||||
"version": 1,
|
||||
"user_cookies": {"cookies": [
|
||||
{"name": AUTH_ID, "value": token(Utc::now().timestamp() + 3600)},
|
||||
{"name": AUTH_REFRESH, "value": "test-refresh"}
|
||||
]},
|
||||
"user_info": {"user_id": TEST_USER_ID, "email": "test@example.invalid", "websocket_token": "old-socket-token"}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
let manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
let session = manager.session().await.unwrap().unwrap();
|
||||
assert_eq!(session.user_id.as_str(), TEST_USER_ID);
|
||||
let contents = std::fs::read_to_string(root.credentials_dir().get_midjourney_credential_path()).unwrap();
|
||||
assert!(contents.contains("midjourney_cookies"));
|
||||
assert!(!contents.contains("old-socket-token"));
|
||||
assert!(legacy.exists(), "migration preserves the old file until logout");
|
||||
let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
assert_eq!(reloaded.session().await.unwrap().unwrap().credential_id, session.credential_id);
|
||||
manager.clear_credentials().await.unwrap();
|
||||
assert!(!legacy.exists());
|
||||
assert!(!root.credentials_dir().get_midjourney_credential_path().exists());
|
||||
assert_eq!(std::fs::read_to_string(artcraft).unwrap(), "existing-artcraft-session");
|
||||
let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
|
||||
assert!(reloaded.session().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
mod refresh {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn refreshes_before_expiry_but_does_not_refresh_a_fresh_token() {
|
||||
assert!(!token_needs_refresh(Some(&token(5000)), 1000));
|
||||
assert!(token_needs_refresh(Some(&token(1120)), 1000));
|
||||
assert!(token_needs_refresh(Some(&token(999)), 1000));
|
||||
assert!(token_needs_refresh(Some("invalid"), 1000));
|
||||
assert!(token_needs_refresh(None, 1000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn temporary_refresh_failures_do_not_force_logout() {
|
||||
assert!(!is_revoked_refresh(&MidjourneyApiError::NetworkError("offline".into()).into()));
|
||||
assert!(!is_revoked_refresh(&MidjourneyApiError::InvalidRequest("API_KEY_INVALID".into()).into()));
|
||||
assert!(is_revoked_refresh(&MidjourneyApiError::InvalidRequest("INVALID_REFRESH_TOKEN".into()).into()));
|
||||
}
|
||||
#[test]
|
||||
fn corrupt_modern_credential_does_not_revive_legacy_state() {
|
||||
let (_dir, root) = root();
|
||||
std::fs::write(root.credentials_dir().get_midjourney_credential_path(), "not valid toml").unwrap();
|
||||
std::fs::write(root.credentials_dir().get_midjourney_state_path(), "{}").unwrap();
|
||||
assert!(load_or_migrate(&root).is_err());
|
||||
}
|
||||
|
||||
fn root() -> (TempDir, AppDataRoot) {
|
||||
@@ -341,10 +130,6 @@ mod tests {
|
||||
(dir, root)
|
||||
}
|
||||
|
||||
fn cookies(exp: i64) -> CookieStore {
|
||||
CookieStore::from_cookie_header(&format!("{AUTH_ID}={}", token(exp)), &Url::parse(MIDJOURNEY_ORIGIN).unwrap())
|
||||
}
|
||||
|
||||
fn token(exp: i64) -> String {
|
||||
let claims = serde_json::json!({"exp": exp, "midjourney_id": TEST_USER_ID});
|
||||
format!("header.{}.signature", BASE64_URL_SAFE_NO_PAD.encode(claims.to_string()))
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
pub (super) mod serializable_midjourney_state;
|
||||
pub mod midjourney_credential_manager;
|
||||
pub mod midjourney_user_info;
|
||||
pub mod midjourney_credential;
|
||||
|
||||
+1
-1
@@ -169,7 +169,7 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::core::commands::generate::generate_image::providers::midjourney::handle_midjourney::handle_midjourney;
|
||||
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
|
||||
use crate::services::midjourney::state::midjourney_credential::MidjourneyCredential;
|
||||
use midjourney_client::credentials::credential_file::MidjourneyCredential;
|
||||
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
|
||||
use artcraft_client::utils::api_host::ApiHost;
|
||||
|
||||
-62
@@ -1,62 +0,0 @@
|
||||
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
|
||||
use base64::Engine;
|
||||
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use serde_json::Value;
|
||||
|
||||
/// The Midjourney auth id-token cookie. Its (unverified) JWT payload carries
|
||||
/// the Midjourney user id under `midjourney_id`.
|
||||
const AUTH_ID_TOKEN_COOKIE: &str = "__Host-Midjourney.AuthUserTokenV3_i";
|
||||
|
||||
/// The JWT claim holding the Midjourney user id (the `singleplayer_{...}`
|
||||
/// channel id is built from it).
|
||||
const MIDJOURNEY_ID_CLAIM: &str = "midjourney_id";
|
||||
|
||||
/// Recover the Midjourney user id from the auth cookie's JWT, WITHOUT hitting
|
||||
/// the Cloudflare-gated index page. `cookie_header` is a raw `Cookie:` header
|
||||
/// string (`"a=1; b=2"`).
|
||||
pub fn extract_midjourney_user_id_from_cookie_header(cookie_header: &str) -> Option<MidjourneyUserId> {
|
||||
let token = cookie_value(cookie_header, AUTH_ID_TOKEN_COOKIE)?;
|
||||
let claims = decode_jwt_payload(token)?;
|
||||
let id = claims.get(MIDJOURNEY_ID_CLAIM).and_then(Value::as_str)?;
|
||||
if id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(MidjourneyUserId::from_str(id))
|
||||
}
|
||||
|
||||
fn cookie_value<'a>(cookie_header: &'a str, name: &str) -> Option<&'a str> {
|
||||
cookie_header.split(';').filter_map(|pair| pair.trim().split_once('=')).find(|(cookie_name, _)| *cookie_name == name).map(|(_, value)| value)
|
||||
}
|
||||
|
||||
/// Decode the (unverified) payload of a `header.payload.signature` JWT.
|
||||
fn decode_jwt_payload(token: &str) -> Option<Value> {
|
||||
let mut parts = token.split('.');
|
||||
let _header = parts.next()?;
|
||||
let payload = parts.next()?;
|
||||
let _signature = parts.next()?;
|
||||
if parts.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
let bytes = BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?;
|
||||
serde_json::from_slice(&bytes).ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Payload: {"midjourney_id":"26c9d38d-1f71-49c7-a356-b29dac58b54c","name":"u1"}
|
||||
const JWT: &str = "eyJhbGciOiJIUzI1NiJ9.eyJtaWRqb3VybmV5X2lkIjoiMjZjOWQzOGQtMWY3MS00OWM3LWEzNTYtYjI5ZGFjNThiNTRjIiwibmFtZSI6InUxIn0.sig";
|
||||
|
||||
#[test]
|
||||
fn extracts_user_id_from_auth_cookie() {
|
||||
let header = format!("_fbp=abc; {}={}; other=x", AUTH_ID_TOKEN_COOKIE, JWT);
|
||||
let user_id = extract_midjourney_user_id_from_cookie_header(&header).unwrap();
|
||||
assert_eq!(user_id.as_str(), "26c9d38d-1f71-49c7-a356-b29dac58b54c");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_without_auth_cookie() {
|
||||
assert!(extract_midjourney_user_id_from_cookie_header("_fbp=abc; other=x").is_none());
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1 @@
|
||||
pub mod download_midjourney_image;
|
||||
pub mod midjourney_browser_profile;
|
||||
pub mod extract_midjourney_user_id_from_cookies;
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
use crate::services::midjourney::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
|
||||
use midjourney_client::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::utils::clear_all_webview_cookies::clear_all_webview_cookies;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
|
||||
Reference in New Issue
Block a user