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
+1
View File
@@ -34,6 +34,7 @@ codex-app-server-test-client = { workspace = true }
codex-arg0 = { workspace = true }
codex-build-info = { workspace = true }
codex-api = { workspace = true }
codex-aws-auth = { workspace = true }
codex-chatgpt = { workspace = true }
codex-cloud-config = { workspace = true }
codex-cloud-tasks = { path = "../cloud-tasks" }
+265
View File
@@ -0,0 +1,265 @@
use clap::error::ErrorKind;
use pretty_assertions::assert_eq;
use super::*;
fn exec_server_from_args(args: &[&str]) -> ExecServerCommand {
let cli = MultitoolCli::try_parse_from(
["codex", "exec-server"]
.into_iter()
.chain(args.iter().copied()),
)
.expect("parse executor arguments");
let Some(Subcommand::ExecServer(command)) = cli.subcommand else {
panic!("expected executor command");
};
command
}
#[test]
fn exec_server_help_documents_remote_options() {
let command = MultitoolCli::command()
.term_width(80)
.mut_subcommand("exec-server", |command| {
command.mut_arg("exit_on_stdin_close", |arg| arg.hide_env_values(true))
});
let help = command
.try_get_matches_from(["codex", "exec-server", "--help"])
.expect_err("help should exit before running the executor");
assert_eq!(help.kind(), ErrorKind::DisplayHelp);
let help_text = help
.to_string()
.lines()
.map(str::trim_end)
.collect::<Vec<_>>()
.join("\n");
insta::assert_snapshot!(help_text);
}
#[test]
fn exec_server_defaults_preserve_local_and_noise_modes() {
for args in [
vec![],
vec!["--listen", "stdio"],
vec![
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
],
vec![
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
"--remote-transport",
"noise",
"--use-agent-identity-auth",
],
] {
let command = exec_server_from_args(&args);
assert_eq!(command.remote_transport, ExecServerRemoteTransport::Noise);
assert!(!command.aws_sigv4);
command
.validate_remote_transport()
.expect("valid existing mode");
}
}
#[test]
fn exec_server_parses_direct_aws_options() {
for (options, expected) in [
(vec![], (None, None, "execute-api")),
(
vec![
"--aws-profile",
"development",
"--aws-region",
"us-west-2",
"--aws-service",
"bedrock-mantle",
],
(Some("development"), Some("us-west-2"), "bedrock-mantle"),
),
] {
let mut args = vec![
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
"--remote-transport",
"direct",
"--aws-sigv4",
];
args.extend(options);
let command = exec_server_from_args(&args);
assert_eq!(
command.remote.as_deref(),
Some("https://registry.example.com")
);
assert_eq!(command.environment_id.as_deref(), Some("env-1"));
assert_eq!(command.remote_transport, ExecServerRemoteTransport::Direct);
assert!(command.aws_sigv4);
assert_eq!(
(
command.aws_profile.as_deref(),
command.aws_region.as_deref(),
command.aws_service.as_str()
),
expected,
);
command
.validate_remote_transport()
.expect("valid Direct mode");
}
}
#[test]
fn exec_server_rejects_missing_or_conflicting_remote_options() {
for (options, expected) in [
(
vec!["--remote-transport", "direct"],
ErrorKind::MissingRequiredArgument,
),
(
vec!["--aws-profile", "development"],
ErrorKind::MissingRequiredArgument,
),
(
vec!["--aws-region", "us-west-2"],
ErrorKind::MissingRequiredArgument,
),
(
vec!["--aws-service", "execute-api"],
ErrorKind::MissingRequiredArgument,
),
(
vec![
"--remote-transport",
"direct",
"--aws-sigv4",
"--use-agent-identity-auth",
],
ErrorKind::ArgumentConflict,
),
(
vec!["--remote-transport", "unsupported"],
ErrorKind::InvalidValue,
),
] {
let error = MultitoolCli::try_parse_from(
[
"codex",
"exec-server",
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
]
.into_iter()
.chain(options),
)
.expect_err("reject invalid remote arguments");
assert_eq!(error.kind(), expected);
}
}
#[test]
fn exec_server_transport_and_aws_options_require_registration_arguments() {
for options in [
vec!["--remote-transport", "noise"],
vec!["--remote-transport", "direct", "--aws-sigv4"],
vec!["--aws-sigv4"],
vec![
"--remote",
"https://registry.example.com",
"--remote-transport",
"direct",
"--aws-sigv4",
],
] {
let error =
MultitoolCli::try_parse_from(["codex", "exec-server"].into_iter().chain(options))
.expect_err("require remote URL and environment ID");
assert_eq!(error.kind(), ErrorKind::MissingRequiredArgument);
}
}
#[tokio::test]
async fn exec_server_sigv4_does_not_enable_aws_auth_for_noise() {
for options in [vec![], vec!["--remote-transport", "noise"]] {
let mut args = vec![
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
"--aws-sigv4",
];
args.extend(options);
let command = exec_server_from_args(&args);
// Unset runtime paths prove validation runs before startup or credential loading.
let error = run_exec_server_command(
command,
&Arg0DispatchPaths::default(),
&CliConfigOverrides::default(),
/*strict_config*/ false,
)
.await
.expect_err("Noise auth is unchanged");
assert_eq!(
error.to_string(),
"--aws-sigv4 requires --remote-transport direct"
);
}
}
#[tokio::test]
async fn exec_server_direct_forwarding_remains_rejected() {
let command = exec_server_from_args(&[
"forward",
"--connect",
"ws://127.0.0.1:8765",
"--remote",
"https://registry.example.com",
"--environment-id",
"env-1",
"--remote-transport",
"direct",
"--aws-sigv4",
"--aws-profile",
"development",
"--aws-region",
"us-west-2",
"--aws-service",
"bedrock-mantle",
]);
assert_eq!(
(
command.remote_transport,
command.aws_sigv4,
command.aws_profile.as_deref(),
command.aws_region.as_deref(),
command.aws_service.as_str()
),
(
ExecServerRemoteTransport::Direct,
true,
Some("development"),
Some("us-west-2"),
"bedrock-mantle"
),
);
let error = run_exec_server_command(
command,
&Arg0DispatchPaths::default(),
&CliConfigOverrides::default(),
/*strict_config*/ false,
)
.await
.expect_err("Direct forwarding is unsupported before startup");
assert_eq!(
error.to_string(),
"direct exec-server transport does not support forwarding"
);
}
+77
View File
@@ -0,0 +1,77 @@
//! CLI-only glue between Codex authentication and generic AWS signing.
//!
//! Keeping this adapter here leaves `codex-api` and `codex-aws-auth` independent.
use std::sync::Arc;
use codex_api::AuthError;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_aws_auth::AwsAuthConfig;
use codex_aws_auth::AwsAuthContext;
use codex_aws_auth::AwsAuthError;
use codex_aws_auth::AwsRequestToSign;
use codex_http_client::Request;
use codex_http_client::RequestBody;
use codex_http_client::RequestCompression;
use http::HeaderMap;
/// Creates a SigV4 provider, preferring an explicit profile over the default credential chain.
pub(super) async fn aws_sigv4_auth_provider(
mut config: AwsAuthConfig,
) -> Result<SharedAuthProvider, AwsAuthError> {
config.profile = config
.profile
.map(|profile| profile.trim().to_string())
.filter(|profile| !profile.is_empty());
config.region = config
.region
.map(|region| region.trim().to_string())
.filter(|region| !region.is_empty());
let context = if config.profile.is_some() {
AwsAuthContext::load_profile(config).await
} else {
AwsAuthContext::load(config).await
}?;
Ok(Arc::new(AwsSigV4AuthProvider { context }))
}
#[derive(Debug)]
struct AwsSigV4AuthProvider {
context: AwsAuthContext,
}
impl AuthProvider for AwsSigV4AuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
fn apply_auth(&self, mut request: Request) -> codex_api::AuthProviderFuture<'_> {
Box::pin(async move {
let prepared = request.prepare_body_for_send().map_err(AuthError::Build)?;
let signed = self
.context
.sign(AwsRequestToSign {
method: request.method.clone(),
url: request.url.clone(),
headers: prepared.headers.clone(),
body: prepared.body_bytes(),
})
.await
.map_err(|error| {
if error.is_retryable() {
AuthError::Transient(error.to_string())
} else {
AuthError::Build(error.to_string())
}
})?;
request.url = signed.url;
request.headers = signed.headers;
request.body = prepared.body.map(RequestBody::Raw);
request.compression = RequestCompression::None;
Ok(request)
})
}
}
#[cfg(test)]
#[path = "exec_server_auth_tests.rs"]
mod tests;
@@ -0,0 +1,96 @@
use std::time::Duration;
use codex_aws_auth::AwsAccessKeys;
use http::HeaderValue;
use http::Method;
use pretty_assertions::assert_eq;
use super::*;
async fn test_provider() -> AwsSigV4AuthProvider {
let context = AwsAuthContext::load_with_access_keys(
AwsAuthConfig {
profile: None,
region: Some("us-east-1".to_string()),
service: "execute-api".to_string(),
},
AwsAccessKeys {
access_key_id: "test-access-key".to_string(),
secret_access_key: "test-secret-key".to_string(),
session_token: Some("test-session-token".to_string()),
},
)
.await
.expect("load fixture signing context");
AwsSigV4AuthProvider { context }
}
#[tokio::test]
async fn signs_requests_without_changing_payload_or_metadata() {
let provider = test_provider().await;
let url = "https://executor.example.com/connect?environment_id=environment-1";
for mut request in [
Request::new(Method::GET, url.to_string()),
Request::new(Method::POST, url.to_string())
.with_json(&serde_json::json!({"transport": "direct_jsonrpc_v1"})),
Request::new(Method::POST, url.to_string())
.with_json(&serde_json::json!({"transport": "direct_jsonrpc_v1"}))
.with_compression(RequestCompression::Zstd),
] {
request
.headers
.insert("session_id", HeaderValue::from_static("session-1"));
request
.headers
.insert("x-custom-header", HeaderValue::from_static("preserved"));
request.timeout = Some(Duration::from_secs(3));
let method = request.method.clone();
let expected = request
.prepare_body_for_send()
.expect("prepare fixture body");
let signed = provider
.apply_auth(request)
.await
.expect("sign fixture request");
assert_eq!(signed.method, method);
assert_eq!(signed.url, url);
assert_eq!(signed.timeout, Some(Duration::from_secs(3)));
assert_eq!(signed.body, expected.body.clone().map(RequestBody::Raw));
assert_eq!(signed.compression, RequestCompression::None);
for (name, value) in &expected.headers {
assert_eq!(signed.headers.get(name), Some(value));
}
assert_eq!(signed.headers["x-amz-security-token"], "test-session-token");
assert!(signed.headers.contains_key("x-amz-date"));
let authorization = signed.headers[http::header::AUTHORIZATION]
.to_str()
.unwrap();
assert!(authorization.starts_with("AWS4-HMAC-SHA256 "));
assert!(authorization.contains("/us-east-1/execute-api/aws4_request"));
assert_eq!(signed.prepare_body_for_send().unwrap().body, expected.body);
}
}
#[tokio::test]
async fn invalid_signing_request_is_a_permanent_auth_error() {
let provider = test_provider().await;
let error = provider
.apply_auth(Request::new(Method::GET, "not a URL".to_string()))
.await
.expect_err("invalid URL should fail signing");
assert!(matches!(error, AuthError::Build(_)));
}
#[tokio::test]
async fn invalid_signing_configuration_is_rejected() {
let error = aws_sigv4_auth_provider(AwsAuthConfig {
profile: None,
region: Some("us-east-1".to_string()),
service: " ".to_string(),
})
.await
.err()
.expect("empty service should fail configuration");
assert!(matches!(error, AwsAuthError::EmptyService));
}
+141 -69
View File
@@ -24,8 +24,6 @@ use codex_exec::Cli as ExecCli;
use codex_exec::Command as ExecCommand;
use codex_exec::ReviewArgs;
use codex_exec_server::ExecServerRuntimePaths;
use codex_exec_server::ExecServerTelemetry;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_execpolicy::ExecPolicyCheckCommand;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
@@ -54,6 +52,10 @@ mod cloud_config;
#[cfg(any(target_os = "macos", target_os = "windows"))]
mod desktop_app;
mod doctor;
#[cfg(test)]
#[path = "exec_server_args_tests.rs"]
mod exec_server_args_tests;
mod exec_server_auth;
mod exec_server_telemetry;
mod marketplace_cmd;
mod mcp_cmd;
@@ -639,6 +641,17 @@ struct ExecServerCommand {
)]
remote: Option<String>,
/// Transport used for the remote executor connection.
#[arg(
long = "remote-transport",
value_enum,
default_value_t = ExecServerRemoteTransport::Noise,
requires = "exec_server_remote",
requires_if("direct", "aws_sigv4"),
global = true
)]
remote_transport: ExecServerRemoteTransport,
/// Environment id to attach to when registering remotely.
#[arg(long = "environment-id", value_name = "ID", global = true)]
environment_id: Option<String>,
@@ -651,10 +664,43 @@ struct ExecServerCommand {
#[arg(
long = "use-agent-identity-auth",
requires = "exec_server_remote",
conflicts_with = "aws_sigv4",
global = true
)]
use_agent_identity_auth: bool,
/// Sign Direct registration and WebSocket handshake requests with AWS SigV4.
#[arg(long = "aws-sigv4", requires = "exec_server_remote", global = true)]
aws_sigv4: bool,
/// AWS profile used for SigV4 authentication.
#[arg(
long = "aws-profile",
value_name = "PROFILE",
requires = "aws_sigv4",
global = true
)]
aws_profile: Option<String>,
/// AWS signing region. Uses the SDK region chain when omitted.
#[arg(
long = "aws-region",
value_name = "REGION",
requires = "aws_sigv4",
global = true
)]
aws_region: Option<String>,
/// AWS signing service.
#[arg(
long = "aws-service",
value_name = "SERVICE",
default_value = "execute-api",
requires = "aws_sigv4",
global = true
)]
aws_service: String,
/// Exit when the parent-owned standard-input pipe closes.
#[arg(
long = "exit-on-stdin-close",
@@ -665,6 +711,37 @@ struct ExecServerCommand {
exit_on_stdin_close: bool,
}
impl ExecServerCommand {
fn validate_remote_transport(&self) -> anyhow::Result<()> {
match (self.remote_transport, self.aws_sigv4) {
(ExecServerRemoteTransport::Noise, true) => {
anyhow::bail!("--aws-sigv4 requires --remote-transport direct");
}
(ExecServerRemoteTransport::Direct, false) => {
anyhow::bail!("--remote-transport direct requires --aws-sigv4");
}
(ExecServerRemoteTransport::Noise, false)
| (ExecServerRemoteTransport::Direct, true) => {}
}
if self.remote_transport == ExecServerRemoteTransport::Direct
&& matches!(
self.command.as_ref(),
Some(ExecServerSubcommand::Forward { .. })
)
{
anyhow::bail!("direct exec-server transport does not support forwarding");
}
Ok(())
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, clap::ValueEnum)]
enum ExecServerRemoteTransport {
#[default]
Noise,
Direct,
}
#[derive(Debug, clap::Subcommand)]
enum ExecServerSubcommand {
/// Register an existing WebSocket exec-server as a remote environment.
@@ -1863,6 +1940,7 @@ async fn run_exec_server_command(
root_config_overrides: &CliConfigOverrides,
strict_config: bool,
) -> anyhow::Result<()> {
cmd.validate_remote_transport()?;
let codex_self_exe = arg0_paths
.codex_self_exe
.clone()
@@ -1880,7 +1958,41 @@ async fn run_exec_server_command(
/*enable_workload_identity*/ true,
)
.await?;
let direct_transport = cmd.remote_transport == ExecServerRemoteTransport::Direct;
let (_otel, telemetry) = exec_server_telemetry::init(Some(&config));
let auth_provider = if cmd.aws_sigv4 {
exec_server_auth::aws_sigv4_auth_provider(codex_aws_auth::AwsAuthConfig {
profile: cmd.aws_profile,
region: cmd.aws_region,
service: cmd.aws_service,
})
.await?
} else {
load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth)
.await?
};
let mut remote_config = codex_exec_server::RemoteEnvironmentConfig::new_with_transport(
base_url,
environment_id,
if direct_transport {
codex_exec_server::RemoteEnvironmentTransport::Direct
} else {
codex_exec_server::RemoteEnvironmentTransport::Noise
},
auth_provider,
config.http_client_factory(),
)?;
if let Some(name) = cmd.name {
remote_config.name = name;
}
remote_config.request_dispatch_mode = cmd.request_dispatch_mode;
let remote_config = remote_config.with_telemetry(telemetry);
let parent_lifetime = if cmd.exit_on_stdin_close {
exec_server_telemetry::ParentLifetime::StdinPipe
} else {
exec_server_telemetry::ParentLifetime::Independent
};
let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel();
#[cfg(target_os = "macos")]
let runtime_paths = runtime_paths.with_allowed_symlinked_codex_home(
codex_config::allowed_symlinked_codex_home(
@@ -1888,13 +2000,33 @@ async fn run_exec_server_command(
&config.codex_home,
),
);
run_remote_exec_server(
cmd,
base_url,
environment_id,
&config,
runtime_paths,
telemetry,
exec_server_telemetry::run_until_shutdown(
async move {
let shutdown = async move {
let _ = shutdown_receiver.await;
};
match cmd.command {
Some(ExecServerSubcommand::Forward { connect }) => {
codex_exec_server::run_remote_environment_forward_until_shutdown(
remote_config,
connect,
shutdown,
)
.await
}
None => {
codex_exec_server::run_remote_environment_until_shutdown(
remote_config,
runtime_paths,
shutdown,
)
.await
}
}
.map_err(anyhow::Error::new)
},
parent_lifetime,
exec_server_telemetry::ShutdownBehavior::Graceful(shutdown_sender),
)
.await
} else {
@@ -1940,66 +2072,6 @@ async fn run_exec_server_command(
}
}
async fn run_remote_exec_server(
cmd: ExecServerCommand,
base_url: String,
environment_id: String,
config: &Config,
runtime_paths: ExecServerRuntimePaths,
telemetry: ExecServerTelemetry,
) -> anyhow::Result<()> {
let auth_provider =
load_exec_server_remote_auth_provider(config, &base_url, cmd.use_agent_identity_auth)
.await?;
let mut remote_config = RemoteEnvironmentConfig::new(
base_url,
environment_id,
auth_provider,
config.http_client_factory(),
)?;
if let Some(name) = cmd.name {
remote_config.name = name;
}
remote_config.request_dispatch_mode = cmd.request_dispatch_mode;
let remote_config = remote_config.with_telemetry(telemetry);
let parent_lifetime = if cmd.exit_on_stdin_close {
exec_server_telemetry::ParentLifetime::StdinPipe
} else {
exec_server_telemetry::ParentLifetime::Independent
};
let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel();
let shutdown = async move {
let _ = shutdown_receiver.await;
};
let run = async move {
match cmd.command {
Some(ExecServerSubcommand::Forward { connect }) => {
codex_exec_server::run_remote_environment_forward_until_shutdown(
remote_config,
connect,
shutdown,
)
.await
}
None => {
codex_exec_server::run_remote_environment_until_shutdown(
remote_config,
runtime_paths,
shutdown,
)
.await
}
}
};
exec_server_telemetry::run_until_shutdown(
run,
parent_lifetime,
exec_server_telemetry::ShutdownBehavior::Graceful(shutdown_sender),
)
.await?;
Ok(())
}
async fn load_exec_server_remote_auth_provider(
config: &codex_core::config::Config,
base_url: &str,
@@ -0,0 +1,85 @@
---
source: cli/src/exec_server_args_tests.rs
expression: help_text
---
[EXPERIMENTAL] Run the standalone exec-server service
Usage: codex exec-server [OPTIONS] [COMMAND]
Commands:
forward Register an existing WebSocket exec-server as a remote environment
help Print this message or the help of the given subcommand(s)
Options:
-c, --config <key=value>
Override a configuration value that would otherwise be loaded from
`~/.codex/config.toml`. Use a dotted path (`foo.bar.baz`) to override
nested values. The `value` portion is parsed as TOML. If it fails to
parse as TOML, the raw string is used as a literal.
Examples: - `-c model="o3"` - `-c
'sandbox_permissions=["disk-full-read-access"]'` - `-c
shell_environment_policy.inherit=all`
--enable <FEATURE>
Enable a feature (repeatable). Equivalent to `-c features.<name>=true`
--strict-config
Error out when config.toml contains fields that are not recognized by
this version of Codex
--concurrent-requests <COUNT>
Maximum number of requests to process concurrently on each connection
[default: 1]
--disable <FEATURE>
Disable a feature (repeatable). Equivalent to `-c
features.<name>=false`
--listen <URL>
Transport endpoint URL. Supported values: `ws://IP:PORT` (default),
`stdio`, `stdio://`
--remote <URL>
Register this exec-server as a remote environment using the given base
URL
--remote-transport <REMOTE_TRANSPORT>
Transport used for the remote executor connection
[default: noise]
[possible values: noise, direct]
--environment-id <ID>
Environment id to attach to when registering remotely
--name <NAME>
Human-readable environment name
--use-agent-identity-auth
Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote
registration
--aws-sigv4
Sign Direct registration and WebSocket handshake requests with AWS
SigV4
--aws-profile <PROFILE>
AWS profile used for SigV4 authentication
--aws-region <REGION>
AWS signing region. Uses the SDK region chain when omitted
--aws-service <SERVICE>
AWS signing service
[default: execute-api]
--exit-on-stdin-close
Exit when the parent-owned standard-input pipe closes
[env: CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE]
-h, --help
Print help (see a summary with '-h')