Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: add bounded binary reads for ViewImage

Browse Source
This commit is contained in:
Hare
2026-09-30 07:44:59 +09:00
parent cafbb73a22
commit 05d6da45fb
14 changed files with 816 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/fs-operation/src/lib.rs
+99 -1
View File
@@ -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};
crates/fs-operation/src/local.rs
+134 -2
View File
@@ -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<ReadBytesResult, FsError> {
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<ReadBytesResult, FsError> {
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<BoundedReadLimits>,
) -> Result<ReadBytesResult, FsError> {
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,
crates/fs-operation/src/operation.rs
+30
View File
@@ -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<ContentHash>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ReadBytesResult {
pub path: FsPath,
pub bytes: Vec<u8>,
#[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 {
crates/tools/src/view_image.rs
+92 -21
View File
@@ -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<Vec<u8>, 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))
}
crates/tools/tests/integration.rs
+218
View File
@@ -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<std::path::PathBuf>,
}
#[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<workdir::StatResult, workdir::WorkdirError> {
self.inner.stat(request).await
}
async fn read(
&self,
request: workdir::ReadRequest,
) -> Result<workdir::ReadResult, workdir::WorkdirError> {
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<workdir::ReadBytesResult, workdir::WorkdirError> {
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<workdir::WriteResult, workdir::WorkdirError> {
self.inner.write(request).await
}
async fn edit(
&self,
request: workdir::EditRequest,
) -> Result<workdir::EditResult, workdir::WorkdirError> {
self.inner.edit(request).await
}
async fn list(
&self,
request: workdir::ListRequest,
) -> Result<workdir::ListResult, workdir::WorkdirError> {
self.inner.list(request).await
}
async fn glob(
&self,
request: workdir::GlobRequest,
) -> Result<workdir::GlobResult, workdir::WorkdirError> {
self.inner.glob(request).await
}
async fn grep(
&self,
request: workdir::GrepRequest,
) -> Result<workdir::GrepResult, workdir::WorkdirError> {
self.inner.grep(request).await
}
async fn start_command(
&self,
request: workdir::CommandRequest,
) -> Result<workdir::CommandHandle, workdir::WorkdirError> {
self.inner.start_command(request).await
}
async fn command_status(
&self,
handle: workdir::CommandHandle,
) -> Result<workdir::CommandStatus, workdir::WorkdirError> {
self.inner.command_status(handle).await
}
async fn command_output(
&self,
request: workdir::CommandOutputRequest,
) -> Result<workdir::CommandOutput, workdir::WorkdirError> {
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();
crates/workdir/src/external.rs
+61 -5
View File
@@ -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::<ExternalWorkdirServerFrame>(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,
crates/workdir/src/http.rs
+22 -3
View File
@@ -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<ReadBytesResult, WorkdirError> {
match self
.operate(WorkdirSessionOperation::ReadBytes(request))
.await?
{
WorkdirSessionOperationResult::ReadBytes(result) => Ok(result),
_ => Err(Self::mismatch("read_bytes")),
}
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
match self
.operate(WorkdirSessionOperation::Write(request))
crates/workdir/src/lib.rs
+10 -1
View File
@@ -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<StatResult, WorkdirError>;
async fn read(&self, request: ReadRequest) -> Result<ReadResult, WorkdirError>;
async fn read_bytes(
&self,
_request: ReadBytesRequest,
) -> Result<ReadBytesResult, WorkdirError> {
Err(WorkdirError::UnsupportedOperation(
"bounded binary reads are not supported by this provider".to_string(),
))
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError>;
async fn edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError>;
async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError>;
crates/workdir/src/local.rs
+19 -4
View File
@@ -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<ReadBytesResult, WorkdirError> {
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<WriteResult, WorkdirError> {
self.ensure_capability(WorkdirSessionCapability::Write)?;
let logical = request.path.clone();
crates/workdir/src/router.rs
+9 -3
View File
@@ -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<ReadBytesResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.read_bytes(request).await
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.write(request).await
crates/workdir/src/scope.rs
+35 -1
View File
@@ -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<ReadBytesResult, WorkdirError> {
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<WriteResult, WorkdirError> {
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<ReadBytesResult, WorkdirError> {
self.inner.read_bytes(request).await
}
async fn write(&self, _request: WriteRequest) -> Result<WriteResult, WorkdirError> {
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(_))
crates/worker/src/feature/builtin/manage_workdir.rs
+9 -2
View File
@@ -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<ReadBytesResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::ReadBytes(request))? {
WorkdirSessionOperationResult::ReadBytes(result) => Ok(result),
_ => Err(Self::mismatch("read_bytes")),
}
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Write(request))? {
WorkdirSessionOperationResult::Write(result) => Ok(result),
crates/workspace-server/src/server.rs
+77 -1
View File
@@ -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<workdir::ReadBytesResult, workdir::WorkdirError> {
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::<ExternalWorkdirServerFrame>(&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;
crates/yoi/src/workdir_share.rs
+1 -1
View File
@@ -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,
}),
},