mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Bound exec-server stdio JSON-RPC messages (#32123)
## Why Newline-delimited stdio input could buffer an unterminated JSON-RPC message without a per-message limit. Apply the same 64 MiB ceiling used by the other exec-server transports. ## What changed - Read stdio messages with bounded lookahead and disconnect when a message exceeds the limit. - Preserve LF and CRLF framing, including messages whose payload is exactly at the limit. ## Testing Added tests for accepting a limit-sized CRLF message and rejecting an unterminated overlong message. GitOrigin-RevId: 73eaf883324a7777960725037cc5b9c34720084f
This commit is contained in:
@@ -24,11 +24,15 @@ use tracing::debug;
|
||||
use tracing::warn;
|
||||
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::io::BufWriter;
|
||||
|
||||
pub(crate) const CHANNEL_CAPACITY: usize = 128;
|
||||
// Match the existing serialized JSON-RPC message ceiling used by Noise and
|
||||
// WebSocket transports so stdio has the same per-message bound.
|
||||
const MAX_STDIO_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
|
||||
const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2);
|
||||
#[cfg(test)]
|
||||
pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25);
|
||||
@@ -230,6 +234,24 @@ pub(crate) struct JsonRpcConnection {
|
||||
|
||||
impl JsonRpcConnection {
|
||||
pub(crate) fn from_stdio<R, W>(reader: R, writer: W, connection_label: String) -> Self
|
||||
where
|
||||
R: AsyncRead + Unpin + Send + 'static,
|
||||
W: AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
Self::from_stdio_with_max_message_len(
|
||||
reader,
|
||||
writer,
|
||||
connection_label,
|
||||
MAX_STDIO_JSONRPC_MESSAGE_LEN,
|
||||
)
|
||||
}
|
||||
|
||||
fn from_stdio_with_max_message_len<R, W>(
|
||||
reader: R,
|
||||
writer: W,
|
||||
connection_label: String,
|
||||
max_message_len: usize,
|
||||
) -> Self
|
||||
where
|
||||
R: AsyncRead + Unpin + Send + 'static,
|
||||
W: AsyncWrite + Unpin + Send + 'static,
|
||||
@@ -241,11 +263,63 @@ impl JsonRpcConnection {
|
||||
let reader_label = connection_label.clone();
|
||||
let incoming_tx_for_reader = incoming_tx.clone();
|
||||
let disconnected_tx_for_reader = disconnected_tx.clone();
|
||||
// Read one byte past the payload limit so an unterminated oversized
|
||||
// message fails promptly. A trailing CR gets one more byte of lookahead
|
||||
// because it may be the first half of a valid CRLF terminator.
|
||||
let read_limit = u64::try_from(max_message_len.saturating_add(1)).unwrap_or(u64::MAX);
|
||||
let reader_task = tokio::spawn(async move {
|
||||
let mut lines = BufReader::new(reader).lines();
|
||||
let mut reader = BufReader::new(reader);
|
||||
let mut line = String::new();
|
||||
loop {
|
||||
match lines.next_line().await {
|
||||
Ok(Some(line)) => {
|
||||
line.clear();
|
||||
let read_result = (&mut reader).take(read_limit).read_line(&mut line).await;
|
||||
match read_result {
|
||||
Ok(0) => {
|
||||
send_disconnected(
|
||||
&incoming_tx_for_reader,
|
||||
&disconnected_tx_for_reader,
|
||||
/*reason*/ None,
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(_) => {
|
||||
if line.ends_with('\n') {
|
||||
line.pop();
|
||||
if line.ends_with('\r') {
|
||||
line.pop();
|
||||
}
|
||||
} else if line.len() > max_message_len && line.ends_with('\r') {
|
||||
match reader.read_u8().await {
|
||||
Ok(b'\n') => {
|
||||
line.pop();
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => {}
|
||||
Err(err) => {
|
||||
send_disconnected(
|
||||
&incoming_tx_for_reader,
|
||||
&disconnected_tx_for_reader,
|
||||
Some(format!(
|
||||
"failed to read JSON-RPC message from {reader_label}: {err}"
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if line.len() > max_message_len {
|
||||
send_disconnected(
|
||||
&incoming_tx_for_reader,
|
||||
&disconnected_tx_for_reader,
|
||||
Some(format!(
|
||||
"JSON-RPC message from {reader_label} exceeds maximum length of {max_message_len} bytes"
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
@@ -270,15 +344,6 @@ impl JsonRpcConnection {
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
send_disconnected(
|
||||
&incoming_tx_for_reader,
|
||||
&disconnected_tx_for_reader,
|
||||
/*reason*/ None,
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Err(err) => {
|
||||
send_disconnected(
|
||||
&incoming_tx_for_reader,
|
||||
@@ -601,6 +666,7 @@ mod tests {
|
||||
use codex_exec_server_protocol::RequestId;
|
||||
use futures::channel::mpsc as futures_mpsc;
|
||||
use futures::task::AtomicWaker;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::accept_async;
|
||||
@@ -608,6 +674,66 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn stdio_connection_accepts_message_at_size_limit() -> anyhow::Result<()> {
|
||||
let message = test_jsonrpc_message();
|
||||
let encoded = serde_json::to_string(&message)?;
|
||||
let max_message_len = encoded.len();
|
||||
let (reader, mut peer) = tokio::io::duplex(max_message_len.saturating_add(1));
|
||||
let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
|
||||
reader,
|
||||
tokio::io::sink(),
|
||||
"test stdio peer".to_string(),
|
||||
max_message_len,
|
||||
);
|
||||
|
||||
peer.write_all(encoded.as_bytes()).await?;
|
||||
peer.write_all(b"\r\n").await?;
|
||||
let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
|
||||
.await?
|
||||
.expect("stdio connection should report the message");
|
||||
match event {
|
||||
JsonRpcConnectionEvent::Message(actual) => assert_eq!(actual, message),
|
||||
event => anyhow::bail!("expected JSON-RPC message, got {event:?}"),
|
||||
}
|
||||
|
||||
drop(peer);
|
||||
drop(connection);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stdio_connection_rejects_overlong_unterminated_message() -> anyhow::Result<()> {
|
||||
let max_message_len: usize = 32;
|
||||
let (reader, mut peer) = tokio::io::duplex(max_message_len.saturating_add(1));
|
||||
let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
|
||||
reader,
|
||||
tokio::io::sink(),
|
||||
"hostile stdio peer".to_string(),
|
||||
max_message_len,
|
||||
);
|
||||
let overlong_message = vec![b'x'; max_message_len + 1];
|
||||
|
||||
peer.write_all(&overlong_message).await?;
|
||||
let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
|
||||
.await?
|
||||
.expect("stdio connection should report the framing violation");
|
||||
match event {
|
||||
JsonRpcConnectionEvent::Disconnected { reason } => assert_eq!(
|
||||
reason,
|
||||
Some(
|
||||
"JSON-RPC message from hostile stdio peer exceeds maximum length of 32 bytes"
|
||||
.to_string()
|
||||
)
|
||||
),
|
||||
event => anyhow::bail!("expected stdio disconnect, got {event:?}"),
|
||||
}
|
||||
|
||||
drop(peer);
|
||||
drop(connection);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> {
|
||||
let (client_websocket, mut server_websocket) = websocket_pair().await?;
|
||||
|
||||
Reference in New Issue
Block a user