diff --git a/codex-rs/.cargo/audit.toml b/codex-rs/.cargo/audit.toml index acfe09b22f..b3808b45a5 100644 --- a/codex-rs/.cargo/audit.toml +++ b/codex-rs/.cargo/audit.toml @@ -12,4 +12,7 @@ ignore = [ "RUSTSEC-2026-0119", # hickory-proto via rama-dns/rama-tcp; remove when rama updates to hickory 0.26.1 or hickory-net "RUSTSEC-2026-0194", # quick-xml via plist/syntect and wayland-scanner; trusted inputs only; remove when rust-plist#191 and wayland-rs#938 are released "RUSTSEC-2026-0195", # quick-xml via plist/syntect and wayland-scanner; trusted inputs only; remove when rust-plist#191 and wayland-rs#938 are released + "RUSTSEC-2026-0370", # zvariant 4.2.0 via zbus 4.4.0 in Linux TUI and keyring/secret-service; remove when those paths migrate off the retired standalone crates + "RUSTSEC-2026-0371", # zbus_names 3.0.0 via zbus 4.4.0 in Linux TUI and keyring/secret-service; remove when those paths migrate off the retired standalone crates + "RUSTSEC-2026-0372", # zvariant_derive 4.2.0 via zbus 4.4.0 in Linux TUI and keyring/secret-service; remove when those paths migrate off the retired standalone crates ] diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f95d5963e2..20eb7fd607 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2356,6 +2356,7 @@ dependencies = [ "tracing", "url", "uuid", + "zstd", ] [[package]] diff --git a/codex-rs/app-server-transport/Cargo.toml b/codex-rs/app-server-transport/Cargo.toml index f78bfebe55..95967581f0 100644 --- a/codex-rs/app-server-transport/Cargo.toml +++ b/codex-rs/app-server-transport/Cargo.toml @@ -53,6 +53,7 @@ tokio-util = { workspace = true } tracing = { workspace = true, features = ["log"] } url = { workspace = true } uuid = { workspace = true, features = ["serde", "v7"] } +zstd = { workspace = true } [target.'cfg(unix)'.dependencies] libc = { workspace = true } diff --git a/codex-rs/app-server-transport/src/lib.rs b/codex-rs/app-server-transport/src/lib.rs index 1f08704ad1..895f65a0fd 100644 --- a/codex-rs/app-server-transport/src/lib.rs +++ b/codex-rs/app-server-transport/src/lib.rs @@ -10,6 +10,7 @@ pub const DAEMON_SHUTDOWN_SOCKET_ENV: &str = "CODEX_DAEMON_SHUTDOWN_SOCKET"; mod connection_auth; mod outgoing_message; mod transport; +mod websocket_compression; pub use connection_auth::ConnectionAuth; pub use outgoing_message::ConnectionId; @@ -44,3 +45,8 @@ pub use transport::start_stdio_connection; pub use transport::start_websocket_acceptor; pub use transport::take_remote_control_disabled_env; pub use transport::verified_local_daemon_socket; +pub use websocket_compression::WEBSOCKET_COMPRESSION_HEADER; +pub use websocket_compression::WebSocketCompression; +pub use websocket_compression::ZSTD_STREAM_PROTOCOL; +pub use websocket_compression::ZstdDecoder; +pub use websocket_compression::ZstdEncoder; diff --git a/codex-rs/app-server-transport/src/transport/unix_socket.rs b/codex-rs/app-server-transport/src/transport/unix_socket.rs index ae53c12db5..70af70c07a 100644 --- a/codex-rs/app-server-transport/src/transport/unix_socket.rs +++ b/codex-rs/app-server-transport/src/transport/unix_socket.rs @@ -7,6 +7,9 @@ use std::path::Path; use super::TransportEvent; use crate::transport::websocket::run_websocket_connection; +use crate::websocket_compression::WEBSOCKET_COMPRESSION_HEADER; +use crate::websocket_compression::WebSocketCompression; +use crate::websocket_compression::ZSTD_STREAM_PROTOCOL; use codex_uds::UnixListener; use codex_uds::UnixStream; use codex_utils_absolute_path::AbsolutePathBuf; @@ -175,6 +178,7 @@ async fn run_control_socket_acceptor( let transport_event_tx = transport_event_tx.clone(); tokio::spawn(async move { let mut shutdown_request = false; + let mut compression = WebSocketCompression::Plain; let websocket_config = WebSocketConfig::default(); let max_unfragmented_message_bytes = [ websocket_config.max_frame_size, @@ -194,6 +198,16 @@ async fn run_control_socket_acceptor( return Err(rejection); } shutdown_request = true; + } else if request + .headers() + .get(WEBSOCKET_COMPRESSION_HEADER) + .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) + { + compression = WebSocketCompression::Zstd; + response.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), + ); } if let Some(max_bytes) = max_unfragmented_message_bytes { response.headers_mut().insert( @@ -218,7 +232,13 @@ async fn run_control_socket_acceptor( return; } let (websocket_writer, websocket_reader) = websocket_stream.split(); - run_websocket_connection(websocket_writer, websocket_reader, transport_event_tx).await; + run_websocket_connection( + websocket_writer, + websocket_reader, + transport_event_tx, + compression, + ) + .await; }); } info!("control socket acceptor shutting down"); diff --git a/codex-rs/app-server-transport/src/transport/websocket.rs b/codex-rs/app-server-transport/src/transport/websocket.rs index a28bdec2ef..b4f3cc97d6 100644 --- a/codex-rs/app-server-transport/src/transport/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/websocket.rs @@ -6,6 +6,11 @@ use super::next_connection_id; use super::serialize_outgoing_message; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::QueuedOutgoingMessage; +use crate::websocket_compression::WEBSOCKET_COMPRESSION_HEADER; +use crate::websocket_compression::WebSocketCompression; +use crate::websocket_compression::ZSTD_STREAM_PROTOCOL; +use crate::websocket_compression::ZstdDecoder; +use crate::websocket_compression::ZstdEncoder; use axum::Router; use axum::body::Body; use axum::body::Bytes; @@ -14,6 +19,7 @@ use axum::extract::State; use axum::extract::ws::Message as AxumWebSocketMessage; use axum::extract::ws::WebSocketUpgrade; use axum::http::HeaderMap; +use axum::http::HeaderValue; use axum::http::Request; use axum::http::StatusCode; use axum::http::header::ORIGIN; @@ -47,6 +53,8 @@ use tracing::warn; /// writer task is healthy, so give them more headroom than internal channels. const WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY: usize = 32 * 1024; const _: () = assert!(WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY > CHANNEL_CAPACITY); +// Match the existing Axum and Tungstenite default WebSocket message limit. +pub(super) const MAX_DECOMPRESSED_MESSAGE_SIZE: usize = 64 << 20; fn colorize(text: &str, style: Style) -> String { text.if_supports_color(Stream::Stderr, |value| value.style(style)) @@ -117,13 +125,33 @@ async fn websocket_upgrade_handler( return (err.status_code(), err.message()).into_response(); } info!(%peer_addr, "websocket client connected"); - websocket + let compression = if headers + .get(WEBSOCKET_COMPRESSION_HEADER) + .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) + { + WebSocketCompression::Zstd + } else { + WebSocketCompression::Plain + }; + let mut response = websocket .on_upgrade(move |stream| async move { let (websocket_writer, websocket_reader) = stream.split(); - run_websocket_connection(websocket_writer, websocket_reader, state.transport_event_tx) - .await; + run_websocket_connection( + websocket_writer, + websocket_reader, + state.transport_event_tx, + compression, + ) + .await; }) - .into_response() + .into_response(); + if compression == WebSocketCompression::Zstd { + response.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), + ); + } + response } pub async fn start_websocket_acceptor( @@ -173,11 +201,25 @@ pub(crate) async fn run_websocket_connection( websocket_writer: impl futures::sink::Sink + Send + 'static, websocket_reader: impl futures::stream::Stream> + Send + 'static, transport_event_tx: mpsc::Sender, + compression: WebSocketCompression, ) where M: AppServerWebSocketMessage + Send + 'static, SinkError: Send + 'static, StreamError: std::fmt::Display + Send + 'static, { + let (encoder, decoder) = match compression { + WebSocketCompression::Plain => (None, None), + WebSocketCompression::Zstd => match ( + ZstdEncoder::new(), + ZstdDecoder::new(MAX_DECOMPRESSED_MESSAGE_SIZE), + ) { + (Ok(encoder), Ok(decoder)) => (Some(encoder), Some(decoder)), + _ => { + warn!("failed to initialize websocket compression"); + return; + } + }, + }; let connection_id = next_connection_id(); let (writer_tx, writer_rx) = mpsc::channel::(WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY); @@ -203,6 +245,7 @@ pub(crate) async fn run_websocket_connection( writer_rx, writer_control_rx, disconnect_token.clone(), + encoder, )); let mut inbound_task = tokio::spawn(run_websocket_inbound_loop( websocket_reader, @@ -211,6 +254,7 @@ pub(crate) async fn run_websocket_connection( writer_control_tx, connection_id, disconnect_token.clone(), + decoder, )); tokio::select! { @@ -231,7 +275,7 @@ pub(crate) async fn run_websocket_connection( pub(crate) enum IncomingWebSocketMessage { Text(String), - Binary, + Binary(Bytes), Ping(Bytes), Pong, Close, @@ -242,6 +286,7 @@ pub(crate) enum IncomingWebSocketMessage { /// sends directly. pub(crate) trait AppServerWebSocketMessage: Sized { fn text(text: String) -> Self; + fn binary(payload: Bytes) -> Self; fn pong(payload: Bytes) -> Self; fn into_incoming(self) -> Option; } @@ -251,6 +296,10 @@ impl AppServerWebSocketMessage for AxumWebSocketMessage { Self::Text(text.into()) } + fn binary(payload: Bytes) -> Self { + Self::Binary(payload) + } + fn pong(payload: Bytes) -> Self { Self::Pong(payload) } @@ -258,7 +307,7 @@ impl AppServerWebSocketMessage for AxumWebSocketMessage { fn into_incoming(self) -> Option { Some(match self { Self::Text(text) => IncomingWebSocketMessage::Text(text.to_string()), - Self::Binary(_) => IncomingWebSocketMessage::Binary, + Self::Binary(payload) => IncomingWebSocketMessage::Binary(payload), Self::Ping(payload) => IncomingWebSocketMessage::Ping(payload), Self::Pong(_) => IncomingWebSocketMessage::Pong, Self::Close(_) => IncomingWebSocketMessage::Close, @@ -271,6 +320,10 @@ impl AppServerWebSocketMessage for TungsteniteWebSocketMessage { Self::Text(text.into()) } + fn binary(payload: Bytes) -> Self { + Self::Binary(payload) + } + fn pong(payload: Bytes) -> Self { Self::Pong(payload) } @@ -278,7 +331,7 @@ impl AppServerWebSocketMessage for TungsteniteWebSocketMessage { fn into_incoming(self) -> Option { Some(match self { Self::Text(text) => IncomingWebSocketMessage::Text(text.to_string()), - Self::Binary(_) => IncomingWebSocketMessage::Binary, + Self::Binary(payload) => IncomingWebSocketMessage::Binary(payload), Self::Ping(payload) => IncomingWebSocketMessage::Ping(payload), Self::Pong(_) => IncomingWebSocketMessage::Pong, Self::Close(_) => IncomingWebSocketMessage::Close, @@ -292,6 +345,7 @@ async fn run_websocket_outbound_loop( mut writer_rx: mpsc::Receiver, mut writer_control_rx: mpsc::Receiver, disconnect_token: CancellationToken, + mut encoder: Option, ) where M: AppServerWebSocketMessage + Send + 'static, SinkError: Send + 'static, @@ -317,7 +371,18 @@ async fn run_websocket_outbound_loop( let Some(json) = serialize_outgoing_message(queued_message.message) else { continue; }; - if websocket_writer.send(M::text(json)).await.is_err() { + let message = if let Some(encoder) = &mut encoder { + match encoder.encode(&json) { + Ok(payload) => M::binary(payload.into()), + Err(err) => { + warn!("websocket compression error: {err}"); + break; + } + } + } else { + M::text(json) + }; + if websocket_writer.send(message).await.is_err() { break; } if let Some(write_complete_tx) = queued_message.write_complete_tx { @@ -335,55 +400,76 @@ async fn run_websocket_inbound_loop( writer_control_tx: mpsc::Sender, connection_id: ConnectionId, disconnect_token: CancellationToken, + mut decoder: Option, ) where M: AppServerWebSocketMessage + Send + 'static, StreamError: std::fmt::Display + Send + 'static, { tokio::pin!(websocket_reader); loop { - tokio::select! { - _ = disconnect_token.cancelled() => { + let incoming_message = tokio::select! { + _ = disconnect_token.cancelled() => break, + incoming_message = websocket_reader.next() => incoming_message, + }; + let message = match incoming_message { + Some(Ok(message)) => message, + None => break, + Some(Err(err)) => { + warn!("websocket receive error: {err}"); break; } - incoming_message = websocket_reader.next() => { - match incoming_message { - Some(Ok(message)) => match message.into_incoming() { - Some(IncomingWebSocketMessage::Text(text)) - if !forward_incoming_message( - &transport_event_tx, - &writer_tx_for_reader, - connection_id, - &text, - ) - .await - => { - break; - } - Some(IncomingWebSocketMessage::Text(_)) => {} - Some(IncomingWebSocketMessage::Ping(payload)) => { - match writer_control_tx.try_send(M::pong(payload)) { - Ok(()) => {} - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - warn!("websocket control queue full while replying to ping; closing connection"); - break; - } - } - } - Some(IncomingWebSocketMessage::Pong) => {} - Some(IncomingWebSocketMessage::Close) => break, - Some(IncomingWebSocketMessage::Binary) => { - warn!("dropping unsupported binary websocket message"); - } - None => {} - }, - None => break, - Some(Err(err)) => { - warn!("websocket receive error: {err}"); + }; + let text = match message.into_incoming() { + Some(IncomingWebSocketMessage::Text(text)) => { + if decoder.is_some() { + warn!("unexpected text on compressed websocket connection"); + break; + } + text + } + Some(IncomingWebSocketMessage::Binary(payload)) => { + let Some(decoder) = &mut decoder else { + warn!("dropping unsupported binary websocket message"); + continue; + }; + match decoder.decode(&payload) { + Ok(text) => text, + Err(err) => { + warn!("websocket decompression error: {err}"); break; } } } + Some(IncomingWebSocketMessage::Ping(payload)) => { + match writer_control_tx.try_send(M::pong(payload)) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break, + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + warn!( + "websocket control queue full while replying to ping; closing connection" + ); + break; + } + } + continue; + } + Some(IncomingWebSocketMessage::Pong) => continue, + Some(IncomingWebSocketMessage::Close) => break, + None => continue, + }; + if !forward_incoming_message( + &transport_event_tx, + &writer_tx_for_reader, + connection_id, + &text, + ) + .await + { + break; } } } + +#[cfg(test)] +#[path = "websocket_compression_tests.rs"] +mod compression_tests; diff --git a/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs b/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs new file mode 100644 index 0000000000..f2897cdfcc --- /dev/null +++ b/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs @@ -0,0 +1,178 @@ +//! Exercise negotiated and legacy peers through the TCP and local socket listeners. + +use super::*; +use crate::OutgoingError; +use crate::OutgoingMessage; +use crate::transport::DaemonShutdownAccess; +use crate::transport::start_control_socket_acceptor; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::RequestId; +use codex_uds::UnixStream; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; +use std::error::Error; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::time::Duration; +use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::client_async; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; + +#[tokio::test] +async fn tcp_compression_negotiates_and_retains_history() -> Result<(), Box> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); + let router = Router::new() + .fallback(any(websocket_upgrade_handler)) + .with_state(WebSocketListenerState { + transport_event_tx, + auth_policy: Arc::new(WebsocketAuthPolicy::default()), + }); + let server = tokio::spawn(async move { + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await + }); + for offer in [None, Some(ZSTD_STREAM_PROTOCOL), Some("future-protocol")] { + let mut request = format!("ws://{address}").into_client_request()?; + if let Some(offer) = offer { + request.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(offer), + ); + } + let (websocket, response) = connect_async(request).await?; + let expected = offer.filter(|offer| *offer == ZSTD_STREAM_PROTOCOL); + assert_eq!( + response + .headers() + .get(WEBSOCKET_COMPRESSION_HEADER) + .map(HeaderValue::to_str) + .transpose()?, + expected + ); + timeout( + Duration::from_secs(/*secs*/ 5), + roundtrip( + websocket, + &mut events, + WebSocketCompression::from_response_header(expected)?, + ), + ) + .await??; + } + server.abort(); + Ok(()) +} + +#[tokio::test] +async fn local_socket_compression_negotiates_and_retains_history() -> Result<(), Box> { + let temp_dir = tempfile::TempDir::new()?; + let socket_path = AbsolutePathBuf::from_absolute_path( + temp_dir.path().join("app-server-control").join("rpc.sock"), + )?; + let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); + let shutdown = CancellationToken::new(); + let server = start_control_socket_acceptor( + socket_path.clone(), + transport_event_tx, + shutdown.clone(), + DaemonShutdownAccess::Disabled, + ) + .await?; + let stream = UnixStream::connect(socket_path.as_path()).await?; + let mut request = "ws://localhost/rpc".into_client_request()?; + request.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), + ); + let (websocket, response) = client_async(request, stream).await?; + assert_eq!( + response.headers().get(WEBSOCKET_COMPRESSION_HEADER), + Some(&HeaderValue::from_static(ZSTD_STREAM_PROTOCOL)) + ); + timeout( + Duration::from_secs(/*secs*/ 5), + roundtrip(websocket, &mut events, WebSocketCompression::Zstd), + ) + .await??; + shutdown.cancel(); + server.await?; + Ok(()) +} + +async fn roundtrip( + mut websocket: WebSocketStream, + events: &mut mpsc::Receiver, + compression: WebSocketCompression, +) -> Result<(), Box> { + let Some(TransportEvent::ConnectionOpened { writer, .. }) = events.recv().await else { + return Err("expected connection opened".into()); + }; + let mut encoder = ZstdEncoder::new()?; + let mut decoder = ZstdDecoder::new(/*max_message_size*/ 1024)?; + for index in 0..2 { + let request = format!(r#"{{"method":"example","id":{index}}}"#); + let message = match compression { + WebSocketCompression::Plain => { + TungsteniteWebSocketMessage::Text(request.clone().into()) + } + WebSocketCompression::Zstd => { + TungsteniteWebSocketMessage::Binary(encoder.encode(&request)?.into()) + } + }; + websocket.send(message).await?; + let Some(TransportEvent::IncomingMessage { message, .. }) = events.recv().await else { + return Err("expected incoming RPC".into()); + }; + assert_eq!(message, serde_json::from_str::(&request)?); + let response = OutgoingMessage::Error(OutgoingError { + id: RequestId::Integer(index), + error: JSONRPCErrorError { + code: -32601, + message: "unknown method".to_string(), + data: None, + }, + }); + let expected = serde_json::to_string(&response)?; + writer.send(QueuedOutgoingMessage::new(response)).await?; + let text = loop { + let response = websocket.next().await.ok_or("expected response")??; + match (compression, response) { + (WebSocketCompression::Plain, TungsteniteWebSocketMessage::Text(text)) => { + break text.to_string(); + } + (WebSocketCompression::Zstd, TungsteniteWebSocketMessage::Binary(bytes)) => { + break decoder.decode(&bytes)?; + } + ( + _, + TungsteniteWebSocketMessage::Ping(_) | TungsteniteWebSocketMessage::Pong(_), + ) => {} + _ => return Err("response has incorrect WebSocket message type".into()), + } + }; + assert_eq!(text, expected); + websocket + .send(TungsteniteWebSocketMessage::Ping(Bytes::from_static( + b"check", + ))) + .await?; + assert_eq!( + websocket.next().await.ok_or("expected pong")??, + TungsteniteWebSocketMessage::Pong(Bytes::from_static(b"check")) + ); + } + websocket.close(/*msg*/ None).await?; + assert!(matches!( + events.recv().await, + Some(TransportEvent::ConnectionClosed { .. }) + )); + Ok(()) +} diff --git a/codex-rs/app-server-transport/src/websocket_compression.rs b/codex-rs/app-server-transport/src/websocket_compression.rs new file mode 100644 index 0000000000..17ac36e0b0 --- /dev/null +++ b/codex-rs/app-server-transport/src/websocket_compression.rs @@ -0,0 +1,123 @@ +//! Negotiated WebSocket payload compression with one persistent zstd stream per direction. +//! +//! The upgrade request and response must both carry +//! `x-codex-app-server-compression: zstd-stream-v1` to enable compression. +//! Without an acknowledgment, messages remain JSON text. +//! Each binary message contains one UTF-8 JSON message flushed with `ZSTD_e_flush`. +//! Keep separate streams per direction and start fresh on reconnect. Never finish +//! the zstd frame between messages. + +use std::io; +use zstd::stream::raw::CParameter; +use zstd::stream::raw::DParameter; +use zstd::stream::raw::Decoder; +use zstd::stream::raw::Encoder; +use zstd::stream::raw::Operation; +use zstd::stream::raw::OutBuffer; + +pub const WEBSOCKET_COMPRESSION_HEADER: &str = "x-codex-app-server-compression"; +pub const ZSTD_STREAM_PROTOCOL: &str = "zstd-stream-v1"; +const COMPRESSION_LEVEL: i32 = -7; +const WINDOW_LOG: u32 = 20; +const BUFFER_SIZE: usize = 32 * 1024; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WebSocketCompression { + Plain, + Zstd, +} + +impl WebSocketCompression { + /// A missing echo selects the existing text protocol for older peers. + pub fn from_response_header(value: Option<&str>) -> io::Result { + match value { + None => Ok(Self::Plain), + Some(ZSTD_STREAM_PROTOCOL) => Ok(Self::Zstd), + Some(_) => Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported app-server WebSocket compression protocol", + )), + } + } +} + +pub struct ZstdEncoder { + encoder: Encoder<'static>, +} + +impl ZstdEncoder { + pub fn new() -> io::Result { + let mut encoder = Encoder::new(COMPRESSION_LEVEL)?; + encoder.set_parameter(CParameter::WindowLog(WINDOW_LOG))?; + Ok(Self { encoder }) + } + + /// Flush one message without ending the frame or resetting its history. + pub fn encode(&mut self, text: &str) -> io::Result> { + let mut input = text.as_bytes(); + let mut encoded = Vec::new(); + let mut buffer = [0; BUFFER_SIZE]; + while !input.is_empty() { + let status = self.encoder.run_on_buffers(input, &mut buffer)?; + encoded.extend_from_slice(&buffer[..status.bytes_written]); + input = &input[status.bytes_read..]; + } + loop { + let mut output = OutBuffer::around(&mut buffer[..]); + let remaining = self.encoder.flush(&mut output)?; + encoded.extend_from_slice(output.as_slice()); + if remaining == 0 { + return Ok(encoded); + } + } + } +} + +pub struct ZstdDecoder { + decoder: Decoder<'static>, + max_message_size: usize, +} + +impl ZstdDecoder { + pub fn new(max_message_size: usize) -> io::Result { + let mut decoder = Decoder::new()?; + decoder.set_parameter(DParameter::WindowLogMax(WINDOW_LOG))?; + Ok(Self { + decoder, + max_message_size, + }) + } + + /// Decode one flushed message, retaining history and bounding its decoded size. + /// Any error invalidates the stream; the caller must close the connection. + pub fn decode(&mut self, mut input: &[u8]) -> io::Result { + let mut decoded = Vec::new(); + let mut buffer = [0; BUFFER_SIZE]; + loop { + let status = self.decoder.run_on_buffers(input, &mut buffer)?; + if status.bytes_written > self.max_message_size.saturating_sub(decoded.len()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "decompressed app-server WebSocket message exceeds size limit", + )); + } + decoded.extend_from_slice(&buffer[..status.bytes_written]); + input = &input[status.bytes_read..]; + if status.bytes_read == 0 && status.bytes_written == 0 { + if !input.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid app-server WebSocket compressed message", + )); + } + break; + } + } + String::from_utf8(decoded) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + } +} + +#[cfg(test)] +#[path = "websocket_compression_tests.rs"] +mod tests; diff --git a/codex-rs/app-server-transport/src/websocket_compression_tests.rs b/codex-rs/app-server-transport/src/websocket_compression_tests.rs new file mode 100644 index 0000000000..91a7c58ec4 --- /dev/null +++ b/codex-rs/app-server-transport/src/websocket_compression_tests.rs @@ -0,0 +1,73 @@ +//! Verify persistent history, peer fallback, and bounded decompression. + +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn response_header_selects_compression_or_plain_fallback() -> io::Result<()> { + assert_eq!( + WebSocketCompression::from_response_header(/*value*/ None)?, + WebSocketCompression::Plain + ); + assert_eq!( + WebSocketCompression::from_response_header(Some(ZSTD_STREAM_PROTOCOL))?, + WebSocketCompression::Zstd + ); + assert!(WebSocketCompression::from_response_header(Some("unknown")).is_err()); + Ok(()) +} + +#[test] +fn messages_keep_history_and_have_independent_size_limits() -> io::Result<()> { + let text = serde_json::to_string( + &(0..1_000) + .map(|index| format!("thread-{index}: individual message body {index}")) + .collect::>(), + )?; + let mut encoder = ZstdEncoder::new()?; + let first = encoder.encode(&text)?; + let second = encoder.encode(&text)?; + let mut decoder = ZstdDecoder::new(text.len())?; + assert_eq!(decoder.decode(&first)?, text); + assert_eq!(decoder.decode(&second)?, text); + assert!(second.len() < first.len()); + assert!(ZstdDecoder::new(text.len())?.decode(&second).is_err()); + Ok(()) +} + +#[test] +fn decoded_size_limit_covers_output_larger_than_scratch_buffer() -> io::Result<()> { + let text = "abcdefgh".repeat(BUFFER_SIZE); + let encoded = ZstdEncoder::new()?.encode(&text)?; + assert_eq!(ZstdDecoder::new(text.len())?.decode(&encoded)?, text); + let error = ZstdDecoder::new(text.len() - 1)? + .decode(&encoded) + .expect_err("oversize message must fail"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + Ok(()) +} + +#[test] +fn malformed_compressed_data_fails() -> io::Result<()> { + assert!( + ZstdDecoder::new(/*max_message_size*/ 1024)? + .decode(b"not zstd") + .is_err() + ); + Ok(()) +} + +#[test] +fn larger_compression_window_is_rejected() -> io::Result<()> { + let mut encoder = ZstdEncoder::new()?; + encoder + .encoder + .set_parameter(CParameter::WindowLog(WINDOW_LOG + 1))?; + let encoded = encoder.encode("a message with an unsupported compression window")?; + assert!( + ZstdDecoder::new(/*max_message_size*/ 1024)? + .decode(&encoded) + .is_err() + ); + Ok(()) +} diff --git a/codex-rs/deny.toml b/codex-rs/deny.toml index 31c5517af0..2ae097e209 100644 --- a/codex-rs/deny.toml +++ b/codex-rs/deny.toml @@ -84,6 +84,9 @@ ignore = [ { id = "RUSTSEC-2026-0119", reason = "hickory-proto v0.25.2 is pulled in via rama-dns/rama-tcp used by codex-network-proxy; no fixed rama release is available yet; remove when rama updates to hickory 0.26.1 or hickory-net" }, { id = "RUSTSEC-2026-0194", reason = "quick-xml v0.39.4 remains via plist/syntect and wayland-scanner; plist does not exercise the affected APIs and wayland-scanner parses trusted build-time XML; remove when rust-plist#191 and wayland-rs#938 are released" }, { id = "RUSTSEC-2026-0195", reason = "quick-xml v0.39.4 remains via plist/syntect and wayland-scanner; plist does not exercise the affected APIs and wayland-scanner parses trusted build-time XML; remove when rust-plist#191 and wayland-rs#938 are released" }, + { id = "RUSTSEC-2026-0370", reason = "zvariant 4.2.0 is unmaintained; pulled in via zbus 4.4.0 by Linux TUI accessibility integration and keyring/secret-service; remove when those paths migrate off the retired standalone crates" }, + { id = "RUSTSEC-2026-0371", reason = "zbus_names 3.0.0 is unmaintained; pulled in via zbus 4.4.0 by Linux TUI accessibility integration and keyring/secret-service; remove when those paths migrate off the retired standalone crates" }, + { id = "RUSTSEC-2026-0372", reason = "zvariant_derive 4.2.0 is unmaintained; pulled in via zbus 4.4.0 by Linux TUI accessibility integration and keyring/secret-service; remove when those paths migrate off the retired standalone crates" }, ] # If this is true, then cargo deny will use the git executable to fetch advisory database. # If this is false, then it uses a built-in git library.