From 05d6da45fbf7fd8e1300e9f296fed4a413d20a4e Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 30 Sep 2026 07:44:59 +0900 Subject: [PATCH] fix: add bounded binary reads for ViewImage --- crates/fs-operation/src/lib.rs | 100 +++++++- crates/fs-operation/src/local.rs | 136 ++++++++++- crates/fs-operation/src/operation.rs | 30 +++ crates/tools/src/view_image.rs | 113 +++++++-- crates/tools/tests/integration.rs | 218 ++++++++++++++++++ crates/workdir/src/external.rs | 66 +++++- crates/workdir/src/http.rs | 25 +- crates/workdir/src/lib.rs | 11 +- crates/workdir/src/local.rs | 23 +- crates/workdir/src/router.rs | 12 +- crates/workdir/src/scope.rs | 36 ++- .../src/feature/builtin/manage_workdir.rs | 11 +- crates/workspace-server/src/server.rs | 78 ++++++- crates/yoi/src/workdir_share.rs | 2 +- 14 files changed, 816 insertions(+), 45 deletions(-) diff --git a/crates/fs-operation/src/lib.rs b/crates/fs-operation/src/lib.rs index 465e3b6b9..7fa897df8 100644 --- a/crates/fs-operation/src/lib.rs +++ b/crates/fs-operation/src/lib.rs @@ -15,7 +15,8 @@ use thiserror::Error; pub use glob::run_glob; pub use local::{ - resolve_access_path, run_edit, run_list, run_read, run_read_bounded, run_stat, run_write, + resolve_access_path, run_edit, run_list, run_read, run_read_bounded, run_read_bytes, + run_read_bytes_bounded, run_stat, run_write, }; pub use operation::*; pub use search::run_grep; @@ -569,6 +570,103 @@ mod tests { } } + #[test] + fn bounded_byte_reads_preserve_binary_ranges_and_fence_versions() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap(); + let original = b"\x89PNG\r\n\x1a\n\0line\n\xfftail"; + std::fs::write(root.join("image.bin"), original).unwrap(); + let access = RootAccess(root.clone()); + let limits = BoundedReadLimits::new(1024, 8).unwrap(); + + let first = run_read_bytes_bounded( + &root, + ReadBytesRequest { + path: FsPath::new("image.bin").unwrap(), + offset: 0, + max_bytes: 8, + expected_hash: None, + }, + &access, + limits, + ) + .unwrap(); + assert_eq!(first.bytes, original[..8]); + assert_eq!(first.total_bytes, original.len() as u64); + assert!(!first.eof); + + let second = run_read_bytes_bounded( + &root, + ReadBytesRequest { + path: FsPath::new("image.bin").unwrap(), + offset: 8, + max_bytes: 8, + expected_hash: Some(first.content_hash), + }, + &access, + limits, + ) + .unwrap(); + assert_eq!(second.bytes, original[8..16]); + assert_eq!(second.content_hash, first.content_hash); + + std::fs::write(root.join("image.bin"), b"different bytes long!").unwrap(); + let error = run_read_bytes_bounded( + &root, + ReadBytesRequest { + path: FsPath::new("image.bin").unwrap(), + offset: 16, + max_bytes: 8, + expected_hash: Some(first.content_hash), + }, + &access, + limits, + ) + .expect_err("a later chunk must not cross file versions"); + assert!(matches!(error, FsError::Conflict(_))); + } + + #[test] + fn bounded_byte_reads_reject_request_and_source_limit_overflow() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap(); + std::fs::write(root.join("image.bin"), [0_u8; 17]).unwrap(); + let access = RootAccess(root.clone()); + let limits = BoundedReadLimits::new(16, 8).unwrap(); + + let response_error = run_read_bytes_bounded( + &root, + ReadBytesRequest { + path: FsPath::new("image.bin").unwrap(), + offset: 0, + max_bytes: 9, + expected_hash: None, + }, + &access, + limits, + ) + .expect_err("requests above the response ceiling must be rejected"); + assert!( + matches!(response_error, FsError::InvalidArgument(message) if message.contains("8 bytes")) + ); + + let source_error = run_read_bytes_bounded( + &root, + ReadBytesRequest { + path: FsPath::new("image.bin").unwrap(), + offset: 0, + max_bytes: 8, + expected_hash: None, + }, + &access, + limits, + ) + .expect_err("sources above the provider ceiling must be rejected"); + assert!( + matches!(source_error, FsError::InvalidArgument(message) if message.contains("provider limit 16")) + ); + } + #[test] fn cancellation_stops_a_read_already_in_progress() { use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; diff --git a/crates/fs-operation/src/local.rs b/crates/fs-operation/src/local.rs index aaff61910..34f60b92d 100644 --- a/crates/fs-operation/src/local.rs +++ b/crates/fs-operation/src/local.rs @@ -7,8 +7,9 @@ use sha2::{Digest, Sha256}; use crate::{ AtomicWriteMode, BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, - FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadRequest, ReadResult, - StatRequest, StatResult, WriteRequest, WriteResult, direct_symlink, + FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadBytesRequest, + ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, + direct_symlink, }; const READ_BUFFER_BYTES: usize = 16 * 1024; @@ -195,6 +196,137 @@ fn run_read_with_limits( }) } +pub fn run_read_bytes( + root: &Path, + request: ReadBytesRequest, + access: &dyn FsAccessPolicy, +) -> Result { + run_read_bytes_with_limits(root, request, access, None) +} + +pub fn run_read_bytes_bounded( + root: &Path, + request: ReadBytesRequest, + access: &dyn FsAccessPolicy, + limits: BoundedReadLimits, +) -> Result { + run_read_bytes_with_limits(root, request, access, Some(limits)) +} + +fn run_read_bytes_with_limits( + root: &Path, + request: ReadBytesRequest, + access: &dyn FsAccessPolicy, + limits: Option, +) -> Result { + let logical = request.path; + if request.max_bytes == 0 { + return Err(FsError::InvalidArgument( + "byte read max_bytes must be greater than zero".to_string(), + )); + } + if let Some(limits) = limits + && request.max_bytes > limits.max_response_bytes + { + return Err(FsError::InvalidArgument(format!( + "byte read response limit is {} bytes", + limits.max_response_bytes + ))); + } + + access + .check_cancelled() + .map_err(|error| map_io(&logical, error))?; + let path = resolve(root, &logical)?; + let target = require_access(&path, &logical, access, false, false)?; + let file = access + .open_read_file(&path, &target) + .map_err(|error| map_io(&logical, error))?; + let metadata = file.metadata().map_err(|error| map_io(&logical, error))?; + if metadata.is_dir() { + return Err(FsError::IsDirectory(PathBuf::from(logical.as_str()))); + } + if !metadata.is_file() { + return Err(FsError::InvalidArgument( + "byte read source must be a regular file".to_string(), + )); + } + if let Some(limits) = limits + && metadata.len() > limits.max_source_bytes + { + return Err(FsError::InvalidArgument(format!( + "byte read source exceeds provider limit {}", + limits.max_source_bytes + ))); + } + + let mut bytes = Vec::with_capacity(request.max_bytes.min(READ_BUFFER_BYTES)); + let mut source_bytes = 0_u64; + let mut hasher = Sha256::new(); + let mut reader = BufReader::with_capacity(READ_BUFFER_BYTES, file); + let mut buffer = [0_u8; READ_BUFFER_BYTES]; + loop { + access + .check_cancelled() + .map_err(|error| map_io(&logical, error))?; + let read_limit = limits + .map(|limits| { + limits + .max_source_bytes + .saturating_sub(source_bytes) + .saturating_add(1) + .min(READ_BUFFER_BYTES as u64) as usize + }) + .unwrap_or(READ_BUFFER_BYTES); + let read = reader + .read(&mut buffer[..read_limit]) + .map_err(|error| map_io(&logical, error))?; + if read == 0 { + break; + } + let chunk_start = source_bytes; + source_bytes = source_bytes.saturating_add(read as u64); + if let Some(limits) = limits + && source_bytes > limits.max_source_bytes + { + return Err(FsError::InvalidArgument(format!( + "byte read source exceeds provider limit {}", + limits.max_source_bytes + ))); + } + hasher.update(&buffer[..read]); + + if bytes.len() < request.max_bytes && source_bytes > request.offset { + let start = request.offset.saturating_sub(chunk_start).min(read as u64) as usize; + let retained = (request.max_bytes - bytes.len()).min(read - start); + bytes.extend_from_slice(&buffer[start..start + retained]); + } + } + + let content_hash: ContentHash = hasher.finalize().into(); + if request + .expected_hash + .is_some_and(|expected| expected != content_hash) + { + return Err(FsError::Conflict(logical.as_str().to_string())); + } + if request.offset > source_bytes { + return Err(FsError::InvalidArgument(format!( + "byte offset {} exceeds file size {source_bytes}", + request.offset + ))); + } + let eof = request.offset.saturating_add(bytes.len() as u64) == source_bytes; + Ok(ReadBytesResult { + path: logical, + bytes, + offset: request.offset, + total_bytes: source_bytes, + content_hash, + eof, + }) +} + pub fn run_write( root: &Path, request: WriteRequest, diff --git a/crates/fs-operation/src/operation.rs b/crates/fs-operation/src/operation.rs index f5765b6f6..ddb309a8d 100644 --- a/crates/fs-operation/src/operation.rs +++ b/crates/fs-operation/src/operation.rs @@ -207,6 +207,36 @@ pub struct ReadResult { pub truncated: bool, } +/// Binary-safe byte range read with an optional whole-file version fence. +/// +/// Unlike [`ReadRequest`], offsets and limits are bytes rather than lines. The +/// provider still hashes the complete bounded source so callers can reject +/// chunks from different file versions. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ReadBytesRequest { + pub path: FsPath, + #[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))] + pub offset: u64, + #[schemars(range(min = 1, max = 9_007_199_254_740_991_usize))] + pub max_bytes: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expected_hash: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ReadBytesResult { + pub path: FsPath, + pub bytes: Vec, + #[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))] + pub offset: u64, + #[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))] + pub total_bytes: u64, + pub content_hash: ContentHash, + pub eof: bool, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct WriteRequest { diff --git a/crates/tools/src/view_image.rs b/crates/tools/src/view_image.rs index 1fc86ca44..255695b9a 100644 --- a/crates/tools/src/view_image.rs +++ b/crates/tools/src/view_image.rs @@ -7,12 +7,14 @@ use agen::tool::{ }; use async_trait::async_trait; use serde::Deserialize; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; +use workdir::{ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::error::ToolsError; /// Maximum image body accepted for one model request. pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024; +/// One binary response remains within the External Workdir transport bound. +const IMAGE_READ_CHUNK_BYTES: usize = 1024 * 1024; const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \ The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \ @@ -47,42 +49,111 @@ impl Tool for ViewImageTool { workdir::WorkdirSessionCapability::Read, )?; let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?; - let result = selected - .session - .read(ReadRequest { - path: path.clone(), - offset: 0, - limit: usize::MAX, - // The scoped provider enforces this cap while reading, rather - // than allocating an unbounded binary body first. - max_bytes: MAX_IMAGE_BYTES + 1, - }) + let bytes = read_image_bytes(selected.session.as_ref(), path.clone()) .await .map_err(ToolsError::from)?; - - if result.truncated || result.bytes.len() > MAX_IMAGE_BYTES { - return Err(ToolError::InvalidArgument(format!( - "image exceeds the {MAX_IMAGE_BYTES}-byte limit" - ))); - } - let mime_type = detect_image_mime(&result.bytes).ok_or_else(|| { + let mime_type = detect_image_mime(&bytes).ok_or_else(|| { ToolError::InvalidArgument( "unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(), ) })?; - let bytes = result.bytes.len(); + let byte_count = bytes.len(); Ok(ToolOutput { - summary: format!("Attached image {path} ({mime_type}, {bytes} bytes)"), + summary: format!("Attached image {path} ({mime_type}, {byte_count} bytes)"), content: None, attachments: vec![Attachment::Image(ImageAttachment::new( mime_type, - Arc::<[u8]>::from(result.bytes), + Arc::<[u8]>::from(bytes), ))], }) } } +async fn read_image_bytes( + session: &dyn workdir::WorkdirSession, + path: WorkdirPath, +) -> Result, workdir::WorkdirError> { + let mut bytes = Vec::new(); + let mut offset = 0_u64; + let mut expected_hash = None; + let mut expected_total = None; + + loop { + let remaining = expected_total + .map(|total: u64| total.saturating_sub(offset) as usize) + .unwrap_or(IMAGE_READ_CHUNK_BYTES); + let request_bytes = remaining.min(IMAGE_READ_CHUNK_BYTES).max(1); + let result = session + .read_bytes(ReadBytesRequest { + path: path.clone(), + offset, + max_bytes: request_bytes, + expected_hash, + }) + .await?; + + if result.path != path || result.offset != offset { + return Err(workdir::WorkdirError::Transport( + "bounded binary read returned a mismatched path or offset".to_string(), + )); + } + if result.bytes.len() > request_bytes { + return Err(workdir::WorkdirError::Transport( + "bounded binary read exceeded the requested response size".to_string(), + )); + } + let end = result + .offset + .checked_add(result.bytes.len() as u64) + .ok_or_else(|| { + workdir::WorkdirError::Transport( + "bounded binary read returned an invalid byte range".to_string(), + ) + })?; + if end > result.total_bytes || result.eof != (end == result.total_bytes) { + return Err(workdir::WorkdirError::Transport( + "bounded binary read returned a truncated or invalid byte range".to_string(), + )); + } + if result.total_bytes > MAX_IMAGE_BYTES as u64 { + return Err(workdir::WorkdirError::InvalidArgument(format!( + "image exceeds the {MAX_IMAGE_BYTES}-byte limit" + ))); + } + if expected_total.is_some_and(|total| total != result.total_bytes) + || expected_hash.is_some_and(|hash| hash != result.content_hash) + { + return Err(workdir::WorkdirError::Conflict( + "image changed while it was being read; retry ViewImage".to_string(), + )); + } + if expected_total.is_none() { + bytes + .try_reserve_exact(result.total_bytes as usize) + .map_err(|_| workdir::WorkdirError::OperationFailed)?; + expected_total = Some(result.total_bytes); + expected_hash = Some(result.content_hash); + } + + if result.bytes.is_empty() && !result.eof { + return Err(workdir::WorkdirError::Transport( + "bounded binary read made no progress".to_string(), + )); + } + bytes.extend_from_slice(&result.bytes); + offset = end; + if result.eof { + if bytes.len() as u64 != result.total_bytes { + return Err(workdir::WorkdirError::Transport( + "bounded binary read returned a partial image".to_string(), + )); + } + return Ok(bytes); + } + } +} + pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { routed_view_image_tool(crate::routing::singleton_router(session)) } diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index f1d346e60..eb4ca3c7f 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -399,6 +399,224 @@ async fn detach_and_reattach_same_alias_requires_a_fresh_read() { router.close_all().await.unwrap(); } +#[derive(Debug)] +struct ExternalValidatingSession { + inner: LocalWorkdirSession, + truncate_first_binary_response: bool, + mutate_after_first_binary_response: Option, +} + +#[async_trait::async_trait] +impl workdir::WorkdirSession for ExternalValidatingSession { + fn workdir(&self) -> &workdir::Workdir { + self.inner.workdir() + } + + fn capabilities(&self) -> workdir::WorkdirSessionCapabilities { + self.inner.capabilities() + } + + async fn stat( + &self, + request: workdir::StatRequest, + ) -> Result { + self.inner.stat(request).await + } + + async fn read( + &self, + request: workdir::ReadRequest, + ) -> Result { + let operation = workdir::http::WorkdirSessionOperation::Read(request); + let operation = workdir::external::ExternalWorkdirOperation::try_from(operation) + .map_err(workdir::WorkdirError::UnsupportedOperation)?; + let workdir::http::WorkdirSessionOperation::Read(request) = operation.into_inner() else { + unreachable!("validated Read changed operation kind") + }; + self.inner.read(request).await + } + + async fn read_bytes( + &self, + request: workdir::ReadBytesRequest, + ) -> Result { + let operation = workdir::http::WorkdirSessionOperation::ReadBytes(request); + let operation = workdir::external::ExternalWorkdirOperation::try_from(operation) + .map_err(workdir::WorkdirError::UnsupportedOperation)?; + let workdir::http::WorkdirSessionOperation::ReadBytes(request) = operation.into_inner() + else { + unreachable!("validated ReadBytes changed operation kind") + }; + let result = self.inner.read_bytes(request).await?; + let result = workdir::external::ExternalWorkdirOperationResult::try_from( + workdir::http::WorkdirSessionOperationResult::ReadBytes(result), + ) + .map_err(workdir::WorkdirError::UnsupportedOperation)?; + let workdir::http::WorkdirSessionOperationResult::ReadBytes(mut result) = + result.into_inner() + else { + unreachable!("validated ReadBytes result changed operation kind") + }; + if result.offset == 0 { + if let Some(path) = &self.mutate_after_first_binary_response { + std::fs::write(path, vec![b'z'; result.total_bytes as usize]).unwrap(); + } + if self.truncate_first_binary_response { + result.bytes.pop(); + } + } + Ok(result) + } + + async fn write( + &self, + request: workdir::WriteRequest, + ) -> Result { + self.inner.write(request).await + } + + async fn edit( + &self, + request: workdir::EditRequest, + ) -> Result { + self.inner.edit(request).await + } + + async fn list( + &self, + request: workdir::ListRequest, + ) -> Result { + self.inner.list(request).await + } + + async fn glob( + &self, + request: workdir::GlobRequest, + ) -> Result { + self.inner.glob(request).await + } + + async fn grep( + &self, + request: workdir::GrepRequest, + ) -> Result { + self.inner.grep(request).await + } + + async fn start_command( + &self, + request: workdir::CommandRequest, + ) -> Result { + self.inner.start_command(request).await + } + + async fn command_status( + &self, + handle: workdir::CommandHandle, + ) -> Result { + self.inner.command_status(handle).await + } + + async fn command_output( + &self, + request: workdir::CommandOutputRequest, + ) -> Result { + self.inner.command_output(request).await + } + + async fn cancel_command( + &self, + handle: workdir::CommandHandle, + ) -> Result<(), workdir::WorkdirError> { + self.inner.cancel_command(handle).await + } + + async fn close(&self) -> Result<(), workdir::WorkdirError> { + self.inner.close().await + } +} + +#[tokio::test] +async fn external_view_image_uses_a_protocol_bounded_request() { + let dir = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let scope = scope_with_spill(dir.path(), spill.path()); + let session: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { + inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()), + truncate_first_binary_response: false, + mutate_after_first_binary_response: None, + }); + let png = b"\x89PNG\r\n\x1a\nexternal\0image\nbody\xff"; + std::fs::write(dir.path().join("image.png"), png).unwrap(); + let (_meta, tool) = view_image_tool(session)(); + + let output = call(&tool, json!({ "path": "image.png" })).await; + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.mime_type(), "image/png"); + assert_eq!(image.data(), png); +} + +#[tokio::test] +async fn external_view_image_rejects_truncated_and_cross_version_chunks() { + let dir = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let image_path = dir.path().join("image.png"); + let scope = scope_with_spill(dir.path(), spill.path()); + std::fs::write(&image_path, b"\x89PNG\r\n\x1a\npartial").unwrap(); + let truncated: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { + inner: LocalWorkdirSession::new(scope.clone(), dir.path().to_path_buf()), + truncate_first_binary_response: true, + mutate_after_first_binary_response: None, + }); + let (_meta, tool) = view_image_tool(truncated)(); + let error = call_err(&tool, json!({ "path": "image.png" })) + .await + .to_string(); + assert!( + error.contains("truncated") || error.contains("partial"), + "{error}" + ); + + let mut large = vec![b'x'; 1024 * 1024 + 32]; + large[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); + std::fs::write(&image_path, large).unwrap(); + let changing: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { + inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()), + truncate_first_binary_response: false, + mutate_after_first_binary_response: Some(image_path), + }); + let (_meta, tool) = view_image_tool(changing)(); + let error = call_err(&tool, json!({ "path": "image.png" })) + .await + .to_string(); + assert!(error.contains("changed"), "{error}"); +} + +#[tokio::test] +async fn view_image_enforces_the_ten_mib_boundary_without_partial_attachments() { + const LIMIT: usize = 10 * 1024 * 1024; + let dir = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let scope = scope_with_spill(dir.path(), spill.path()); + let session: WorkdirSessionHandle = + Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); + let (_meta, tool) = view_image_tool(session)(); + let mut png = vec![0xff; LIMIT]; + png[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); + std::fs::write(dir.path().join("boundary.png"), &png).unwrap(); + + let output = call(&tool, json!({ "path": "boundary.png" })).await; + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.data(), png); + + png.push(0); + std::fs::write(dir.path().join("boundary.png"), png).unwrap(); + let error = call_err(&tool, json!({ "path": "boundary.png" })) + .await + .to_string(); + assert!(error.contains("10485760-byte limit"), "{error}"); +} + #[tokio::test] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); diff --git a/crates/workdir/src/external.rs b/crates/workdir/src/external.rs index 5557231ff..c94582867 100644 --- a/crates/workdir/src/external.rs +++ b/crates/workdir/src/external.rs @@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize}; use crate::http::{WorkdirSessionOperation, WorkdirSessionOperationResult, WorkdirTransportError}; use crate::{BoundedReadLimits, WorkdirId, WorkdirPath, WorkdirSessionCapabilities}; -pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 1; +pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 2; /// Maximum serialized Backend-to-provider command frame size. pub const MAX_EXTERNAL_SERVER_FRAME_BYTES: usize = 1024 * 1024; /// Maximum serialized provider-to-Backend result frame size. Read bytes are @@ -155,6 +155,12 @@ fn validate_external_operation(operation: &WorkdirSessionOperation) -> Result<() && request.limit <= 1_000_000 && request.max_bytes <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes } + WorkdirSessionOperation::ReadBytes(request) => { + is_root_relative(&request.path) + && request.offset <= BoundedReadLimits::EXTERNAL_DEFAULT.max_source_bytes + && (1..=BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes) + .contains(&request.max_bytes) + } WorkdirSessionOperation::List(request) => { is_root_relative(&request.path) && request.limit <= MAX_EXTERNAL_RESULT_ITEMS } @@ -275,6 +281,15 @@ fn validate_external_result(result: &WorkdirSessionOperationResult) -> Result<() is_root_relative(&result.path) && result.bytes.len() <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes } + WorkdirSessionOperationResult::ReadBytes(result) => { + let returned_end = result.offset.checked_add(result.bytes.len() as u64); + is_root_relative(&result.path) + && result.bytes.len() <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes + && result.total_bytes <= BoundedReadLimits::EXTERNAL_DEFAULT.max_source_bytes + && returned_end.is_some_and(|end| { + end <= result.total_bytes && result.eof == (end == result.total_bytes) + }) + } WorkdirSessionOperationResult::List(result) => { result.entries.len() <= MAX_EXTERNAL_RESULT_ITEMS && result @@ -482,8 +497,8 @@ mod tests { use super::*; use crate::{ - CommandRequest, EditRequest, GrepResult, ListRequest, ReadRequest, Workdir, WorkdirPath, - WriteRequest, + CommandRequest, EditRequest, GrepResult, ListRequest, ReadBytesRequest, ReadBytesResult, + ReadRequest, Workdir, WorkdirPath, WriteRequest, }; fn assert_no_provider_authority(value: &Value) { @@ -614,7 +629,7 @@ mod tests { #[test] fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() { let future = r#"{ - "version": 2, + "version": 3, "message": {"kind": "heartbeat", "generation": 1, "sequence": 1} }"#; assert!(serde_json::from_str::(future).is_err()); @@ -726,10 +741,51 @@ mod tests { ); } + #[test] + fn protocol_bounds_binary_read_requests_and_rejects_partial_results() { + let path = WorkdirPath::new("image.png").unwrap(); + let bounded = WorkdirSessionOperation::ReadBytes(ReadBytesRequest { + path: path.clone(), + offset: 0, + max_bytes: BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes, + expected_hash: None, + }); + assert!(ExternalWorkdirOperation::try_from(bounded).is_ok()); + + let oversized = WorkdirSessionOperation::ReadBytes(ReadBytesRequest { + path: path.clone(), + offset: 0, + max_bytes: BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes + 1, + expected_hash: None, + }); + assert!(ExternalWorkdirOperation::try_from(oversized).is_err()); + + let hash = [7; 32]; + let complete = WorkdirSessionOperationResult::ReadBytes(ReadBytesResult { + path: path.clone(), + bytes: vec![1, 2, 3], + offset: 0, + total_bytes: 3, + content_hash: hash, + eof: true, + }); + assert!(ExternalWorkdirOperationResult::try_from(complete).is_ok()); + + let partial_marked_complete = WorkdirSessionOperationResult::ReadBytes(ReadBytesResult { + path, + bytes: vec![1, 2], + offset: 0, + total_bytes: 3, + content_hash: hash, + eof: true, + }); + assert!(ExternalWorkdirOperationResult::try_from(partial_marked_complete).is_err()); + } + #[test] fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() { let unknown = r#"{ - "version": 1, + "version": 2, "message": { "kind": "heartbeat", "generation": 1, diff --git a/crates/workdir/src/http.rs b/crates/workdir/src/http.rs index 0a3f7e8aa..4017056a0 100644 --- a/crates/workdir/src/http.rs +++ b/crates/workdir/src/http.rs @@ -11,9 +11,9 @@ use serde::{Deserialize, Serialize}; use crate::{ CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, WorkdirError, WorkdirId, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WriteRequest, WriteResult, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, + WorkdirError, WorkdirId, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, + WorkdirSession, WorkdirSessionCapabilities, WriteRequest, WriteResult, }; /// Opaque Runtime-owned identifier for one ephemeral Workdir session. @@ -66,6 +66,7 @@ pub enum WorkdirSessionOperation { ScopeRulesOverlap(WorkdirScopeOverlapRequest), Stat(StatRequest), Read(ReadRequest), + ReadBytes(ReadBytesRequest), Write(WriteRequest), Edit(EditRequest), List(ListRequest), @@ -97,6 +98,7 @@ pub enum WorkdirSessionOperationResult { ScopeRulesOverlap { overlaps: bool }, Stat(StatResult), Read(ReadResult), + ReadBytes(ReadBytesResult), Write(WriteResult), Edit(EditResult), List(ListResult), @@ -135,6 +137,10 @@ pub async fn dispatch_workdir_session_operation( .read(request) .await .map(WorkdirSessionOperationResult::Read), + WorkdirSessionOperation::ReadBytes(request) => session + .read_bytes(request) + .await + .map(WorkdirSessionOperationResult::ReadBytes), WorkdirSessionOperation::Write(request) => session .write(request) .await @@ -572,6 +578,19 @@ mod client { } } + async fn read_bytes( + &self, + request: ReadBytesRequest, + ) -> Result { + match self + .operate(WorkdirSessionOperation::ReadBytes(request)) + .await? + { + WorkdirSessionOperationResult::ReadBytes(result) => Ok(result), + _ => Err(Self::mismatch("read_bytes")), + } + } + async fn write(&self, request: WriteRequest) -> Result { match self .operate(WorkdirSessionOperation::Write(request)) diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 806a1431a..a390b1084 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -27,7 +27,8 @@ use tokio::sync::broadcast; pub use fs_operation::{ BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath, GlobRequest, GlobResult, GrepOutputMode, GrepRequest, GrepResult, ListEntry, ListRequest, - ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, + ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, + StatResult, WriteRequest, WriteResult, }; pub use http::dispatch_workdir_session_operation; pub use local::{ @@ -220,6 +221,14 @@ pub trait WorkdirSession: std::fmt::Debug + Send + Sync { async fn stat(&self, request: StatRequest) -> Result; async fn read(&self, request: ReadRequest) -> Result; + async fn read_bytes( + &self, + _request: ReadBytesRequest, + ) -> Result { + Err(WorkdirError::UnsupportedOperation( + "bounded binary reads are not supported by this provider".to_string(), + )) + } async fn write(&self, request: WriteRequest) -> Result; async fn edit(&self, request: EditRequest) -> Result; async fn list(&self, request: ListRequest) -> Result; diff --git a/crates/workdir/src/local.rs b/crates/workdir/src/local.rs index 4a4a3b61e..33850729b 100644 --- a/crates/workdir/src/local.rs +++ b/crates/workdir/src/local.rs @@ -30,10 +30,10 @@ use crate::{ BoundedReadLimits, CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirPath, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirToolScopePermission, WriteRequest, - WriteResult, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, + WorkdirError, WorkdirPath, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, + WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionCapability, + WorkdirToolScopePermission, WriteRequest, WriteResult, }; #[cfg(test)] use crate::{EntryKind, WriteOutcome}; @@ -1261,6 +1261,21 @@ impl WorkdirSession for LocalWorkdirSession { .map_err(|error| sanitize_error(error, &logical)) } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + self.ensure_capability(WorkdirSessionCapability::Read)?; + let logical = request.path.clone(); + self.validate_operation_path(&request.path)?; + let access = self.scope_access(); + match self.inner.read_limits { + Some(limits) => { + fs_operation::run_read_bytes_bounded(&self.inner.root, request, &access, limits) + } + None => fs_operation::run_read_bytes(&self.inner.root, request, &access), + } + .map_err(WorkdirError::from) + .map_err(|error| sanitize_error(error, &logical)) + } + async fn write(&self, request: WriteRequest) -> Result { self.ensure_capability(WorkdirSessionCapability::Write)?; let logical = request.path.clone(); diff --git a/crates/workdir/src/router.rs b/crates/workdir/src/router.rs index eed9a9e50..270f0ec90 100644 --- a/crates/workdir/src/router.rs +++ b/crates/workdir/src/router.rs @@ -16,9 +16,10 @@ use tokio::sync::broadcast; use crate::{ CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, - GrepResult, ListRequest, ListResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, - WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, + GrepResult, ListRequest, ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, + ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirScopeAuthorizationRequest, + WorkdirScopeOverlapRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, + WriteRequest, WriteResult, }; /// Stable Worker-local routing key for one attached Workdir. @@ -257,6 +258,11 @@ impl WorkdirSession for RoutedWorkdirSession { self.state.session.read(request).await } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.read_bytes(request).await + } + async fn write(&self, request: WriteRequest) -> Result { let _active = self.state.enter()?; self.state.session.write(request).await diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index e76f59125..efe773d9f 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -6,7 +6,8 @@ use std::sync::{Arc, Mutex, Weak}; use async_trait::async_trait; use fs_operation::{ EditRequest, EditResult, FsPath, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, - ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, + ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, + StatResult, WriteRequest, WriteResult, }; use manifest::SymlinkPolicy; use schemars::JsonSchema; @@ -1011,6 +1012,18 @@ impl WorkdirSession for ScopedWorkdirSession { self.source.read(request).await } + async fn read_bytes( + &self, + mut request: ReadBytesRequest, + ) -> Result { + let path = self + .resolve_operation_path(&request.path, WorkdirToolScopePermission::Read) + .await?; + self.ensure_read(&path, WorkdirSessionCapability::Read)?; + request.path = path; + self.source.read_bytes(request).await + } + async fn write(&self, mut request: WriteRequest) -> Result { let _scope_guard = self.scope_lock.lock().await; let path = self @@ -1276,6 +1289,10 @@ impl WorkdirSession for ReadOnlyWorkdirSession { self.inner.read(request).await } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + self.inner.read_bytes(request).await + } + async fn write(&self, _request: WriteRequest) -> Result { Err(WorkdirError::Denied("read-only workdir session".into())) } @@ -1656,6 +1673,15 @@ mod tests { } } + fn read_bytes(path: &str) -> ReadBytesRequest { + ReadBytesRequest { + path: fs_path(path), + offset: 0, + max_bytes: 1024, + expected_hash: None, + } + } + fn write(path: &str, content: &str) -> WriteRequest { WriteRequest { path: fs_path(path), @@ -1887,6 +1913,14 @@ mod tests { child.read(read("readme.md")).await.unwrap().bytes, b"visible" ); + assert_eq!( + child + .read_bytes(read_bytes("readme.md")) + .await + .unwrap() + .bytes, + b"visible" + ); assert!(matches!( child.write(write("new.md", "no")).await, Err(WorkdirError::Denied(_)) diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 517f7c236..54514dd75 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -18,8 +18,8 @@ use workdir::workspace::WorkspaceWorkdirSessionOperationRequest; use workdir::{ CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, + WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, }; @@ -357,6 +357,13 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { } } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + match self.operate(WorkdirSessionOperation::ReadBytes(request))? { + WorkdirSessionOperationResult::ReadBytes(result) => Ok(result), + _ => Err(Self::mismatch("read_bytes")), + } + } + async fn write(&self, request: WriteRequest) -> Result { match self.operate(WorkdirSessionOperation::Write(request))? { WorkdirSessionOperationResult::Write(result) => Ok(result), diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 31649fcb3..35f37069b 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -1074,6 +1074,7 @@ fn external_workdir_operation_kind(operation: &WorkdirSessionOperation) -> &'sta WorkdirSessionOperation::ScopeRulesOverlap(_) => "scope_rules_overlap", WorkdirSessionOperation::Stat(_) => "stat", WorkdirSessionOperation::Read(_) => "read", + WorkdirSessionOperation::ReadBytes(_) => "read_bytes", WorkdirSessionOperation::Write(_) => "write", WorkdirSessionOperation::Edit(_) => "edit", WorkdirSessionOperation::List(_) => "list", @@ -1141,6 +1142,19 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { _ => Err(Self::mismatch("read")), } } + async fn read_bytes( + &self, + request: workdir::ReadBytesRequest, + ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; + match self + .operate(WorkdirSessionOperation::ReadBytes(request)) + .await? + { + WorkdirSessionOperationResult::ReadBytes(value) => Ok(value), + _ => Err(Self::mismatch("read_bytes")), + } + } async fn write( &self, request: workdir::WriteRequest, @@ -14824,6 +14838,7 @@ async fn scoped_execute_current_worker_workdir_operation( | WorkdirSessionOperation::ScopeRulesOverlap(_) | WorkdirSessionOperation::Stat(_) | WorkdirSessionOperation::Read(_) + | WorkdirSessionOperation::ReadBytes(_) | WorkdirSessionOperation::Write(_) | WorkdirSessionOperation::Edit(_) | WorkdirSessionOperation::List(_) @@ -14890,6 +14905,10 @@ async fn execute_workdir_session_operation( .read(request) .await .map(WorkdirSessionOperationResult::Read), + WorkdirSessionOperation::ReadBytes(request) => session + .read_bytes(request) + .await + .map(WorkdirSessionOperationResult::ReadBytes), WorkdirSessionOperation::Write(request) => session .write(request) .await @@ -17642,7 +17661,7 @@ async fn serve_external_workdir_provider( Ok(operation) => operation, Err(message) => { let _ = response.send(Err(WorkdirTransportError { - code: workdir::http::WorkdirTransportErrorCode::Unsupported, + code: workdir::http::WorkdirTransportErrorCode::InvalidRequest, message, })); continue; @@ -28109,6 +28128,8 @@ mod tests { ); let provider_root = tempfile::tempdir().unwrap(); fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); + let image_bytes = b"\x89PNG\r\n\x1a\nexternal\0image\n\xff"; + fs::write(provider_root.path().join("image.png"), image_bytes).unwrap(); let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( workdir::Workdir::new(&grant.working_directory_id), @@ -28178,6 +28199,61 @@ mod tests { let read_result = read.await.unwrap().unwrap(); assert_eq!(read_result.bytes, b"one\ntwo\n"); + let read_bytes = { + let broker_session = broker_session.clone(); + tokio::spawn(async move { + broker_session + .read_bytes(workdir::ReadBytesRequest { + path: workdir::WorkdirPath::new("image.png").unwrap(), + offset: 0, + max_bytes: 1024, + expected_hash: None, + }) + .await + }) + }; + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + assert!(matches!( + operation.as_inner(), + WorkdirSessionOperation::ReadBytes(_) + )); + let result = workdir::http::dispatch_workdir_session_operation( + &local_session, + operation.into_inner(), + ) + .await + .unwrap() + .try_into() + .unwrap(); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 2, + operation_id, + outcome: ExternalWorkdirOperationOutcome::Completed { result }, + }, + )) + .unwrap() + .into(), + )) + .await + .unwrap(); + let read_bytes_result = read_bytes.await.unwrap().unwrap(); + assert_eq!(read_bytes_result.bytes, image_bytes); + assert!(read_bytes_result.eof); + let edit = { let broker_session = broker_session.clone(); let expected_hash = read_result.content_hash; diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index 2330f6419..a2d3bc6b2 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -903,7 +903,7 @@ async fn execute_operation( Ok(result) => ExternalWorkdirOperationOutcome::Completed { result }, Err(message) => ExternalWorkdirOperationOutcome::Failed { error: ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError { - code: workdir::http::WorkdirTransportErrorCode::Unsupported, + code: workdir::http::WorkdirTransportErrorCode::InvalidRequest, message, }), },