From 92fbfc03d219544c16e8f7db47a547f9ff402c27 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 5 Oct 2026 16:50:03 +0200 Subject: [PATCH] fix: bound s3 ai attachment downloads before encoding (#11534) --- backend/windmill-ai/src/image_handler.rs | 10 +- backend/windmill-common/src/client.rs | 108 +++++++++++++++++++- backend/windmill-worker/src/sql_s3_input.rs | 2 +- 3 files changed, 117 insertions(+), 3 deletions(-) diff --git a/backend/windmill-ai/src/image_handler.rs b/backend/windmill-ai/src/image_handler.rs index 5bd9794da5..a4f360cdb4 100644 --- a/backend/windmill-ai/src/image_handler.rs +++ b/backend/windmill-ai/src/image_handler.rs @@ -47,6 +47,9 @@ pub async fn upload_image_to_s3( }) } +// Bounds each attachment's download before base64 expansion; provider limits may be lower. +const MAX_S3_MEDIA_BYTES: usize = 30 * 1024 * 1024; + /// Download an S3 image and convert it to a base64 data URL. /// /// The caller must provide an AuthedClient authorized for `workspace_id`. @@ -57,7 +60,12 @@ pub async fn download_and_encode_s3_image( ) -> Result<(String, String), Error> { // Download the image from S3 let image_bytes = client - .download_s3_file(workspace_id, &image.s3, image.storage.clone()) + .download_s3_file( + workspace_id, + &image.s3, + image.storage.clone(), + Some(MAX_S3_MEDIA_BYTES), + ) .await .map_err(|e| Error::internal_err(format!("Failed to download S3 image: {}", e)))?; diff --git a/backend/windmill-common/src/client.rs b/backend/windmill-common/src/client.rs index 48f3e0205b..985d3bbbfc 100644 --- a/backend/windmill-common/src/client.rs +++ b/backend/windmill-common/src/client.rs @@ -307,12 +307,13 @@ impl AuthedClient { workspace_id: &str, file_key: &str, storage: Option, + max_bytes: Option, ) -> anyhow::Result { let mut query = vec![("file_key", file_key.to_string())]; if let Some(storage) = storage { query.push(("storage", storage)); } - let response = self + let mut response = self .force_client .as_ref() .unwrap_or(&HTTP_CLIENT) @@ -331,6 +332,39 @@ impl AuthedClient { .context("Failed to send download_s3_file request") .map_err(|e| anyhow::anyhow!(e.to_string()))?; + if let Some(max_bytes) = max_bytes { + let size_error = || { + anyhow::anyhow!( + "S3 file '{file_key}' exceeds maximum supported download size of {max_bytes} bytes" + ) + }; + if response + .content_length() + .is_some_and(|len| len > max_bytes as u64) + { + return Err(size_error()); + } + + let status = response.status(); + let mut bytes = Vec::new(); + // Content-Length may be absent (including for decoded compressed responses). + while let Some(chunk) = response + .chunk() + .await + .context("Failed to read response bytes")? + { + if chunk.len() > max_bytes - bytes.len() { + return Err(size_error()); + } + bytes.extend_from_slice(&chunk); + } + return if status == reqwest::StatusCode::OK { + Ok(bytes.into()) + } else { + Err(anyhow::anyhow!(String::from_utf8_lossy(&bytes).into_owned())) + }; + } + match response.status().as_u16() { 200u16 => Ok(response .bytes() @@ -341,6 +375,78 @@ impl AuthedClient { } } +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + + async fn download_response( + response: &'static str, + max_bytes: Option, + ) -> anyhow::Result { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let client = AuthedClient::new( + format!("http://{}", listener.local_addr().unwrap()), + "test_ws".to_string(), + "test_token".to_string(), + Some(reqwest::Client::builder().no_proxy().build().unwrap()), + ); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = BufReader::new(socket); + loop { + let mut line = String::new(); + assert!(socket.read_line(&mut line).await.unwrap() > 0); + if line == "\r\n" { + break; + } + } + socket + .get_mut() + .write_all(response.as_bytes()) + .await + .unwrap(); + // Keep incomplete bodies open so rejection must happen before EOF. + std::future::pending::<()>().await; + }); + let result = tokio::time::timeout( + Duration::from_secs(5), + client.download_s3_file("test_ws", "image.png", None, max_bytes), + ) + .await; + server.abort(); + result.expect("download waited for the oversized body to finish") + } + + #[tokio::test] + async fn download_s3_file_rejects_oversized_body_before_eof() { + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n1\r\ne\r\n", + "HTTP/1.1 500 Internal Server Error\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nerror\r\n", + ] { + let error = download_response(response, Some(4)).await.unwrap_err(); + assert!(error.to_string().contains( + "S3 file 'image.png' exceeds maximum supported download size of 4 bytes" + )); + } + } + + #[tokio::test] + async fn download_s3_file_accepts_body_at_limit() { + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\nabcd", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n0\r\n\r\n", + ] { + assert_eq!( + download_response(response, Some(4)).await.unwrap(), + &b"abcd"[..] + ); + } + } +} + #[inline] fn query_from_json_path(json_path: Option) -> Vec<(&'static str, String)> { json_path diff --git a/backend/windmill-worker/src/sql_s3_input.rs b/backend/windmill-worker/src/sql_s3_input.rs index e3a5b25411..ecf2c53dec 100644 --- a/backend/windmill-worker/src/sql_s3_input.rs +++ b/backend/windmill-worker/src/sql_s3_input.rs @@ -56,7 +56,7 @@ pub async fn fetch_s3object_as_json_text( crate::worker_utils::JobPingHeartbeat::start(conn, job_id, "s3object materialization"); let bytes = client - .download_s3_file(workspace_id, &obj.s3, obj.storage.clone()) + .download_s3_file(workspace_id, &obj.s3, obj.storage.clone(), None) .await .with_context(|| format!("Failed to download S3 object `{}`", obj.s3))?;