From 03d54d008b07968e2830b6cd0957f4aa68b322b4 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 3 Jun 2026 00:44:44 -0700 Subject: [PATCH] fix: make secure exec relay opt in Co-authored-by: Codex noreply@openai.com --- codex-rs/cli/Cargo.toml | 2 +- codex-rs/cli/src/main.rs | 35 ++++ codex-rs/exec-server/README.md | 10 +- codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/relay.rs | 199 +++++++++++++++++++++++ codex-rs/exec-server/src/remote.rs | 164 +++++++++++++++---- codex-rs/exec-server/src/remote_tests.rs | 61 +++++-- codex-rs/exec-server/tests/relay.rs | 83 +++++++++- 8 files changed, 507 insertions(+), 48 deletions(-) diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 6a4d78fbc7..d66273b9b2 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -19,7 +19,7 @@ workspace = true [dependencies] anyhow = { workspace = true } -clap = { workspace = true, features = ["derive"] } +clap = { workspace = true, features = ["derive", "env"] } clap_complete = { workspace = true } codex-app-server = { workspace = true } codex-app-server-daemon = { workspace = true } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 2a5ce171b7..5fb065dce9 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -542,6 +542,14 @@ struct ExecServerCommand { /// Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote registration. #[arg(long = "use-agent-identity-auth", requires = "remote")] use_agent_identity_auth: bool, + + /// Opt into Noise-encrypted remote relay communication. + #[arg( + long = "enable-noise", + env = "CODEX_EXEC_SERVER_ENABLE_NOISE", + requires = "remote" + )] + enable_noise: bool, } #[derive(Debug, clap::Subcommand)] @@ -1611,6 +1619,9 @@ async fn run_exec_server_command( if let Some(name) = cmd.name { remote_config.name = name; } + if cmd.enable_noise { + remote_config.relay_security = codex_exec_server::RemoteRelaySecurity::Noise; + } codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; Ok(()) } else { @@ -3222,6 +3233,30 @@ mod tests { .expect("exec-server should support root --strict-config"); } + #[test] + fn exec_server_enable_noise_is_explicit_remote_opt_in() { + let cli = MultitoolCli::try_parse_from([ + "codex", + "exec-server", + "--remote", + "https://registry.example", + "--environment-id", + "environment-1", + "--enable-noise", + ]) + .expect("parse"); + + assert_matches!( + cli.subcommand, + Some(Subcommand::ExecServer(ExecServerCommand { + enable_noise: true, + .. + })) + ); + MultitoolCli::try_parse_from(["codex", "exec-server", "--enable-noise"]) + .expect_err("--enable-noise should require --remote"); + } + #[test] fn root_strict_config_is_rejected_for_unsupported_subcommands() { let cli = MultitoolCli::try_parse_from(["codex", "--strict-config", "mcp", "list"]) diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 8fa1a9eb75..5f93e8bda9 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -22,10 +22,15 @@ the wire. The CLI entrypoint supports: - `ws://IP:PORT` (default) -- `--remote URL --environment-id ID [--name NAME]` +- `--remote URL --environment-id ID [--name NAME] [--enable-noise]` Remote mode registers the local exec-server with the environment registry, then reconnects to the service-provided rendezvous websocket as the environment. +It uses the legacy relay contract by default for compatibility with existing +registries and harnesses. `--enable-noise` explicitly opts into Noise-encrypted +relay registration and requires both the registry and harness to support the +secure relay contract. Container entrypoints can make the same explicit choice +by setting `CODEX_EXEC_SERVER_ENABLE_NOISE=true`. It uses the standard Codex ChatGPT sign-in state; run `codex login` first when remote registration needs authentication. Containerized callers that receive an Agent Identity JWT in `CODEX_ACCESS_TOKEN` can opt into that auth path with @@ -45,7 +50,8 @@ codex exec-server \ Wire framing: - local websocket: one JSON-RPC message per websocket frame -- remote websocket: binary protobuf relay frames carrying JSON-RPC payloads +- legacy remote websocket: binary protobuf relay frames carrying JSON-RPC payloads +- Noise remote websocket: binary protobuf relay frames carrying encrypted payloads ## Remote Relay Message Format diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index beaa3be56d..5a1bed176d 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -107,6 +107,7 @@ pub use protocol::WriteParams; pub use protocol::WriteResponse; pub use protocol::WriteStatus; pub use remote::RemoteEnvironmentConfig; +pub use remote::RemoteRelaySecurity; pub use remote::run_remote_environment; pub use runtime_paths::ExecServerRuntimePaths; pub use secure_channel::SecureChannelError; diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index e42f520b58..e9574921b0 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,12 +1,18 @@ +use std::collections::HashMap; + use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; use futures::SinkExt; use futures::Stream; use futures::StreamExt; use prost::Message as ProstMessage; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; use tokio::sync::mpsc; use tokio::sync::watch; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; +use tracing::debug; use tracing::warn; use uuid::Uuid; @@ -21,6 +27,7 @@ use crate::relay_proto::RelayMessageFrame; use crate::relay_proto::RelayReset; use crate::relay_proto::RelayResume; use crate::relay_proto::relay_message_frame; +use crate::server::ConnectionProcessor; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const MAX_RELAY_RESET_REASON_BYTES: usize = 256; @@ -365,6 +372,198 @@ where } } +/// Runs the backwards-compatible, cleartext multiplexed relay protocol. +/// +/// The physical websocket carries protobuf [`RelayMessageFrame`] values. Each +/// frame's `stream_id` selects an independent logical JSON-RPC connection, so a +/// single registered exec-server can serve multiple orchestrator sessions. +/// This runner intentionally preserves the protocol used before secure relay +/// support was introduced. Callers must select the separate secure runner when +/// Noise protection was explicitly negotiated during registration. +/// +/// Relay framing is not an authentication boundary. Every frame is validated +/// before its routing metadata or payload is used, and malformed frames are +/// dropped without affecting the other virtual streams on the connection. +pub(crate) async fn run_multiplexed_environment( + stream: WebSocketStream, + processor: ConnectionProcessor, +) where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let mut websocket = stream; + let (physical_outgoing_tx, mut physical_outgoing_rx) = + mpsc::channel::>(CHANNEL_CAPACITY); + + let mut streams: HashMap = HashMap::new(); + loop { + // Serialize all logical-stream writes through this task so only one + // owner writes to the physical websocket. Reads remain in the same + // select loop, which also keeps disconnect handling deterministic. + let frame = tokio::select! { + maybe_encoded = physical_outgoing_rx.recv() => { + let Some(encoded) = maybe_encoded else { + break; + }; + if websocket.send(Message::Binary(encoded.into())).await.is_err() { + break; + } + continue; + } + incoming_message = websocket.next() => match incoming_message { + Some(Ok(Message::Binary(payload))) => { + match decode_relay_message_frame(payload.as_ref()) { + Ok(frame) => frame, + Err(err) => { + warn!("dropping malformed relay message frame from harness: {err}"); + continue; + } + } + } + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, + Some(Ok(Message::Text(_))) => { + warn!("dropping non-binary relay message frame from harness"); + continue; + } + Some(Err(err)) => { + debug!("multiplexed environment websocket read failed: {err}"); + break; + } + } + }; + + let kind = match frame.validate() { + Ok(kind) => kind, + Err(err) => { + warn!("dropping invalid relay message frame: {err}"); + continue; + } + }; + + match kind { + RelayFrameBodyKind::Data => { + let stream_id = frame.stream_id.clone(); + let message = match frame.into_jsonrpc_message() { + Ok(message) => message, + Err(err) => { + warn!("dropping malformed relay data message frame: {err}"); + continue; + } + }; + + // A logical connection is created lazily on its first data + // frame. The connection processor owns the JSON-RPC lifecycle; + // this task only translates between relay and connection + // events. + let stream = streams.entry(stream_id.clone()).or_insert_with(|| { + spawn_virtual_stream( + stream_id.clone(), + processor.clone(), + physical_outgoing_tx.clone(), + ) + }); + if stream + .incoming_tx + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + streams.remove(&stream_id); + } + } + RelayFrameBodyKind::Reset => { + if let Some(stream) = streams.remove(&frame.stream_id) { + stream.disconnect(frame.into_reset_reason()).await; + } + } + // These control frames are meaningful to other relay protocol + // variants. The legacy protocol has no resume or handshake state, + // so ignoring them preserves its existing wire behavior. + RelayFrameBodyKind::Ack + | RelayFrameBodyKind::Resume + | RelayFrameBodyKind::Heartbeat + | RelayFrameBodyKind::Handshake => {} + } + } + + // A physical disconnect ends every virtual connection. Notify each + // processor explicitly so requests do not remain live after the relay + // websocket has disappeared. + for (_stream_id, stream) in streams { + stream.disconnect(/*reason*/ None).await; + } + drop(physical_outgoing_tx); +} + +/// The exec-server-facing half of one logical connection on the shared relay. +struct VirtualStream { + incoming_tx: mpsc::Sender, + disconnected_tx: watch::Sender, +} + +impl VirtualStream { + /// Marks the logical connection disconnected and supplies the peer's reset + /// reason, when one was provided. + async fn disconnect(self, reason: Option) { + let _ = self.disconnected_tx.send(true); + let _ = self + .incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { reason }) + .await; + } +} + +/// Creates a logical JSON-RPC connection and forwards its outgoing messages +/// back to the physical relay writer as legacy cleartext data frames. +fn spawn_virtual_stream( + stream_id: String, + processor: ConnectionProcessor, + physical_outgoing_tx: mpsc::Sender>, +) -> VirtualStream { + let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (disconnected_tx, disconnected_rx) = watch::channel(false); + + let writer_stream_id = stream_id; + let writer_task = tokio::spawn(async move { + let mut next_seq = 0u32; + while let Some(message) = json_outgoing_rx.recv().await { + let payload = match jsonrpc_payload(&message) { + Ok(payload) => payload, + Err(err) => { + warn!("failed to serialize virtual stream JSON-RPC payload: {err}"); + break; + } + }; + let frame = RelayMessageFrame::data(writer_stream_id.clone(), next_seq, payload); + next_seq = next_seq.wrapping_add(1); + if physical_outgoing_tx + .send(encode_relay_message_frame(&frame)) + .await + .is_err() + { + break; + } + } + }); + + let connection = JsonRpcConnection { + outgoing_tx: json_outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![writer_task], + transport: JsonRpcTransport::External, + }; + tokio::spawn(async move { + processor.run_connection(connection).await; + }); + + VirtualStream { + incoming_tx, + disconnected_tx, + } +} + #[cfg(test)] mod tests { use std::pin::Pin; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 5ebf464dc5..dd83234a14 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -16,6 +16,7 @@ use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::SecureChannelIdentity; use crate::SecureChannelPublicKey; +use crate::relay::run_multiplexed_environment; use crate::secure_relay::HarnessKeyValidator; use crate::secure_relay::run_secure_multiplexed_environment; use crate::secure_relay::secure_relay_websocket_config; @@ -55,11 +56,36 @@ impl EnvironmentRegistryClient { }) } - async fn register_environment( + /// Register using the original body-less registry contract. + /// + /// This remains the default so a new exec-server can still connect through + /// registries and harnesses that have not rolled out secure relay support. + async fn register_legacy_environment( + &self, + environment_id: &str, + ) -> Result { + let response = self + .http + .post(endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/register"), + )) + .headers(self.auth_provider.to_auth_headers()) + .send() + .await?; + self.parse_json_response(response).await + } + + /// Register the exec-server's static Noise identity with a secure-aware registry. + /// + /// Supplying this request body is the protocol-level opt in. Registries can + /// therefore distinguish secure registrations from the legacy body-less + /// contract without guessing based on binary version or rollout state. + async fn register_noise_environment( &self, environment_id: &str, executor_public_key: &SecureChannelPublicKey, - ) -> Result { + ) -> Result { let response = self .http .post(endpoint_url( @@ -135,6 +161,12 @@ struct EnvironmentRegistryRegistrationRequest<'a> { struct EnvironmentRegistryRegistrationResponse { environment_id: String, url: String, +} + +#[derive(Debug, Clone, Eq, PartialEq, Deserialize)] +struct EnvironmentRegistryNoiseRegistrationResponse { + environment_id: String, + url: String, security_profile: String, executor_registration_id: String, } @@ -173,12 +205,54 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { } } +/// Security mode used for an exec-server's registered remote relay. +/// +/// Legacy is intentionally the default during rollout. Noise must be selected +/// explicitly so mixed-version deployments keep the original registry and relay +/// contract until both endpoints are ready for the secure protocol. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum RemoteRelaySecurity { + #[default] + Legacy, + Noise, +} + +enum RemoteRelaySecurityState { + Legacy, + Noise(SecureChannelIdentity), +} + +enum RegisteredRemoteEnvironment { + Legacy(EnvironmentRegistryRegistrationResponse), + Noise { + response: EnvironmentRegistryNoiseRegistrationResponse, + identity: SecureChannelIdentity, + }, +} + +impl RegisteredRemoteEnvironment { + fn environment_id(&self) -> &str { + match self { + Self::Legacy(response) => &response.environment_id, + Self::Noise { response, .. } => &response.environment_id, + } + } + + fn websocket_url(&self) -> &str { + match self { + Self::Legacy(response) => &response.url, + Self::Noise { response, .. } => &response.url, + } + } +} + /// Configuration for registering an exec-server for remote use. #[derive(Clone)] pub struct RemoteEnvironmentConfig { pub base_url: String, pub environment_id: String, pub name: String, + pub relay_security: RemoteRelaySecurity, auth_provider: SharedAuthProvider, } @@ -188,6 +262,7 @@ impl std::fmt::Debug for RemoteEnvironmentConfig { .field("base_url", &self.base_url) .field("environment_id", &self.environment_id) .field("name", &self.name) + .field("relay_security", &self.relay_security) .field("auth_provider", &"") .finish() } @@ -204,6 +279,7 @@ impl RemoteEnvironmentConfig { base_url, environment_id, name: "codex-exec-server".to_string(), + relay_security: RemoteRelaySecurity::Legacy, auth_provider, }) } @@ -219,35 +295,56 @@ pub async fn run_remote_environment( let client = EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; let processor = ConnectionProcessor::new(runtime_paths); - let identity = SecureChannelIdentity::generate()?; - let executor_public_key = identity.public_key(); + let security_state = match config.relay_security { + RemoteRelaySecurity::Legacy => RemoteRelaySecurityState::Legacy, + RemoteRelaySecurity::Noise => { + RemoteRelaySecurityState::Noise(SecureChannelIdentity::generate()?) + } + }; let mut backoff = Duration::from_secs(1); loop { - let response = client - .register_environment(&config.environment_id, &executor_public_key) - .await?; - if response.security_profile != SECURE_RELAY_SECURITY_PROFILE { - return Err(ExecServerError::Protocol(format!( - "environment registry returned unsupported security profile `{}`", - response.security_profile - ))); - } - if response.environment_id != config.environment_id { + let registration = match &security_state { + RemoteRelaySecurityState::Legacy => RegisteredRemoteEnvironment::Legacy( + client + .register_legacy_environment(&config.environment_id) + .await?, + ), + RemoteRelaySecurityState::Noise(identity) => { + let response = client + .register_noise_environment(&config.environment_id, &identity.public_key()) + .await?; + if response.security_profile != SECURE_RELAY_SECURITY_PROFILE { + return Err(ExecServerError::Protocol(format!( + "environment registry returned unsupported security profile `{}`", + response.security_profile + ))); + } + RegisteredRemoteEnvironment::Noise { + response, + identity: identity.clone(), + } + } + }; + if registration.environment_id() != config.environment_id { return Err(ExecServerError::Protocol( "environment registry returned a different environment id".to_string(), )); } - let environment_id = &response.environment_id; + let environment_id = registration.environment_id(); info!( "codex exec-server remote environment registered with environment_id {environment_id}" ); + let websocket_config = match ®istration { + RegisteredRemoteEnvironment::Legacy(_) => None, + RegisteredRemoteEnvironment::Noise { .. } => Some(secure_relay_websocket_config()), + }; match timeout( REMOTE_RENDEZVOUS_CONNECT_TIMEOUT, connect_async_with_config( - response.url.as_str(), - Some(secure_relay_websocket_config()), + registration.websocket_url(), + websocket_config, /*disable_nagle*/ false, ), ) @@ -255,19 +352,26 @@ pub async fn run_remote_environment( { Ok(Ok((websocket, _))) => { backoff = Duration::from_secs(1); - run_secure_multiplexed_environment( - websocket, - processor.clone(), - response.environment_id, - response.executor_registration_id.clone(), - identity.clone(), - RegistryHarnessKeyValidator { - client: client.clone(), - environment_id: config.environment_id.clone(), - executor_registration_id: response.executor_registration_id, - }, - ) - .await; + match registration { + RegisteredRemoteEnvironment::Legacy(_) => { + run_multiplexed_environment(websocket, processor.clone()).await; + } + RegisteredRemoteEnvironment::Noise { response, identity } => { + run_secure_multiplexed_environment( + websocket, + processor.clone(), + response.environment_id, + response.executor_registration_id.clone(), + identity, + RegistryHarnessKeyValidator { + client: client.clone(), + environment_id: config.environment_id.clone(), + executor_registration_id: response.executor_registration_id, + }, + ) + .await; + } + } } Ok(Err(err)) => { warn!("failed to connect remote exec-server websocket: {err}"); diff --git a/codex-rs/exec-server/src/remote_tests.rs b/codex-rs/exec-server/src/remote_tests.rs index 9f303a2f42..1b77285ecc 100644 --- a/codex-rs/exec-server/src/remote_tests.rs +++ b/codex-rs/exec-server/src/remote_tests.rs @@ -8,6 +8,8 @@ use reqwest::StatusCode; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; +use wiremock::matchers::body_bytes; +use wiremock::matchers::body_partial_json; use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -47,9 +49,45 @@ async fn register_environment_posts_with_auth_provider_headers() { .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) + .and(body_bytes(Vec::::new())) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "environment_id": "env-1", "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc", + }))) + .mount(&server) + .await; + let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) + .expect("client"); + let response = client + .register_legacy_environment(&config.environment_id) + .await + .expect("register environment"); + + assert_eq!( + response, + EnvironmentRegistryRegistrationResponse { + environment_id: "env-1".to_string(), + url: "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc".to_string(), + } + ); +} + +#[tokio::test] +async fn register_noise_environment_posts_security_profile_and_public_key() { + let server = MockServer::start().await; + let executor_public_key = SecureChannelIdentity::generate() + .expect("identity") + .public_key(); + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-requested/register")) + .and(header("authorization", "Bearer registry-token")) + .and(body_partial_json(serde_json::json!({ + "security_profile": SECURE_RELAY_SECURITY_PROFILE, + "executor_public_key": executor_public_key.clone(), + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "url": "wss://rendezvous.test/secure", "security_profile": SECURE_RELAY_SECURITY_PROFILE, "executor_registration_id": "registration-1", }))) @@ -57,20 +95,17 @@ async fn register_environment_posts_with_auth_provider_headers() { .await; let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) .expect("client"); - let executor_public_key = SecureChannelIdentity::generate() - .expect("identity") - .public_key(); let response = client - .register_environment(&config.environment_id, &executor_public_key) + .register_noise_environment("environment-requested", &executor_public_key) .await - .expect("register environment"); + .expect("register secure environment"); assert_eq!( response, - EnvironmentRegistryRegistrationResponse { - environment_id: "env-1".to_string(), - url: "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc".to_string(), + EnvironmentRegistryNoiseRegistrationResponse { + environment_id: "environment-requested".to_string(), + url: "wss://rendezvous.test/secure".to_string(), security_profile: SECURE_RELAY_SECURITY_PROFILE.to_string(), executor_registration_id: "registration-1".to_string(), } @@ -78,7 +113,7 @@ async fn register_environment_posts_with_auth_provider_headers() { } #[tokio::test] -async fn register_environment_does_not_follow_redirects_with_auth_headers() { +async fn register_legacy_environment_does_not_follow_redirects_with_auth_headers() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/cloud/environment/environment-requested/register")) @@ -97,12 +132,8 @@ async fn register_environment_does_not_follow_redirects_with_auth_headers() { .await; let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) .expect("client"); - let executor_public_key = SecureChannelIdentity::generate() - .expect("identity") - .public_key(); - let error = client - .register_environment("environment-requested", &executor_public_key) + .register_legacy_environment("environment-requested") .await .expect_err("redirect response should not be followed"); @@ -128,6 +159,7 @@ fn debug_output_redacts_auth_provider() { assert!(debug.contains("")); assert!(!debug.contains("workspace-123")); + assert!(debug.contains("Legacy")); } #[test] @@ -142,6 +174,7 @@ fn remote_environment_config_accepts_cloud_environment_id() { .expect("config"); assert_eq!(config.environment_id, environment_id); + assert_eq!(config.relay_security, RemoteRelaySecurity::Legacy); } #[test] diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index fb04dd9de3..1e9107b40a 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -24,6 +24,7 @@ use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::FsReadFileParams; use codex_exec_server::ProcessId; use codex_exec_server::RemoteEnvironmentConfig; +use codex_exec_server::RemoteRelaySecurity; use codex_exec_server::SecureChannelIdentity; use codex_exec_server::SecureChannelPublicKey; use codex_exec_server::SecureRendezvousConnectArgs; @@ -38,6 +39,7 @@ use http::HeaderMap; use http::HeaderValue; use pretty_assertions::assert_eq; use prost::Message as ProstMessage; +use relay_proto::RelayData; use relay_proto::RelayMessageFrame; use relay_proto::relay_message_frame; use tempfile::TempDir; @@ -124,6 +126,84 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { Arc::new(StaticRegistryAuthProvider) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_environment_uses_legacy_relay_by_default() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/register" + ))) + .and(header("authorization", format!("Bearer {REGISTRY_TOKEN}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": ENVIRONMENT_ID, + "url": rendezvous_url, + }))) + .mount(®istry) + .await; + + let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; + let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; + let config = RemoteEnvironmentConfig::new( + registry.uri(), + ENVIRONMENT_ID.to_string(), + static_registry_auth_provider(), + )?; + assert_eq!(config.relay_security, RemoteRelaySecurity::Legacy); + let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment( + config, + runtime_paths, + )); + + let mut websocket = accept_websocket(&listener, "legacy environment").await?; + let initialize = serde_json::json!({ + "id": 1, + "method": "initialize", + "params": { + "clientName": "legacy-default-test", + "resumeSessionId": null, + }, + }); + websocket + .send(Message::Binary( + RelayMessageFrame { + version: 1, + stream_id: "legacy-stream".to_string(), + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Data(RelayData { + seq: 0, + segment_index: 0, + segment_count: 1, + payload: serde_json::to_vec(&initialize)?, + })), + } + .encode_to_vec() + .into(), + )) + .await?; + + let response = timeout(TEST_TIMEOUT, websocket.next()) + .await + .context("legacy relay should return initialize response")? + .context("legacy relay websocket should remain open")??; + let Message::Binary(response) = response else { + anyhow::bail!("legacy relay should return a binary protobuf frame"); + }; + let response = RelayMessageFrame::decode(response.as_ref())?; + let Some(relay_message_frame::Body::Data(response)) = response.body else { + anyhow::bail!("legacy relay should return a data frame"); + }; + let response: serde_json::Value = serde_json::from_slice(&response.payload)?; + assert_eq!(response["id"], 1); + assert!(response["result"].is_object()); + + remote_environment.abort(); + let _ = remote_environment.await; + Ok(()) +} + #[tokio::test] async fn secure_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> { let attempts = Arc::new(AtomicUsize::new(0)); @@ -228,11 +308,12 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; - let config = RemoteEnvironmentConfig::new( + let mut config = RemoteEnvironmentConfig::new( registry.uri(), ENVIRONMENT_ID.to_string(), static_registry_auth_provider(), )?; + config.relay_security = RemoteRelaySecurity::Noise; let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment( config, runtime_paths,