From 0bb8d84568bbe6093f3a79f71ab174439ec76502 Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Tue, 9 Jun 2026 14:40:17 -0700 Subject: [PATCH] Harden remote image materialization --- codex-rs/core/src/image_url_materializer.rs | 194 +++++++++++--------- 1 file changed, 104 insertions(+), 90 deletions(-) diff --git a/codex-rs/core/src/image_url_materializer.rs b/codex-rs/core/src/image_url_materializer.rs index 3a3c5381a7..afa5556445 100644 --- a/codex-rs/core/src/image_url_materializer.rs +++ b/codex-rs/core/src/image_url_materializer.rs @@ -1,55 +1,28 @@ use std::borrow::Cow; use std::path::Path; use std::sync::LazyLock; +use std::time::Duration; +use anyhow::Context; +use anyhow::Result; +use anyhow::bail; +use codex_login::default_client::build_reqwest_client; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ResponseItem; -use codex_utils_image::ImageProcessingError; use codex_utils_image::PromptImageMode; use codex_utils_image::load_for_prompt_bytes; use reqwest::Client; -use thiserror::Error; +use reqwest::Response; use tracing::warn; use url::Url; +const IMAGE_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(60); +const MAX_IMAGE_DOWNLOAD_BYTES: u64 = 50 * 1024 * 1024; const IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER: &str = "image content omitted because it could not be downloaded or processed"; -static IMAGE_URL_CLIENT: LazyLock = LazyLock::new(Client::new); - -#[derive(Debug, Error)] -pub(crate) enum ImageUrlMaterializationError { - #[error("failed to download image: {0}")] - Download(#[from] reqwest::Error), - #[error("failed to process downloaded image: {0}")] - Processing(#[from] ImageProcessingError), -} - -/// Downloads an HTTP(S) image and returns it as an inline data URL. -/// -/// Non-HTTP(S) values, including existing data URLs, are returned unchanged. -pub(crate) async fn materialize_http_image_url<'a>( - client: &Client, - image_url: &'a str, -) -> Result, ImageUrlMaterializationError> { - let Ok(url) = Url::parse(image_url) else { - return Ok(Cow::Borrowed(image_url)); - }; - if !matches!(url.scheme(), "http" | "https") { - return Ok(Cow::Borrowed(image_url)); - } - - let response = client.get(url.as_str()).send().await?.error_for_status()?; - let bytes = response.bytes().await?; - let image = load_for_prompt_bytes( - Path::new(url.path()), - bytes.to_vec(), - PromptImageMode::Original, - )?; - - Ok(Cow::Owned(image.into_data_url())) -} +static IMAGE_URL_CLIENT: LazyLock = LazyLock::new(build_reqwest_client); /// Materializes HTTP(S) image URLs in a newly recorded batch. /// @@ -65,11 +38,40 @@ pub(crate) async fn materialize_conversation_item_images<'a>( let mut items = items.to_vec(); for item in &mut items { match item { - ResponseItem::Message { content, .. } => materialize_content_items(content).await, + ResponseItem::Message { content, .. } => { + for content_item in content { + let result = match content_item { + ContentItem::InputImage { image_url, .. } => { + materialize_image_url(image_url).await + } + ContentItem::InputText { .. } | ContentItem::OutputText { .. } => continue, + }; + if let Err(err) = result { + warn!(error = %err, "failed to materialize remote message image"); + *content_item = ContentItem::InputText { + text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(), + }; + } + } + } ResponseItem::FunctionCallOutput { output, .. } | ResponseItem::CustomToolCallOutput { output, .. } => { if let Some(content_items) = output.content_items_mut() { - materialize_function_call_output_items(content_items).await; + for content_item in content_items { + let result = match content_item { + FunctionCallOutputContentItem::InputImage { image_url, .. } => { + materialize_image_url(image_url).await + } + FunctionCallOutputContentItem::InputText { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => continue, + }; + if let Err(err) = result { + warn!(error = %err, "failed to materialize remote tool output image"); + *content_item = FunctionCallOutputContentItem::InputText { + text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(), + }; + } + } } } ResponseItem::Reasoning { .. } @@ -94,13 +96,21 @@ pub(crate) async fn materialize_conversation_item_images<'a>( fn response_item_has_http_image(item: &ResponseItem) -> bool { match item { ResponseItem::Message { content, .. } => content.iter().any(|item| match item { - ContentItem::InputImage { image_url, .. } => is_http_url(image_url), + ContentItem::InputImage { image_url, .. } => parse_http_url(image_url).is_some(), ContentItem::InputText { .. } | ContentItem::OutputText { .. } => false, }), ResponseItem::FunctionCallOutput { output, .. } - | ResponseItem::CustomToolCallOutput { output, .. } => output - .content_items() - .is_some_and(|items| items.iter().any(function_call_output_item_has_http_image)), + | ResponseItem::CustomToolCallOutput { output, .. } => { + output.content_items().is_some_and(|items| { + items.iter().any(|item| match item { + FunctionCallOutputContentItem::InputImage { image_url, .. } => { + parse_http_url(image_url).is_some() + } + FunctionCallOutputContentItem::InputText { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => false, + }) + }) + } ResponseItem::Reasoning { .. } | ResponseItem::AgentMessage { .. } | ResponseItem::LocalShellCall { .. } @@ -117,56 +127,60 @@ fn response_item_has_http_image(item: &ResponseItem) -> bool { } } -fn function_call_output_item_has_http_image(item: &FunctionCallOutputContentItem) -> bool { - match item { - FunctionCallOutputContentItem::InputImage { image_url, .. } => is_http_url(image_url), - FunctionCallOutputContentItem::InputText { .. } - | FunctionCallOutputContentItem::EncryptedContent { .. } => false, +fn parse_http_url(value: &str) -> Option { + // Avoid parsing potentially large inline data URLs on the common no-op path. + let (scheme, _) = value.split_once(':')?; + if !scheme.eq_ignore_ascii_case("http") && !scheme.eq_ignore_ascii_case("https") { + return None; } + Url::parse(value).ok() } -fn is_http_url(value: &str) -> bool { - Url::parse(value).is_ok_and(|url| matches!(url.scheme(), "http" | "https")) -} - -async fn materialize_content_items(items: &mut [ContentItem]) { - for item in items { - let result = match item { - ContentItem::InputImage { image_url, .. } => materialize_image_url(image_url).await, - ContentItem::InputText { .. } | ContentItem::OutputText { .. } => continue, - }; - if let Err(err) = result { - warn!(error = %err, "failed to materialize remote message image"); - *item = ContentItem::InputText { - text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(), - }; - } - } -} - -async fn materialize_function_call_output_items(items: &mut [FunctionCallOutputContentItem]) { - for item in items { - let result = match item { - FunctionCallOutputContentItem::InputImage { image_url, .. } => { - materialize_image_url(image_url).await - } - FunctionCallOutputContentItem::InputText { .. } - | FunctionCallOutputContentItem::EncryptedContent { .. } => continue, - }; - if let Err(err) = result { - warn!(error = %err, "failed to materialize remote tool output image"); - *item = FunctionCallOutputContentItem::InputText { - text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(), - }; - } - } -} - -async fn materialize_image_url(image_url: &mut String) -> Result<(), ImageUrlMaterializationError> { - if let Cow::Owned(materialized_url) = - materialize_http_image_url(&IMAGE_URL_CLIENT, image_url).await? - { - *image_url = materialized_url; - } +async fn materialize_image_url(image_url: &mut String) -> Result<()> { + let Some(url) = parse_http_url(image_url) else { + return Ok(()); + }; + let response = IMAGE_URL_CLIENT + .get(url) + .timeout(IMAGE_DOWNLOAD_TIMEOUT) + .send() + .await + .map_err(reqwest::Error::without_url) + .context("failed to download image")? + .error_for_status() + .map_err(reqwest::Error::without_url) + .context("failed to download image")?; + let bytes = read_response_body_with_limit(response).await?; + let image = load_for_prompt_bytes( + Path::new(""), + bytes, + PromptImageMode::Original, + ) + .context("failed to process downloaded image")?; + *image_url = image.into_data_url(); Ok(()) } + +async fn read_response_body_with_limit(mut response: Response) -> Result> { + if response + .content_length() + .is_some_and(|length| length > MAX_IMAGE_DOWNLOAD_BYTES) + { + bail!("downloaded image exceeded the maximum size of {MAX_IMAGE_DOWNLOAD_BYTES} bytes"); + } + + let mut body = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .map_err(reqwest::Error::without_url) + .context("failed to download image")? + { + let next_length = body.len() as u64 + chunk.len() as u64; + if next_length > MAX_IMAGE_DOWNLOAD_BYTES { + bail!("downloaded image exceeded the maximum size of {MAX_IMAGE_DOWNLOAD_BYTES} bytes"); + } + body.extend_from_slice(&chunk); + } + Ok(body) +}