From 4338ef75e35e1f92fa89e013a2770ca22cdd5867 Mon Sep 17 00:00:00 2001
From: "Dylan M. Taylor" <dylan@dylanmtaylor.com>
Date: Tue, 3 Nov 2026 17:28:38 -0500
Subject: [PATCH] MCP: render_preview and window captures decode the app's PNG
 under size limits (#5445)

---
 crates/mcp/src/tools/dispatch.rs |  3 ++-
 crates/mcp/src/tools/handlers.rs | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/crates/mcp/src/tools/dispatch.rs b/crates/mcp/src/tools/dispatch.rs
index 7c91c9c8b..9e3af3dba 100644
--- a/crates/mcp/src/tools/dispatch.rs
+++ b/crates/mcp/src/tools/dispatch.rs
@@ -452,7 +452,8 @@ fn render_preview(b: &mut dyn Backend, a: &Args) -> Result<ToolResult, String> {
     let r = b.call("ui.render", p)?;
     let encoded = r.get("pngBase64").and_then(Value::as_str).ok_or("renderer returned no image")?;
     let png = vectorcraft_format::base64_decode(encoded).ok_or("renderer returned bad base64")?;
-    let img = image::load_from_memory_with_format(&png, image::ImageFormat::Png).map_err(|e| format!("read preview: {e}"))?;
+    // The render was asked for at most 4096 a side.
+    let img = super::handlers::decode_png(&png, 8192).map_err(|e| format!("read preview: {e}"))?;
     let img = if img.width() > max || img.height() > max { img.thumbnail(max, max) } else { img };
     let mut out = std::io::Cursor::new(Vec::new());
     img.write_to(&mut out, image::ImageFormat::Png).map_err(|e| format!("encode preview: {e}"))?;
diff --git a/crates/mcp/src/tools/handlers.rs b/crates/mcp/src/tools/handlers.rs
index c9d203c6b..ea2c6ec70 100644
--- a/crates/mcp/src/tools/handlers.rs
+++ b/crates/mcp/src/tools/handlers.rs
@@ -442,12 +442,27 @@ pub(super) fn max_side(a: &Args) -> Result<Option<u32>, String> {
     }
 }
 
+/// The longest side of a window capture that is decoded (a screen is never wider).
+const MAX_CAPTURE_SIDE: u32 = 16_384;
+
+/// `png` decoded, refused when a side is longer than `side` or the decode needs more than 256 MiB:
+/// the picture is the app's reply, and its header names the size it is decoded at.
+pub(super) fn decode_png(png: &[u8], side: u32) -> Result<image::DynamicImage, image::ImageError> {
+    let mut reader = image::ImageReader::with_format(std::io::Cursor::new(png), image::ImageFormat::Png);
+    let mut limits = image::Limits::default();
+    limits.max_image_width = Some(side);
+    limits.max_image_height = Some(side);
+    limits.max_alloc = Some(256 << 20);
+    reader.limits(limits);
+    reader.decode()
+}
+
 /// A window capture with `max_side`: scaled down when its longer side is larger, with its own size
 /// in `capturedWidth` / `capturedHeight` and the returned image's in `width` / `height`. Without
 /// `max_side` the capture and `info` stay as they are.
 fn fit_capture(png: Vec<u8>, max: Option<u32>, info: &mut Value) -> Result<Vec<u8>, String> {
     let Some(max) = max else { return Ok(png) };
-    let img = image::load_from_memory_with_format(&png, image::ImageFormat::Png).map_err(|e| format!("read capture: {e}"))?;
+    let img = decode_png(&png, MAX_CAPTURE_SIDE).map_err(|e| format!("read capture: {e}"))?;
     info["capturedWidth"] = json!(img.width());
     info["capturedHeight"] = json!(img.height());
     if img.width().max(img.height()) <= max {
@@ -472,3 +487,17 @@ pub(super) fn image_result(png: &[u8], info: Value) -> ToolResult {
         is_error: false,
     }
 }
+
+#[cfg(test)]
+mod decode_tests {
+    use super::decode_png;
+
+    #[test]
+    fn a_picture_with_a_side_past_the_limit_is_not_decoded() {
+        let mut png = std::io::Cursor::new(Vec::new());
+        image::RgbaImage::new(200, 10).write_to(&mut png, image::ImageFormat::Png).unwrap();
+        let png = png.into_inner();
+        assert!(decode_png(&png, 100).is_err());
+        assert_eq!(decode_png(&png, 200).unwrap().width(), 200);
+    }
+}