From fd62a249359eaff578aaf7dba7312523f74fdc45 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 3 Jun 2026 13:39:58 -0700 Subject: [PATCH] refactor: keep Noise foundation independently usable Co-authored-by: Codex noreply@openai.com --- codex-rs/exec-server/src/noise_channel.rs | 217 ----------------- .../exec-server/src/noise_channel_tests.rs | 222 ------------------ 2 files changed, 439 deletions(-) diff --git a/codex-rs/exec-server/src/noise_channel.rs b/codex-rs/exec-server/src/noise_channel.rs index 036129b1fd..f174808e4f 100644 --- a/codex-rs/exec-server/src/noise_channel.rs +++ b/codex-rs/exec-server/src/noise_channel.rs @@ -9,20 +9,11 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD; -use clatter::HybridHandshake; -use clatter::HybridHandshakeParams; use clatter::KeyPair; use clatter::bytearray::ByteArray; -use clatter::constants::MAX_MESSAGE_LEN; -use clatter::crypto::cipher::AesGcm; use clatter::crypto::dh::X25519; -use clatter::crypto::hash::Sha256; -use clatter::handshakepattern::noise_hybrid_ik; -use clatter::traits::Cipher; use clatter::traits::Dh; -use clatter::traits::Handshaker; use clatter::traits::Kem; -use clatter::transportstate::TransportState; use serde::Deserialize; use serde::Serialize; @@ -32,11 +23,7 @@ use crate::aws_lc_ml_kem::PUBLIC_KEY_LEN as MLKEM768_PUBLIC_KEY_LEN; pub const NOISE_CHANNEL_SUITE: &str = "Noise_hybridIK_X25519+MLKEM768_AESGCM_SHA256"; const X25519_PUBLIC_KEY_LEN: usize = 32; -const MAX_TRANSPORT_RECORDS_PER_DIRECTION: u64 = u32::MAX as u64; -const PROLOGUE_DOMAIN: &[u8] = b"codex-exec-server-relay-noise/v1"; -type Handshake = HybridHandshake; -type Transport = TransportState; type DhKeyPair = KeyPair<::PubKey, ::PrivateKey>; type KemKeyPair = KeyPair<::PubKey, ::SecretKey>; @@ -78,14 +65,6 @@ impl NoiseChannelPublicKey { } } - fn from_raw(dh: &::PubKey, kem: &::PubKey) -> Self { - Self { - suite: NOISE_CHANNEL_SUITE.to_string(), - x25519_public_key: STANDARD.encode(dh), - mlkem768_public_key: STANDARD.encode(kem.as_slice()), - } - } - fn decode( &self, ) -> Result<(::PubKey, ::PubKey), NoiseChannelError> { @@ -150,202 +129,6 @@ impl NoiseChannelIdentity { } } -pub(crate) struct InitiatorHandshake { - handshake: Handshake, -} - -impl InitiatorHandshake { - /// Start hybrid IK while pinning the expected responder static key. - /// - /// `payload` is encrypted by the first IK message. The relay uses it for - /// the short-lived registry authorization naming this harness public key. - pub(crate) fn start( - identity: &NoiseChannelIdentity, - responder_public_key: &NoiseChannelPublicKey, - prologue: &[u8], - payload: &[u8], - ) -> Result<(Self, Vec), NoiseChannelError> { - let (responder_dh, responder_kem) = responder_public_key.decode()?; - - // IK authenticates both static identities. Supplying both responder - // components here is what makes a misrouted or impersonating - // exec-server fail before any JSON-RPC plaintext is released. - let params = HybridHandshakeParams::new(noise_hybrid_ik(), true) - .with_prologue(prologue) - .with_s(identity.dh.clone()) - .with_s_kem(identity.kem.clone()) - .with_rs(responder_dh) - .with_rs_kem(responder_kem); - let mut handshake = Handshake::new(params)?; - let mut output = [0u8; MAX_MESSAGE_LEN]; - let output_len = handshake.write_message(payload, &mut output)?; - Ok((Self { handshake }, output[..output_len].to_vec())) - } - - /// Consume the responder message and enter transport mode. - /// - /// The responder message carries no application payload in v1. Rejecting - /// one keeps future protocol additions from becoming an implicit channel. - pub(crate) fn finish(mut self, response: &[u8]) -> Result { - ensure_noise_frame_len(response.len(), "handshake response is too large")?; - let mut payload = [0u8; MAX_MESSAGE_LEN]; - let payload_len = self.handshake.read_message(response, &mut payload)?; - if payload_len != 0 { - return Err(NoiseChannelError::InvalidMessage( - "handshake response payload must be empty", - )); - } - Ok(NoiseTransport { - transport: self.handshake.finalize()?, - }) - } -} - -pub(crate) struct PendingResponderHandshake { - handshake: Handshake, - initiator_public_key: NoiseChannelPublicKey, - payload: Vec, -} - -impl PendingResponderHandshake { - /// Authenticate and parse the first IK message without completing it. - /// - /// This split is intentional: callers must authorize `initiator_public_key` - /// with the registry before calling [`Self::complete`]. - pub(crate) fn read_request( - identity: &NoiseChannelIdentity, - prologue: &[u8], - request: &[u8], - ) -> Result { - ensure_noise_frame_len(request.len(), "handshake request is too large")?; - let params = HybridHandshakeParams::new(noise_hybrid_ik(), false) - .with_prologue(prologue) - .with_s(identity.dh.clone()) - .with_s_kem(identity.kem.clone()); - let mut handshake = Handshake::new(params)?; - let mut payload = [0u8; MAX_MESSAGE_LEN]; - let payload_len = handshake.read_message(request, &mut payload)?; - // Clatter exposes the initiator static key only after the first IK - // message authenticates and decrypts successfully. - let remote = handshake - .get_remote_static() - .ok_or(NoiseChannelError::InvalidMessage( - "handshake request is missing initiator static key", - ))?; - Ok(Self { - handshake, - initiator_public_key: NoiseChannelPublicKey::from_raw(remote.dh(), remote.kem()), - payload: payload[..payload_len].to_vec(), - }) - } - - pub(crate) fn initiator_public_key(&self) -> &NoiseChannelPublicKey { - &self.initiator_public_key - } - - /// Move the authenticated first-message payload out of pending state. - /// - /// The v1 payload is a short-lived registry authorization and is not - /// needed to complete the handshake. Moving it avoids retaining a second - /// copy while external authorization is in flight. - pub(crate) fn take_payload(&mut self) -> Vec { - std::mem::take(&mut self.payload) - } - - /// Finish the responder handshake after external harness authorization. - pub(crate) fn complete(mut self) -> Result<(NoiseTransport, Vec), NoiseChannelError> { - let mut response = [0u8; MAX_MESSAGE_LEN]; - let response_len = self.handshake.write_message(&[], &mut response)?; - Ok(( - NoiseTransport { - transport: self.handshake.finalize()?, - }, - response[..response_len].to_vec(), - )) - } -} - -pub(crate) struct NoiseTransport { - transport: Transport, -} - -impl NoiseTransport { - /// Encrypt exactly one ordered transport record. - /// - /// The caller owns relay sequence assignment and must never encrypt the - /// same logical record twice under different transport nonces. - pub(crate) fn encrypt(&mut self, plaintext: &[u8]) -> Result, NoiseChannelError> { - if self.transport.sending_nonce() >= MAX_TRANSPORT_RECORDS_PER_DIRECTION { - return Err(NoiseChannelError::InvalidState( - "transport record nonce exhausted", - )); - } - let frame_len = plaintext.len().checked_add(AesGcm::tag_len()).ok_or( - NoiseChannelError::InvalidMessage("transport plaintext is too large"), - )?; - ensure_noise_frame_len(frame_len, "transport plaintext is too large")?; - Ok(self.transport.send_vec(plaintext)?) - } - - /// Decrypt exactly the next ordered transport record. - /// - /// Clatter advances the receiving nonce during this call, so callers must - /// reorder and deduplicate relay frames before invoking it. - pub(crate) fn decrypt(&mut self, ciphertext: &[u8]) -> Result, NoiseChannelError> { - if self.transport.receiving_nonce() >= MAX_TRANSPORT_RECORDS_PER_DIRECTION { - return Err(NoiseChannelError::InvalidState( - "transport record nonce exhausted", - )); - } - if ciphertext.len() < AesGcm::tag_len() { - return Err(NoiseChannelError::InvalidMessage( - "transport ciphertext is too short", - )); - } - ensure_noise_frame_len(ciphertext.len(), "transport ciphertext is too large")?; - Ok(self.transport.receive_vec(ciphertext)?) - } -} - -/// Build the transcript prologue that binds cryptographic identity to routing. -/// -/// A handshake captured from one environment, exec-server registration, or -/// relay stream cannot be replayed into another because every participant must -/// construct the same prologue before the first Noise message is processed. -pub(crate) fn noise_channel_prologue( - environment_id: &str, - executor_registration_id: &str, - stream_id: &str, -) -> Result, NoiseChannelError> { - let mut prologue = Vec::new(); - append_prologue_part(&mut prologue, PROLOGUE_DOMAIN)?; - append_prologue_part(&mut prologue, environment_id.as_bytes())?; - append_prologue_part(&mut prologue, executor_registration_id.as_bytes())?; - append_prologue_part(&mut prologue, stream_id.as_bytes())?; - Ok(prologue) -} - -fn append_prologue_part(prologue: &mut Vec, part: &[u8]) -> Result<(), NoiseChannelError> { - // Length prefixes make component boundaries unambiguous. Raw concatenation - // would allow different identifier tuples to produce the same prologue. - let len = u32::try_from(part.len()).map_err(|_| { - NoiseChannelError::InvalidMessage("Noise channel prologue part is too large") - })?; - prologue.extend_from_slice(&len.to_be_bytes()); - prologue.extend_from_slice(part); - Ok(()) -} - -fn ensure_noise_frame_len( - frame_len: usize, - message: &'static str, -) -> Result<(), NoiseChannelError> { - if frame_len > MAX_MESSAGE_LEN { - return Err(NoiseChannelError::InvalidMessage(message)); - } - Ok(()) -} - #[derive(Debug, thiserror::Error)] pub enum NoiseChannelError { #[error("Noise channel key generation failed: {0}")] diff --git a/codex-rs/exec-server/src/noise_channel_tests.rs b/codex-rs/exec-server/src/noise_channel_tests.rs index 5034420c2c..260d64c333 100644 --- a/codex-rs/exec-server/src/noise_channel_tests.rs +++ b/codex-rs/exec-server/src/noise_channel_tests.rs @@ -1,230 +1,8 @@ use pretty_assertions::assert_eq; -use super::InitiatorHandshake; -use super::MAX_TRANSPORT_RECORDS_PER_DIRECTION; use super::NOISE_CHANNEL_SUITE; -use super::NoiseChannelError; use super::NoiseChannelIdentity; use super::NoiseChannelPublicKey; -use super::PendingResponderHandshake; -use super::noise_channel_prologue; - -#[test] -fn hybrid_ik_roundtrip_authenticates_both_endpoints() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - let authorization = b"harness-key-authorization"; - - let (initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &responder.public_key(), - &prologue, - authorization, - ) - .expect("start initiator handshake"); - let mut responder_handshake = - PendingResponderHandshake::read_request(&responder, &prologue, &request) - .expect("read responder handshake"); - - assert_eq!( - responder_handshake.initiator_public_key(), - &initiator.public_key() - ); - assert_eq!(responder_handshake.take_payload(), authorization); - - let (mut responder_transport, response) = responder_handshake - .complete() - .expect("complete responder handshake"); - let mut initiator_transport = initiator_handshake - .finish(&response) - .expect("complete initiator handshake"); - - let request_ciphertext = initiator_transport - .encrypt(b"request") - .expect("encrypt request"); - assert_ne!(request_ciphertext, b"request"); - assert_eq!( - responder_transport - .decrypt(&request_ciphertext) - .expect("decrypt request"), - b"request" - ); - - let response_ciphertext = responder_transport - .encrypt(b"response") - .expect("encrypt response"); - assert_ne!(response_ciphertext, b"response"); - assert_eq!( - initiator_transport - .decrypt(&response_ciphertext) - .expect("decrypt response"), - b"response" - ); -} - -#[test] -fn initiator_rejects_wrong_responder_key() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let expected_responder = NoiseChannelIdentity::generate().expect("generate expected identity"); - let actual_responder = NoiseChannelIdentity::generate().expect("generate actual identity"); - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - - let (_initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &expected_responder.public_key(), - &prologue, - b"authorization", - ) - .expect("start initiator handshake"); - - assert!( - PendingResponderHandshake::read_request(&actual_responder, &prologue, &request).is_err() - ); -} - -#[test] -fn responder_rejects_mismatched_prologue() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); - let initiator_prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - let responder_prologue = - noise_channel_prologue("env-1", "registration-1", "stream-2").expect("build prologue"); - let (_initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &responder.public_key(), - &initiator_prologue, - b"authorization", - ) - .expect("start initiator handshake"); - - assert!( - PendingResponderHandshake::read_request(&responder, &responder_prologue, &request).is_err() - ); -} - -#[test] -fn prologue_encoding_is_stable_and_unambiguous() { - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - - assert_eq!( - prologue, - b"\x00\x00\x00\x20codex-exec-server-relay-noise/v1\ - \x00\x00\x00\x05env-1\ - \x00\x00\x00\x0eregistration-1\ - \x00\x00\x00\x08stream-1" - .to_vec() - ); -} - -#[test] -fn transport_rejects_tampered_ciphertext() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - let (initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &responder.public_key(), - &prologue, - b"authorization", - ) - .expect("start initiator handshake"); - let responder_handshake = - PendingResponderHandshake::read_request(&responder, &prologue, &request) - .expect("read responder handshake"); - let (mut responder_transport, response) = responder_handshake - .complete() - .expect("complete responder handshake"); - let mut initiator_transport = initiator_handshake - .finish(&response) - .expect("complete initiator handshake"); - let mut ciphertext = initiator_transport - .encrypt(b"request") - .expect("encrypt request"); - ciphertext[0] ^= 1; - - assert!(responder_transport.decrypt(&ciphertext).is_err()); -} - -#[test] -fn transport_rejects_exhausted_receiving_nonce_before_decryption() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - let (initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &responder.public_key(), - &prologue, - b"authorization", - ) - .expect("start initiator handshake"); - let responder_handshake = - PendingResponderHandshake::read_request(&responder, &prologue, &request) - .expect("read responder handshake"); - let (mut responder_transport, response) = responder_handshake - .complete() - .expect("complete responder handshake"); - let mut initiator_transport = initiator_handshake - .finish(&response) - .expect("complete initiator handshake"); - let ciphertext = initiator_transport - .encrypt(b"request") - .expect("encrypt request"); - responder_transport - .transport - .set_receiving_nonce(MAX_TRANSPORT_RECORDS_PER_DIRECTION); - - assert!(matches!( - responder_transport.decrypt(&ciphertext), - Err(NoiseChannelError::InvalidState( - "transport record nonce exhausted" - )) - )); -} - -#[test] -fn transport_rejects_replayed_ciphertext() { - let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); - let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); - let prologue = - noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue"); - let (initiator_handshake, request) = InitiatorHandshake::start( - &initiator, - &responder.public_key(), - &prologue, - b"authorization", - ) - .expect("start initiator handshake"); - let responder_handshake = - PendingResponderHandshake::read_request(&responder, &prologue, &request) - .expect("read responder handshake"); - let (mut responder_transport, response) = responder_handshake - .complete() - .expect("complete responder handshake"); - let mut initiator_transport = initiator_handshake - .finish(&response) - .expect("complete initiator handshake"); - let ciphertext = initiator_transport - .encrypt(b"request") - .expect("encrypt request"); - - assert_eq!( - responder_transport - .decrypt(&ciphertext) - .expect("decrypt request"), - b"request" - ); - assert!(matches!( - responder_transport.decrypt(&ciphertext), - Err(NoiseChannelError::Transport(_)) - )); -} #[test] fn public_key_validation_rejects_unknown_suite() {