diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index a0ce5bd555a9..0b69f46b486d 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -86,3 +86,7 @@ impl JsonRpcMessageDecoder { Ok(Some(message_len)) } } + +#[cfg(test)] +#[path = "message_framing_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs new file mode 100644 index 000000000000..c3df14bf5622 --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs @@ -0,0 +1,52 @@ +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_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; + +#[test] +fn fragments_and_reassembles_large_jsonrpc_message() { + let message = JSONRPCMessage::Notification(JSONRPCNotification { + method: "large/test".to_string(), + params: Some(serde_json::json!({ + "data": "x".repeat(128 * 1024), + })), + }); + let framed = frame_jsonrpc_message(&message).unwrap(); + assert!(framed.len() > 128 * 1024); + + let mut decoder = JsonRpcMessageDecoder::default(); + let mut decoded = Vec::new(); + for record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) { + decoded.extend(decoder.push(record).unwrap()); + } + + assert_eq!(decoded, vec![message]); +} + +#[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(); + + assert!(matches!( + decoder.push(&declared_len), + Err(ExecServerError::Protocol(message)) + if message == "Noise relay JSON-RPC message has invalid length" + )); +} + +#[test] +fn rejects_oversized_plaintext_record() { + let mut decoder = JsonRpcMessageDecoder::default(); + + assert!(matches!( + decoder.push(&vec![0; NOISE_RECORD_PLAINTEXT_LEN + 1]), + Err(ExecServerError::Protocol(message)) + if message == "Noise relay plaintext record exceeds maximum length" + )); +} diff --git a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs index dea332ec06fd..f122d4970743 100644 --- a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs +++ b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs @@ -79,3 +79,7 @@ impl OrderedCiphertextFrames { Ok(()) } } + +#[cfg(test)] +#[path = "ordered_ciphertext_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs new file mode 100644 index 000000000000..6aa86fcdf055 --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs @@ -0,0 +1,52 @@ +use pretty_assertions::assert_eq; + +use super::MAX_PENDING_BYTES; +use super::OrderedCiphertextFrames; + +#[test] +fn releases_ciphertexts_only_in_nonce_order() { + let mut frames = OrderedCiphertextFrames::default(); + + assert_eq!( + frames.push(/*seq*/ 1, b"second".to_vec()).unwrap(), + Vec::>::new() + ); + assert_eq!( + frames.push(/*seq*/ 0, b"first".to_vec()).unwrap(), + vec![b"first".to_vec(), b"second".to_vec()] + ); +} + +#[test] +fn ignores_duplicate_ciphertexts_without_replacing_buffered_record() { + let mut frames = OrderedCiphertextFrames::default(); + + assert_eq!( + frames.push(/*seq*/ 1, b"first copy".to_vec()).unwrap(), + Vec::>::new() + ); + assert_eq!( + frames.push(/*seq*/ 1, b"replacement".to_vec()).unwrap(), + Vec::>::new() + ); + assert_eq!( + frames.push(/*seq*/ 0, b"zero".to_vec()).unwrap(), + vec![b"zero".to_vec(), b"first copy".to_vec()] + ); + assert_eq!( + frames.push(/*seq*/ 0, b"duplicate".to_vec()).unwrap(), + Vec::>::new() + ); +} + +#[test] +fn rejects_unbounded_reordering() { + let mut frames = OrderedCiphertextFrames::default(); + + assert!(frames.push(/*seq*/ 65, Vec::new()).is_err()); + assert!( + frames + .push(/*seq*/ 1, vec![0; MAX_PENDING_BYTES + 1]) + .is_err() + ); +}