exec-server/src/noise_relay/message_framing.rs
85 lines
use codex_app_server_protocol::JSONRPCMessage;use crate::ExecServerError;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.////// Clatter limits an individual Noise message to 65,535 bytes, while valid/// exec-server responses can be much larger. A four-byte authenticated length/// prefix lets the caller split this byte stream into bounded Noise records and/// lets the receiver reconstruct exact JSON-RPC message boundaries.pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> { 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 { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC message exceeds maximum length".to_string(), )); } framed[..LENGTH_PREFIX_BYTES].copy_from_slice(&(message_len as u32).to_be_bytes()); Ok(framed)}/// Incrementally reconstructs authenticated JSON-RPC messages from Noise records.////// The length prefix is encrypted along with the message. It is still bounded/// here so a bad authenticated peer cannot grow the reassembly buffer forever.#[derive(Default)]pub(crate) struct JsonRpcMessageDecoder { buffered: Vec<u8>,}impl JsonRpcMessageDecoder { /// Append one decrypted record and return all complete framed messages. pub(crate) fn push( &mut self, plaintext_record: &[u8], ) -> Result<Vec<JSONRPCMessage>, ExecServerError> { if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN { return Err(ExecServerError::Protocol( "Noise relay plaintext record exceeds maximum length".to_string(), )); } self.buffered.extend_from_slice(plaintext_record); // One record can finish multiple messages, and one message can span // multiple records. Parse only after the authenticated length prefix // and the full declared payload are present. let mut messages = Vec::new(); while let Some(prefix) = self.buffered.get(..LENGTH_PREFIX_BYTES) { 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 { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC message has invalid length".to_string(), )); } let framed_len = LENGTH_PREFIX_BYTES + message_len; if self.buffered.len() < framed_len { break; } messages.push(serde_json::from_slice( &self.buffered[LENGTH_PREFIX_BYTES..framed_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 { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC reassembly buffer exceeds maximum length".to_string(), )); } Ok(messages) }}#[cfg(test)]#[path = "message_framing_tests.rs"]mod tests;