From 51bfb5f3b115f4905af891e2c09fd1641dc9f8aa Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 12 May 2026 18:40:53 -0700 Subject: [PATCH] Restore app-server websocket listener with auth guard (#22404) ## Why PR #21843 removed the TCP websocket app-server listener, but that also removed functionality that still needs to exist. Restoring it as-is would reopen the old remote exposure problem, so this keeps the restored listener while making remote and non-loopback usage require explicit auth. ## What Changed - Mostly reverts #21843 and reapplies the small merge-conflict resolutions needed on top of current main. - Restores ws://IP:PORT parsing, the app-server TCP websocket acceptor, websocket auth CLI flags, and the associated tests. - The only intentional behavior change from the restored code is that non-loopback websocket listeners now fail startup unless --ws-auth capability-token or --ws-auth signed-bearer-token is configured. Loopback listeners remain available for local and SSH-forwarding workflows. ## Reviewer Focus Please focus review on the small auth-enforcement delta layered on top of the revert: - codex-rs/app-server-transport/src/transport/websocket.rs: start_websocket_acceptor now rejects unauthenticated non-loopback websocket binds before accepting connections. - codex-rs/app-server-transport/src/transport/auth.rs: helper logic classifies unauthenticated non-loopback listeners. - codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs: tests cover unauthenticated ws://0.0.0.0 startup rejection and authenticated non-loopback capability-token startup. Everything else is intended to be revert/merge-conflict restoration rather than new product behavior. ## Verification - Manually verified that TUI remoting is restored and that auth is enforced for non-localhost urls. --- codex-rs/Cargo.lock | 13 +- codex-rs/Cargo.toml | 1 + codex-rs/app-server-transport/Cargo.toml | 14 + codex-rs/app-server-transport/src/lib.rs | 2 + .../src/transport/auth.rs | 751 ++++++++++++++++++ .../app-server-transport/src/transport/mod.rs | 21 +- .../src/transport/remote_control/enroll.rs | 2 +- .../src/transport/remote_control/websocket.rs | 2 +- .../src/transport/unix_socket_tests.rs | 8 - .../src/transport/websocket.rs | 174 +++- codex-rs/app-server/Cargo.toml | 5 +- codex-rs/app-server/README.md | 9 + codex-rs/app-server/src/app_server_tracing.rs | 3 +- codex-rs/app-server/src/lib.rs | 27 +- codex-rs/app-server/src/main.rs | 8 +- codex-rs/app-server/src/transport.rs | 2 + .../app-server/tests/suite/v2/command_exec.rs | 8 +- .../suite/v2/connection_handling_websocket.rs | 524 ++++++++++-- .../v2/connection_handling_websocket_unix.rs | 25 +- .../tests/suite/v2/thread_name_websocket.rs | 16 +- codex-rs/cli/src/main.rs | 81 +- 21 files changed, 1578 insertions(+), 118 deletions(-) create mode 100644 codex-rs/app-server-transport/src/transport/auth.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 75c1ea772..dbd754435 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1922,7 +1922,6 @@ dependencies = [ "codex-state", "codex-thread-store", "codex-tools", - "codex-uds", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-cli", @@ -1931,13 +1930,16 @@ dependencies = [ "core_test_support", "flate2", "futures", + "hmac", "opentelemetry", "opentelemetry_sdk", "pretty_assertions", + "reqwest", "rmcp", "serde", "serde_json", "serial_test", + "sha2", "shlex", "tar", "tempfile", @@ -2054,8 +2056,11 @@ dependencies = [ name = "codex-app-server-transport" version = "0.0.0" dependencies = [ + "anyhow", + "axum", "base64 0.22.1", "chrono", + "clap", "codex-api", "codex-app-server-protocol", "codex-config", @@ -2066,12 +2071,18 @@ dependencies = [ "codex-uds", "codex-utils-absolute-path", "codex-utils-rustls-provider", + "constant_time_eq 0.3.1", "futures", "gethostname", + "hmac", + "jsonwebtoken", + "owo-colors", "pretty_assertions", "serde", "serde_json", + "sha2", "tempfile", + "time", "tokio", "tokio-tungstenite", "tokio-util", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 3b9a0f0d7..d8251598d 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -265,6 +265,7 @@ chrono = "0.4.43" clap = "4" clap_complete = "4" color-eyre = "0.6.3" +constant_time_eq = "0.3.1" crossbeam-channel = "0.5.15" crypto_box = { version = "0.9.1", features = ["seal"] } crossterm = "0.28.1" diff --git a/codex-rs/app-server-transport/Cargo.toml b/codex-rs/app-server-transport/Cargo.toml index 523c6ada5..175890962 100644 --- a/codex-rs/app-server-transport/Cargo.toml +++ b/codex-rs/app-server-transport/Cargo.toml @@ -13,7 +13,15 @@ doctest = false workspace = true [dependencies] +anyhow = { workspace = true } +axum = { workspace = true, default-features = false, features = [ + "http1", + "json", + "tokio", + "ws", +] } base64 = { workspace = true } +clap = { workspace = true, features = ["derive"] } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-core = { workspace = true } @@ -23,10 +31,16 @@ codex-state = { workspace = true } codex-uds = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-rustls-provider = { workspace = true } +constant_time_eq = { workspace = true } futures = { workspace = true } gethostname = { workspace = true } +hmac = { workspace = true } +jsonwebtoken = { workspace = true } +owo-colors = { workspace = true, features = ["supports-colors"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } +time = { workspace = true } tokio = { workspace = true, features = [ "io-std", "macros", diff --git a/codex-rs/app-server-transport/src/lib.rs b/codex-rs/app-server-transport/src/lib.rs index 5d0adb642..e3f39f60c 100644 --- a/codex-rs/app-server-transport/src/lib.rs +++ b/codex-rs/app-server-transport/src/lib.rs @@ -14,6 +14,8 @@ pub use transport::RemoteControlHandle; pub use transport::RemoteControlStartConfig; pub use transport::TransportEvent; pub use transport::app_server_control_socket_path; +pub use transport::auth; pub use transport::start_control_socket_acceptor; pub use transport::start_remote_control; pub use transport::start_stdio_connection; +pub use transport::start_websocket_acceptor; diff --git a/codex-rs/app-server-transport/src/transport/auth.rs b/codex-rs/app-server-transport/src/transport/auth.rs new file mode 100644 index 000000000..eeccf21ed --- /dev/null +++ b/codex-rs/app-server-transport/src/transport/auth.rs @@ -0,0 +1,751 @@ +use anyhow::Context; +use axum::http::HeaderMap; +use axum::http::StatusCode; +use axum::http::header::AUTHORIZATION; +use clap::Args; +use clap::ValueEnum; +use codex_utils_absolute_path::AbsolutePathBuf; +use constant_time_eq::constant_time_eq_32; +use jsonwebtoken::Algorithm; +use jsonwebtoken::DecodingKey; +use jsonwebtoken::Validation; +use jsonwebtoken::decode; +use serde::Deserialize; +use sha2::Digest; +use sha2::Sha256; +use std::io; +use std::io::ErrorKind; +use std::net::SocketAddr; +use std::path::Path; +use std::path::PathBuf; +use time::OffsetDateTime; + +const DEFAULT_MAX_CLOCK_SKEW_SECONDS: u64 = 30; +const MIN_SIGNED_BEARER_SECRET_BYTES: usize = 32; +const INVALID_AUTHORIZATION_HEADER_MESSAGE: &str = "invalid authorization header"; + +#[derive(Debug, Clone, Default, PartialEq, Eq, Args)] +pub struct AppServerWebsocketAuthArgs { + /// Websocket auth mode for non-loopback listeners. + #[arg(long = "ws-auth", value_name = "MODE", value_enum)] + pub ws_auth: Option, + + /// Absolute path to the capability-token file. + #[arg(long = "ws-token-file", value_name = "PATH")] + pub ws_token_file: Option, + + /// Hex-encoded SHA-256 digest of the capability token. + #[arg(long = "ws-token-sha256", value_name = "HEX")] + pub ws_token_sha256: Option, + + /// Absolute path to the shared secret file for signed JWT bearer tokens. + #[arg(long = "ws-shared-secret-file", value_name = "PATH")] + pub ws_shared_secret_file: Option, + + /// Expected issuer for signed JWT bearer tokens. + #[arg(long = "ws-issuer", value_name = "ISSUER")] + pub ws_issuer: Option, + + /// Expected audience for signed JWT bearer tokens. + #[arg(long = "ws-audience", value_name = "AUDIENCE")] + pub ws_audience: Option, + + /// Maximum clock skew when validating signed JWT bearer tokens. + #[arg(long = "ws-max-clock-skew-seconds", value_name = "SECONDS")] + pub ws_max_clock_skew_seconds: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum WebsocketAuthCliMode { + CapabilityToken, + SignedBearerToken, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AppServerWebsocketAuthSettings { + pub config: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AppServerWebsocketAuthConfig { + CapabilityToken { + source: AppServerWebsocketCapabilityTokenSource, + }, + SignedBearerToken { + shared_secret_file: AbsolutePathBuf, + issuer: Option, + audience: Option, + max_clock_skew_seconds: u64, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AppServerWebsocketCapabilityTokenSource { + TokenFile { token_file: AbsolutePathBuf }, + TokenSha256 { token_sha256: [u8; 32] }, +} + +#[derive(Clone, Debug, Default)] +pub struct WebsocketAuthPolicy { + pub(crate) mode: Option, +} + +#[derive(Clone, Debug)] +pub(crate) enum WebsocketAuthMode { + CapabilityToken { + token_sha256: [u8; 32], + }, + SignedBearerToken { + shared_secret: Vec, + issuer: Option, + audience: Option, + max_clock_skew_seconds: i64, + }, +} + +#[derive(Debug)] +pub(crate) struct WebsocketAuthError { + status_code: StatusCode, + message: &'static str, +} + +#[derive(Deserialize)] +struct JwtClaims { + exp: i64, + nbf: Option, + iss: Option, + aud: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum JwtAudienceClaim { + Single(String), + Multiple(Vec), +} + +impl WebsocketAuthError { + pub(crate) fn status_code(&self) -> StatusCode { + self.status_code + } + + pub(crate) fn message(&self) -> &'static str { + self.message + } +} + +impl AppServerWebsocketAuthArgs { + pub fn try_into_settings(self) -> anyhow::Result { + let normalize = |value: Option| { + value.and_then(|value| { + let trimmed = value.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + }) + }; + + let config = match self.ws_auth { + Some(WebsocketAuthCliMode::CapabilityToken) => { + if self.ws_shared_secret_file.is_some() + || self.ws_issuer.is_some() + || self.ws_audience.is_some() + || self.ws_max_clock_skew_seconds.is_some() + { + anyhow::bail!( + "`--ws-shared-secret-file`, `--ws-issuer`, `--ws-audience`, and `--ws-max-clock-skew-seconds` require `--ws-auth signed-bearer-token`" + ); + } + let source = match (self.ws_token_file, self.ws_token_sha256) { + (Some(_), Some(_)) => { + anyhow::bail!( + "`--ws-token-file` and `--ws-token-sha256` are mutually exclusive" + ); + } + (Some(token_file), None) => { + AppServerWebsocketCapabilityTokenSource::TokenFile { + token_file: absolute_path_arg("--ws-token-file", token_file)?, + } + } + (None, Some(token_sha256)) => { + AppServerWebsocketCapabilityTokenSource::TokenSha256 { + token_sha256: sha256_digest_arg("--ws-token-sha256", &token_sha256)?, + } + } + (None, None) => { + anyhow::bail!( + "`--ws-token-file` or `--ws-token-sha256` is required when `--ws-auth capability-token` is set" + ); + } + }; + Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) + } + Some(WebsocketAuthCliMode::SignedBearerToken) => { + if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() { + anyhow::bail!( + "`--ws-token-file` and `--ws-token-sha256` require `--ws-auth capability-token`, not `signed-bearer-token`" + ); + } + let shared_secret_file = self.ws_shared_secret_file.context( + "`--ws-shared-secret-file` is required when `--ws-auth signed-bearer-token` is set", + )?; + Some(AppServerWebsocketAuthConfig::SignedBearerToken { + shared_secret_file: absolute_path_arg( + "--ws-shared-secret-file", + shared_secret_file, + )?, + issuer: normalize(self.ws_issuer), + audience: normalize(self.ws_audience), + max_clock_skew_seconds: self + .ws_max_clock_skew_seconds + .unwrap_or(DEFAULT_MAX_CLOCK_SKEW_SECONDS), + }) + } + None => { + if self.ws_token_file.is_some() + || self.ws_token_sha256.is_some() + || self.ws_shared_secret_file.is_some() + || self.ws_issuer.is_some() + || self.ws_audience.is_some() + || self.ws_max_clock_skew_seconds.is_some() + { + anyhow::bail!( + "websocket auth flags require `--ws-auth capability-token` or `--ws-auth signed-bearer-token`" + ); + } + None + } + }; + + Ok(AppServerWebsocketAuthSettings { config }) + } +} + +pub fn policy_from_settings( + settings: &AppServerWebsocketAuthSettings, +) -> io::Result { + let mode = match settings.config.as_ref() { + Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source { + AppServerWebsocketCapabilityTokenSource::TokenFile { token_file } => { + let token = read_trimmed_secret(token_file.as_ref())?; + Some(WebsocketAuthMode::CapabilityToken { + token_sha256: sha256_digest(token.as_bytes()), + }) + } + AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => { + Some(WebsocketAuthMode::CapabilityToken { + token_sha256: *token_sha256, + }) + } + }, + Some(AppServerWebsocketAuthConfig::SignedBearerToken { + shared_secret_file, + issuer, + audience, + max_clock_skew_seconds, + }) => { + let shared_secret = read_trimmed_secret(shared_secret_file.as_ref())?.into_bytes(); + validate_signed_bearer_secret(shared_secret_file.as_ref(), &shared_secret)?; + let max_clock_skew_seconds = i64::try_from(*max_clock_skew_seconds).map_err(|_| { + io::Error::new( + ErrorKind::InvalidInput, + "websocket auth clock skew must fit in a signed 64-bit integer", + ) + })?; + Some(WebsocketAuthMode::SignedBearerToken { + shared_secret, + issuer: issuer.clone(), + audience: audience.clone(), + max_clock_skew_seconds, + }) + } + None => None, + }; + + Ok(WebsocketAuthPolicy { mode }) +} + +pub(crate) fn is_unauthenticated_non_loopback_listener( + bind_address: SocketAddr, + policy: &WebsocketAuthPolicy, +) -> bool { + !bind_address.ip().is_loopback() && policy.mode.is_none() +} + +pub(crate) fn authorize_upgrade( + headers: &HeaderMap, + policy: &WebsocketAuthPolicy, +) -> Result<(), WebsocketAuthError> { + let Some(mode) = policy.mode.as_ref() else { + return Ok(()); + }; + + let token = bearer_token_from_headers(headers)?; + match mode { + WebsocketAuthMode::CapabilityToken { token_sha256 } => { + let actual_sha256 = sha256_digest(token.as_bytes()); + if constant_time_eq_32(token_sha256, &actual_sha256) { + Ok(()) + } else { + Err(unauthorized("invalid websocket bearer token")) + } + } + WebsocketAuthMode::SignedBearerToken { + shared_secret, + issuer, + audience, + max_clock_skew_seconds, + } => verify_signed_bearer_token( + token, + shared_secret, + issuer.as_deref(), + audience.as_deref(), + *max_clock_skew_seconds, + ), + } +} + +fn verify_signed_bearer_token( + token: &str, + shared_secret: &[u8], + issuer: Option<&str>, + audience: Option<&str>, + max_clock_skew_seconds: i64, +) -> Result<(), WebsocketAuthError> { + let claims = decode_jwt_claims(token, shared_secret)?; + validate_jwt_claims(&claims, issuer, audience, max_clock_skew_seconds) +} + +fn decode_jwt_claims(token: &str, shared_secret: &[u8]) -> Result { + let mut validation = Validation::new(Algorithm::HS256); + validation.required_spec_claims.clear(); + validation.validate_exp = false; + validation.validate_nbf = false; + validation.validate_aud = false; + + decode::(token, &DecodingKey::from_secret(shared_secret), &validation) + .map(|token_data| token_data.claims) + .map_err(|_| unauthorized("invalid websocket jwt")) +} + +fn validate_jwt_claims( + claims: &JwtClaims, + issuer: Option<&str>, + audience: Option<&str>, + max_clock_skew_seconds: i64, +) -> Result<(), WebsocketAuthError> { + let now = OffsetDateTime::now_utc().unix_timestamp(); + if now > claims.exp.saturating_add(max_clock_skew_seconds) { + return Err(unauthorized("expired websocket jwt")); + } + if let Some(nbf) = claims.nbf + && now < nbf.saturating_sub(max_clock_skew_seconds) + { + return Err(unauthorized("websocket jwt is not valid yet")); + } + if let Some(expected_issuer) = issuer + && claims.iss.as_deref() != Some(expected_issuer) + { + return Err(unauthorized("websocket jwt issuer mismatch")); + } + if let Some(expected_audience) = audience + && !audience_matches(claims.aud.as_ref(), expected_audience) + { + return Err(unauthorized("websocket jwt audience mismatch")); + } + + Ok(()) +} + +fn audience_matches(audience: Option<&JwtAudienceClaim>, expected_audience: &str) -> bool { + match audience { + Some(JwtAudienceClaim::Single(actual)) => actual == expected_audience, + Some(JwtAudienceClaim::Multiple(actual)) => { + actual.iter().any(|audience| audience == expected_audience) + } + None => false, + } +} + +fn bearer_token_from_headers(headers: &HeaderMap) -> Result<&str, WebsocketAuthError> { + let raw_header = headers + .get(AUTHORIZATION) + .ok_or_else(|| unauthorized("missing websocket bearer token"))?; + let header = raw_header + .to_str() + .map_err(|_| unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE))?; + let Some((scheme, token)) = header.split_once(' ') else { + return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE)); + }; + if !scheme.eq_ignore_ascii_case("Bearer") { + return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE)); + } + let token = token.trim(); + if token.is_empty() { + return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE)); + } + Ok(token) +} + +fn validate_signed_bearer_secret(path: &Path, shared_secret: &[u8]) -> io::Result<()> { + if shared_secret.len() < MIN_SIGNED_BEARER_SECRET_BYTES { + return Err(io::Error::new( + ErrorKind::InvalidInput, + format!( + "signed websocket bearer secret {} must be at least {MIN_SIGNED_BEARER_SECRET_BYTES} bytes", + path.display() + ), + )); + } + Ok(()) +} + +fn read_trimmed_secret(path: &std::path::Path) -> io::Result { + let raw = std::fs::read_to_string(path).map_err(|err| { + io::Error::new( + err.kind(), + format!( + "failed to read websocket auth secret {}: {err}", + path.display() + ), + ) + })?; + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(io::Error::new( + ErrorKind::InvalidInput, + format!("websocket auth secret {} must not be empty", path.display()), + )); + } + Ok(trimmed.to_string()) +} + +fn absolute_path_arg(flag_name: &str, path: PathBuf) -> anyhow::Result { + AbsolutePathBuf::try_from(path).with_context(|| format!("{flag_name} must be an absolute path")) +} + +fn sha256_digest_arg(flag_name: &str, value: &str) -> anyhow::Result<[u8; 32]> { + let trimmed = value.trim(); + if trimmed.len() != 64 { + anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest"); + } + + let mut digest = [0u8; 32]; + for (index, pair) in trimmed.as_bytes().chunks_exact(2).enumerate() { + let high = hex_nibble(flag_name, pair[0])?; + let low = hex_nibble(flag_name, pair[1])?; + digest[index] = (high << 4) | low; + } + Ok(digest) +} + +fn hex_nibble(flag_name: &str, byte: u8) -> anyhow::Result { + match byte { + b'0'..=b'9' => Ok(byte - b'0'), + b'a'..=b'f' => Ok(byte - b'a' + 10), + b'A'..=b'F' => Ok(byte - b'A' + 10), + _ => anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest"), + } +} + +fn sha256_digest(input: &[u8]) -> [u8; 32] { + let mut digest = [0u8; 32]; + digest.copy_from_slice(&Sha256::digest(input)); + digest +} + +fn unauthorized(message: &'static str) -> WebsocketAuthError { + WebsocketAuthError { + status_code: StatusCode::UNAUTHORIZED, + message, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::HeaderValue; + use base64::Engine; + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use hmac::Hmac; + use hmac::Mac; + use serde_json::json; + + type HmacSha256 = Hmac; + + fn signed_token(shared_secret: &[u8], claims: serde_json::Value) -> String { + let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#); + let claims_segment = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap()); + let payload = format!("{header}.{claims_segment}"); + let mut mac = HmacSha256::new_from_slice(shared_secret).unwrap(); + mac.update(payload.as_bytes()); + let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()); + format!("{payload}.{signature}") + } + + #[test] + fn detects_unauthenticated_non_loopback_listener() { + let policy = WebsocketAuthPolicy::default(); + assert!(is_unauthenticated_non_loopback_listener( + "0.0.0.0:8765".parse().unwrap(), + &policy, + )); + assert!(!is_unauthenticated_non_loopback_listener( + "127.0.0.1:8765".parse().unwrap(), + &policy, + )); + assert!(!is_unauthenticated_non_loopback_listener( + "0.0.0.0:8765".parse().unwrap(), + &WebsocketAuthPolicy { + mode: Some(WebsocketAuthMode::CapabilityToken { + token_sha256: [0u8; 32], + }), + }, + )); + } + + #[test] + fn capability_token_args_require_token_file_or_hash() { + let err = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::CapabilityToken), + ..Default::default() + } + .try_into_settings() + .expect_err("capability-token mode should require a token source"); + assert!( + err.to_string().contains("--ws-token-file") + && err.to_string().contains("--ws-token-sha256"), + "unexpected error: {err}" + ); + } + + #[test] + fn capability_token_args_accept_token_hash() { + let settings = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::CapabilityToken), + ws_token_sha256: Some("ab".repeat(32)), + ..Default::default() + } + .try_into_settings() + .expect("capability-token hash args should parse"); + + assert_eq!( + settings, + AppServerWebsocketAuthSettings { + config: Some(AppServerWebsocketAuthConfig::CapabilityToken { + source: AppServerWebsocketCapabilityTokenSource::TokenSha256 { + token_sha256: [0xab; 32], + }, + }), + } + ); + } + + #[test] + fn capability_token_args_reject_multiple_token_sources() { + let err = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::CapabilityToken), + ws_token_file: Some(PathBuf::from("/tmp/token")), + ws_token_sha256: Some("ab".repeat(32)), + ..Default::default() + } + .try_into_settings() + .expect_err("capability-token mode should reject multiple token sources"); + assert!( + err.to_string().contains("mutually exclusive"), + "unexpected error: {err}" + ); + } + + #[test] + fn capability_token_args_reject_malformed_token_hash() { + let err = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::CapabilityToken), + ws_token_sha256: Some("not-a-sha256".to_string()), + ..Default::default() + } + .try_into_settings() + .expect_err("capability-token mode should reject malformed token hashes"); + assert!( + err.to_string().contains("64-character hex"), + "unexpected error: {err}" + ); + } + + #[test] + fn capability_token_hash_policy_authorizes_matching_bearer_token() { + let settings = AppServerWebsocketAuthSettings { + config: Some(AppServerWebsocketAuthConfig::CapabilityToken { + source: AppServerWebsocketCapabilityTokenSource::TokenSha256 { + token_sha256: sha256_digest(b"super-secret-token"), + }, + }), + }; + let policy = policy_from_settings(&settings).expect("hash policy should build"); + let mut headers = HeaderMap::new(); + headers.insert( + AUTHORIZATION, + HeaderValue::from_static("Bearer super-secret-token"), + ); + authorize_upgrade(&headers, &policy).expect("matching token should authorize"); + + headers.insert( + AUTHORIZATION, + HeaderValue::from_static("Bearer wrong-token"), + ); + let err = authorize_upgrade(&headers, &policy).expect_err("wrong token should fail"); + assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); + } + + #[test] + fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() { + let err = AppServerWebsocketAuthArgs { + ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")), + ..Default::default() + } + .try_into_settings() + .expect_err("mode-specific flags should require --ws-auth"); + assert!( + err.to_string().contains("websocket auth flags require"), + "unexpected error: {err}" + ); + } + + #[test] + fn signed_bearer_args_default_clock_skew_and_trim_optional_claims() { + let settings = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::SignedBearerToken), + ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")), + ws_issuer: Some(" issuer ".to_string()), + ws_audience: Some(" ".to_string()), + ..Default::default() + } + .try_into_settings() + .expect("signed bearer args should parse"); + + assert_eq!( + settings, + AppServerWebsocketAuthSettings { + config: Some(AppServerWebsocketAuthConfig::SignedBearerToken { + shared_secret_file: AbsolutePathBuf::from_absolute_path("/tmp/secret") + .expect("absolute path"), + issuer: Some("issuer".to_string()), + audience: None, + max_clock_skew_seconds: DEFAULT_MAX_CLOCK_SKEW_SECONDS, + }), + } + ); + } + + #[test] + fn signed_bearer_token_verification_rejects_tampering() { + let shared_secret = b"0123456789abcdef0123456789abcdef"; + let token = signed_token( + shared_secret, + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + }), + ); + let tampered = token.replace(".eyJleHAi", ".eyJleHBi"); + let err = verify_signed_bearer_token( + &tampered, + shared_secret, + /*issuer*/ None, + /*audience*/ None, + /*max_clock_skew_seconds*/ 30, + ) + .expect_err("tampered jwt should fail"); + assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); + } + + #[test] + fn signed_bearer_token_verification_accepts_valid_token() { + let shared_secret = b"0123456789abcdef0123456789abcdef"; + let token = signed_token( + shared_secret, + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "iss": "issuer", + "aud": "audience", + }), + ); + verify_signed_bearer_token( + &token, + shared_secret, + Some("issuer"), + Some("audience"), + /*max_clock_skew_seconds*/ 30, + ) + .expect("valid signed token should verify"); + } + + #[test] + fn signed_bearer_token_verification_accepts_multiple_audiences() { + let shared_secret = b"0123456789abcdef0123456789abcdef"; + let token = signed_token( + shared_secret, + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "aud": ["other-audience", "audience"], + }), + ); + verify_signed_bearer_token( + &token, + shared_secret, + /*issuer*/ None, + Some("audience"), + /*max_clock_skew_seconds*/ 30, + ) + .expect("jwt audience arrays should verify"); + } + + #[test] + fn signed_bearer_token_verification_rejects_alg_none_tokens() { + let claims_segment = URL_SAFE_NO_PAD.encode( + serde_json::to_vec(&json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + })) + .unwrap(), + ); + let header_segment = URL_SAFE_NO_PAD.encode(br#"{"alg":"none","typ":"JWT"}"#); + let token = format!("{header_segment}.{claims_segment}."); + let err = verify_signed_bearer_token( + &token, + b"0123456789abcdef0123456789abcdef", + /*issuer*/ None, + /*audience*/ None, + /*max_clock_skew_seconds*/ 30, + ) + .expect_err("alg=none jwt should be rejected"); + assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); + } + + #[test] + fn signed_bearer_token_verification_rejects_missing_exp() { + let shared_secret = b"0123456789abcdef0123456789abcdef"; + let token = signed_token( + shared_secret, + json!({ + "iss": "issuer", + }), + ); + let err = verify_signed_bearer_token( + &token, + shared_secret, + /*issuer*/ None, + /*audience*/ None, + /*max_clock_skew_seconds*/ 30, + ) + .expect_err("jwt without exp should be rejected"); + assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); + } + + #[test] + fn validate_signed_bearer_secret_rejects_short_secret() { + let err = validate_signed_bearer_secret(Path::new("/tmp/secret"), b"too-short") + .expect_err("short shared secret should be rejected"); + assert_eq!(err.kind(), ErrorKind::InvalidInput); + assert!( + err.to_string().contains("must be at least 32 bytes"), + "unexpected error: {err}" + ); + } +} diff --git a/codex-rs/app-server-transport/src/transport/mod.rs b/codex-rs/app-server-transport/src/transport/mod.rs index dfd1d8fae..3e76069f4 100644 --- a/codex-rs/app-server-transport/src/transport/mod.rs +++ b/codex-rs/app-server-transport/src/transport/mod.rs @@ -1,3 +1,5 @@ +pub mod auth; + use crate::outgoing_message::ConnectionId; use crate::outgoing_message::OutgoingError; use crate::outgoing_message::OutgoingMessage; @@ -6,6 +8,7 @@ use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCMessage; use codex_core::config::find_codex_home; use codex_utils_absolute_path::AbsolutePathBuf; +use std::net::SocketAddr; use std::path::Path; use std::str::FromStr; use std::sync::atomic::AtomicU64; @@ -32,6 +35,7 @@ pub use remote_control::RemoteControlStartConfig; pub use remote_control::start_remote_control; pub use stdio::start_stdio_connection; pub use unix_socket::start_control_socket_acceptor; +pub use websocket::start_websocket_acceptor; const OVERLOADED_ERROR_CODE: i64 = -32001; @@ -50,6 +54,7 @@ pub fn app_server_control_socket_path(codex_home: &Path) -> std::io::Result write!( f, - "unsupported --listen URL `{listen_url}`; expected `stdio://`, `unix://`, `unix://PATH`, or `off`" + "unsupported --listen URL `{listen_url}`; expected `stdio://`, `unix://`, `unix://PATH`, `ws://IP:PORT`, or `off`" ), AppServerTransportParseError::InvalidUnixSocketPath { listen_url, @@ -73,6 +79,10 @@ impl std::fmt::Display for AppServerTransportParseError { f, "invalid unix socket --listen URL `{listen_url}`; failed to resolve socket path: {message}" ), + AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!( + f, + "invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`" + ), } } } @@ -116,6 +126,13 @@ impl AppServerTransport { return Ok(Self::Off); } + if let Some(socket_addr) = listen_url.strip_prefix("ws://") { + let bind_address = socket_addr.parse::().map_err(|_| { + AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string()) + })?; + return Ok(Self::WebSocket { bind_address }); + } + Err(AppServerTransportParseError::UnsupportedListenUrl( listen_url.to_string(), )) @@ -151,7 +168,7 @@ pub enum TransportEvent { pub enum ConnectionOrigin { Stdio, InProcess, - UnixSocket, + WebSocket, RemoteControl, } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs index a44c8522d..60dfc3d84 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs @@ -1,6 +1,7 @@ use super::protocol::EnrollRemoteServerRequest; use super::protocol::EnrollRemoteServerResponse; use super::protocol::RemoteControlTarget; +use axum::http::HeaderMap; use codex_api::SharedAuthProvider; use codex_login::default_client::build_reqwest_client; use codex_state::RemoteControlEnrollmentRecord; @@ -8,7 +9,6 @@ use codex_state::StateRuntime; use gethostname::gethostname; use std::io; use std::io::ErrorKind; -use tokio_tungstenite::tungstenite::http::HeaderMap; use tracing::info; use tracing::warn; diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index c3111fbc7..472639bc6 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -19,6 +19,7 @@ use super::segment::ClientSegmentObservation; use super::segment::ClientSegmentReassembler; use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES; use super::segment::split_server_envelope_for_transport; +use axum::http::HeaderValue; use base64::Engine; use codex_app_server_protocol::RemoteControlConnectionStatus; use codex_app_server_protocol::RemoteControlStatusChangedNotification; @@ -47,7 +48,6 @@ use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite; use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_util::sync::CancellationToken; use tracing::error; use tracing::info; diff --git a/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs b/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs index 4c77d8488..0b7dec0a2 100644 --- a/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs +++ b/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs @@ -52,14 +52,6 @@ fn listen_unix_socket_accepts_relative_custom_path() { ); } -#[test] -fn listen_websocket_url_is_not_supported() { - assert!(matches!( - AppServerTransport::from_listen_url("ws://127.0.0.1:0"), - Err(super::AppServerTransportParseError::UnsupportedListenUrl(_)) - )); -} - #[tokio::test] async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_and_pings() { let temp_dir = tempfile::TempDir::new().expect("temp dir"); diff --git a/codex-rs/app-server-transport/src/transport/websocket.rs b/codex-rs/app-server-transport/src/transport/websocket.rs index 086f76d7a..177e37890 100644 --- a/codex-rs/app-server-transport/src/transport/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/websocket.rs @@ -1,17 +1,46 @@ use super::CHANNEL_CAPACITY; use super::ConnectionOrigin; use super::TransportEvent; +use super::auth::WebsocketAuthPolicy; +use super::auth::authorize_upgrade; +use super::auth::is_unauthenticated_non_loopback_listener; use super::forward_incoming_message; use super::next_connection_id; use super::serialize_outgoing_message; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::QueuedOutgoingMessage; +use axum::Router; +use axum::body::Body; +use axum::body::Bytes; +use axum::extract::ConnectInfo; +use axum::extract::State; +use axum::extract::ws::Message as AxumWebSocketMessage; +use axum::extract::ws::WebSocketUpgrade; +use axum::http::HeaderMap; +use axum::http::Request; +use axum::http::StatusCode; +use axum::http::header::ORIGIN; +use axum::middleware; +use axum::middleware::Next; +use axum::response::IntoResponse; +use axum::response::Response; +use axum::routing::any; +use axum::routing::get; use futures::SinkExt; use futures::StreamExt; +use owo_colors::OwoColorize; +use owo_colors::Stream; +use owo_colors::Style; +use std::io::Result as IoResult; +use std::net::SocketAddr; +use std::sync::Arc; +use tokio::net::TcpListener; use tokio::sync::mpsc; -use tokio_tungstenite::tungstenite::Bytes; +use tokio::task::JoinHandle; use tokio_tungstenite::tungstenite::Message as TungsteniteWebSocketMessage; use tokio_util::sync::CancellationToken; +use tracing::error; +use tracing::info; use tracing::warn; /// WebSocket clients can briefly lag behind normal turn output bursts while the @@ -19,6 +48,127 @@ use tracing::warn; const WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY: usize = 32 * 1024; const _: () = assert!(WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY > CHANNEL_CAPACITY); +fn colorize(text: &str, style: Style) -> String { + text.if_supports_color(Stream::Stderr, |value| value.style(style)) + .to_string() +} + +#[allow(clippy::print_stderr)] +fn print_websocket_startup_banner(addr: SocketAddr) { + let title = colorize("codex app-server (WebSockets)", Style::new().bold().cyan()); + let listening_label = colorize("listening on:", Style::new().dimmed()); + let listen_url = colorize(&format!("ws://{addr}"), Style::new().green()); + let ready_label = colorize("readyz:", Style::new().dimmed()); + let ready_url = colorize(&format!("http://{addr}/readyz"), Style::new().green()); + let health_label = colorize("healthz:", Style::new().dimmed()); + let health_url = colorize(&format!("http://{addr}/healthz"), Style::new().green()); + let note_label = colorize("note:", Style::new().dimmed()); + eprintln!("{title}"); + eprintln!(" {listening_label} {listen_url}"); + eprintln!(" {ready_label} {ready_url}"); + eprintln!(" {health_label} {health_url}"); + if addr.ip().is_loopback() { + eprintln!( + " {note_label} binds localhost only (use SSH port-forwarding for remote access)" + ); + } else { + eprintln!(" {note_label} websocket auth is required for non-localhost listeners"); + } +} + +#[derive(Clone)] +struct WebSocketListenerState { + transport_event_tx: mpsc::Sender, + auth_policy: Arc, +} + +async fn health_check_handler() -> StatusCode { + StatusCode::OK +} + +async fn reject_requests_with_origin_header( + request: Request, + next: Next, +) -> Result { + if request.headers().contains_key(ORIGIN) { + warn!( + method = %request.method(), + uri = %request.uri(), + "rejecting websocket listener request with Origin header" + ); + Err(StatusCode::FORBIDDEN) + } else { + Ok(next.run(request).await) + } +} + +async fn websocket_upgrade_handler( + websocket: WebSocketUpgrade, + ConnectInfo(peer_addr): ConnectInfo, + State(state): State, + headers: HeaderMap, +) -> impl IntoResponse { + if let Err(err) = authorize_upgrade(&headers, state.auth_policy.as_ref()) { + warn!( + %peer_addr, + message = err.message(), + "rejecting websocket client during upgrade" + ); + return (err.status_code(), err.message()).into_response(); + } + info!(%peer_addr, "websocket client connected"); + websocket + .on_upgrade(move |stream| async move { + let (websocket_writer, websocket_reader) = stream.split(); + run_websocket_connection(websocket_writer, websocket_reader, state.transport_event_tx) + .await; + }) + .into_response() +} + +pub async fn start_websocket_acceptor( + bind_address: SocketAddr, + transport_event_tx: mpsc::Sender, + shutdown_token: CancellationToken, + auth_policy: WebsocketAuthPolicy, +) -> IoResult> { + if is_unauthenticated_non_loopback_listener(bind_address, &auth_policy) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "refusing to start non-loopback websocket listener {bind_address} without auth; configure `--ws-auth capability-token` or `--ws-auth signed-bearer-token`" + ), + )); + } + let listener = TcpListener::bind(bind_address).await?; + let local_addr = listener.local_addr()?; + print_websocket_startup_banner(local_addr); + info!("app-server websocket listening on ws://{local_addr}"); + + let router = Router::new() + .route("/readyz", get(health_check_handler)) + .route("/healthz", get(health_check_handler)) + .fallback(any(websocket_upgrade_handler)) + .layer(middleware::from_fn(reject_requests_with_origin_header)) + .with_state(WebSocketListenerState { + transport_event_tx, + auth_policy: Arc::new(auth_policy), + }); + let server = axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async move { + shutdown_token.cancelled().await; + }); + Ok(tokio::spawn(async move { + if let Err(err) = server.await { + error!("websocket acceptor failed: {err}"); + } + info!("websocket acceptor shutting down"); + })) +} + pub(crate) async fn run_websocket_connection( websocket_writer: impl futures::sink::Sink + Send + 'static, websocket_reader: impl futures::stream::Stream> + Send + 'static, @@ -36,7 +186,7 @@ pub(crate) async fn run_websocket_connection( if transport_event_tx .send(TransportEvent::ConnectionOpened { connection_id, - origin: ConnectionOrigin::UnixSocket, + origin: ConnectionOrigin::WebSocket, writer: writer_tx, disconnect_sender: Some(disconnect_token.clone()), }) @@ -95,6 +245,26 @@ pub(crate) trait AppServerWebSocketMessage: Sized { fn into_incoming(self) -> Option; } +impl AppServerWebSocketMessage for AxumWebSocketMessage { + fn text(text: String) -> Self { + Self::Text(text.into()) + } + + fn pong(payload: Bytes) -> Self { + Self::Pong(payload) + } + + fn into_incoming(self) -> Option { + Some(match self { + Self::Text(text) => IncomingWebSocketMessage::Text(text.to_string()), + Self::Binary(_) => IncomingWebSocketMessage::Binary, + Self::Ping(payload) => IncomingWebSocketMessage::Ping(payload), + Self::Pong(_) => IncomingWebSocketMessage::Pong, + Self::Close(_) => IncomingWebSocketMessage::Close, + }) + } +} + impl AppServerWebSocketMessage for TungsteniteWebSocketMessage { fn text(text: String) -> Self { Self::Text(text.into()) diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index fcf731aa9..ecc8d1c25 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -28,6 +28,7 @@ axum = { workspace = true, default-features = false, features = [ "http1", "json", "tokio", + "ws", ] } codex-analytics = { workspace = true } codex-arg0 = { workspace = true } @@ -100,20 +101,22 @@ axum = { workspace = true, default-features = false, features = [ "tokio", ] } base64 = { workspace = true } -codex-uds = { workspace = true } codex-model-provider-info = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } flate2 = { workspace = true } +hmac = { workspace = true } opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } pretty_assertions = { workspace = true } +reqwest = { workspace = true, features = ["rustls-tls"] } rmcp = { workspace = true, default-features = false, features = [ "elicitation", "server", "transport-streamable-http-server", ] } serial_test = { workspace = true } +sha2 = { workspace = true } shlex = { workspace = true } tar = { workspace = true } tokio-tungstenite = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index c16762725..2be4694bf 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -24,9 +24,18 @@ Similar to [MCP](https://modelcontextprotocol.io/), `codex app-server` supports Supported transports: - stdio (`--listen stdio://`, default): newline-delimited JSON (JSONL) +- websocket (`--listen ws://IP:PORT`): one JSON-RPC message per websocket text frame (**experimental / unsupported**) - unix socket (`--listen unix://` or `--listen unix://PATH`): websocket connections over `$CODEX_HOME/app-server-control/app-server-control.sock` or a custom socket path, using the standard HTTP Upgrade handshake - off (`--listen off`): do not expose a local transport +When running with `--listen ws://IP:PORT`, the same listener also serves basic HTTP health probes: + +- `GET /readyz` returns `200 OK` once the listener is accepting new connections. +- `GET /healthz` returns `200 OK` when no `Origin` header is present. +- Any request carrying an `Origin` header is rejected with `403 Forbidden`. + +Websocket transport is currently experimental and unsupported. Do not rely on it for production workloads. + The unix socket transport is intended for local app-server control-plane clients. `codex app-server proxy` opens exactly one raw stream connection to `$CODEX_HOME/app-server-control/app-server-control.sock` by default, or to `--sock PATH` when provided, and proxies bytes between that socket and stdin/stdout. diff --git a/codex-rs/app-server/src/app_server_tracing.rs b/codex-rs/app-server/src/app_server_tracing.rs index 21acf1c86..6e8133740 100644 --- a/codex-rs/app-server/src/app_server_tracing.rs +++ b/codex-rs/app-server/src/app_server_tracing.rs @@ -1,7 +1,7 @@ //! Tracing helpers shared by socket and in-process app-server entry points. //! //! The in-process path intentionally reuses the same span shape as JSON-RPC -//! transports so request telemetry stays comparable across stdio, unix socket, +//! transports so request telemetry stays comparable across stdio, websocket, //! and embedded callers. [`typed_request_span`] is the in-process counterpart //! of [`request_span`] and stamps `rpc.transport` as `"in-process"` while //! deriving client identity from the typed [`ClientRequest`] rather than @@ -86,6 +86,7 @@ fn transport_name(transport: &AppServerTransport) -> &'static str { match transport { AppServerTransport::Stdio => "stdio", AppServerTransport::UnixSocket { .. } => "unix_socket", + AppServerTransport::WebSocket { .. } => "websocket", AppServerTransport::Off => "off", } } diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 1538b45ce..baff7c8d4 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -31,10 +31,12 @@ use crate::transport::ConnectionState; use crate::transport::OutboundConnectionState; use crate::transport::RemoteControlStartConfig; use crate::transport::TransportEvent; +use crate::transport::auth::policy_from_settings; use crate::transport::route_outgoing_envelope; use crate::transport::start_control_socket_acceptor; use crate::transport::start_remote_control; use crate::transport::start_stdio_connection; +use crate::transport::start_websocket_acceptor; use codex_analytics::AppServerRpcTransport; use codex_app_server_protocol::ConfigLayerSource; use codex_app_server_protocol::ConfigWarningNotification; @@ -101,6 +103,9 @@ pub use crate::error_code::INPUT_TOO_LARGE_ERROR_CODE; pub use crate::error_code::INVALID_PARAMS_ERROR_CODE; pub use crate::transport::AppServerTransport; pub use crate::transport::app_server_control_socket_path; +pub use crate::transport::auth::AppServerWebsocketAuthArgs; +pub use crate::transport::auth::AppServerWebsocketAuthSettings; +pub use crate::transport::auth::WebsocketAuthCliMode; const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT"; const OTEL_SERVICE_NAME: &str = "codex-app-server"; @@ -378,6 +383,7 @@ pub async fn run_main( default_analytics_enabled, AppServerTransport::Stdio, SessionSource::VSCode, + AppServerWebsocketAuthSettings::default(), ) .await } @@ -410,6 +416,7 @@ pub async fn run_main_with_transport( default_analytics_enabled: bool, transport: AppServerTransport, session_source: SessionSource, + auth: AppServerWebsocketAuthSettings, ) -> IoResult<()> { run_main_with_transport_options( arg0_paths, @@ -418,6 +425,7 @@ pub async fn run_main_with_transport( default_analytics_enabled, transport, session_source, + auth, AppServerRuntimeOptions::default(), ) .await @@ -431,6 +439,7 @@ pub async fn run_main_with_transport_options( default_analytics_enabled: bool, transport: AppServerTransport, session_source: SessionSource, + auth: AppServerWebsocketAuthSettings, runtime_options: AppServerRuntimeOptions, ) -> IoResult<()> { let (transport_event_tx, mut transport_event_rx) = @@ -670,6 +679,16 @@ pub async fn run_main_with_transport_options( .await?; transport_accept_handles.push(accept_handle); } + AppServerTransport::WebSocket { bind_address } => { + let accept_handle = start_websocket_acceptor( + *bind_address, + transport_event_tx.clone(), + transport_shutdown_token.clone(), + policy_from_settings(&auth)?, + ) + .await?; + transport_accept_handles.push(accept_handle); + } AppServerTransport::Off => {} } @@ -741,7 +760,7 @@ pub async fn run_main_with_transport_options( } OutboundControlEvent::DisconnectAll => { info!( - "disconnecting {} outbound connection(s) for graceful restart", + "disconnecting {} outbound websocket connection(s) for graceful restart", outbound_connections.len() ); for connection_state in outbound_connections.values() { @@ -1068,9 +1087,9 @@ pub async fn run_main_with_transport_options( fn analytics_rpc_transport(transport: &AppServerTransport) -> AppServerRpcTransport { match transport { AppServerTransport::Stdio => AppServerRpcTransport::Stdio, - AppServerTransport::UnixSocket { .. } | AppServerTransport::Off => { - AppServerRpcTransport::Websocket - } + AppServerTransport::UnixSocket { .. } + | AppServerTransport::WebSocket { .. } + | AppServerTransport::Off => AppServerRpcTransport::Websocket, } } diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index acf5655e2..31a017d17 100644 --- a/codex-rs/app-server/src/main.rs +++ b/codex-rs/app-server/src/main.rs @@ -1,6 +1,7 @@ use clap::Parser; use codex_app_server::AppServerRuntimeOptions; use codex_app_server::AppServerTransport; +use codex_app_server::AppServerWebsocketAuthArgs; use codex_app_server::PluginStartupTasks; use codex_app_server::run_main_with_transport_options; use codex_arg0::Arg0DispatchPaths; @@ -18,7 +19,7 @@ const DISABLE_MANAGED_CONFIG_ENV_VAR: &str = "CODEX_APP_SERVER_DISABLE_MANAGED_C #[derive(Debug, Parser)] struct AppServerArgs { /// Transport endpoint URL. Supported values: `stdio://` (default), - /// `unix://`, `unix://PATH`, `off`. + /// `unix://`, `unix://PATH`, `ws://IP:PORT`, `off`. #[arg( long = "listen", value_name = "URL", @@ -35,6 +36,9 @@ struct AppServerArgs { )] session_source: SessionSource, + #[command(flatten)] + auth: AppServerWebsocketAuthArgs, + /// Hidden debug-only test hook used by integration tests that spawn the /// production app-server binary. #[cfg(debug_assertions)] @@ -58,6 +62,7 @@ fn main() -> anyhow::Result<()> { }; let transport = args.listen; let session_source = args.session_source; + let auth = args.auth.try_into_settings()?; let mut runtime_options = AppServerRuntimeOptions::default(); #[cfg(debug_assertions)] if args.disable_plugin_startup_tasks_for_tests { @@ -72,6 +77,7 @@ fn main() -> anyhow::Result<()> { /*default_analytics_enabled*/ false, transport, session_source, + auth, runtime_options, ) .await?; diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 100eed716..3d18cf321 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -21,9 +21,11 @@ pub(crate) use codex_app_server_transport::QueuedOutgoingMessage; pub(crate) use codex_app_server_transport::RemoteControlStartConfig; pub(crate) use codex_app_server_transport::TransportEvent; pub use codex_app_server_transport::app_server_control_socket_path; +pub use codex_app_server_transport::auth; pub(crate) use codex_app_server_transport::start_control_socket_acceptor; pub(crate) use codex_app_server_transport::start_remote_control; pub(crate) use codex_app_server_transport::start_stdio_connection; +pub(crate) use codex_app_server_transport::start_websocket_acceptor; pub(crate) struct ConnectionState { pub(crate) outbound_initialized: Arc, diff --git a/codex-rs/app-server/tests/suite/v2/command_exec.rs b/codex-rs/app-server/tests/suite/v2/command_exec.rs index f65503aa9..211cec935 100644 --- a/codex-rs/app-server/tests/suite/v2/command_exec.rs +++ b/codex-rs/app-server/tests/suite/v2/command_exec.rs @@ -861,10 +861,10 @@ async fn command_exec_process_ids_are_connection_scoped_and_disconnect_terminate .as_nanos() ); - let (mut process, socket_path, _socket_dir) = spawn_websocket_server(codex_home.path()).await?; + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; - let mut ws1 = connect_websocket(&socket_path).await?; - let mut ws2 = connect_websocket(&socket_path).await?; + let mut ws1 = connect_websocket(bind_addr).await?; + let mut ws2 = connect_websocket(bind_addr).await?; send_initialize_request(&mut ws1, /*id*/ 1, "ws_client_one").await?; read_initialize_response(&mut ws1, /*request_id*/ 1).await?; @@ -929,7 +929,7 @@ async fn command_exec_process_ids_are_connection_scoped_and_disconnect_terminate process .kill() .await - .context("failed to stop app-server process")?; + .context("failed to stop websocket app-server process")?; Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index ac322cbdc..cf76b9d57 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -4,6 +4,8 @@ use anyhow::bail; use app_test_support::DISABLE_PLUGIN_STARTUP_TASKS_ARG; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::JSONRPCError; @@ -16,14 +18,18 @@ use codex_app_server_protocol::ThreadLoadedListParams; use codex_app_server_protocol::ThreadLoadedListResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; -use codex_uds::UnixStream; use futures::SinkExt; use futures::StreamExt; +use hmac::Hmac; +use hmac::Mac; +use reqwest::StatusCode; use serde_json::json; +use sha2::Sha256; +use std::net::SocketAddr; use std::path::Path; -use std::path::PathBuf; use std::process::Stdio; use tempfile::TempDir; +use time::OffsetDateTime; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Child; @@ -32,29 +38,37 @@ use tokio::time::Duration; use tokio::time::Instant; use tokio::time::sleep; use tokio::time::timeout; +use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; -use tokio_tungstenite::client_async; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Error as WsError; use tokio_tungstenite::tungstenite::Message as WebSocketMessage; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::tungstenite::http::header::ORIGIN; // macOS and Windows CI can spend tens of seconds starting the app-server test -// binary under Bazel before it accepts JSON-RPC over the control socket. +// binary under Bazel before it accepts JSON-RPC or reports its websocket bind +// address. #[cfg(any(target_os = "macos", windows))] pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60); #[cfg(not(any(target_os = "macos", windows)))] pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); -pub(super) type WsClient = WebSocketStream; +pub(super) type WsClient = WebSocketStream>; +type HmacSha256 = Hmac; #[tokio::test] -async fn unix_socket_transport_routes_per_connection_handshake_and_responses() -> Result<()> { +async fn websocket_transport_routes_per_connection_handshake_and_responses() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), &server.uri(), "never")?; - let (mut process, socket_path, _socket_dir) = spawn_websocket_server(codex_home.path()).await?; + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; - let mut ws1 = connect_websocket(&socket_path).await?; - let mut ws2 = connect_websocket(&socket_path).await?; + let mut ws1 = connect_websocket(bind_addr).await?; + let mut ws2 = connect_websocket(bind_addr).await?; send_initialize_request(&mut ws1, /*id*/ 1, "ws_client_one").await?; let first_init = read_response_for_id(&mut ws1, /*id*/ 1).await?; @@ -86,20 +100,259 @@ async fn unix_socket_transport_routes_per_connection_handshake_and_responses() - process .kill() .await - .context("failed to stop app-server process")?; + .context("failed to stop websocket app-server process")?; Ok(()) } #[tokio::test] -async fn unix_socket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_timeout() --> Result<()> { +async fn websocket_transport_serves_health_endpoints_on_same_listener() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), &server.uri(), "never")?; - let (mut process, socket_path, _socket_dir) = spawn_websocket_server(codex_home.path()).await?; + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; + let client = reqwest::Client::new(); - let mut ws1 = connect_websocket(&socket_path).await?; + let readyz = http_get(&client, bind_addr, "/readyz").await?; + assert_eq!(readyz.status(), StatusCode::OK); + + let healthz = http_get(&client, bind_addr, "/healthz").await?; + assert_eq!(healthz.status(), StatusCode::OK); + + let mut ws = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_health_client").await?; + let init = read_response_for_id(&mut ws, /*id*/ 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_rejects_browser_origin_without_auth() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; + + let mut ws = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_loopback_client").await?; + let init = read_response_for_id(&mut ws, /*id*/ 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + drop(ws); + + assert_websocket_connect_rejected_with_headers( + bind_addr, + /*bearer_token*/ None, + Some("https://evil.example"), + StatusCode::FORBIDDEN, + ) + .await?; + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_rejects_missing_and_invalid_capability_tokens() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + let token_file = codex_home.path().join("app-server-token"); + std::fs::write(&token_file, "super-secret-token\n")?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + let auth_args = vec![ + "--ws-auth".to_string(), + "capability-token".to_string(), + "--ws-token-file".to_string(), + token_file.display().to_string(), + ]; + + let (mut process, bind_addr) = + spawn_websocket_server_with_args(codex_home.path(), "ws://0.0.0.0:0", &auth_args).await?; + + assert_websocket_connect_rejected(bind_addr, /*bearer_token*/ None).await?; + assert_websocket_connect_rejected(bind_addr, Some("wrong-token")).await?; + + let mut ws = connect_websocket_with_bearer(bind_addr, Some("super-secret-token")).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_auth_client").await?; + let init = read_response_for_id(&mut ws, /*id*/ 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_verifies_signed_short_lived_bearer_tokens() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + let shared_secret_file = codex_home.path().join("app-server-signing-secret"); + let shared_secret = "0123456789abcdef0123456789abcdef"; + std::fs::write(&shared_secret_file, format!("{shared_secret}\n"))?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + let auth_args = vec![ + "--ws-auth".to_string(), + "signed-bearer-token".to_string(), + "--ws-shared-secret-file".to_string(), + shared_secret_file.display().to_string(), + "--ws-issuer".to_string(), + "codex-enroller".to_string(), + "--ws-audience".to_string(), + "codex-app-server".to_string(), + "--ws-max-clock-skew-seconds".to_string(), + "1".to_string(), + ]; + + let (mut process, bind_addr) = + spawn_websocket_server_with_args(codex_home.path(), "ws://127.0.0.1:0", &auth_args).await?; + let expired_token = signed_bearer_token( + shared_secret.as_bytes(), + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() - 30, + "iss": "codex-enroller", + "aud": "codex-app-server", + }), + )?; + assert_websocket_connect_rejected(bind_addr, Some(expired_token.as_str())).await?; + + let malformed_token = "not-a-jwt"; + assert_websocket_connect_rejected(bind_addr, Some(malformed_token)).await?; + + let not_yet_valid_token = signed_bearer_token( + shared_secret.as_bytes(), + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "nbf": OffsetDateTime::now_utc().unix_timestamp() + 30, + "iss": "codex-enroller", + "aud": "codex-app-server", + }), + )?; + assert_websocket_connect_rejected(bind_addr, Some(not_yet_valid_token.as_str())).await?; + + let wrong_issuer_token = signed_bearer_token( + shared_secret.as_bytes(), + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "iss": "someone-else", + "aud": "codex-app-server", + }), + )?; + assert_websocket_connect_rejected(bind_addr, Some(wrong_issuer_token.as_str())).await?; + + let wrong_audience_token = signed_bearer_token( + shared_secret.as_bytes(), + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "iss": "codex-enroller", + "aud": "wrong-audience", + }), + )?; + assert_websocket_connect_rejected(bind_addr, Some(wrong_audience_token.as_str())).await?; + + let wrong_signature_token = signed_bearer_token( + b"fedcba9876543210fedcba9876543210", + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "iss": "codex-enroller", + "aud": "codex-app-server", + }), + )?; + assert_websocket_connect_rejected(bind_addr, Some(wrong_signature_token.as_str())).await?; + + let valid_token = signed_bearer_token( + shared_secret.as_bytes(), + json!({ + "exp": OffsetDateTime::now_utc().unix_timestamp() + 60, + "iss": "codex-enroller", + "aud": "codex-app-server", + }), + )?; + let mut ws = connect_websocket_with_bearer(bind_addr, Some(valid_token.as_str())).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_signed_auth_client").await?; + let init = read_response_for_id(&mut ws, /*id*/ 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_rejects_short_signed_bearer_secret_configuration() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + let shared_secret_file = codex_home.path().join("app-server-signing-secret"); + std::fs::write(&shared_secret_file, "too-short\n")?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let output = run_websocket_server_to_completion_with_args( + codex_home.path(), + "ws://127.0.0.1:0", + &[ + "--ws-auth".to_string(), + "signed-bearer-token".to_string(), + "--ws-shared-secret-file".to_string(), + shared_secret_file.display().to_string(), + ], + ) + .await?; + assert!( + !output.status.success(), + "short shared secret should fail websocket server startup" + ); + let stderr = String::from_utf8(output.stderr).context("stderr should be valid utf-8")?; + assert!( + stderr.contains("must be at least 32 bytes"), + "unexpected stderr: {stderr}" + ); + + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_rejects_unauthenticated_non_loopback_startup() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let output = + run_websocket_server_to_completion_with_args(codex_home.path(), "ws://0.0.0.0:0", &[]) + .await?; + assert!( + !output.status.success(), + "unauthenticated non-loopback listener should fail websocket server startup" + ); + let stderr = String::from_utf8(output.stderr).context("stderr should be valid utf-8")?; + assert!( + stderr.contains("refusing to start non-loopback websocket listener"), + "unexpected stderr: {stderr}" + ); + + Ok(()) +} + +#[tokio::test] +async fn websocket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_timeout() -> Result<()> +{ + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; + + let mut ws1 = connect_websocket(bind_addr).await?; send_initialize_request(&mut ws1, /*id*/ 1, "ws_thread_owner").await?; read_response_for_id(&mut ws1, /*id*/ 1).await?; @@ -109,7 +362,7 @@ async fn unix_socket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_t ws1.close(None).await.context("failed to close websocket")?; drop(ws1); - let mut ws2 = connect_websocket(&socket_path).await?; + let mut ws2 = connect_websocket(bind_addr).await?; send_initialize_request(&mut ws2, /*id*/ 4, "ws_reconnect_client").await?; read_response_for_id(&mut ws2, /*id*/ 4).await?; @@ -118,29 +371,26 @@ async fn unix_socket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_t process .kill() .await - .context("failed to stop app-server process")?; + .context("failed to stop websocket app-server process")?; Ok(()) } -pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, PathBuf, TempDir)> { +pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, SocketAddr)> { + spawn_websocket_server_with_args(codex_home, "ws://127.0.0.1:0", &[]).await +} + +pub(super) async fn spawn_websocket_server_with_args( + codex_home: &Path, + listen_url: &str, + extra_args: &[String], +) -> Result<(Child, SocketAddr)> { let program = codex_utils_cargo_bin::cargo_bin("codex-app-server") .context("should find app-server binary")?; - #[cfg(unix)] - let socket_dir = tempfile::Builder::new() - .prefix("cxs-") - .tempdir_in("/tmp") - .context("failed to create short app-server socket temp dir")?; - #[cfg(not(unix))] - let socket_dir = tempfile::Builder::new() - .prefix("cxs-") - .tempdir() - .context("failed to create app-server socket temp dir")?; - let socket_path = socket_dir.path().join("c.sock"); - let listen_url = format!("unix://{}", socket_path.display()); let mut cmd = Command::new(program); cmd.arg("--listen") - .arg(&listen_url) + .arg(listen_url) .arg(DISABLE_PLUGIN_STARTUP_TASKS_ARG) + .args(extra_args) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::piped()) @@ -149,62 +399,194 @@ pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, let mut process = cmd .kill_on_drop(true) .spawn() - .context("failed to spawn app-server process")?; + .context("failed to spawn websocket app-server process")?; let stderr = process .stderr .take() - .context("failed to capture app-server stderr")?; + .context("failed to capture websocket app-server stderr")?; let mut stderr_reader = BufReader::new(stderr).lines(); + let deadline = Instant::now() + DEFAULT_READ_TIMEOUT; + let bind_addr = loop { + let line = timeout( + deadline.saturating_duration_since(Instant::now()), + stderr_reader.next_line(), + ) + .await + .context("timed out waiting for websocket app-server to report bound websocket address")? + .context("failed to read websocket app-server stderr")? + .context("websocket app-server exited before reporting bound websocket address")?; + eprintln!("[websocket app-server stderr] {line}"); + + let stripped_line = { + let mut stripped = String::with_capacity(line.len()); + let mut chars = line.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '\u{1b}' && matches!(chars.peek(), Some(&'[')) { + chars.next(); + for next in chars.by_ref() { + if ('@'..='~').contains(&next) { + break; + } + } + continue; + } + stripped.push(ch); + } + stripped + }; + + if let Some(bind_addr) = stripped_line + .split_whitespace() + .find_map(|token| token.strip_prefix("ws://")) + .and_then(|addr| addr.parse::().ok()) + { + break bind_addr; + } + }; + tokio::spawn(async move { while let Ok(Some(line)) = stderr_reader.next_line().await { - eprintln!("[app-server stderr] {line}"); + eprintln!("[websocket app-server stderr] {line}"); } }); + Ok((process, bind_addr)) +} + +pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result { + connect_websocket_with_bearer(bind_addr, /*bearer_token*/ None).await +} + +pub(super) async fn connect_websocket_with_bearer( + bind_addr: SocketAddr, + bearer_token: Option<&str>, +) -> Result { + let url = format!("ws://{}", connectable_bind_addr(bind_addr)); + let request = websocket_request(url.as_str(), bearer_token, /*origin*/ None)?; let deadline = Instant::now() + DEFAULT_READ_TIMEOUT; loop { - if socket_path.exists() { - return Ok((process, socket_path, socket_dir)); + match connect_async(request.clone()).await { + Ok((stream, _response)) => return Ok(stream), + Err(err) => { + if Instant::now() >= deadline { + bail!("failed to connect websocket to {url}: {err}"); + } + sleep(Duration::from_millis(50)).await; + } } - if let Some(status) = process.try_wait()? { - bail!("app-server exited before creating control socket: {status}"); - } - if Instant::now() >= deadline { - bail!( - "timed out waiting for app-server control socket at {}", - socket_path.display() - ); - } - sleep(Duration::from_millis(50)).await; } } -pub(super) async fn connect_websocket(socket_path: &Path) -> Result { +async fn assert_websocket_connect_rejected( + bind_addr: SocketAddr, + bearer_token: Option<&str>, +) -> Result<()> { + assert_websocket_connect_rejected_with_headers( + bind_addr, + bearer_token, + /*origin*/ None, + StatusCode::UNAUTHORIZED, + ) + .await +} + +async fn assert_websocket_connect_rejected_with_headers( + bind_addr: SocketAddr, + bearer_token: Option<&str>, + origin: Option<&str>, + expected_status: StatusCode, +) -> Result<()> { + let url = format!("ws://{}", connectable_bind_addr(bind_addr)); + let request = websocket_request(url.as_str(), bearer_token, origin)?; + + match connect_async(request).await { + Ok((_stream, response)) => { + bail!( + "expected websocket handshake rejection, got {}", + response.status() + ) + } + Err(WsError::Http(response)) => { + assert_eq!(response.status(), expected_status); + Ok(()) + } + Err(err) => bail!("expected http rejection during websocket handshake: {err}"), + } +} + +async fn run_websocket_server_to_completion_with_args( + codex_home: &Path, + listen_url: &str, + extra_args: &[String], +) -> Result { + let program = codex_utils_cargo_bin::cargo_bin("codex-app-server") + .context("should find app-server binary")?; + let mut cmd = Command::new(program); + cmd.arg("--listen") + .arg(listen_url) + .arg(DISABLE_PLUGIN_STARTUP_TASKS_ARG) + .args(extra_args) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .env("CODEX_HOME", codex_home) + .env("RUST_LOG", "warn"); + timeout(DEFAULT_READ_TIMEOUT, cmd.output()) + .await + .context("timed out waiting for websocket app-server to exit")? + .context("failed to run websocket app-server") +} + +async fn http_get( + client: &reqwest::Client, + bind_addr: SocketAddr, + path: &str, +) -> Result { + let connectable_bind_addr = connectable_bind_addr(bind_addr); let deadline = Instant::now() + DEFAULT_READ_TIMEOUT; loop { - match UnixStream::connect(socket_path).await { - Ok(stream) => match client_async("ws://localhost/rpc", stream).await { - Ok((websocket, _response)) => return Ok(websocket), - Err(err) => { - if Instant::now() >= deadline { - bail!( - "failed to upgrade websocket over {}: {err}", - socket_path.display() - ); - } - } - }, + match client + .get(format!("http://{connectable_bind_addr}{path}")) + .send() + .await + .with_context(|| format!("failed to GET http://{connectable_bind_addr}{path}")) + { + Ok(response) => return Ok(response), Err(err) => { if Instant::now() >= deadline { - bail!("failed to connect to {}: {err}", socket_path.display()); + bail!("failed to GET http://{connectable_bind_addr}{path}: {err}"); } + sleep(Duration::from_millis(50)).await; } } - sleep(Duration::from_millis(50)).await; } } +fn websocket_request( + url: &str, + bearer_token: Option<&str>, + origin: Option<&str>, +) -> Result> { + let mut request = url + .into_client_request() + .context("failed to create websocket request")?; + if let Some(bearer_token) = bearer_token { + request.headers_mut().insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {bearer_token}")) + .context("invalid bearer token header")?, + ); + } + if let Some(origin) = origin { + request.headers_mut().insert( + ORIGIN, + HeaderValue::from_str(origin).context("invalid origin header")?, + ); + } + Ok(request) +} + pub(super) async fn send_initialize_request( stream: &mut WsClient, id: i64, @@ -467,3 +849,25 @@ stream_max_retries = 0 ), ) } + +fn connectable_bind_addr(bind_addr: SocketAddr) -> SocketAddr { + match bind_addr { + SocketAddr::V4(addr) if addr.ip().is_unspecified() => { + SocketAddr::from(([127, 0, 0, 1], addr.port())) + } + SocketAddr::V6(addr) if addr.ip().is_unspecified() => { + SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], addr.port())) + } + _ => bind_addr, + } +} + +fn signed_bearer_token(shared_secret: &[u8], claims: serde_json::Value) -> Result { + let header_segment = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#); + let claims_segment = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims)?); + let payload = format!("{header_segment}.{claims_segment}"); + let mut mac = HmacSha256::new_from_slice(shared_secret).context("failed to create hmac")?; + mac.update(payload.as_bytes()); + let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()); + Ok(format!("{payload}.{signature}")) +} diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs index dcc7b6016..4b708d4ab 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs @@ -32,13 +32,12 @@ use wiremock::matchers::method; use wiremock::matchers::path_regex; #[tokio::test] -async fn unix_socket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Result<()> { +async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Result<()> { let GracefulCtrlCFixture { _codex_home, _server, mut process, mut ws, - .. } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; send_sigint(&process)?; @@ -58,13 +57,12 @@ async fn unix_socket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Re } #[tokio::test] -async fn unix_socket_transport_second_ctrl_c_forces_exit_while_turn_running() -> Result<()> { +async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> Result<()> { let GracefulCtrlCFixture { _codex_home, _server, mut process, mut ws, - .. } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; send_sigint(&process)?; @@ -85,13 +83,12 @@ async fn unix_socket_transport_second_ctrl_c_forces_exit_while_turn_running() -> } #[tokio::test] -async fn unix_socket_transport_sigterm_waits_for_running_turn_before_exit() -> Result<()> { +async fn websocket_transport_sigterm_waits_for_running_turn_before_exit() -> Result<()> { let GracefulCtrlCFixture { _codex_home, _server, mut process, mut ws, - .. } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; send_sigterm(&process)?; @@ -111,13 +108,12 @@ async fn unix_socket_transport_sigterm_waits_for_running_turn_before_exit() -> R } #[tokio::test] -async fn unix_socket_transport_second_sigterm_forces_exit_while_turn_running() -> Result<()> { +async fn websocket_transport_second_sigterm_forces_exit_while_turn_running() -> Result<()> { let GracefulCtrlCFixture { _codex_home, _server, mut process, mut ws, - .. } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; send_sigterm(&process)?; @@ -144,7 +140,6 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() -> _server, mut process, mut ws, - .. } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; send_sighup(&process)?; @@ -168,7 +163,6 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() -> struct GracefulCtrlCFixture { _codex_home: TempDir, - _socket_dir: TempDir, _server: wiremock::MockServer, process: Child, ws: WsClient, @@ -187,8 +181,8 @@ async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result Result Result<()> { } fn send_signal(process: &Child, signal: &str) -> Result<()> { - let pid = process.id().context("app-server process has no pid")?; + let pid = process + .id() + .context("websocket app-server process has no pid")?; let status = StdCommand::new("kill") .arg(signal) .arg(pid.to_string()) @@ -304,7 +299,7 @@ async fn wait_for_process_exit_within( timeout(window, process.wait()) .await .context(timeout_context)? - .context("failed waiting for app-server process exit") + .context("failed waiting for websocket app-server process exit") } async fn expect_websocket_disconnect(stream: &mut WsClient) -> Result<()> { diff --git a/codex-rs/app-server/tests/suite/v2/thread_name_websocket.rs b/codex-rs/app-server/tests/suite/v2/thread_name_websocket.rs index 2382d8fe3..951e4d74e 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_name_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_name_websocket.rs @@ -36,11 +36,11 @@ async fn thread_name_updated_broadcasts_for_loaded_threads() -> Result<()> { create_config_toml(codex_home.path(), &server.uri(), "never")?; let conversation_id = create_rollout(codex_home.path(), "2025-01-05T12-00-00")?; - let (mut process, socket_path, _socket_dir) = spawn_websocket_server(codex_home.path()).await?; + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; let result = async { - let mut ws1 = connect_websocket(&socket_path).await?; - let mut ws2 = connect_websocket(&socket_path).await?; + let mut ws1 = connect_websocket(bind_addr).await?; + let mut ws2 = connect_websocket(bind_addr).await?; initialize_both_clients(&mut ws1, &mut ws2).await?; send_request( @@ -91,7 +91,7 @@ async fn thread_name_updated_broadcasts_for_loaded_threads() -> Result<()> { process .kill() .await - .context("failed to stop app-server process")?; + .context("failed to stop websocket app-server process")?; result } @@ -102,11 +102,11 @@ async fn thread_name_updated_broadcasts_for_not_loaded_threads() -> Result<()> { create_config_toml(codex_home.path(), &server.uri(), "never")?; let conversation_id = create_rollout(codex_home.path(), "2025-01-05T12-05-00")?; - let (mut process, socket_path, _socket_dir) = spawn_websocket_server(codex_home.path()).await?; + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; let result = async { - let mut ws1 = connect_websocket(&socket_path).await?; - let mut ws2 = connect_websocket(&socket_path).await?; + let mut ws1 = connect_websocket(bind_addr).await?; + let mut ws2 = connect_websocket(bind_addr).await?; initialize_both_clients(&mut ws1, &mut ws2).await?; let renamed = "Stored rename"; @@ -143,7 +143,7 @@ async fn thread_name_updated_broadcasts_for_not_loaded_threads() -> Result<()> { process .kill() .await - .context("failed to stop app-server process")?; + .context("failed to stop websocket app-server process")?; result } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index e5dc7deb3..aa841e5a2 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -420,7 +420,7 @@ struct AppServerCommand { subcommand: Option, /// Transport endpoint URL. Supported values: `stdio://` (default), - /// `unix://`, `unix://PATH`, `off`. + /// `unix://`, `unix://PATH`, `ws://IP:PORT`, `off`. #[arg( long = "listen", value_name = "URL", @@ -449,6 +449,9 @@ struct AppServerCommand { /// See https://developers.openai.com/codex/config-advanced/#metrics for more details. #[arg(long = "analytics-default-enabled")] analytics_default_enabled: bool, + + #[command(flatten)] + auth: codex_app_server::AppServerWebsocketAuthArgs, } #[derive(Debug, Parser)] @@ -894,6 +897,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { listen, remote_control, analytics_default_enabled, + auth, } = app_server_cli; reject_remote_mode_for_app_server_subcommand( root_remote.as_deref(), @@ -903,6 +907,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { match subcommand { None => { let transport = listen; + let auth = auth.try_into_settings()?; let runtime_options = codex_app_server::AppServerRuntimeOptions { remote_control_enabled: remote_control, ..Default::default() @@ -914,6 +919,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { analytics_default_enabled, transport, codex_protocol::protocol::SessionSource::VSCode, + auth, runtime_options, ) .await?; @@ -2532,14 +2538,16 @@ mod tests { } #[test] - fn app_server_listen_websocket_url_fails_to_parse() { - let parse_result = MultitoolCli::try_parse_from([ - "codex", - "app-server", - "--listen", - "ws://127.0.0.1:4500", - ]); - assert!(parse_result.is_err()); + fn app_server_listen_websocket_url_parses() { + let app_server = app_server_from_args( + ["codex", "app-server", "--listen", "ws://127.0.0.1:4500"].as_ref(), + ); + assert_eq!( + app_server.listen, + codex_app_server::AppServerTransport::WebSocket { + bind_address: "127.0.0.1:4500".parse().expect("valid socket address"), + } + ); } #[test] @@ -2708,6 +2716,61 @@ mod tests { assert!(err.to_string().contains("app-server daemon version")); } + #[test] + fn app_server_capability_token_flags_parse() { + let app_server = app_server_from_args( + [ + "codex", + "app-server", + "--ws-auth", + "capability-token", + "--ws-token-file", + "/tmp/codex-token", + ] + .as_ref(), + ); + assert_eq!( + app_server.auth.ws_auth, + Some(codex_app_server::WebsocketAuthCliMode::CapabilityToken) + ); + assert_eq!( + app_server.auth.ws_token_file, + Some(PathBuf::from("/tmp/codex-token")) + ); + } + + #[test] + fn app_server_signed_bearer_flags_parse() { + let app_server = app_server_from_args( + [ + "codex", + "app-server", + "--ws-auth", + "signed-bearer-token", + "--ws-shared-secret-file", + "/tmp/codex-secret", + "--ws-issuer", + "issuer", + "--ws-audience", + "audience", + "--ws-max-clock-skew-seconds", + "9", + ] + .as_ref(), + ); + assert_eq!( + app_server.auth.ws_auth, + Some(codex_app_server::WebsocketAuthCliMode::SignedBearerToken) + ); + assert_eq!( + app_server.auth.ws_shared_secret_file, + Some(PathBuf::from("/tmp/codex-secret")) + ); + assert_eq!(app_server.auth.ws_issuer.as_deref(), Some("issuer")); + assert_eq!(app_server.auth.ws_audience.as_deref(), Some("audience")); + assert_eq!(app_server.auth.ws_max_clock_skew_seconds, Some(9)); + } + #[test] fn app_server_rejects_removed_insecure_non_loopback_flag() { let parse_result = MultitoolCli::try_parse_from([