From ab0fa25897d412947392cb445de762bc3c32e940 Mon Sep 17 00:00:00 2001 From: Gh0stly <26194374+Gh0stlyKn1ght@users.noreply.github.com> Date: Mon, 5 Oct 2026 01:02:40 -0400 Subject: [PATCH] security: bound automation previews and reply resources (#93) --- apps/photocraft/src/control_server.rs | 58 +++++++- book/src/automation/control-protocol.md | 7 +- book/src/automation/security-model.md | 5 +- book/src/security/automation-security.md | 21 ++- book/src/security/overview.md | 9 +- crates/automation/src/bridge.rs | 20 ++- crates/automation/src/budgets.rs | 169 +++++++++++++++++++++++ crates/automation/src/headless.rs | 10 +- crates/automation/src/lib.rs | 1 + crates/automation/src/rpc.rs | 120 +++++++++++++--- crates/automation/src/server.rs | 103 +++++++++++--- crates/automation/tests/mcp.rs | 50 +++++++ docs/control-protocol.md | 26 +++- 13 files changed, 544 insertions(+), 55 deletions(-) create mode 100644 crates/automation/src/budgets.rs diff --git a/apps/photocraft/src/control_server.rs b/apps/photocraft/src/control_server.rs index d5b747f2..38501178 100644 --- a/apps/photocraft/src/control_server.rs +++ b/apps/photocraft/src/control_server.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::mpsc::{Receiver, Sender, channel}; use std::time::Duration; +use photocraft_automation::budgets::write_reply; use photocraft_automation::security::{ ConnectionLimiter, LineRead, MAX_CONNECTIONS, MAX_REQUEST_BYTES, authentication_reply, configure_stream, read_bounded_line, }; @@ -62,7 +63,7 @@ fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui:: "ok": false, "error": format!("request exceeds {MAX_REQUEST_BYTES} bytes"), }); - let _ = writeln!(out, "{reply}"); + let _ = write_reply(&mut out, &reply); let _ = out.flush(); break; } @@ -72,7 +73,7 @@ fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui:: if !authenticated { let (reply, ok) = authentication_reply(&line, token); authenticated = ok; - if writeln!(out, "{reply}").is_err() || out.flush().is_err() || !authenticated { + if write_reply(&mut out, &reply).is_err() || out.flush().is_err() || !authenticated { break; } continue; @@ -95,7 +96,7 @@ fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui:: } Err(e) => json!({"ok": false, "error": format!("bad JSON: {e}")}), }; - if writeln!(out, "{reply}").is_err() { + if write_reply(&mut out, &reply).is_err() { break; } if out.flush().is_err() { @@ -103,3 +104,54 @@ fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui:: } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::io::BufRead; + + #[test] + fn oversized_reply_preserves_framing_id_and_the_next_control_request() { + use photocraft_automation::budgets::MAX_RESPONSE_BYTES; + const TOKEN: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, rx) = channel::(); + let server = std::thread::spawn(move || { + let (stream, _) = listener.accept().unwrap(); + serve(stream, TOKEN, tx, egui::Context::default()); + }); + let handler = std::thread::spawn(move || { + let first = rx.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!(first.method, "test.large"); + first.reply.send(json!({"ok": true, "result": "x".repeat(MAX_RESPONSE_BYTES)})).unwrap(); + let second = rx.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!(second.method, "test.small"); + second.reply.send(json!({"ok": true, "result": "still serving"})).unwrap(); + }); + let mut stream = TcpStream::connect(addr).unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + writeln!(stream, "{}", json!({"id": 1, "method": "auth", "params": {"token": TOKEN}})).unwrap(); + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + assert_eq!(serde_json::from_str::(&line).unwrap()["ok"], true); + writeln!(stream, "{}", json!({"id": 2, "method": "test.large"})).unwrap(); + line.clear(); + reader.read_line(&mut line).unwrap(); + let rejected: Value = serde_json::from_str(&line).unwrap(); + assert_eq!(rejected["id"], 2); + assert_eq!(rejected["ok"], false); + assert!(rejected["error"].as_str().unwrap().contains("operation may have completed")); + writeln!(stream, "{}", json!({"id": 3, "method": "test.small"})).unwrap(); + line.clear(); + reader.read_line(&mut line).unwrap(); + let accepted: Value = serde_json::from_str(&line).unwrap(); + assert_eq!(accepted["id"], 3); + assert_eq!(accepted["result"], "still serving"); + drop(reader); + drop(stream); + handler.join().unwrap(); + server.join().unwrap(); + } +} diff --git a/book/src/automation/control-protocol.md b/book/src/automation/control-protocol.md index 84baef48..c81c2058 100644 --- a/book/src/automation/control-protocol.md +++ b/book/src/automation/control-protocol.md @@ -21,13 +21,18 @@ Implementation is split across: before dispatch; the bridge accepts only loopback-style addresses and authenticates on connection; encoded request lines are limited to 1 MiB; active connections are limited to 16; socket I/O and handler waits use timeouts; MCP and headless batches are limited to 256 steps. +JSON-lines replies and encoded MCP tool results are capped at 8 MiB, with aggregate batch reply +budgets. Headless previews have source-pixel, requested-edge, and encoded-PNG ceilings. The +bridge bounds incoming reply frames and screenshot decoding. Automation file methods require separately configured read and write roots. Request paths are relative, and absolute, traversal, alternate-separator, device-name, and link-escape paths are rejected before file effects. **Known limitations:** the token is a bearer credential and grants the complete non-filesystem control surface. There is no encryption, client identity, general per-method capability check, -JSON-depth limit, or operation/resource budget. The listener creates one thread per authenticated +explicit JSON-depth policy, aggregate session-memory accounting, or command-duration/cancellation +budget. Desktop capture/encoding and import/export still need their own operation budgets. +The listener creates one thread per authenticated active connection, within the connection cap. Trusted one-shot CLI operations and interactive desktop file pickers retain the launching user's normal filesystem authority. diff --git a/book/src/automation/security-model.md b/book/src/automation/security-model.md index 2d0f181c..27c15eed 100644 --- a/book/src/automation/security-model.md +++ b/book/src/automation/security-model.md @@ -13,6 +13,8 @@ Automation requests are untrusted input with potential effects beyond the active | Symlink-safe capability filesystem | **Implemented:** relative operations use held directory capabilities and reject link escapes | | Request-byte and JSON-depth limits | **Partial:** 1 MiB request-line limit; no explicit JSON-depth policy | | Batch-step limit | **Implemented:** 256 steps for headless and MCP batches | +| Reply-size limit | **Implemented:** 8 MiB encoded JSON/MCP result; aggregate batch reply budget | +| Headless preview budgets | **Implemented:** 2048-pixel requested edge, 67,108,864 source pixels, 5 MiB encoded PNG | | Connection/worker limit | **Implemented:** 16 active TCP connections; one worker thread per accepted active connection | | Security audit events | **Proposed** | @@ -36,5 +38,6 @@ MCP or control client The current token proves possession of a secret but does not provide general method authorization. Filesystem handles and a defensive command-path policy are implemented; explicit non-filesystem -capabilities, broader resource budgets, and audit events remain. Private token handling, process +capabilities, session-memory/command-duration budgets, and audit events remain. Preview and +reply ceilings do not bound all document operations or desktop screenshot capture. Private token handling, process isolation, and least-privileged execution are still practical containment measures. diff --git a/book/src/security/automation-security.md b/book/src/security/automation-security.md index 3fbfb959..e4f7cd29 100644 --- a/book/src/security/automation-security.md +++ b/book/src/security/automation-security.md @@ -8,6 +8,15 @@ Automation is a privilege boundary because requests can cause filesystem access, - Headless TCP refuses a successfully bound non-loopback address. - Every desktop-control and headless-TCP connection must authenticate with a 256-bit bearer token before method dispatch. - Encoded request lines are limited to 1 MiB, active TCP connections to 16, and headless/MCP batches to 256 steps. +- JSON-lines replies and encoded MCP tool results are limited to 8 MiB; retained batch replies + have an aggregate budget and stop later steps when exhausted. The headless stdio JSON-lines + transport enforces the same request/reply ceilings as TCP. +- Headless previews reject requested edges above 2048 pixels and source documents above + 67,108,864 pixels before compositing. Full-size requests must fit the edge ceiling. Rendered + PNGs are limited to 5 MiB before base64 encoding or file writes. Trusted-local CLI rendering + retains its existing behavior. +- The MCP bridge bounds incoming replies and does not retry an operation after an oversized + reply. Screenshot decoding has separate dimension, pixel, and allocation ceilings. - MCP normally uses stdio and can optionally bridge to loopback control TCP. - The bridge and UI request paths use timeouts. - Engine commands are expected to reject invalid parameters without panicking. @@ -23,10 +32,18 @@ Automation is a privilege boundary because requests can cause filesystem access, - no per-client or per-tool capabilities; - no general per-client or per-tool capability model beyond filesystem read/write authority; -- no explicit JSON-depth, response-size, render, document-memory, or command-duration budget; +- no explicit JSON-depth policy, aggregate document/session-memory accounting, compositor + scratch-space accounting, or command-duration/cancellation budget; - one thread per accepted TCP connection; - no structured security audit event stream; -- full-size renders and expensive engine commands are not charged to a session budget. +- desktop screenshot capture/encoding and document import/export have no automation-specific + operation budgets; the desktop transport limits replies after their creation; +- expensive engine commands are not charged to a session budget. + +Reply rejection can happen after a command has changed state. Such errors report that the +operation may have completed. Exhausting a batch reply budget stops subsequent steps, including +when normal command errors would allow the batch to continue; it does not roll back prior edits. +See [Control protocol](../automation/control-protocol.md) for the authoritative protocol reference. ## Proposed session model diff --git a/book/src/security/overview.md b/book/src/security/overview.md index e9545076..21ce8e85 100644 --- a/book/src/security/overview.md +++ b/book/src/security/overview.md @@ -10,8 +10,8 @@ PhotoCraft processes complex attacker-controlled files and exposes automation th | Raster codecs | Default dimension/pixel/allocation limits; decoder-specific budgets; malformed/property tests; fuzz targets | Continuous fuzzing and more format-specific semantic cases | | `.pcraft` | Manifest/blob/total-byte limits; bounded decompression; CRC/hash and version checks; fuzz target | Tighter complexity/entry limits and directory-bundle path-race analysis | | Engine commands | Error-return contract, focused bad-parameter tests, adversarial `panic_hunt` test | Continue eliminating input-derived panic and exhaustion paths | -| Control TCP | Loopback binding; 256-bit bearer-token authentication; 1 MiB request lines; 16 active connections; I/O timeouts | No capabilities, filesystem policy, encryption, client identity, JSON-depth/operation budgets, or audit trail | -| MCP | Stdio transport; typed tool schemas; headless/bridge separation; authenticated bridge; 256-step batch limit | No authorization scopes or filesystem policy | +| Control TCP | Loopback binding; bearer-token authentication; rooted read/write authority; 1 MiB request lines; 8 MiB replies; 16 active connections; I/O timeouts | No general method capabilities, encryption, client identity, explicit JSON-depth policy, aggregate session-memory/command-duration budgets, or audit trail | +| MCP | Stdio; typed tools; authenticated and bounded bridge; rooted read/write authority; 256-step batches with aggregate reply budgets; headless preview/output ceilings | No general tool authorization scopes, aggregate session-memory accounting, or command cancellation; transport reply caps do not bound desktop screenshot capture | | Dependencies | Cargo lockfile and normal CI build/test/clippy | No `cargo-audit` or `cargo-deny` job/configuration confirmed in current tree | | Releases | Checksums and platform signing/notarization paths when secrets are configured | Signing can be absent; no SBOM generation confirmed in current tree | @@ -20,8 +20,9 @@ The table is a snapshot of the current repository, not a guarantee about deploye ## Security priorities 1. Add capability-scoped automation sessions on top of the authenticated transport. -2. Restrict automation file access through explicit workspace handles. -3. Add JSON-complexity, response-size, render, document-memory, and command-duration budgets. +2. Maintain the implemented workspace-handle boundary as file-bearing commands evolve. +3. Extend the preview/reply budgets to JSON complexity, aggregate document/session memory, + desktop capture/import/export, and command duration/cancellation. 4. Expand PSD/PSB, codec, and `.pcraft` fuzzing with regression corpora. 5. Add dependency advisories/license policy checks and release SBOMs. diff --git a/crates/automation/src/bridge.rs b/crates/automation/src/bridge.rs index 77efa392..210ac511 100644 --- a/crates/automation/src/bridge.rs +++ b/crates/automation/src/bridge.rs @@ -6,11 +6,12 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use serde_json::{Value, json}; -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; use tokio::net::TcpStream; use tokio::sync::Mutex; use crate::AutomationError; +use crate::budgets::MAX_RESPONSE_BYTES; use crate::security::{AUTH_METHOD, MAX_REQUEST_BYTES, validate_token}; type Conn = (BufReader, tokio::net::tcp::OwnedWriteHalf); @@ -75,6 +76,12 @@ impl BridgeClient { return Err(AutomationError::Bridge(format!("not connected to {}", self.addr))); }; match tokio::time::timeout(self.timeout, exchange(conn, id, method, ¶ms)).await { + Ok(Ok(Err(error @ AutomationError::BadRequest(_)))) => { + // An oversized frame leaves unread bytes. Drop this connection and + // report the budget failure without retrying a possibly completed edit. + *guard = None; + return Err(error); + } Ok(Ok(v)) => return v, Ok(Err(e)) if attempt == 0 => { *guard = None; @@ -104,14 +111,19 @@ async fn exchange(conn: &mut Conn, id: u64, method: &str, params: &Value) -> Res let io = |e: std::io::Error| AutomationError::Bridge(e.to_string()); conn.1.write_all(line.as_bytes()).await.map_err(io)?; conn.1.flush().await.map_err(io)?; - let mut buf = String::new(); + let mut buf = Vec::new(); loop { buf.clear(); - let n = conn.0.read_line(&mut buf).await.map_err(io)?; + // Read raw bytes so truncation inside a UTF-8 character still reports the + // budget error instead of a retryable decoding/transport error. + let n = (&mut conn.0).take((MAX_RESPONSE_BYTES + 1) as u64).read_until(b'\n', &mut buf).await.map_err(io)?; + if n > MAX_RESPONSE_BYTES { + return Ok(Err(AutomationError::BadRequest(format!("bridge response exceeds {MAX_RESPONSE_BYTES} bytes; operation may have completed")))); + } if n == 0 { return Err(AutomationError::Bridge("connection closed by the app".into())); } - let Ok(v) = serde_json::from_str::(buf.trim()) else { + let Ok(v) = serde_json::from_slice::(&buf) else { continue; }; if v.get("id").and_then(Value::as_u64) != Some(id) { diff --git a/crates/automation/src/budgets.rs b/crates/automation/src/budgets.rs new file mode 100644 index 00000000..37f6e024 --- /dev/null +++ b/crates/automation/src/budgets.rs @@ -0,0 +1,169 @@ +//! Fixed ceilings for automation previews and outbound JSON. These bound individual +//! outputs, not total document memory or the duration of an engine command. + +use std::io::{self, Write}; + +use serde::Serialize; +use serde_json::{Value, json}; + +use crate::AutomationError; +use crate::security::MAX_REQUEST_BYTES; + +/// Maximum encoded JSON reply, including its newline for JSON-lines transports. +pub const MAX_RESPONSE_BYTES: usize = 8 << 20; +/// Maximum PNG bytes before base64 expansion (leaves room for JSON/MCP metadata). +pub const MAX_PNG_BYTES: usize = 5 << 20; +/// Maximum requested preview edge. Zero still requests full size, within this ceiling. +pub const MAX_PREVIEW_SIDE: u32 = 2048; +/// Maximum source pixels processed by a headless automation preview. +pub const MAX_RENDER_SOURCE_PIXELS: u64 = 64 << 20; + +fn bad(message: impl Into) -> AutomationError { + AutomationError::BadRequest(message.into()) +} + +/// Check before the compositor runs. Oversized requests fail rather than silently downscale. +pub fn check_preview(width: u32, height: u32, max_side: u32) -> Result<(), AutomationError> { + if width == 0 || height == 0 { + return Err(bad("automation preview requires nonzero document dimensions")); + } + if u64::from(width) * u64::from(height) > MAX_RENDER_SOURCE_PIXELS { + return Err(bad(format!("automation preview source exceeds {MAX_RENDER_SOURCE_PIXELS} pixels"))); + } + let side = if max_side == 0 { width.max(height) } else { max_side }; + if side > MAX_PREVIEW_SIDE { + return Err(bad(format!("automation preview side exceeds {MAX_PREVIEW_SIDE} pixels"))); + } + Ok(()) +} + +/// Check before base64 encoding or writing a rendered PNG. +pub fn check_png(bytes: usize) -> Result<(), AutomationError> { + if bytes > MAX_PNG_BYTES { + return Err(bad(format!("automation PNG exceeds {MAX_PNG_BYTES} bytes"))); + } + Ok(()) +} + +struct LimitedWriter { + bytes: Vec, + maximum: usize, +} + +impl Write for LimitedWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + if buf.len() > self.maximum.saturating_sub(self.bytes.len()) { + return Err(io::Error::new(io::ErrorKind::InvalidData, format!("response exceeds {} bytes", self.maximum))); + } + self.bytes.try_reserve(buf.len()).map_err(|error| io::Error::other(format!("response allocation failed: {error}")))?; + self.bytes.extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +/// Serialize into a bounded buffer. No partial response is sent to the transport. +pub fn json_bytes(value: &impl Serialize) -> Result, AutomationError> { + encode_with_limit(value, MAX_RESPONSE_BYTES - 1) +} + +fn encode_with_limit(value: &impl Serialize, maximum: usize) -> Result, AutomationError> { + let mut writer = LimitedWriter { bytes: Vec::new(), maximum }; + serde_json::to_writer(&mut writer, value).map_err(|error| bad(format!("response encoding failed: {error}")))?; + Ok(writer.bytes) +} + +/// Encode the entire envelope before writing. A rejected reply is replaced by a small +/// error with the original ID; the operation may already have completed. +pub fn write_reply(out: &mut impl Write, reply: &Value) -> io::Result<()> { + let encoded = match json_bytes(reply) { + Ok(bytes) => bytes, + Err(_) => { + let error = json!({ + "id": reply.get("id").cloned().unwrap_or(Value::Null), + "ok": false, + "error": format!("response exceeds {MAX_RESPONSE_BYTES} bytes; operation may have completed"), + }); + // Request IDs arriving over the wire fit in MAX_REQUEST_BYTES. Bound even + // callers that construct a reply directly instead of using a transport. + match json_bytes(&error) { + Ok(bytes) => bytes, + Err(_) => b"{\"id\":null,\"ok\":false,\"error\":\"response budget exceeded\"}".to_vec(), + } + } + }; + out.write_all(&encoded)?; + out.write_all(b"\n") +} + +/// Reserve envelope/ID space and bound retained batch results incrementally, so +/// many individually valid replies cannot accumulate into an enormous batch. +pub struct BatchReplyBudget { + remaining: usize, +} + +impl Default for BatchReplyBudget { + fn default() -> Self { + Self { remaining: MAX_RESPONSE_BYTES - MAX_REQUEST_BYTES - 4096 } + } +} + +impl BatchReplyBudget { + /// Must be called before retaining the result or dispatching the next step. + pub fn charge(&mut self, result: &Value) -> Result<(), AutomationError> { + let bytes = encode_with_limit(result, self.remaining.saturating_sub(1))?; + self.remaining = self.remaining.saturating_sub(bytes.len() + 1); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preview_checks_full_size_requested_edge_and_source_before_rendering() { + assert!(check_preview(64, 32, 0).is_ok()); + assert!(check_preview(MAX_PREVIEW_SIDE, MAX_PREVIEW_SIDE, 0).is_ok()); + assert!(check_preview(MAX_PREVIEW_SIDE + 1, 1, 0).is_err()); + assert!(check_preview(64, 32, MAX_PREVIEW_SIDE + 1).is_err()); + assert!(check_preview(8192, 8192, 1024).is_ok()); + assert!(check_preview(8193, 8192, 1024).is_err()); + assert!(check_preview(u32::MAX, u32::MAX, 1).is_err()); + assert!(check_preview(0, 32, 1).is_err()); + } + + #[test] + fn encoded_bytes_include_json_escaping_and_accept_exact_boundary() { + assert_eq!(encode_with_limit(&json!("abc"), 5).unwrap(), b"\"abc\""); + assert!(encode_with_limit(&json!("abc"), 4).is_err()); + assert!(encode_with_limit(&json!("\n\n"), 5).is_err()); + assert!(encode_with_limit(&Value::Null, 0).is_err()); + assert!(check_png(MAX_PNG_BYTES).is_ok()); + assert!(check_png(MAX_PNG_BYTES + 1).is_err()); + } + + #[test] + fn oversized_reply_is_one_complete_error_with_matching_id() { + let reply = json!({"id": 7, "ok": true, "result": "x".repeat(MAX_RESPONSE_BYTES)}); + let mut out = Vec::new(); + write_reply(&mut out, &reply).unwrap(); + assert!(out.len() < 1024); + assert_eq!(out.iter().filter(|&&byte| byte == b'\n').count(), 1); + let error: Value = serde_json::from_slice(&out).unwrap(); + assert_eq!(error["id"], 7); + assert_eq!(error["ok"], false); + assert!(error["error"].as_str().unwrap().contains("operation may have completed")); + } + + #[test] + fn batch_cannot_retain_more_than_the_aggregate_reply_budget() { + let mut budget = BatchReplyBudget { remaining: 10 }; + budget.charge(&json!("abc")).unwrap(); + assert!(budget.charge(&json!("abc")).is_err()); + assert_eq!(budget.remaining, 4); + } +} diff --git a/crates/automation/src/headless.rs b/crates/automation/src/headless.rs index d73cac6c..995529de 100644 --- a/crates/automation/src/headless.rs +++ b/crates/automation/src/headless.rs @@ -139,7 +139,15 @@ impl Headless { pub fn render_png(&self, index: Option, max_side: u32) -> Result, AutomationError> { let i = self.doc_index(index)?; - files::render_png(&self.session.documents()[i].doc, max_side) + let doc = &self.session.documents()[i].doc; + if !matches!(&self.filesystem, Filesystem::TrustedLocal) { + crate::budgets::check_preview(doc.size.width, doc.size.height, max_side)?; + } + let png = files::render_png(doc, max_side)?; + if !matches!(&self.filesystem, Filesystem::TrustedLocal) { + crate::budgets::check_png(png.len())?; + } + Ok(png) } pub fn session_list(&self) -> Value { diff --git a/crates/automation/src/lib.rs b/crates/automation/src/lib.rs index 644b736d..f1e7049d 100644 --- a/crates/automation/src/lib.rs +++ b/crates/automation/src/lib.rs @@ -17,6 +17,7 @@ #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] pub mod bridge; +pub mod budgets; pub mod files; pub mod headless; pub mod rpc; diff --git a/crates/automation/src/rpc.rs b/crates/automation/src/rpc.rs index e1ff4924..d52c364c 100644 --- a/crates/automation/src/rpc.rs +++ b/crates/automation/src/rpc.rs @@ -21,6 +21,7 @@ use std::sync::{Arc, Mutex, PoisonError}; use base64::Engine as _; use serde_json::{Value, json}; +use crate::budgets::{BatchReplyBudget, write_reply}; use crate::security::{ ConnectionLimiter, LineRead, MAX_BATCH_STEPS, MAX_CONNECTIONS, MAX_REQUEST_BYTES, authentication_reply, configure_stream, read_bounded_line, }; @@ -96,7 +97,12 @@ impl Headless { } "doc.inspect" => self.inspect(index_of(&p)), "doc.render" => { - let max = p.get("maxSide").and_then(Value::as_u64).unwrap_or(1024) as u32; + let max = match p.get("maxSide") { + None => 1024, + Some(value) => { + value.as_u64().and_then(|value| u32::try_from(value).ok()).ok_or_else(|| bad("maxSide must be an unsigned 32-bit integer"))? + } + }; let png = self.render_png(index_of(&p), max)?; match str_of(&p, "path") { Some(path) => { @@ -131,6 +137,7 @@ impl Headless { let stop = p.get("stopOnError").and_then(Value::as_bool).unwrap_or(true); let mut results = Vec::with_capacity(steps.len()); let mut failed = 0usize; + let mut reply_budget = BatchReplyBudget::default(); for (i, s) in steps.iter().enumerate() { let params = s.get("params").cloned().unwrap_or(Value::Null); let r = if let Some(c) = str_of(s, "command") { @@ -140,14 +147,21 @@ impl Headless { } else { Err(bad(format!("step {i} needs `command` or `method`"))) }; - match r { - Ok(v) => results.push(json!({"ok": true, "result": v})), - Err(e) => { - failed += 1; - results.push(json!({"ok": false, "error": e.to_string()})); - if stop { - break; - } + let was_error = r.is_err(); + let result = match r { + Ok(v) => json!({"ok": true, "result": v}), + Err(e) => json!({"ok": false, "error": e.to_string()}), + }; + if let Err(error) = reply_budget.charge(&result) { + failed += 1; + results.push(json!({"ok": false, "error": format!("batch response budget exceeded; current step may have completed: {error}")})); + break; + } + results.push(result); + if was_error { + failed += 1; + if stop { + break; } } } @@ -179,17 +193,25 @@ pub fn respond(h: &Mutex, line: &str) -> Value { } /// Serve JSON lines from `r` to `w` until EOF. Blank lines are ignored. -pub fn serve_lines(h: &Mutex, r: impl BufRead, mut w: impl Write) -> std::io::Result<()> { - for line in r.lines() { - let line = line?; +pub fn serve_lines(h: &Mutex, mut r: impl BufRead, mut w: impl Write) -> std::io::Result<()> { + let mut line = String::new(); + loop { + match read_bounded_line(&mut r, &mut line)? { + LineRead::Eof => return Ok(()), + LineRead::TooLong => { + write_reply(&mut w, &json!({"id": null, "ok": false, "error": format!("request exceeds {MAX_REQUEST_BYTES} bytes")}))?; + w.flush()?; + return Ok(()); + } + LineRead::Line => {} + } if line.trim().is_empty() { continue; } let reply = respond(h, &line); - writeln!(w, "{reply}")?; + write_reply(&mut w, &reply)?; w.flush()?; } - Ok(()) } fn serve_tcp_connection(h: &Mutex, stream: std::net::TcpStream, token: &str) -> std::io::Result<()> { @@ -208,7 +230,7 @@ fn serve_tcp_connection(h: &Mutex, stream: std::net::TcpStream, token: "ok": false, "error": format!("request exceeds {MAX_REQUEST_BYTES} bytes"), }); - writeln!(out, "{reply}")?; + write_reply(&mut out, &reply)?; out.flush()?; return Ok(()); } @@ -222,7 +244,7 @@ fn serve_tcp_connection(h: &Mutex, stream: std::net::TcpStream, token: authenticated = ok; reply }; - writeln!(out, "{reply}")?; + write_reply(&mut out, &reply)?; out.flush()?; if !authenticated { return Ok(()); @@ -356,6 +378,72 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn automation_preview_rejects_oversized_full_size_and_numeric_wraparound() { + let mut h = Headless::new(); + h.handle("doc.new", json!({"width": 2049, "height": 1})).unwrap(); + for max in [json!(0), json!(2049), json!(u64::from(u32::MAX) + 1), json!(-1), json!(1.5)] { + assert!(h.handle("doc.render", json!({"maxSide": max})).is_err()); + } + // A smaller preview still works, and rejection leaves the source document intact. + let rendered = h.handle("doc.render", json!({"maxSide": 32})).unwrap(); + let png = base64::engine::general_purpose::STANDARD.decode(rendered["base64"].as_str().unwrap()).unwrap(); + assert_eq!(photocraft_codecs::decode(&png).unwrap().dimensions(), (32, 1)); + assert_eq!(h.session.active().unwrap().doc.size.width, 2049); + } + + #[test] + fn stdio_rejects_oversized_frames_before_dispatch() { + let h = Mutex::new(Headless::new()); + let input = " ".repeat(MAX_REQUEST_BYTES + 1); + let mut out = Vec::new(); + serve_lines(&h, input.as_bytes(), &mut out).unwrap(); + let reply: Value = serde_json::from_slice(&out).unwrap(); + assert_eq!(reply["ok"], false); + assert!(reply["error"].as_str().unwrap().contains("request exceeds")); + assert!(h.lock().unwrap().session.documents().is_empty()); + } + + #[test] + fn preview_rejection_precedes_output_file_creation() { + let root = std::env::temp_dir().join(format!("pc-preview-budget-{}", std::process::id())); + std::fs::create_dir_all(&root).unwrap(); + let workspace = crate::AuthorizedWorkspace::new(None, Some(&root)).unwrap(); + let mut h = Headless::with_workspace(workspace); + h.handle("doc.new", json!({"width": 16, "height": 8})).unwrap(); + let result = h.handle("doc.render", json!({"maxSide": 2049, "path": "rejected.png"})); + assert!(result.unwrap_err().to_string().contains("preview side exceeds")); + assert!(!root.join("rejected.png").exists()); + assert!(h.handle("doc.render", json!({"maxSide": 16, "path": "accepted.png"})).is_ok()); + assert!(root.join("accepted.png").exists()); + // Release directory capabilities before cleanup on Windows. + drop(h); + std::fs::remove_file(root.join("accepted.png")).unwrap(); + std::fs::remove_dir(root).unwrap(); + } + + #[test] + fn trusted_local_render_keeps_full_size_behavior() { + let mut h = Headless::trusted_local(); + h.handle("doc.new", json!({"width": 2049, "height": 1})).unwrap(); + let rendered = h.handle("doc.render", json!({"maxSide": 0})).unwrap(); + let png = base64::engine::general_purpose::STANDARD.decode(rendered["base64"].as_str().unwrap()).unwrap(); + assert_eq!(photocraft_codecs::decode(&png).unwrap().dimensions(), (2049, 1)); + } + + #[test] + fn batch_response_budget_stops_later_steps_even_when_stop_on_error_is_false() { + let mut h = Headless::new(); + let mut steps = vec![json!({"command": "command.list"}); MAX_BATCH_STEPS - 1]; + steps.push(json!({"command": "file.new", "params": {"width": 8, "height": 8}})); + let result = h.batch(&json!({"steps": steps, "stopOnError": false})).unwrap(); + assert_eq!(result["failed"], 1); + assert!(result["results"].as_array().unwrap().last().unwrap()["error"].as_str().unwrap().contains("batch response budget exceeded")); + assert!(h.session.documents().is_empty(), "later edit must not execute after the quota is exhausted"); + assert!(crate::budgets::json_bytes(&result).is_ok()); + assert!(h.handle("methods", Value::Null).is_ok()); + } + #[test] fn tcp_serves_loopback() { use std::io::{BufReader, Write as _}; diff --git a/crates/automation/src/server.rs b/crates/automation/src/server.rs index b10ce788..0e7ae568 100644 --- a/crates/automation/src/server.rs +++ b/crates/automation/src/server.rs @@ -14,6 +14,7 @@ use serde::Deserialize; use serde_json::{Value, json}; use crate::bridge::BridgeClient; +use crate::budgets::{BatchReplyBudget, MAX_PREVIEW_SIDE, check_png, json_bytes}; use crate::headless::Headless; use crate::security::MAX_BATCH_STEPS; use crate::{AuthorizedWorkspace, AutomationError, files}; @@ -88,7 +89,8 @@ pub struct SaveParams { pub struct PreviewParams { #[serde(default)] pub index: Option, - /// Longest side of the returned PNG in pixels (default 1024, 0 = full size). + /// Longest side of a headless preview (default 1024, maximum 2048). + /// Zero requests full size within that ceiling. Bridge mode returns the window screenshot. #[serde(default)] pub max_side: Option, } @@ -160,11 +162,26 @@ pub struct ControlParams { // --------------------------------------------------------------------------- fn ok_json(v: &Value) -> CallToolResult { - CallToolResult::success(vec![Content::text(serde_json::to_string_pretty(v).unwrap_or_default())]) + let text = match json_bytes(v).and_then(|bytes| String::from_utf8(bytes).map_err(|error| AutomationError::Other(error.to_string()))) { + Ok(text) => text, + Err(error) => return fail(format!("{error}; operation may have completed")), + }; + bounded_tool_result(CallToolResult::success(vec![Content::text(text)])) +} + +fn bounded_tool_result(result: CallToolResult) -> CallToolResult { + match json_bytes(&result) { + Ok(_) => result, + Err(error) => fail(format!("{error}; operation may have completed")), + } } fn fail(e: impl std::fmt::Display) -> CallToolResult { - CallToolResult::error(vec![Content::text(e.to_string())]) + let result = CallToolResult::error(vec![Content::text(e.to_string())]); + if json_bytes(&result).is_err() { + return CallToolResult::error(vec![Content::text("response budget exceeded; operation may have completed")]); + } + result } /// Neither backend took the call. The backend is always headless or a bridge, so this only @@ -174,8 +191,11 @@ fn no_backend() -> CallToolResult { } fn png_result(png: &[u8], note: String) -> CallToolResult { + if let Err(error) = check_png(png.len()) { + return fail(error); + } let b64 = base64::engine::general_purpose::STANDARD.encode(png); - CallToolResult::success(vec![Content::image(b64, "image/png"), Content::text(note)]) + bounded_tool_result(CallToolResult::success(vec![Content::image(b64, "image/png"), Content::text(note)])) } fn to_result(r: Result) -> Result { @@ -247,6 +267,9 @@ impl PhotocraftMcp { } async fn screenshot(&self, b: &BridgeClient, max_side: Option) -> Result { + if max_side.is_some_and(|side| side > MAX_PREVIEW_SIDE) { + return Ok(fail(format!("automation preview side exceeds {MAX_PREVIEW_SIDE} pixels"))); + } let response = match b.call("ui.screenshot", json!({})).await { Ok(response) => response, Err(error) => return Ok(fail(error)), @@ -258,19 +281,26 @@ impl PhotocraftMcp { Ok(bytes) => bytes, Err(error) => return Ok(fail(format!("app screenshot data: {error}"))), }; - let bytes = match max_side { - Some(m) if m > 0 => downscale_png(&bytes, m).unwrap_or(bytes), - _ => bytes, + if let Err(error) = check_png(bytes.len()) { + return Ok(fail(error)); + } + let bytes = match downscale_png(&bytes, max_side.unwrap_or(0)) { + Ok(Some(small)) => small, + Ok(None) => bytes, + Err(error) => return Ok(fail(error)), }; Ok(png_result(&bytes, "screenshot of the live app window".into())) } } -fn downscale_png(png: &[u8], max_side: u32) -> Option> { - let img = photocraft_codecs::decode(png).ok()?; +fn downscale_png(png: &[u8], max_side: u32) -> Result>, AutomationError> { + // A compressed reply's byte limit does not bound its decoded pixel allocation. + let opts = + photocraft_codecs::DecodeOptions { limits: photocraft_codecs::Limits { max_width: 8192, max_height: 8192, max_pixels: 16 << 20, max_alloc: 64 << 20 } }; + let img = photocraft_codecs::decode_as_with(photocraft_codecs::Format::Png, png, &opts).map_err(|error| AutomationError::Other(error.to_string()))?; let (w, h) = img.dimensions(); - if w.max(h) <= max_side { - return None; + if max_side == 0 || w.max(h) <= max_side { + return Ok(None); } let s = max_side as f32 / w.max(h) as f32; let (nw, nh) = (((w as f32 * s) as u32).max(1), ((h as f32 * s) as u32).max(1)); @@ -285,8 +315,11 @@ fn downscale_png(png: &[u8], max_side: u32) -> Option> { out[di..di + 4].copy_from_slice(&src[si..si + 4]); } } - let small = photocraft_codecs::Image::from_u8(nw, nh, photocraft_codecs::ChannelLayout::Rgba, out).ok()?; - photocraft_codecs::encode(&small, photocraft_codecs::Format::Png, &Default::default()).ok() + let small = + photocraft_codecs::Image::from_u8(nw, nh, photocraft_codecs::ChannelLayout::Rgba, out).map_err(|error| AutomationError::Other(error.to_string()))?; + let bytes = + photocraft_codecs::encode(&small, photocraft_codecs::Format::Png, &Default::default()).map_err(|error| AutomationError::Other(error.to_string()))?; + Ok(Some(bytes)) } // --------------------------------------------------------------------------- @@ -450,15 +483,24 @@ impl PhotocraftMcp { }; let mut results = Vec::new(); let mut failed = 0; + let mut reply_budget = BatchReplyBudget::default(); for s in &steps { - match b.call("engine.execute", s.clone()).await { - Ok(v) => results.push(json!({"ok": true, "result": v})), - Err(e) => { - failed += 1; - results.push(json!({"ok": false, "error": e.to_string()})); - if stop { - break; - } + let response = b.call("engine.execute", s.clone()).await; + let was_error = response.is_err(); + let result = match response { + Ok(value) => json!({"ok": true, "result": value}), + Err(error) => json!({"ok": false, "error": error.to_string()}), + }; + if let Err(error) = reply_budget.charge(&result) { + failed += 1; + results.push(json!({"ok": false, "error": format!("batch response budget exceeded; current step may have completed: {error}")})); + break; + } + results.push(result); + if was_error { + failed += 1; + if stop { + break; } } } @@ -574,6 +616,25 @@ pub fn render_document_png(doc: &photocraft_doc::Document, max_side: u32) -> Res mod tests { use super::*; + #[test] + fn mcp_text_checks_the_encoded_tool_envelope() { + // JSON escaping in the outer text-content envelope also consumes the budget. + let result = ok_json(&json!("\n".repeat(crate::budgets::MAX_RESPONSE_BYTES / 3))); + assert_eq!(result.is_error, Some(true)); + assert!(json_bytes(&result).is_ok()); + let error = fail("x".repeat(crate::budgets::MAX_RESPONSE_BYTES)); + assert_eq!(error.is_error, Some(true)); + assert!(json_bytes(&error).is_ok()); + } + + #[test] + fn bridge_screenshot_decode_limits_apply_even_without_downscaling() { + let image = photocraft_codecs::Image::from_u8(8193, 1, photocraft_codecs::ChannelLayout::Rgba, vec![0; 8193 * 4]).unwrap(); + let png = photocraft_codecs::encode(&image, photocraft_codecs::Format::Png, &Default::default()).unwrap(); + assert!(downscale_png(&png, 0).is_err()); + assert!(downscale_png(&png, 1024).is_err()); + } + #[tokio::test(flavor = "multi_thread")] async fn a_panicking_command_does_not_wedge_the_session() { let mcp = PhotocraftMcp::headless(); diff --git a/crates/automation/tests/mcp.rs b/crates/automation/tests/mcp.rs index f4179bad..51220611 100644 --- a/crates/automation/tests/mcp.rs +++ b/crates/automation/tests/mcp.rs @@ -189,6 +189,56 @@ async fn errors_are_tool_errors_not_crashes() { client.cancel().await.unwrap(); } +#[tokio::test(flavor = "multi_thread")] +async fn preview_budget_failure_preserves_the_mcp_session() { + let client = connect(PhotocraftMcp::headless()).await; + json_of(&call(&client, "doc_new", json!({"width": 2049, "height": 1})).await); + for side in [0, 2049] { + let reply = call(&client, "doc_render_preview", json!({"max_side": side})).await; + assert_eq!(reply.is_error, Some(true)); + assert!(text(&reply).contains("preview side exceeds")); + } + let preview = call(&client, "doc_render_preview", json!({"max_side": 32})).await; + assert!(preview.content.iter().any(|content| content.as_image().is_some())); + let inspected = json_of(&call(&client, "doc_inspect", json!({})).await); + assert_eq!(inspected["width"], 2049); + client.cancel().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn bridge_response_budget_drops_connection_without_retrying_the_operation() { + use photocraft_automation::{BridgeClient, budgets::MAX_RESPONSE_BYTES}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let app = tokio::spawn(async move { + let (sock, _) = listener.accept().await.unwrap(); + let (read, mut write) = sock.into_split(); + let mut reader = BufReader::new(read); + let mut line = String::new(); + reader.read_line(&mut line).await.unwrap(); + let auth: Value = serde_json::from_str(&line).unwrap(); + let reply = json!({"id": auth["id"], "ok": true, "result": {"authenticated": true}}); + write.write_all(format!("{reply}\n").as_bytes()).await.unwrap(); + line.clear(); + reader.read_line(&mut line).await.unwrap(); + let request: Value = serde_json::from_str(&line).unwrap(); + assert_eq!(request["method"], "engine.execute"); + let mut oversized = vec![b'x'; MAX_RESPONSE_BYTES]; + // The read limit cuts this multi-byte character in half. Classification must + // still be a non-retryable budget error, not a UTF-8 transport failure. + oversized.extend_from_slice("é".as_bytes()); + write.write_all(&oversized).await.unwrap(); + // A retry would connect to this still-live listener. The budget error must return + // directly instead, without waiting for another authentication exchange. + let next = tokio::time::timeout(std::time::Duration::from_millis(250), listener.accept()).await; + assert!(next.is_err(), "oversized reply must not retry the operation"); + }); + let bridge = BridgeClient::new(&addr, CONTROL_TOKEN).unwrap(); + let result = tokio::time::timeout(std::time::Duration::from_secs(5), bridge.call("engine.execute", json!({"command": "command.list"}))).await.unwrap(); + assert!(result.unwrap_err().to_string().contains("bridge response exceeds")); + app.await.unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn open_png_and_inspect() { let dir = tmp("open"); diff --git a/docs/control-protocol.md b/docs/control-protocol.md index 1a2dc8a6..640c9592 100644 --- a/docs/control-protocol.md +++ b/docs/control-protocol.md @@ -209,10 +209,32 @@ in headless and bridge mode. The desktop and headless TCP listeners currently enforce: - a 1 MiB maximum encoded request line; +- an 8 MiB maximum encoded JSON reply, including the newline; - at most 16 simultaneously serviced connections per listener; - a 30-second socket read/write timeout; - at most 256 steps in a headless `batch` or MCP `command_batch` request. An oversized line, excess connection, unauthenticated request, or unauthorized filesystem path is -rejected before command dispatch or file effects. These limits do not impose JSON-depth, -response-size, render, document-memory, command-duration, or general per-method capability budgets. +rejected before command dispatch or file effects. The headless JSON-lines stdio server also +enforces the request and reply byte ceilings. MCP tool results are checked as encoded JSON, +including the text/image content envelope, and the MCP bridge bounds incoming desktop replies. + +Headless automation previews (`doc.render` / MCP `doc_render_preview`) allow a maximum requested +edge of 2048 pixels and a source document of at most 67,108,864 pixels. `maxSide: 0` (MCP +`max_side: 0`) still means full size, but fails if the document's longest edge exceeds 2048. +PNG results are capped at 5 MiB before base64 encoding or writing a rendered file. Preview +dimension/source checks run before compositing; the encoded-PNG check runs after encoding. +Explicit trusted-local CLI rendering keeps its existing behavior. + +Batch replies have an aggregate byte budget with space reserved for the outer reply and ID. +Exhausting it stops later steps even when `stopOnError` is false. A result can exceed the reply +budget after an edit has run: the error says the operation may have completed. Earlier steps +are not rolled back; inspect state before retrying. The MCP bridge drops an oversized incoming +reply without automatically retrying the operation. Bridged screenshots are decoded with +8192-pixel edge, 16,777,216-pixel, and 64 MiB allocation ceilings, even without downscaling. + +These ceilings do not implement explicit JSON-depth policy, total session/document-memory +accounting, compositor scratch-space accounting, command cancellation/duration limits, or +general per-method capabilities. Desktop screenshot capture/encoding and document import/export +still need their own operation budgets; the desktop reply ceiling applies after the UI creates +its response. A bounded output does not imply bounded command cost.