From d5814d841a11540fa130455047c0ffb838a3b981 Mon Sep 17 00:00:00 2001 From: Uday Chandra Date: Tue, 5 May 2026 15:32:47 +0000 Subject: [PATCH] Harden agent attachment drops --- crates/agent_ui/src/mention_set.rs | 107 ++++++++++++++++++++++---- crates/agent_ui/src/message_editor.rs | 46 +++++++---- 2 files changed, 122 insertions(+), 31 deletions(-) diff --git a/crates/agent_ui/src/mention_set.rs b/crates/agent_ui/src/mention_set.rs index 808111f1eb..e8a6036df5 100644 --- a/crates/agent_ui/src/mention_set.rs +++ b/crates/agent_ui/src/mention_set.rs @@ -663,7 +663,7 @@ mod tests { use semver::Version; use serde_json::json; use settings::SettingsStore; - use std::path::Path; + use std::path::{Path, PathBuf}; use theme; use util::path; @@ -757,6 +757,49 @@ mod tests { assert!(!is_raster_image_path(Path::new("/tmp/notes.txt"))); assert!(!is_raster_image_path(Path::new("/tmp/README"))); } + + #[test] + fn test_load_external_image_rejects_invalid_image_extension() { + let path = temp_path("png"); + std::fs::write(&path, b"not an image").expect("write temp image"); + + let result = load_external_image_from_path(&path, &"Image".into()); + std::fs::remove_file(&path).expect("remove temp image"); + + assert!(result.is_err()); + } + + #[test] + fn test_load_external_image_ignores_non_image_extension() { + let path = temp_path("txt"); + std::fs::write(&path, b"plain text").expect("write temp text"); + + let result = load_external_image_from_path(&path, &"Image".into()); + std::fs::remove_file(&path).expect("remove temp text"); + + assert!(matches!(result, Ok(None))); + } + + #[test] + fn test_load_external_image_rejects_oversized_image_extension() { + let path = temp_path("png"); + let file = std::fs::File::create(&path).expect("create temp image"); + file.set_len(MAX_EXTERNAL_ATTACHMENT_BYTES + 1) + .expect("resize temp image"); + + let result = load_external_image_from_path(&path, &"Image".into()); + std::fs::remove_file(&path).expect("remove temp image"); + + assert!(result.is_err()); + } + + fn temp_path(extension: &str) -> PathBuf { + std::env::temp_dir().join(format!( + "mast-agent-ui-test-{}.{}", + uuid::Uuid::new_v4(), + extension + )) + } } /// Inserts a list of images into the editor as context mentions. @@ -882,7 +925,7 @@ fn image_format_from_external_content(format: image::ImageFormat) -> Option bool { +pub(crate) fn is_raster_image_path(path: &Path) -> bool { let Some(extension) = path.extension().and_then(OsStr::to_str) else { return false; }; @@ -905,18 +948,41 @@ pub(crate) fn is_pdf_path(path: &Path) -> bool { pub(crate) fn load_external_image_from_path( path: &Path, default_name: &SharedString, -) -> Option<(Image, SharedString)> { - let content = std::fs::read(path).ok()?; - let format = image::guess_format(&content) +) -> Result> { + let has_raster_extension = is_raster_image_path(path); + if path.extension().is_some() && !has_raster_extension { + return Ok(None); + } + + let size = path.metadata()?.len(); + if size > MAX_EXTERNAL_ATTACHMENT_BYTES { + if has_raster_extension { + return Err(anyhow!( + "{} is too large to attach. The limit is {} MB.", + path.display(), + MAX_EXTERNAL_ATTACHMENT_BYTES / 1024 / 1024 + )); + } + return Ok(None); + } + + let content = std::fs::read(path)?; + let Some(format) = image::guess_format(&content) .ok() - .and_then(image_format_from_external_content)?; + .and_then(image_format_from_external_content) + else { + if has_raster_extension { + return Err(anyhow!("{} is not a supported image file.", path.display())); + } + return Ok(None); + }; let name = path .file_name() .and_then(|name| name.to_str()) .map(|name| SharedString::from(name.to_owned())) .unwrap_or_else(|| default_name.clone()); - Some((Image::from_bytes(format, content), name)) + Ok(Some((Image::from_bytes(format, content), name))) } pub(crate) fn load_external_pdf_from_path(path: &Path) -> Result> { @@ -1044,16 +1110,25 @@ pub(crate) fn paste_images_as_context( .partition_map::, Vec<_>, _, _, _>(std::convert::identity); if !paths.is_empty() { - images.extend( - cx.background_spawn(async move { - paths - .into_iter() - .flat_map(|paths| paths.paths().to_owned()) - .filter_map(|path| load_external_image_from_path(&path, &default_name)) - .collect::>() + let (loaded_images, errors) = cx + .background_spawn(async move { + let mut images = Vec::new(); + let mut errors = Vec::new(); + for path in paths.into_iter().flat_map(|paths| paths.paths().to_owned()) { + match load_external_image_from_path(&path, &default_name) { + Ok(Some(image)) => images.push(image), + Ok(None) => {} + Err(error) => errors.push(error), + } + } + (images, errors) }) - .await, - ); + .await; + + for error in errors { + Err::<(), _>(error).notify_workspace_async_err(workspace.clone(), &mut cx); + } + images.extend(loaded_images); } if !images.is_empty() { diff --git a/crates/agent_ui/src/message_editor.rs b/crates/agent_ui/src/message_editor.rs index 3f421b0999..6d5f1c830d 100644 --- a/crates/agent_ui/src/message_editor.rs +++ b/crates/agent_ui/src/message_editor.rs @@ -273,8 +273,8 @@ async fn resolve_external_context_items( let default_image_name: SharedString = "Image".into(); for path in paths { - if supports_images - && let Some((image, name)) = cx + if supports_images { + match cx .background_spawn({ let path = path.clone(); let default_image_name = default_image_name.clone(); @@ -286,9 +286,17 @@ async fn resolve_external_context_items( } }) .await - { - items.push(ResolvedPastedContextItem::Image(image, name)); - continue; + { + Ok(Some((image, name))) => { + items.push(ResolvedPastedContextItem::Image(image, name)); + continue; + } + Ok(None) => {} + Err(error) => { + errors.push(error); + continue; + } + } } if crate::mention_set::is_pdf_path(&path) { @@ -1640,20 +1648,28 @@ impl MessageEditor { }; let default_image_name: SharedString = "Image".into(); - let images = cx + let (images, errors) = cx .background_spawn(async move { - paths - .into_iter() - .filter_map(|path| { - crate::mention_set::load_external_image_from_path( - &path, - &default_image_name, - ) - }) - .collect::>() + let mut images = Vec::new(); + let mut errors = Vec::new(); + for path in paths { + match crate::mention_set::load_external_image_from_path( + &path, + &default_image_name, + ) { + Ok(Some(image)) => images.push(image), + Ok(None) => {} + Err(error) => errors.push(error), + } + } + (images, errors) }) .await; + for error in errors { + Err::<(), _>(error).notify_workspace_async_err(workspace.clone(), cx); + } + crate::mention_set::insert_images_as_context( images, editor,