mirror of
https://github.com/openai/codex.git
synced 2026-09-20 12:47:38 +00:00
Advertise the control socket's WebSocket message size limit (#46548)
## Why Clients need to know the incoming size limit for single-frame requests so they can reject oversized requests before the socket closes. ## What changed Add `x-codex-websocket-max-unfragmented-message-bytes` to successful control socket WebSocket handshake responses. Derive the advertised byte limit from the minimum configured frame and message size limits, using the same `WebSocketConfig` for the connection. ## Testing Extend the control socket upgrade test to verify that the header contains a numeric value matching the effective default limit. GitOrigin-RevId: c4f46d12e00e20f56d4c14a3e84f9fe6e2f020fa
This commit is contained in:
@@ -15,10 +15,13 @@ use futures::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio::time::Duration;
|
||||
use tokio_tungstenite::accept_hdr_async;
|
||||
use tokio_tungstenite::accept_hdr_async_with_config;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::handshake::server::Response as HandshakeResponse;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::http::Response;
|
||||
use tokio_tungstenite::tungstenite::http::StatusCode;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::error;
|
||||
use tracing::info;
|
||||
@@ -26,6 +29,10 @@ use tracing::warn;
|
||||
|
||||
#[cfg(unix)]
|
||||
const CONTROL_SOCKET_MODE: u32 = 0o600;
|
||||
// Advertise the effective incoming cap for single-frame messages so clients can
|
||||
// reject oversized requests before the socket closes.
|
||||
const MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER: &str =
|
||||
"x-codex-websocket-max-unfragmented-message-bytes";
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub enum DaemonShutdownAccess {
|
||||
@@ -150,9 +157,18 @@ async fn run_control_socket_acceptor(
|
||||
let transport_event_tx = transport_event_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut shutdown_request = false;
|
||||
let websocket_stream = match accept_hdr_async(
|
||||
let websocket_config = WebSocketConfig::default();
|
||||
let max_unfragmented_message_bytes = [
|
||||
websocket_config.max_frame_size,
|
||||
websocket_config.max_message_size,
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.min();
|
||||
let websocket_stream = match accept_hdr_async_with_config(
|
||||
stream,
|
||||
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
|
||||
|request: &tokio_tungstenite::tungstenite::handshake::server::Request,
|
||||
mut response: HandshakeResponse| {
|
||||
if request.uri().path() == "/daemon/shutdown" {
|
||||
if !matches!(daemon_shutdown_access, DaemonShutdownAccess::Managed) {
|
||||
let mut rejection = Response::new(Some("unmanaged server".to_string()));
|
||||
@@ -161,8 +177,15 @@ async fn run_control_socket_acceptor(
|
||||
}
|
||||
shutdown_request = true;
|
||||
}
|
||||
if let Some(max_bytes) = max_unfragmented_message_bytes {
|
||||
response.headers_mut().insert(
|
||||
MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER,
|
||||
HeaderValue::from(max_bytes),
|
||||
);
|
||||
}
|
||||
Ok(response)
|
||||
},
|
||||
Some(websocket_config),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -21,6 +21,7 @@ use tokio::time::timeout;
|
||||
use tokio_tungstenite::client_async;
|
||||
use tokio_tungstenite::tungstenite::Bytes;
|
||||
use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[test]
|
||||
@@ -77,6 +78,26 @@ async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_a
|
||||
.await
|
||||
.expect("websocket upgrade should complete");
|
||||
assert_eq!(response.status().as_u16(), 101);
|
||||
let advertised_max = response
|
||||
.headers()
|
||||
.get("x-codex-websocket-max-unfragmented-message-bytes")
|
||||
.expect("byte cap header should be advertised")
|
||||
.to_str()
|
||||
.expect("byte cap header should be ASCII")
|
||||
.parse::<usize>()
|
||||
.expect("byte cap header should be a number");
|
||||
let websocket_config = WebSocketConfig::default();
|
||||
assert_eq!(
|
||||
advertised_max,
|
||||
[
|
||||
websocket_config.max_frame_size,
|
||||
websocket_config.max_message_size,
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.min()
|
||||
.expect("default websocket config should have an incoming size limit")
|
||||
);
|
||||
|
||||
let opened = timeout(Duration::from_secs(1), transport_event_rx.recv())
|
||||
.await
|
||||
|
||||
Reference in New Issue
Block a user