From a7862549018316d5b4606cc0dad2c9903e70c7ea Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 13 Jun 2025 22:23:52 -0400 Subject: [PATCH] better base64 handling --- Cargo.lock | 9 +++++++ Cargo.toml | 2 ++ crates/desktop/artcraft/Cargo.toml | 1 + .../utils/save_base64_image_to_temp_dir.rs | 5 ++-- .../fal_background_removal_command.rs | 4 +++ crates/lib/data/web_base64/Cargo.toml | 25 +++++++++++++++++++ crates/lib/data/web_base64/src/lib.rs | 5 ++++ .../data/web_base64/src/web_base64_decode.rs | 14 +++++++++++ 8 files changed, 62 insertions(+), 3 deletions(-) create mode 100644 crates/lib/data/web_base64/Cargo.toml create mode 100644 crates/lib/data/web_base64/src/lib.rs create mode 100644 crates/lib/data/web_base64/src/web_base64_decode.rs diff --git a/Cargo.lock b/Cargo.lock index c5ffc174b0..df6854720c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -688,6 +688,7 @@ dependencies = [ "tokio", "tokio-util", "url", + "web_base64", "workspace-build-acceleration", ] @@ -12210,6 +12211,14 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web_base64" +version = "0.0.1" +dependencies = [ + "base64 0.22.1", + "workspace-build-acceleration", +] + [[package]] name = "web_scrapers" version = "0.0.1" diff --git a/Cargo.toml b/Cargo.toml index 49125c675a..02d95bec34 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ members = [ "crates/lib/data/crockford", "crates/lib/data/datetimes", "crates/lib/data/primitives", + "crates/lib/data/web_base64", "crates/lib/deprecated/config", # TODO(bt, 2023-02-13): Deprecate or repurpose "crates/lib/deprecated/http_server_common", # TODO(bt, 2023-03-20): Move useful pieces to actix_helpers "crates/lib/dimensions", @@ -228,3 +229,4 @@ shared_service_components = { path = "crates/lib/shared_service_components" } storyteller_client = { path = "crates/lib/api_clients/storyteller_client" } subprocess_common = { path = "crates/lib/subprocess_common" } tokens = { path = "crates/schema/public/tokens" } +web_base64= { path = "crates/lib/data/web_base64" } diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index 494d9ec0fe..7953bc8e35 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -35,6 +35,7 @@ openai_sora_client = { workspace = true } primitives = { workspace = true } storyteller_client = { workspace = true } tokens = { workspace = true } +web_base64 = { workspace = true } # External anyhow = { workspace = true } diff --git a/crates/desktop/artcraft/src/core/utils/save_base64_image_to_temp_dir.rs b/crates/desktop/artcraft/src/core/utils/save_base64_image_to_temp_dir.rs index 65a7363f71..8e5138a1d8 100644 --- a/crates/desktop/artcraft/src/core/utils/save_base64_image_to_temp_dir.rs +++ b/crates/desktop/artcraft/src/core/utils/save_base64_image_to_temp_dir.rs @@ -1,13 +1,12 @@ use crate::core::artcraft_error::ArtcraftError; use crate::core::state::data_dir::app_data_root::AppDataRoot; -use base64::prelude::BASE64_STANDARD; -use base64::Engine; use mimetypes::mimetype_info::mimetype_info::MimetypeInfo; use std::io::Write; use tempfile::NamedTempFile; +use web_base64::web_base64_decode::web_base64_decode; pub async fn save_base64_image_to_temp_dir(app_data_root: &AppDataRoot, base64_image: String) -> Result { - let bytes = BASE64_STANDARD.decode(base64_image)?; + let bytes = web_base64_decode(&base64_image)?; let extension = MimetypeInfo::get_for_bytes(&bytes) .map(|info| info.file_extension()) diff --git a/crates/desktop/artcraft/src/services/fal/commands/fal_background_removal_command.rs b/crates/desktop/artcraft/src/services/fal/commands/fal_background_removal_command.rs index 7bb2e64a26..73f8bbf33c 100644 --- a/crates/desktop/artcraft/src/services/fal/commands/fal_background_removal_command.rs +++ b/crates/desktop/artcraft/src/services/fal/commands/fal_background_removal_command.rs @@ -238,9 +238,11 @@ pub async fn remove_background( let mut temp_download; if let Some(media_token) = request.image_media_token { + info!("From media token: {:?}", &media_token); temp_download = download_media_file_to_temp_dir(&app_data_root, &media_token).await?; } else if let Some(base64_bytes) = request.base64_image { + info!("From base64 bytes..."); temp_download = save_base64_image_to_temp_dir(&app_data_root, base64_bytes).await?; } else { return Err(InnerError::AnyhowError(anyhow!("No image media token or base64 image provided"))); @@ -251,6 +253,8 @@ pub async fn remove_background( let filename = temp_download.path().to_path_buf(); let result = remove_background_rembg(filename, &api_key).await?; + + info!("Unpacking FAL result..."); let extension_with_dot = get_url_file_extension(&result.image_url) .map(|ext| format!(".{}", ext)) diff --git a/crates/lib/data/web_base64/Cargo.toml b/crates/lib/data/web_base64/Cargo.toml new file mode 100644 index 0000000000..ff3818572c --- /dev/null +++ b/crates/lib/data/web_base64/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "web_base64" +edition = "2024" +version = "0.0.1" +publish = false + +[package.metadata.cargo-udeps.ignore] +normal = ["workspace-build-acceleration"] + +[lib] +name = "web_base64" +path = "src/lib.rs" + +[dependencies] +# Workspace hack (faster compile times) +workspace-build-acceleration = { workspace = true } + +# Internal +# (None) + +# External +base64 = { workspace = true } + +[dev-dependencies] +# (None) diff --git a/crates/lib/data/web_base64/src/lib.rs b/crates/lib/data/web_base64/src/lib.rs new file mode 100644 index 0000000000..c5c08aeb78 --- /dev/null +++ b/crates/lib/data/web_base64/src/lib.rs @@ -0,0 +1,5 @@ +//! Base64 encoding and decoding library that handles data from the web. +//! Most of the base64 we receive from the web is of a particular encoding +//! type, and we also have to juggle URL prefixes. + +pub mod web_base64_decode; diff --git a/crates/lib/data/web_base64/src/web_base64_decode.rs b/crates/lib/data/web_base64/src/web_base64_decode.rs new file mode 100644 index 0000000000..f80161118c --- /dev/null +++ b/crates/lib/data/web_base64/src/web_base64_decode.rs @@ -0,0 +1,14 @@ +use base64::prelude::BASE64_STANDARD; +use base64::Engine; + +pub fn web_base64_decode(base64_string: &str) -> Result, base64::DecodeError> { + // Remove the data URL prefix if it exists + // eg: "data:image/png;base64,..." + let base64_data = if base64_string.starts_with("data:") { + base64_string.split(',').nth(1).unwrap_or(base64_string) + } else { + base64_string + }; + + BASE64_STANDARD.decode(base64_data) +}