openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-24 18:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bound exec-server ingress

Browse Source
This commit is contained in:
jif-oai
2026-07-08 14:48:23 +01:00
parent 6849549b49
commit e8870ee187
10 changed files with 323 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/client_recovery.rs
+7
View File
@@ -516,6 +516,13 @@ impl ExecServerClient {
inner.request_recovery(rpc_client, disconnected_message(reason.as_deref()));
return;
}
RpcClientEvent::ProtocolError { reason } => {
rpc_client.close_transport().await;
inner
.fail(format!("exec-server protocol violation: {reason}"))
.await;
return;
}
}
}
});
codex-rs/exec-server/src/client_transport.rs
+18 -11
View File
@@ -6,7 +6,6 @@ use tokio::io::AsyncBufReadExt;
use tokio::io::BufReader;
use tokio::process::Command;
use tokio::time::timeout;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::connect_async_with_config;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tracing::debug;
@@ -26,6 +25,7 @@ use crate::client_api::RemoteExecServerConnectArgs;
use crate::client_api::StdioExecServerCommand;
use crate::client_api::StdioExecServerConnectArgs;
use crate::connection::JsonRpcConnection;
use crate::connection::jsonrpc_websocket_config;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_relay::NoiseHarnessConnectionArgs;
use crate::noise_relay::noise_harness_connection_from_websocket;
@@ -191,16 +191,23 @@ impl ExecServerClient {
ensure_rustls_crypto_provider();
let websocket_url = args.websocket_url.clone();
let connect_timeout = args.connect_timeout;
let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str()))
.await
.map_err(|_| ExecServerError::WebSocketConnectTimeout {
url: websocket_url.clone(),
timeout: connect_timeout,
})?
.map_err(|source| ExecServerError::WebSocketConnect {
url: websocket_url.clone(),
source,
})?;
let (stream, _) = timeout(
connect_timeout,
connect_async_with_config(
websocket_url.as_str(),
Some(jsonrpc_websocket_config()),
/*disable_nagle*/ false,
),
)
.await
.map_err(|_| ExecServerError::WebSocketConnectTimeout {
url: websocket_url.clone(),
timeout: connect_timeout,
})?
.map_err(|source| ExecServerError::WebSocketConnect {
url: websocket_url.clone(),
source,
})?;
let connection_label = format!("exec-server websocket {websocket_url}");
let connection = if is_rendezvous_harness_url(&websocket_url) {
codex-rs/exec-server/src/connection.rs
+160 -43
View File
@@ -24,11 +24,14 @@ 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;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
pub(crate) const CHANNEL_CAPACITY: usize = 128;
pub(crate) const MAX_JSONRPC_MESSAGE_BYTES: usize = 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);
@@ -37,9 +40,16 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30
#[derive(Debug)]
pub(crate) enum JsonRpcConnectionEvent {
Message(JSONRPCMessage),
MalformedMessage { reason: String },
Disconnected { reason: Option<String> },
Message {
message: JSONRPCMessage,
encoded_len: usize,
},
MalformedMessage {
reason: String,
},
Disconnected {
reason: Option<String>,
},
}
#[derive(Clone)]
@@ -242,17 +252,51 @@ impl JsonRpcConnection {
let incoming_tx_for_reader = incoming_tx.clone();
let disconnected_tx_for_reader = disconnected_tx.clone();
let reader_task = tokio::spawn(async move {
let mut lines = BufReader::new(reader).lines();
let mut reader = BufReader::new(reader);
loop {
match lines.next_line().await {
Ok(Some(line)) => {
if line.trim().is_empty() {
let mut encoded = Vec::new();
match (&mut reader)
.take((MAX_JSONRPC_MESSAGE_BYTES + 2) as u64)
.read_until(b'\n', &mut encoded)
.await
{
Ok(0) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
/*reason*/ None,
)
.await;
break;
}
Ok(_) => {
if encoded.last() == Some(&b'\n') {
encoded.pop();
if encoded.last() == Some(&b'\r') {
encoded.pop();
}
}
if encoded.len() > MAX_JSONRPC_MESSAGE_BYTES {
send_malformed_message(
&incoming_tx_for_reader,
Some(format!(
"JSON-RPC message from {reader_label} exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes"
)),
)
.await;
break;
}
if encoded.iter().all(u8::is_ascii_whitespace) {
continue;
}
match serde_json::from_str::<JSONRPCMessage>(&line) {
let encoded_len = encoded.len();
match serde_json::from_slice::<JSONRPCMessage>(&encoded) {
Ok(message) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.send(JsonRpcConnectionEvent::Message {
message,
encoded_len,
})
.await
.is_err()
{
@@ -270,15 +314,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,
@@ -392,36 +427,52 @@ impl JsonRpcConnection {
}
incoming_message = websocket.next() => {
match incoming_message {
Some(Ok(message)) => match message.parse_jsonrpc_frame() {
Ok(JsonRpcWebSocketFrame::Message(message)) => {
if incoming_tx
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
break;
}
}
Ok(JsonRpcWebSocketFrame::Close) => {
send_disconnected(
Some(Ok(message)) => {
let encoded_len = message.jsonrpc_payload_len();
if encoded_len.is_some_and(|len| len > MAX_JSONRPC_MESSAGE_BYTES) {
send_malformed_message(
&incoming_tx,
&disconnected_tx,
/*reason*/ None,
Some(format!(
"websocket JSON-RPC message from {connection_label} exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes"
)),
)
.await;
break;
}
Ok(JsonRpcWebSocketFrame::Ignore) => {}
Err(err) => {
send_malformed_message(
&incoming_tx,
Some(format!(
"failed to parse websocket JSON-RPC message from {connection_label}: {err}"
)),
)
.await;
match message.parse_jsonrpc_frame() {
Ok(JsonRpcWebSocketFrame::Message(message)) => {
if incoming_tx
.send(JsonRpcConnectionEvent::Message {
message,
encoded_len: encoded_len.unwrap_or_default(),
})
.await
.is_err()
{
break;
}
}
Ok(JsonRpcWebSocketFrame::Close) => {
send_disconnected(
&incoming_tx,
&disconnected_tx,
/*reason*/ None,
)
.await;
break;
}
Ok(JsonRpcWebSocketFrame::Ignore) => {}
Err(err) => {
send_malformed_message(
&incoming_tx,
Some(format!(
"failed to parse websocket JSON-RPC message from {connection_label}: {err}"
)),
)
.await;
}
}
},
}
Some(Err(err)) => {
send_disconnected(
&incoming_tx,
@@ -470,12 +521,21 @@ enum JsonRpcWebSocketFrame {
}
trait JsonRpcWebSocketMessage: Send + 'static {
fn jsonrpc_payload_len(&self) -> Option<usize>;
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error>;
fn from_text(text: String) -> Self;
fn ping() -> Self;
}
impl JsonRpcWebSocketMessage for Message {
fn jsonrpc_payload_len(&self) -> Option<usize> {
match self {
Message::Text(text) => Some(text.len()),
Message::Binary(bytes) => Some(bytes.len()),
Message::Close(_) | Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => None,
}
}
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
match self {
Message::Text(text) => {
@@ -501,6 +561,16 @@ impl JsonRpcWebSocketMessage for Message {
}
impl JsonRpcWebSocketMessage for AxumWebSocketMessage {
fn jsonrpc_payload_len(&self) -> Option<usize> {
match self {
AxumWebSocketMessage::Text(text) => Some(text.len()),
AxumWebSocketMessage::Binary(bytes) => Some(bytes.len()),
AxumWebSocketMessage::Close(_)
| AxumWebSocketMessage::Ping(_)
| AxumWebSocketMessage::Pong(_) => None,
}
}
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
match self {
AxumWebSocketMessage::Text(text) => {
@@ -525,6 +595,12 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage {
}
}
pub(crate) fn jsonrpc_websocket_config() -> WebSocketConfig {
WebSocketConfig::default()
.max_frame_size(Some(MAX_JSONRPC_MESSAGE_BYTES))
.max_message_size(Some(MAX_JSONRPC_MESSAGE_BYTES))
}
async fn send_disconnected(
incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
disconnected_tx: &watch::Sender<bool>,
@@ -675,13 +751,54 @@ mod tests {
.await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::Message(actual)) if actual == message
Some(JsonRpcConnectionEvent::Message { message: actual, .. }) if actual == message
));
drop(connection);
Ok(())
}
#[tokio::test]
async fn websocket_connection_rejects_oversized_jsonrpc_message() -> anyhow::Result<()> {
let (websocket, control, _outbound_rx) =
ControlledWebSocket::new(/*write_ready*/ true);
let mut connection = JsonRpcConnection::from_websocket_stream(
websocket,
"test".into(),
/*ping_interval*/ None,
);
control.send_inbound(Message::Text(
"x".repeat(MAX_JSONRPC_MESSAGE_BYTES + 1).into(),
))?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::MalformedMessage { reason })
if reason.contains("exceeds")
));
Ok(())
}
#[tokio::test]
async fn stdio_connection_rejects_oversized_jsonrpc_message() -> anyhow::Result<()> {
let (mut server_writer, client_reader) = tokio::io::duplex(MAX_JSONRPC_MESSAGE_BYTES + 2);
let (client_writer, _server_reader) = tokio::io::duplex(64);
let mut connection =
JsonRpcConnection::from_stdio(client_reader, client_writer, "test".into());
let mut oversized_message = vec![b'x'; MAX_JSONRPC_MESSAGE_BYTES + 1];
oversized_message.push(b'\n');
server_writer.write_all(&oversized_message).await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::MalformedMessage { reason })
if reason.contains("exceeds")
));
Ok(())
}
#[tokio::test]
async fn websocket_connection_keeps_outbound_message_while_send_is_backpressured()
-> anyhow::Result<()> {
codex-rs/exec-server/src/noise_relay/executor_stream.rs
+5 -2
View File
@@ -74,9 +74,12 @@ impl NoiseVirtualStream {
ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))
})?
};
for message in self.inbound_decoder.push(&plaintext)? {
for (message, encoded_len) in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.try_send(JsonRpcConnectionEvent::Message(message))
.try_send(JsonRpcConnectionEvent::Message {
message,
encoded_len,
})
.map_err(|_| {
ExecServerError::Protocol(
"Noise virtual stream inbound queue is full or closed".to_string(),
codex-rs/exec-server/src/noise_relay/harness.rs
+5 -2
View File
@@ -582,10 +582,13 @@ async fn receive_data(
// The authenticated byte stream can carry partial or multiple JSON-RPC
// messages; emit only complete, successfully parsed messages.
for message in decoder.push(&plaintext)? {
for (message, encoded_len) in decoder.push(&plaintext)? {
send_incoming_event(
incoming_tx,
JsonRpcConnectionEvent::Message(message),
JsonRpcConnectionEvent::Message {
message,
encoded_len,
},
delivery_deadline,
)
.await?;
codex-rs/exec-server/src/noise_relay/message_framing.rs
+9 -8
View File
@@ -1,9 +1,9 @@
use codex_exec_server_protocol::JSONRPCMessage;
use crate::ExecServerError;
use crate::connection::MAX_JSONRPC_MESSAGE_BYTES;
const LENGTH_PREFIX_BYTES: usize = size_of::<u32>();
const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024;
/// Serialize one JSON-RPC message into the encrypted record byte stream.
@@ -16,7 +16,7 @@ pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>,
let mut framed = vec![0; LENGTH_PREFIX_BYTES];
serde_json::to_writer(&mut framed, message)?;
let message_len = framed.len() - LENGTH_PREFIX_BYTES;
if message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN {
if message_len > MAX_JSONRPC_MESSAGE_BYTES {
return Err(ExecServerError::Protocol(
"Noise relay JSON-RPC message exceeds maximum length".to_string(),
));
@@ -39,7 +39,7 @@ impl JsonRpcMessageDecoder {
pub(crate) fn push(
&mut self,
plaintext_record: &[u8],
) -> Result<Vec<JSONRPCMessage>, ExecServerError> {
) -> Result<Vec<(JSONRPCMessage, usize)>, ExecServerError> {
if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN {
return Err(ExecServerError::Protocol(
"Noise relay plaintext record exceeds maximum length".to_string(),
@@ -55,7 +55,7 @@ impl JsonRpcMessageDecoder {
let message_len =
u32::from_be_bytes([prefix[0], prefix[1], prefix[2], prefix[3]]) as usize;
// Reject the authenticated length before waiting for its payload.
if message_len == 0 || message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN {
if message_len == 0 || message_len > MAX_JSONRPC_MESSAGE_BYTES {
return Err(ExecServerError::Protocol(
"Noise relay JSON-RPC message has invalid length".to_string(),
));
@@ -64,14 +64,15 @@ impl JsonRpcMessageDecoder {
if self.buffered.len() < framed_len {
break;
}
messages.push(serde_json::from_slice(
&self.buffered[LENGTH_PREFIX_BYTES..framed_len],
)?);
messages.push((
serde_json::from_slice(&self.buffered[LENGTH_PREFIX_BYTES..framed_len])?,
message_len,
));
self.buffered.drain(..framed_len);
}
// Even before a message is complete, keep reassembly memory bounded.
if self.buffered.len() > LENGTH_PREFIX_BYTES + MAX_NOISE_JSONRPC_MESSAGE_LEN {
if self.buffered.len() > LENGTH_PREFIX_BYTES + MAX_JSONRPC_MESSAGE_BYTES {
return Err(ExecServerError::Protocol(
"Noise relay JSON-RPC reassembly buffer exceeds maximum length".to_string(),
));
codex-rs/exec-server/src/noise_relay/message_framing_tests.rs
+3 -3
View File
@@ -3,10 +3,10 @@ use codex_exec_server_protocol::JSONRPCNotification;
use pretty_assertions::assert_eq;
use super::JsonRpcMessageDecoder;
use super::MAX_NOISE_JSONRPC_MESSAGE_LEN;
use super::NOISE_RECORD_PLAINTEXT_LEN;
use super::frame_jsonrpc_message;
use crate::ExecServerError;
use crate::connection::MAX_JSONRPC_MESSAGE_BYTES;
#[test]
fn fragments_and_reassembles_large_jsonrpc_message() {
@@ -25,13 +25,13 @@ fn fragments_and_reassembles_large_jsonrpc_message() {
decoded.extend(decoder.push(record).unwrap());
}
assert_eq!(decoded, vec![message]);
assert_eq!(decoded, vec![(message, framed.len() - size_of::<u32>())]);
}
#[test]
fn rejects_declared_message_length_above_limit_without_payload() {
let mut decoder = JsonRpcMessageDecoder::default();
let declared_len = (MAX_NOISE_JSONRPC_MESSAGE_LEN as u32 + 1).to_be_bytes();
let declared_len = (MAX_JSONRPC_MESSAGE_BYTES as u32 + 1).to_be_bytes();
assert!(matches!(
decoder.push(&declared_len),
codex-rs/exec-server/src/relay.rs
+24 -8
View File
@@ -25,6 +25,7 @@ use crate::connection::CHANNEL_CAPACITY;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
use crate::connection::MAX_JSONRPC_MESSAGE_BYTES;
use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_channel::NoiseChannelPublicKey;
@@ -194,9 +195,16 @@ impl RelayMessageFrame {
}
}
fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
fn into_jsonrpc_message(self) -> Result<(JSONRPCMessage, usize), ExecServerError> {
let payload = self.into_data()?.payload;
serde_json::from_slice(&payload).map_err(ExecServerError::Json)
let encoded_len = payload.len();
if encoded_len > MAX_JSONRPC_MESSAGE_BYTES {
return Err(ExecServerError::Protocol(format!(
"relay JSON-RPC message exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes"
)));
}
let message = serde_json::from_slice(&payload).map_err(ExecServerError::Json)?;
Ok((message, encoded_len))
}
pub(crate) fn into_handshake_payload(self) -> Result<Vec<u8>, ExecServerError> {
@@ -365,12 +373,15 @@ where
};
match kind {
RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() {
Ok(message) => {
Ok((message, encoded_len)) => {
match send_event_with_keepalive(
&mut websocket,
&mut keepalive,
&incoming_tx,
JsonRpcConnectionEvent::Message(message),
JsonRpcConnectionEvent::Message {
message,
encoded_len,
},
)
.await
{
@@ -969,7 +980,7 @@ mod tests {
.await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::Message(actual)) if actual == message
Some(JsonRpcConnectionEvent::Message { message: actual, .. }) if actual == message
));
drop(connection);
@@ -1019,6 +1030,7 @@ mod tests {
ControlledWebSocket::new(/*write_ready*/ true);
let (incoming_tx, mut incoming_rx) = mpsc::channel(/*buffer*/ 1);
let message = test_jsonrpc_message();
let encoded_len = serde_json::to_vec(&message)?.len();
let expected_message = message.clone();
incoming_tx
.send(JsonRpcConnectionEvent::MalformedMessage {
@@ -1036,7 +1048,10 @@ mod tests {
&mut websocket,
&mut keepalive,
&incoming_tx,
JsonRpcConnectionEvent::Message(message),
JsonRpcConnectionEvent::Message {
encoded_len,
message,
},
)
.await
});
@@ -1055,7 +1070,8 @@ mod tests {
));
assert!(matches!(
incoming_rx.recv().await,
Some(JsonRpcConnectionEvent::Message(actual)) if actual == expected_message
Some(JsonRpcConnectionEvent::Message { message: actual, .. })
if actual == expected_message
));
Ok(())
}
@@ -1129,7 +1145,7 @@ mod tests {
};
let frame = decode_relay_message_frame(data_payload.as_ref())?;
assert_eq!(frame.stream_id, stream_id);
assert_eq!(frame.into_jsonrpc_message()?, message);
assert_eq!(frame.into_jsonrpc_message()?.0, message);
drop(connection);
Ok(())
}
codex-rs/exec-server/src/rpc.rs
+91 -13
View File
@@ -6,6 +6,7 @@ use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use codex_exec_server_protocol::JSONRPCError;
use codex_exec_server_protocol::JSONRPCErrorError;
@@ -29,6 +30,9 @@ use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010;
const INGRESS_RATE_WINDOW: Duration = Duration::from_secs(1);
const MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW: usize = 8 * 1024 * 1024;
const MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW: usize = 8 * 1024;
#[derive(Debug)]
pub(crate) enum RpcCallError {
@@ -59,6 +63,54 @@ enum RpcCallTimeout {
pub(crate) enum RpcClientEvent {
Notification(JSONRPCNotification),
Disconnected { reason: Option<String> },
ProtocolError { reason: String },
}
enum RpcReaderEnd {
Disconnected { reason: Option<String> },
ProtocolError { reason: String },
}
struct IngressRateWindow {
started_at: Instant,
bytes: usize,
messages: usize,
}
impl IngressRateWindow {
fn new() -> Self {
Self {
started_at: Instant::now(),
bytes: 0,
messages: 0,
}
}
fn record(&mut self, encoded_len: usize) -> Result<(), String> {
let now = Instant::now();
if now.duration_since(self.started_at) >= INGRESS_RATE_WINDOW {
self.started_at = now;
self.bytes = 0;
self.messages = 0;
}
let next_bytes = self.bytes.saturating_add(encoded_len);
let next_messages = self.messages.saturating_add(1);
if next_bytes > MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW {
return Err(format!(
"exec-server sent more than {MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW} JSON-RPC bytes in one second"
));
}
if next_messages > MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW {
return Err(format!(
"exec-server sent more than {MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW} JSON-RPC messages in one second"
));
}
self.bytes = next_bytes;
self.messages = next_messages;
Ok(())
}
}
#[derive(Debug, Clone, PartialEq)]
@@ -263,36 +315,41 @@ impl RpcClient {
let closed_for_reader = Arc::clone(&closed);
let transport_for_reader = transport.clone();
let reader_task = tokio::spawn(async move {
let disconnect_reason = loop {
let mut ingress_rate = IngressRateWindow::new();
let reader_end = loop {
let Some(event) = incoming_rx.recv().await else {
break None;
break RpcReaderEnd::Disconnected { reason: None };
};
match event {
JsonRpcConnectionEvent::Message(message) => {
JsonRpcConnectionEvent::Message {
message,
encoded_len,
} => {
if let Err(reason) = ingress_rate.record(encoded_len) {
break RpcReaderEnd::ProtocolError { reason };
}
if let Err(err) =
handle_server_message(&pending_for_reader, &event_tx, message).await
{
let _ = err;
break None;
break RpcReaderEnd::ProtocolError { reason: err };
}
}
JsonRpcConnectionEvent::MalformedMessage { reason } => {
let _ = reason;
break None;
break RpcReaderEnd::ProtocolError { reason };
}
JsonRpcConnectionEvent::Disconnected { reason } => {
break reason;
break RpcReaderEnd::Disconnected { reason };
}
}
};
closed_for_reader.store(true, Ordering::Release);
drain_pending(&pending_for_reader).await;
let _ = event_tx
.send(RpcClientEvent::Disconnected {
reason: disconnect_reason,
})
.await;
let event = match reader_end {
RpcReaderEnd::Disconnected { reason } => RpcClientEvent::Disconnected { reason },
RpcReaderEnd::ProtocolError { reason } => RpcClientEvent::ProtocolError { reason },
};
let _ = event_tx.send(event).await;
transport_for_reader.terminate();
});
@@ -626,6 +683,9 @@ mod tests {
use tracing_subscriber::filter::filter_fn;
use tracing_subscriber::prelude::*;
use super::IngressRateWindow;
use super::MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW;
use super::MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW;
use super::RpcClient;
use crate::connection::JsonRpcConnection;
@@ -665,6 +725,24 @@ mod tests {
}
}
#[test]
fn ingress_rate_window_bounds_bytes() {
let mut window = IngressRateWindow::new();
assert!(window.record(MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW).is_ok());
assert!(window.record(1).is_err());
}
#[test]
fn ingress_rate_window_bounds_messages() {
let mut window = IngressRateWindow::new();
for _ in 0..MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW {
assert!(window.record(0).is_ok());
}
assert!(window.record(0).is_err());
}
#[tokio::test]
async fn rpc_client_matches_out_of_order_responses_by_request_id() {
let (client_stdin, server_reader) = tokio::io::duplex(4096);
codex-rs/exec-server/src/server/processor.rs
+1 -1
View File
@@ -125,7 +125,7 @@ async fn run_connection(
break;
}
}
JsonRpcConnectionEvent::Message(message) => match message {
JsonRpcConnectionEvent::Message { message, .. } => match message {
codex_exec_server_protocol::JSONRPCMessage::Request(request) => {
let request_started_at = Instant::now();
if let Some((method, route)) = router.request_route(request.method.as_str()) {