diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 4003afe..04018d9 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -66,6 +66,10 @@ ### Fixed +- Bound newline-delimited input framing for `AcpAgent`, `Stdio`, and `ByteStreams` to 16 MiB by + default before JSON-RPC parsing, with explicit overrides for peers that legitimately need a + different finite ceiling. ([#340](https://github.com/agentclientprotocol/rust-sdk/issues/340), + [#342](https://github.com/agentclientprotocol/rust-sdk/issues/342)) - Preserve stable v1 `NewSessionResponse::config_options` on `ActiveSession`, expose them through `ActiveSession::config_options`, and include them in reconstructed and proxied session responses. diff --git a/src/agent-client-protocol/src/acp_agent.rs b/src/agent-client-protocol/src/acp_agent.rs index b43e931..fb5d029 100644 --- a/src/agent-client-protocol/src/acp_agent.rs +++ b/src/agent-client-protocol/src/acp_agent.rs @@ -6,15 +6,15 @@ use std::collections::{BTreeMap, VecDeque}; use std::path::{Path, PathBuf}; +use std::pin::pin; use std::str::FromStr; use std::sync::Arc; use std::time::Duration; use async_process::Child; use serde::{Deserialize, Serialize}; -use std::pin::pin; -use crate::{Client, Conductor, Role}; +use crate::{Client, Conductor, DEFAULT_LINE_LIMIT, Role, line::BoundedLines}; type DebugCallback = Arc; @@ -165,6 +165,7 @@ impl AcpAgentConfig { pub struct AcpAgent { config: AcpAgentConfig, debug_callback: Option, + stdout_line_limit: usize, } impl std::fmt::Debug for AcpAgent { @@ -175,6 +176,7 @@ impl std::fmt::Debug for AcpAgent { "debug_callback", &self.debug_callback.as_ref().map(|_| "..."), ) + .field("stdout_line_limit", &self.stdout_line_limit) .finish() } } @@ -186,6 +188,7 @@ impl AcpAgent { Self { config, debug_callback: None, + stdout_line_limit: DEFAULT_LINE_LIMIT, } } @@ -242,6 +245,19 @@ impl AcpAgent { self } + /// Set the maximum number of bytes accepted before the newline terminating one ACP stdout + /// frame. The default is [`DEFAULT_LINE_LIMIT`]. An optional carriage return counts + /// toward this limit; the newline itself does not. + /// + /// Exceeding the limit fails the protocol transport and terminates the spawned agent process + /// group. Set a larger finite value when an agent legitimately emits unusually large inline + /// content. + #[must_use] + pub fn with_stdout_line_limit(mut self, limit: usize) -> Self { + self.stdout_line_limit = limit; + self + } + /// Spawn the configured process and return its stdio streams and raw child handle. /// /// This is a low-level escape hatch. The caller owns the returned child process and is @@ -633,9 +649,10 @@ impl crate::ConnectTo for Acp self, client: impl crate::ConnectTo, ) -> Result<(), crate::Error> { + use futures::StreamExt; use futures::io::BufReader; - use futures::{AsyncBufReadExt, StreamExt}; + let stdout_line_limit = self.stdout_line_limit; let (child_stdin, child_stdout, child_stderr, child) = self.spawn_process()?; // Create a channel to collect stderr for error reporting @@ -668,13 +685,20 @@ impl crate::ConnectTo for Acp let incoming_lines: std::pin::Pin< Box> + Send>, > = if let Some(callback) = self.debug_callback.clone() { - Box::pin(BufReader::new(child_stdout).lines().inspect(move |result| { - if let Ok(line) = result { - callback(line, LineDirection::Stdout); - } - })) + Box::pin( + BoundedLines::new(BufReader::new(child_stdout), stdout_line_limit).inspect( + move |result| { + if let Ok(line) = result { + callback(line, LineDirection::Stdout); + } + }, + ), + ) } else { - Box::pin(BufReader::new(child_stdout).lines()) + Box::pin(BoundedLines::new( + BufReader::new(child_stdout), + stdout_line_limit, + )) }; // The JSON-RPC transport keeps polling stdout while it drains stdin. @@ -893,6 +917,30 @@ mod tests { (callback, lines) } + #[cfg(unix)] + #[tokio::test] + async fn stdout_line_overflow_fails_transport_without_waiting_for_child_exit() { + let agent = AcpAgent::from_args(["/bin/sh", "-c", "printf '%065d' 0; sleep 30"]) + .unwrap() + .with_stdout_line_limit(64); + + let error = + tokio::time::timeout(Duration::from_secs(5), Client.builder().connect_to(agent)) + .await + .expect("overflow should stop the transport before the child exits") + .expect_err("oversized stdout should fail the transport"); + let detail = error + .data + .as_ref() + .map(serde_json::Value::to_string) + .unwrap_or_default(); + + assert!( + detail.contains("64-byte limit"), + "unexpected error: {error:?}" + ); + } + #[test] fn stderr_tail_keeps_last_bytes() { let initial = vec![b'a'; STDERR_CAPTURE_LIMIT]; diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 9cff5a6..888034b 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -6317,6 +6317,7 @@ where pub struct ByteStreams { outgoing: OB, incoming: IB, + incoming_line_limit: usize, } impl ByteStreams @@ -6326,7 +6327,20 @@ where { /// Create a new byte stream transport. pub fn new(outgoing: OB, incoming: IB) -> Self { - Self { outgoing, incoming } + Self { + outgoing, + incoming, + incoming_line_limit: crate::DEFAULT_LINE_LIMIT, + } + } + + /// Set the maximum number of bytes accepted before the newline terminating one incoming ACP + /// frame. The default is [`crate::DEFAULT_LINE_LIMIT`]. An optional carriage return counts + /// toward this limit; the newline itself does not. + #[must_use] + pub fn with_incoming_line_limit(mut self, limit: usize) -> Self { + self.incoming_line_limit = limit; + self } fn into_lines( @@ -6335,11 +6349,17 @@ where impl futures::Sink + Send + 'static, impl futures::Stream> + Send + 'static, > { - use futures::AsyncBufReadExt; use futures::io::BufReader; - let Self { outgoing, incoming } = self; + let Self { + outgoing, + incoming, + incoming_line_limit, + } = self; - let incoming_lines = Box::pin(BufReader::new(incoming).lines()); + let incoming_lines = Box::pin(crate::line::BoundedLines::new( + BufReader::new(incoming), + incoming_line_limit, + )); let outgoing_lines = futures::sink::unfold(Box::pin(outgoing), async move |mut writer, line: String| { write_line(&mut writer, line).await?; @@ -6509,6 +6529,20 @@ impl ConnectTo for Channel { mod tests { use super::*; + #[tokio::test] + async fn byte_streams_reject_an_incoming_line_over_the_configured_limit() { + let mut input = vec![b'x'; 65]; + input.push(b'\n'); + let mut lines = ByteStreams::new(futures::io::sink(), futures::io::Cursor::new(input)) + .with_incoming_line_limit(64) + .into_lines(); + + let error = lines.incoming.next().await.unwrap().unwrap_err(); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + assert!(error.to_string().contains("64-byte limit")); + } + #[cfg(feature = "unstable_protocol_v2")] fn connection_with_task_receiver() -> ( ConnectionTo, diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 4f71dbd..ff179ea 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -161,6 +161,9 @@ pub use agent_client_protocol_derive::{JsonRpcNotification, JsonRpcRequest, Json mod session; pub use session::*; +mod line; +pub use line::DEFAULT_LINE_LIMIT; + #[cfg(not(target_family = "wasm"))] mod acp_agent; #[cfg(not(target_family = "wasm"))] diff --git a/src/agent-client-protocol/src/line.rs b/src/agent-client-protocol/src/line.rs new file mode 100644 index 0000000..2b1c397 --- /dev/null +++ b/src/agent-client-protocol/src/line.rs @@ -0,0 +1,127 @@ +use std::{ + io, + pin::Pin, + task::{Context, Poll}, +}; + +/// Default maximum size of one newline-delimited ACP frame. +pub const DEFAULT_LINE_LIMIT: usize = 16 * 1024 * 1024; + +pub(crate) struct BoundedLines { + reader: Pin>, + buffer: Vec, + limit: usize, + finished: bool, +} + +impl BoundedLines { + pub(crate) fn new(reader: R, limit: usize) -> Self { + Self { + reader: Box::pin(reader), + buffer: Vec::new(), + limit, + finished: false, + } + } + + fn take_line(&mut self, terminated: bool) -> io::Result { + if terminated && self.buffer.last() == Some(&b'\r') { + self.buffer.pop(); + } + String::from_utf8(std::mem::take(&mut self.buffer)).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + ) + }) + } + + fn overflow(&mut self) -> Poll>> { + self.finished = true; + Poll::Ready(Some(Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("ACP line exceeds configured {}-byte limit", self.limit), + )))) + } +} + +impl futures::Stream for BoundedLines { + type Item = io::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.finished { + return Poll::Ready(None); + } + + loop { + let this = &mut *self; + let available = match this.reader.as_mut().poll_fill_buf(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(error)) => { + this.finished = true; + return Poll::Ready(Some(Err(error))); + } + Poll::Ready(Ok(available)) => available, + }; + if available.is_empty() { + this.finished = true; + return if this.buffer.is_empty() { + Poll::Ready(None) + } else { + Poll::Ready(Some(this.take_line(false))) + }; + } + + let newline = available.iter().position(|byte| *byte == b'\n'); + let payload_bytes = newline.unwrap_or(available.len()); + let Some(payload_len) = this.buffer.len().checked_add(payload_bytes) else { + return this.overflow(); + }; + if payload_len > this.limit { + return this.overflow(); + } + + let consumed = newline.map_or(available.len(), |position| position + 1); + this.buffer.extend_from_slice(&available[..payload_bytes]); + this.reader.as_mut().consume(consumed); + if newline.is_some() { + return Poll::Ready(Some(this.take_line(true))); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt as _; + + #[tokio::test] + async fn accepts_the_exact_limit_and_strips_line_endings() { + let limit = 64; + let mut input = vec![b'x'; limit - 1]; + input.extend_from_slice(b"\r\nnext"); + let reader = futures::io::BufReader::with_capacity(8, futures::io::Cursor::new(input)); + let mut lines = BoundedLines::new(reader, limit); + + assert_eq!(lines.next().await.unwrap().unwrap(), "x".repeat(limit - 1)); + assert_eq!(lines.next().await.unwrap().unwrap(), "next"); + assert!(lines.next().await.is_none()); + } + + #[tokio::test] + async fn rejects_oversize_without_retaining_more_than_the_limit() { + let limit = 64; + let mut input = vec![b'x'; limit + 1]; + input.push(b'\n'); + let reader = futures::io::BufReader::with_capacity(8, futures::io::Cursor::new(input)); + let mut lines = BoundedLines::new(reader, limit); + + let error = lines.next().await.unwrap().unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert!(error.to_string().contains("64-byte limit")); + assert!(lines.buffer.len() <= limit); + assert!(lines.next().await.is_none()); + } +} diff --git a/src/agent-client-protocol/src/stdio.rs b/src/agent-client-protocol/src/stdio.rs index d402ac2..ba8e93d 100644 --- a/src/agent-client-protocol/src/stdio.rs +++ b/src/agent-client-protocol/src/stdio.rs @@ -1,7 +1,7 @@ //! Stdio transport for connecting ACP components via standard input/output. use crate::acp_agent::LineDirection; -use crate::{ByteStreams, ConnectTo, Role}; +use crate::{ByteStreams, ConnectTo, DEFAULT_LINE_LIMIT, Role, line::BoundedLines}; use std::sync::Arc; /// A transport that connects to an ACP peer via standard input/output. @@ -10,11 +10,14 @@ use std::sync::Arc; /// which is the standard transport for MCP and ACP subprocess communication. pub struct Stdio { debug_callback: Option>, + stdin_line_limit: usize, } impl std::fmt::Debug for Stdio { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Stdio").finish_non_exhaustive() + f.debug_struct("Stdio") + .field("stdin_line_limit", &self.stdin_line_limit) + .finish_non_exhaustive() } } @@ -24,6 +27,7 @@ impl Stdio { pub fn new() -> Self { Self { debug_callback: None, + stdin_line_limit: DEFAULT_LINE_LIMIT, } } @@ -36,6 +40,15 @@ impl Stdio { self.debug_callback = Some(Arc::new(callback)); self } + + /// Set the maximum number of bytes accepted before the newline terminating one ACP stdin + /// frame. The default is [`DEFAULT_LINE_LIMIT`]. An optional carriage return counts toward + /// this limit; the newline itself does not. + #[must_use] + pub fn with_stdin_line_limit(mut self, limit: usize) -> Self { + self.stdin_line_limit = limit; + self + } } impl Default for Stdio { @@ -53,15 +66,19 @@ impl ConnectTo for Stdio { let stdout = blocking::Unblock::new(std::io::stdout()); if let Some(callback) = self.debug_callback { + use futures::StreamExt; use futures::io::BufReader; - use futures::{AsyncBufReadExt, StreamExt}; let incoming_callback = callback.clone(); - let incoming_lines = Box::pin(BufReader::new(stdin).lines().inspect(move |result| { - if let Ok(line) = result { - incoming_callback(line, LineDirection::Stdin); - } - })) + let incoming_lines = Box::pin( + BoundedLines::new(BufReader::new(stdin), self.stdin_line_limit).inspect( + move |result| { + if let Ok(line) = result { + incoming_callback(line, LineDirection::Stdin); + } + }, + ), + ) as std::pin::Pin> + Send>>; let outgoing_sink = Box::pin(futures::sink::unfold( @@ -80,7 +97,11 @@ impl ConnectTo for Stdio { ) .await } else { - ConnectTo::::connect_to(ByteStreams::new(stdout, stdin), client).await + ConnectTo::::connect_to( + ByteStreams::new(stdout, stdin).with_incoming_line_limit(self.stdin_line_limit), + client, + ) + .await } } }