Add direct SigV4 transport to exec-server (#42781)

## Why

Allow remote exec servers to connect directly to AWS-hosted registries that
authenticate registry requests and WebSocket handshakes with AWS SigV4.

## What changed

- Add `--remote-transport direct` with SigV4 profile, region, and service
  options while keeping Noise as the default transport.
- Register the `direct_jsonrpc_v1` transport and carry plain exec-server
  JSON-RPC messages over the authenticated WebSocket.
- Reuse direct registrations across transient disconnects, refresh them after
  a `409 Conflict`, and require TLS for non-loopback endpoints.

## Testing

- Cover CLI validation and SigV4 request signing.
- Exercise direct registration, handshake retry behavior, JSON-RPC
  interoperability, and process recovery after reconnecting.

GitOrigin-RevId: 0755df330ba3abe5db0a516fdaa49338d9bbe2d2
This commit is contained in:
ostepanian
2026-09-04 14:49:46 +00:00
committed by copyberry
parent c9fac4dd5a
commit b3f5e45cc1
15 changed files with 1988 additions and 121 deletions
+99 -8
View File
@@ -50,11 +50,25 @@ use crate::server::RequestDispatchMode;
use crate::trace_context::current_rendezvous_headers;
use crate::trace_context::current_trace_context_headers;
#[path = "remote/direct.rs"]
mod direct;
use direct::run_direct_environment;
const ERROR_BODY_PREVIEW_BYTES: usize = 4096;
const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1";
mod registration_retry;
/// Wire transport used after registering a remote exec-server.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
/// The transport used to connect a remote exec-server environment.
pub enum RemoteEnvironmentTransport {
#[default]
Noise,
Direct,
}
#[derive(Clone)]
struct EnvironmentRegistryClient {
base_url: String,
@@ -515,6 +529,7 @@ pub struct RemoteEnvironmentConfig {
pub environment_id: String,
pub name: String,
pub request_dispatch_mode: RequestDispatchMode,
transport: RemoteEnvironmentTransport,
auth_provider: SharedAuthProvider,
telemetry: ExecServerTelemetry,
http_client_factory: HttpClientFactory,
@@ -527,17 +542,36 @@ impl std::fmt::Debug for RemoteEnvironmentConfig {
.field("environment_id", &self.environment_id)
.field("name", &self.name)
.field("request_dispatch_mode", &self.request_dispatch_mode)
.field("transport", &self.transport)
.field("auth_provider", &"<redacted>")
.finish()
}
}
impl RemoteEnvironmentConfig {
/// Creates a remote environment configuration using the default Noise transport.
pub fn new(
base_url: String,
environment_id: String,
auth_provider: SharedAuthProvider,
http_client_factory: HttpClientFactory,
) -> Result<Self, ExecServerError> {
Self::new_with_transport(
base_url,
environment_id,
RemoteEnvironmentTransport::Noise,
auth_provider,
http_client_factory,
)
}
/// Creates a remote environment configuration using an explicit transport.
pub fn new_with_transport(
base_url: String,
environment_id: String,
transport: RemoteEnvironmentTransport,
auth_provider: SharedAuthProvider,
http_client_factory: HttpClientFactory,
) -> Result<Self, ExecServerError> {
let environment_id = normalize_environment_id(environment_id)?;
Ok(Self {
@@ -545,6 +579,7 @@ impl RemoteEnvironmentConfig {
environment_id,
name: "codex-exec-server".to_string(),
request_dispatch_mode: RequestDispatchMode::Inline,
transport,
auth_provider,
telemetry: ExecServerTelemetry::default(),
http_client_factory,
@@ -557,12 +592,15 @@ impl RemoteEnvironmentConfig {
}
}
/// Register an exec-server for remote use and serve requests over Noise.
/// Register an exec-server for remote use and serve requests over its configured transport.
///
/// The executor identity is generated once per process and reused across
/// In Noise mode, the executor identity is generated once per process and reused across
/// reconnects. The registration and rendezvous URL are also reused until
/// rendezvous rejects the URL, at which point the next attempt registers again.
/// The websocket carries cleartext routing metadata and encrypted payloads.
///
/// Direct mode reuses its registration across reconnects. A WebSocket handshake
/// conflict refreshes the registration; other permanent client errors stop the runner.
pub async fn run_remote_environment(
config: RemoteEnvironmentConfig,
runtime_paths: ExecServerRuntimePaths,
@@ -588,7 +626,20 @@ where
config.request_dispatch_mode,
);
let result = run_remote_transport(config, processor.clone(), shutdown).await;
let result = match config.transport {
RemoteEnvironmentTransport::Noise => {
run_remote_transport(config, shutdown, |config, client| {
run_remote_environment_connections(config, client, processor.clone())
})
.await
}
RemoteEnvironmentTransport::Direct => {
run_remote_transport(config, shutdown, |config, client| {
run_direct_environment(config, client, processor.clone())
})
.await
}
};
processor.shutdown().await;
result
}
@@ -602,22 +653,35 @@ pub async fn run_remote_environment_forward_until_shutdown<F>(
where
F: std::future::Future<Output = ()>,
{
// Forwarder implements the Noise stream bridge. Direct forwarding needs a
// separate Direct-compatible bridge, so reject it rather than use the
// Noise-specific path.
// Remove this guard when a Direct-compatible forwarder is added.
if config.transport == RemoteEnvironmentTransport::Direct {
return Err(ExecServerError::EnvironmentRegistryConfig(
"direct exec-server transport does not support forwarding".to_string(),
));
}
let forwarder = Forwarder::new(
websocket_url,
&config.http_client_factory,
config.telemetry.clone(),
)?;
run_remote_transport(config, forwarder, shutdown).await
run_remote_transport(config, shutdown, |config, client| {
run_remote_environment_connections(config, client, forwarder)
})
.await
}
async fn run_remote_transport<F, H>(
async fn run_remote_transport<F, R, T>(
config: RemoteEnvironmentConfig,
handler: H,
shutdown: F,
run_loop: R,
) -> Result<(), ExecServerError>
where
F: std::future::Future<Output = ()>,
H: NoiseStreamHandler,
R: FnOnce(RemoteEnvironmentConfig, EnvironmentRegistryClient) -> T,
T: std::future::Future<Output = Result<(), ExecServerError>>,
{
ensure_rustls_crypto_provider();
let client = EnvironmentRegistryClient::new_with_telemetry(
@@ -626,7 +690,7 @@ where
config.telemetry.clone(),
config.http_client_factory.clone(),
)?;
let run = run_remote_environment_connections(config, client, handler);
let run = run_loop(config, client);
tokio::pin!(run, shutdown);
tokio::select! {
result = &mut run => result,
@@ -1205,6 +1269,33 @@ mod tests {
);
}
#[test]
fn remote_environment_config_new_defaults_to_noise_transport() {
let config = RemoteEnvironmentConfig::new(
"https://registry.example".to_string(),
"env-1".to_string(),
static_registry_auth_provider(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
.expect("config");
assert_eq!(config.transport, RemoteEnvironmentTransport::Noise);
}
#[test]
fn remote_environment_config_new_with_transport_preserves_direct_transport() {
let config = RemoteEnvironmentConfig::new_with_transport(
"https://registry.example".to_string(),
"env-1".to_string(),
RemoteEnvironmentTransport::Direct,
static_registry_auth_provider(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
.expect("config");
assert_eq!(config.transport, RemoteEnvironmentTransport::Direct);
}
#[test]
fn debug_output_redacts_auth_provider() {
let config = RemoteEnvironmentConfig::new(