openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: make Noise environment runtime landable

Browse Source 
Co-authored-by: Codex noreply@openai.com
This commit is contained in:
viyatb-oai
2026-06-09 16:16:47 -07:00
co-authored by Codex noreply@openai.com
parent 3a92723378
commit 2df71e6837
4 changed files with 273 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/lib.rs
+1
View File
@@ -115,6 +115,7 @@ pub use protocol::WriteParams;
pub use protocol::WriteResponse;
pub use protocol::WriteStatus;
pub use remote::RemoteEnvironmentConfig;
pub use remote::RemoteRelayProtocol;
pub use remote::run_remote_environment;
pub use runtime_paths::ExecServerRuntimePaths;
pub use server::DEFAULT_LISTEN_URL;
codex-rs/exec-server/src/noise_relay/environment.rs
+3 -5
View File
@@ -387,7 +387,9 @@ impl NoiseVirtualStream {
.transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.decrypt(&ciphertext)?
transport.decrypt(&ciphertext).map_err(|error| {
ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))
})?
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
@@ -493,7 +495,3 @@ async fn send_reset(physical_outgoing_tx: &mpsc::Sender<Vec<u8>>, stream_id: Str
.send(encode_relay_message_frame(&reset))
.await;
}
#[cfg(test)]
#[path = "environment_tests.rs"]
mod tests;
codex-rs/exec-server/src/noise_relay/environment_tests.rs
-163
View File
@@ -1,163 +0,0 @@
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use anyhow::Result;
use futures::SinkExt;
use futures::StreamExt;
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::time::timeout;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;
use super::HarnessKeyValidator;
use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES;
use super::run_noise_multiplexed_environment;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::noise_channel::InitiatorHandshake;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_channel::NoiseChannelPublicKey;
use crate::noise_channel::noise_channel_prologue;
use crate::relay::RelayFrameBodyKind;
use crate::relay::decode_relay_message_frame;
use crate::relay::encode_relay_message_frame;
use crate::relay_proto::RelayMessageFrame;
use crate::server::ConnectionProcessor;
const ENVIRONMENT_ID: &str = "environment-1";
const EXECUTOR_REGISTRATION_ID: &str = "registration-1";
#[derive(Clone)]
struct BlockingValidator {
calls: Arc<AtomicUsize>,
release: Arc<Notify>,
}
impl HarnessKeyValidator for BlockingValidator {
fn validate_harness_key(
&self,
_harness_public_key: &NoiseChannelPublicKey,
_authorization: &str,
) -> impl std::future::Future<Output = Result<(), ExecServerError>> + Send {
let calls = Arc::clone(&self.calls);
let release = Arc::clone(&self.release);
async move {
calls.fetch_add(1, Ordering::SeqCst);
release.notified().await;
Ok(())
}
}
}
#[tokio::test]
async fn pending_harness_key_validation_does_not_block_new_handshakes() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let harness_connection = tokio::spawn(connect_async(websocket_url));
let (socket, _peer_addr) = listener.accept().await?;
let environment_websocket = accept_async(socket).await?;
let (mut harness_websocket, _response) = harness_connection.await??;
let environment_identity = NoiseChannelIdentity::generate()?;
let harness_identity = NoiseChannelIdentity::generate()?;
let calls = Arc::new(AtomicUsize::new(0));
let environment_task = tokio::spawn(run_noise_multiplexed_environment(
environment_websocket,
ConnectionProcessor::new(ExecServerRuntimePaths::new(
std::env::current_exe()?,
/*codex_linux_sandbox_exe*/ None,
)?),
ENVIRONMENT_ID.to_string(),
EXECUTOR_REGISTRATION_ID.to_string(),
environment_identity.clone(),
BlockingValidator {
calls: Arc::clone(&calls),
release: Arc::new(Notify::new()),
},
));
for stream_id in ["stream-1", "stream-2"] {
let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id)?;
let (_handshake, request) = InitiatorHandshake::start(
&harness_identity,
&environment_identity.public_key(),
&prologue,
b"authorization",
)?;
let frame = RelayMessageFrame::handshake(stream_id.to_string(), request);
harness_websocket
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await?;
}
timeout(Duration::from_secs(1), async {
while calls.load(Ordering::SeqCst) != 2 {
tokio::task::yield_now().await;
}
})
.await?;
harness_websocket.close(None).await?;
timeout(Duration::from_secs(1), environment_task).await??;
Ok(())
}
#[tokio::test]
async fn oversized_harness_authorization_is_rejected_before_validation() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let harness_connection = tokio::spawn(connect_async(websocket_url));
let (socket, _peer_addr) = listener.accept().await?;
let environment_websocket = accept_async(socket).await?;
let (mut harness_websocket, _response) = harness_connection.await??;
let environment_identity = NoiseChannelIdentity::generate()?;
let harness_identity = NoiseChannelIdentity::generate()?;
let calls = Arc::new(AtomicUsize::new(0));
let environment_task = tokio::spawn(run_noise_multiplexed_environment(
environment_websocket,
ConnectionProcessor::new(ExecServerRuntimePaths::new(
std::env::current_exe()?,
/*codex_linux_sandbox_exe*/ None,
)?),
ENVIRONMENT_ID.to_string(),
EXECUTOR_REGISTRATION_ID.to_string(),
environment_identity.clone(),
BlockingValidator {
calls: Arc::clone(&calls),
release: Arc::new(Notify::new()),
},
));
let stream_id = "stream-1";
let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id)?;
let oversized_authorization = vec![b'a'; MAX_HARNESS_KEY_AUTHORIZATION_BYTES + 1];
let (_handshake, request) = InitiatorHandshake::start(
&harness_identity,
&environment_identity.public_key(),
&prologue,
&oversized_authorization,
)?;
let frame = RelayMessageFrame::handshake(stream_id.to_string(), request);
harness_websocket
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await?;
let Message::Binary(payload) = timeout(Duration::from_secs(1), harness_websocket.next())
.await?
.ok_or_else(|| anyhow::anyhow!("environment closed before sending reset"))??
else {
anyhow::bail!("expected binary reset frame");
};
let reset = decode_relay_message_frame(payload.as_ref())?;
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
assert_eq!(calls.load(Ordering::SeqCst), 0);
harness_websocket.close(None).await?;
timeout(Duration::from_secs(1), environment_task).await??;
Ok(())
}
codex-rs/exec-server/src/remote.rs
+269 -138
View File
@@ -3,18 +3,29 @@ use std::time::Duration;
use codex_api::SharedAuthProvider;
use reqwest::StatusCode;
use serde::Deserialize;
use serde::Serialize;
use tokio::time::sleep;
use tokio_tungstenite::connect_async;
use tokio::time::timeout;
use tokio_tungstenite::connect_async_with_config;
use tracing::info;
use tracing::warn;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::NoiseChannelIdentity;
use crate::NoiseChannelPublicKey;
use crate::noise_relay::HarnessKeyValidator;
use crate::noise_relay::noise_relay_websocket_config;
use crate::noise_relay::run_noise_multiplexed_environment;
use crate::relay::run_multiplexed_environment;
use crate::server::ConnectionProcessor;
const ERROR_BODY_PREVIEW_BYTES: usize = 4096;
const ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_REMOTE_ENVIRONMENT_ID_LEN: usize = 256;
const REMOTE_RENDEZVOUS_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(Clone)]
struct EnvironmentRegistryClient {
@@ -40,11 +51,16 @@ impl EnvironmentRegistryClient {
auth_provider,
http: reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT)
.build()?,
})
}
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 Noise relay support.
async fn register_legacy_environment(
&self,
environment_id: &str,
) -> Result<EnvironmentRegistryRegistrationResponse, ExecServerError> {
@@ -60,6 +76,64 @@ impl EnvironmentRegistryClient {
self.parse_json_response(response).await
}
/// Register the exec-server's static Noise identity with a Noise-aware registry.
///
/// Supplying this request body is the protocol-level opt in. Registries can
/// therefore distinguish Noise 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: &NoiseChannelPublicKey,
) -> Result<EnvironmentRegistryNoiseRegistrationResponse, ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/register"),
))
.headers(self.auth_provider.to_auth_headers())
.json(&EnvironmentRegistryRegistrationRequest {
security_profile: NOISE_RELAY_SECURITY_PROFILE,
executor_public_key,
})
.send()
.await?;
self.parse_json_response(response).await
}
async fn validate_harness_key(
&self,
environment_id: &str,
executor_registration_id: &str,
harness_public_key: &NoiseChannelPublicKey,
harness_key_authorization: &str,
) -> Result<(), ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/validate"),
))
.headers(self.auth_provider.to_auth_headers())
.json(&EnvironmentRegistryHarnessKeyValidationRequest {
executor_registration_id,
harness_public_key,
harness_key_authorization,
})
.send()
.await?;
let response = self
.parse_json_response::<EnvironmentRegistryHarnessKeyValidationResponse>(response)
.await?;
if !response.valid {
return Err(ExecServerError::Protocol(
"environment registry rejected Noise relay harness key".to_string(),
));
}
Ok(())
}
async fn parse_json_response<R>(
&self,
response: reqwest::Response,
@@ -81,18 +155,112 @@ impl EnvironmentRegistryClient {
}
}
const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1";
#[derive(Serialize)]
struct EnvironmentRegistryRegistrationRequest<'a> {
security_profile: &'static str,
executor_public_key: &'a NoiseChannelPublicKey,
}
#[derive(Debug, Clone, Eq, PartialEq, Deserialize)]
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,
}
#[derive(Serialize)]
struct EnvironmentRegistryHarnessKeyValidationRequest<'a> {
executor_registration_id: &'a str,
harness_public_key: &'a NoiseChannelPublicKey,
harness_key_authorization: &'a str,
}
#[derive(Deserialize)]
struct EnvironmentRegistryHarnessKeyValidationResponse {
valid: bool,
}
#[derive(Clone)]
struct RegistryHarnessKeyValidator {
client: EnvironmentRegistryClient,
environment_id: String,
executor_registration_id: String,
}
impl HarnessKeyValidator for RegistryHarnessKeyValidator {
async fn validate_harness_key(
&self,
harness_public_key: &NoiseChannelPublicKey,
authorization: &str,
) -> Result<(), ExecServerError> {
self.client
.validate_harness_key(
&self.environment_id,
&self.executor_registration_id,
harness_public_key,
authorization,
)
.await
}
}
/// Protocol 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 Noise.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum RemoteRelayProtocol {
#[default]
Legacy,
Noise,
}
enum RemoteRelayProtocolState {
Legacy,
Noise(NoiseChannelIdentity),
}
enum RegisteredRemoteEnvironment {
Legacy(EnvironmentRegistryRegistrationResponse),
Noise {
response: EnvironmentRegistryNoiseRegistrationResponse,
identity: NoiseChannelIdentity,
},
}
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_protocol: RemoteRelayProtocol,
auth_provider: SharedAuthProvider,
}
@@ -102,6 +270,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_protocol", &self.relay_protocol)
.field("auth_provider", &"<redacted>")
.finish()
}
@@ -118,6 +287,7 @@ impl RemoteEnvironmentConfig {
base_url,
environment_id,
name: "codex-exec-server".to_string(),
relay_protocol: RemoteRelayProtocol::Legacy,
auth_provider,
})
}
@@ -133,23 +303,88 @@ 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 protocol_state = match config.relay_protocol {
RemoteRelayProtocol::Legacy => RemoteRelayProtocolState::Legacy,
RemoteRelayProtocol::Noise => {
RemoteRelayProtocolState::Noise(NoiseChannelIdentity::generate()?)
}
};
let mut backoff = Duration::from_secs(1);
loop {
let response = client.register_environment(&config.environment_id).await?;
eprintln!(
"codex exec-server remote environment registered with environment_id {}",
response.environment_id
);
match connect_async(response.url.as_str()).await {
Ok((websocket, _)) => {
backoff = Duration::from_secs(1);
run_multiplexed_environment(websocket, processor.clone()).await;
let registration = match &protocol_state {
RemoteRelayProtocolState::Legacy => RegisteredRemoteEnvironment::Legacy(
client
.register_legacy_environment(&config.environment_id)
.await?,
),
RemoteRelayProtocolState::Noise(identity) => {
let response = client
.register_noise_environment(&config.environment_id, &identity.public_key())
.await?;
if response.security_profile != NOISE_RELAY_SECURITY_PROFILE {
return Err(ExecServerError::Protocol(format!(
"environment registry returned unsupported security profile `{}`",
response.security_profile
)));
}
RegisteredRemoteEnvironment::Noise {
response,
identity: identity.clone(),
}
}
Err(err) => {
};
if registration.environment_id() != config.environment_id {
return Err(ExecServerError::Protocol(
"environment registry returned a different environment id".to_string(),
));
}
let environment_id = registration.environment_id();
info!(
"codex exec-server remote environment registered with environment_id {environment_id}"
);
let websocket_config = match &registration {
RegisteredRemoteEnvironment::Legacy(_) => None,
RegisteredRemoteEnvironment::Noise { .. } => Some(noise_relay_websocket_config()),
};
match timeout(
REMOTE_RENDEZVOUS_CONNECT_TIMEOUT,
connect_async_with_config(
registration.websocket_url(),
websocket_config,
/*disable_nagle*/ false,
),
)
.await
{
Ok(Ok((websocket, _))) => {
backoff = Duration::from_secs(1);
match registration {
RegisteredRemoteEnvironment::Legacy(_) => {
run_multiplexed_environment(websocket, processor.clone()).await;
}
RegisteredRemoteEnvironment::Noise { response, identity } => {
run_noise_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}");
}
Err(_) => warn!("timed out connecting remote exec-server websocket"),
}
sleep(backoff).await;
@@ -158,12 +393,32 @@ pub async fn run_remote_environment(
}
fn normalize_environment_id(environment_id: String) -> Result<String, ExecServerError> {
let environment_id = environment_id.trim().to_string();
if environment_id.is_empty() {
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id is required for remote exec-server registration".to_string(),
));
}
if environment_id.trim() != environment_id {
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id must not contain surrounding whitespace".to_string(),
));
}
if environment_id.len() > MAX_REMOTE_ENVIRONMENT_ID_LEN {
return Err(ExecServerError::EnvironmentRegistryConfig(format!(
"environment id cannot be longer than {MAX_REMOTE_ENVIRONMENT_ID_LEN} characters"
)));
}
// The ID is interpolated into authenticated registry request paths below.
// Keep it to one URL path component so a caller cannot use a delimiter to
// redirect the exec-server's registration credential to another endpoint.
if !environment_id
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_')
{
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id must contain only ASCII letters, numbers, '-' or '_'".to_string(),
));
}
Ok(environment_id)
}
@@ -240,127 +495,3 @@ fn preview_error_body(body: &str) -> Option<String> {
}
Some(trimmed.chars().take(ERROR_BODY_PREVIEW_BYTES).collect())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use codex_api::AuthProvider;
use http::HeaderMap;
use http::HeaderValue;
use pretty_assertions::assert_eq;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::*;
#[derive(Debug)]
struct StaticRegistryAuthProvider;
impl AuthProvider for StaticRegistryAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
let _ = headers.insert(
http::header::AUTHORIZATION,
HeaderValue::from_static("Bearer registry-token"),
);
let _ = headers.insert(
"ChatGPT-Account-ID",
HeaderValue::from_static("workspace-123"),
);
}
}
fn static_registry_auth_provider() -> SharedAuthProvider {
Arc::new(StaticRegistryAuthProvider)
}
#[tokio::test]
async fn register_environment_posts_with_auth_provider_headers() {
let server = MockServer::start().await;
let config = RemoteEnvironmentConfig::new(
server.uri(),
"environment-requested".to_string(),
static_registry_auth_provider(),
)
.expect("config");
Mock::given(method("POST"))
.and(path("/cloud/environment/environment-requested/register"))
.and(header("authorization", "Bearer registry-token"))
.and(header("chatgpt-account-id", "workspace-123"))
.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_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_environment_does_not_follow_redirects_with_auth_headers() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/cloud/environment/environment-requested/register"))
.and(header("authorization", "Bearer registry-token"))
.respond_with(
ResponseTemplate::new(302)
.insert_header("location", format!("{}/redirect-target", server.uri())),
)
.mount(&server)
.await;
Mock::given(path("/redirect-target"))
.and(header("authorization", "Bearer registry-token"))
.respond_with(ResponseTemplate::new(200))
.expect(0)
.mount(&server)
.await;
let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider())
.expect("client");
let error = client
.register_environment("environment-requested")
.await
.expect_err("redirect response should not be followed");
assert!(matches!(
error,
ExecServerError::EnvironmentRegistryHttp {
status: StatusCode::FOUND,
..
}
));
}
#[test]
fn debug_output_redacts_auth_provider() {
let config = RemoteEnvironmentConfig::new(
"https://registry.example".to_string(),
"env-1".to_string(),
static_registry_auth_provider(),
)
.expect("config");
let debug = format!("{config:?}");
assert!(debug.contains("<redacted>"));
assert!(!debug.contains("workspace-123"));
}
}