diff --git a/codex-rs/app-server-transport/src/transport/auth.rs b/codex-rs/app-server-transport/src/transport/auth.rs index eeccf21ed30a..9a8a8a7d65d4 100644 --- a/codex-rs/app-server-transport/src/transport/auth.rs +++ b/codex-rs/app-server-transport/src/transport/auth.rs @@ -1,7 +1,10 @@ use anyhow::Context; use axum::http::HeaderMap; use axum::http::StatusCode; +use axum::http::Uri; use axum::http::header::AUTHORIZATION; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use clap::Args; use clap::ValueEnum; use codex_utils_absolute_path::AbsolutePathBuf; @@ -10,19 +13,23 @@ use jsonwebtoken::Algorithm; use jsonwebtoken::DecodingKey; use jsonwebtoken::Validation; use jsonwebtoken::decode; +use rand::TryRngCore; +use rand::rngs::OsRng; 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; +use tracing::warn; const DEFAULT_MAX_CLOCK_SKEW_SECONDS: u64 = 30; +const GENERATED_QUERY_TOKEN_BYTES: usize = 32; const MIN_SIGNED_BEARER_SECRET_BYTES: usize = 32; const INVALID_AUTHORIZATION_HEADER_MESSAGE: &str = "invalid authorization header"; +const REMOTE_WS_TOKEN_ENV_VAR: &str = "CODEX_REMOTE_WS_TOKEN"; #[derive(Debug, Clone, Default, PartialEq, Eq, Args)] pub struct AppServerWebsocketAuthArgs { @@ -53,6 +60,13 @@ pub struct AppServerWebsocketAuthArgs { /// 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, + + /// Disable enforcement of the generated websocket query token. + /// + /// A token is still generated and printed. Missing or incorrect tokens are + /// accepted. + #[arg(long = "no-token-check")] + pub no_token_check: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] @@ -61,13 +75,17 @@ pub enum WebsocketAuthCliMode { SignedBearerToken, } -#[derive(Debug, Clone, Default, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct AppServerWebsocketAuthSettings { - pub config: Option, + pub config: AppServerWebsocketAuthConfig, } #[derive(Debug, Clone, PartialEq, Eq)] +#[allow(clippy::enum_variant_names)] pub enum AppServerWebsocketAuthConfig { + QueryToken { + enforce: bool, + }, CapabilityToken { source: AppServerWebsocketCapabilityTokenSource, }, @@ -85,13 +103,26 @@ pub enum AppServerWebsocketCapabilityTokenSource { TokenSha256 { token_sha256: [u8; 32] }, } -#[derive(Clone, Debug, Default)] +impl Default for AppServerWebsocketAuthSettings { + fn default() -> Self { + Self { + config: AppServerWebsocketAuthConfig::QueryToken { enforce: true }, + } + } +} + +#[derive(Clone, Debug)] pub struct WebsocketAuthPolicy { - pub(crate) mode: Option, + pub(crate) mode: WebsocketAuthMode, } #[derive(Clone, Debug)] +#[allow(clippy::enum_variant_names)] pub(crate) enum WebsocketAuthMode { + QueryToken { + token: String, + enforce: bool, + }, CapabilityToken { token_sha256: [u8; 32], }, @@ -176,7 +207,7 @@ impl AppServerWebsocketAuthArgs { ); } }; - Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) + AppServerWebsocketAuthConfig::CapabilityToken { source } } Some(WebsocketAuthCliMode::SignedBearerToken) => { if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() { @@ -187,7 +218,7 @@ impl AppServerWebsocketAuthArgs { 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 { + AppServerWebsocketAuthConfig::SignedBearerToken { shared_secret_file: absolute_path_arg( "--ws-shared-secret-file", shared_secret_file, @@ -197,7 +228,7 @@ impl AppServerWebsocketAuthArgs { 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() @@ -211,7 +242,9 @@ impl AppServerWebsocketAuthArgs { "websocket auth flags require `--ws-auth capability-token` or `--ws-auth signed-bearer-token`" ); } - None + AppServerWebsocketAuthConfig::QueryToken { + enforce: !self.no_token_check, + } } }; @@ -222,26 +255,30 @@ impl AppServerWebsocketAuthArgs { pub fn policy_from_settings( settings: &AppServerWebsocketAuthSettings, ) -> io::Result { - let mode = match settings.config.as_ref() { - Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source { + let mode = match &settings.config { + AppServerWebsocketAuthConfig::QueryToken { enforce } => WebsocketAuthMode::QueryToken { + token: generate_query_token()?, + enforce: *enforce, + }, + AppServerWebsocketAuthConfig::CapabilityToken { source } => match source { AppServerWebsocketCapabilityTokenSource::TokenFile { token_file } => { let token = read_trimmed_secret(token_file.as_ref())?; - Some(WebsocketAuthMode::CapabilityToken { + WebsocketAuthMode::CapabilityToken { token_sha256: sha256_digest(token.as_bytes()), - }) + } } AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => { - Some(WebsocketAuthMode::CapabilityToken { + WebsocketAuthMode::CapabilityToken { token_sha256: *token_sha256, - }) + } } }, - Some(AppServerWebsocketAuthConfig::SignedBearerToken { + 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(|_| { @@ -250,37 +287,57 @@ pub fn policy_from_settings( "websocket auth clock skew must fit in a signed 64-bit integer", ) })?; - Some(WebsocketAuthMode::SignedBearerToken { + 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( + uri: &Uri, 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 { + match &policy.mode { + WebsocketAuthMode::QueryToken { token, enforce } => { + let query_token = match query_token_from_uri(uri) { + Ok(query_token) => query_token, + Err(reason) if !enforce => { + warn!( + %reason, + "allowing websocket client because query token enforcement is disabled" + ); + return Ok(()); + } + Err(reason) => return Err(unauthorized(reason)), + }; + let token_matches = query_token.as_deref() == Some(token.as_str()); + if token_matches { + return Ok(()); + } + let reason = if query_token.is_some() { + "invalid generated websocket query token" + } else { + "missing generated websocket query token" + }; + if *enforce { + Err(unauthorized(reason)) + } else { + warn!( + %reason, + "allowing websocket client because query token enforcement is disabled" + ); + Ok(()) + } + } WebsocketAuthMode::CapabilityToken { token_sha256 } => { + let token = bearer_token_from_headers(headers)?; let actual_sha256 = sha256_digest(token.as_bytes()); if constant_time_eq_32(token_sha256, &actual_sha256) { Ok(()) @@ -293,14 +350,63 @@ pub(crate) fn authorize_upgrade( issuer, audience, max_clock_skew_seconds, - } => verify_signed_bearer_token( - token, - shared_secret, - issuer.as_deref(), - audience.as_deref(), - *max_clock_skew_seconds, - ), + } => { + let token = bearer_token_from_headers(headers)?; + verify_signed_bearer_token( + token, + shared_secret, + issuer.as_deref(), + audience.as_deref(), + *max_clock_skew_seconds, + ) + } + } +} + +impl WebsocketAuthPolicy { + pub(crate) fn generated_query_token(&self) -> Option<&str> { + match &self.mode { + WebsocketAuthMode::QueryToken { token, .. } => Some(token), + WebsocketAuthMode::CapabilityToken { .. } + | WebsocketAuthMode::SignedBearerToken { .. } => None, + } + } +} + +fn generate_query_token() -> io::Result { + match std::env::var(REMOTE_WS_TOKEN_ENV_VAR) { + Ok(token) => return Ok(token), + Err(std::env::VarError::NotPresent) => {} + Err(std::env::VarError::NotUnicode(_)) => { + return Err(io::Error::new( + ErrorKind::InvalidData, + format!("{REMOTE_WS_TOKEN_ENV_VAR} must contain valid Unicode"), + )); + } + } + + let mut bytes = [0_u8; GENERATED_QUERY_TOKEN_BYTES]; + let mut rng = OsRng; + rng.try_fill_bytes(&mut bytes) + .map_err(|err| io::Error::other(format!("failed to generate websocket token: {err}")))?; + Ok(URL_SAFE_NO_PAD.encode(bytes)) +} + +fn query_token_from_uri(uri: &Uri) -> Result, &'static str> { + let Some(query) = uri.query() else { + return Ok(None); + }; + let mut token = None; + for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { + if key != "token" { + continue; + } + if token.is_some() { + return Err("multiple websocket query tokens"); + } + token = Some(value.into_owned()); } + Ok(token) } fn verify_signed_bearer_token( @@ -481,27 +587,6 @@ mod tests { 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 { @@ -530,11 +615,11 @@ mod tests { assert_eq!( settings, AppServerWebsocketAuthSettings { - config: Some(AppServerWebsocketAuthConfig::CapabilityToken { + config: AppServerWebsocketAuthConfig::CapabilityToken { source: AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256: [0xab; 32], }, - }), + }, } ); } @@ -573,28 +658,79 @@ mod tests { #[test] fn capability_token_hash_policy_authorizes_matching_bearer_token() { let settings = AppServerWebsocketAuthSettings { - config: Some(AppServerWebsocketAuthConfig::CapabilityToken { + config: AppServerWebsocketAuthConfig::CapabilityToken { source: AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256: sha256_digest(b"super-secret-token"), }, - }), + }, }; let policy = policy_from_settings(&settings).expect("hash policy should build"); + assert_eq!(policy.generated_query_token(), None); + let uri = Uri::from_static("/"); let mut headers = HeaderMap::new(); headers.insert( AUTHORIZATION, HeaderValue::from_static("Bearer super-secret-token"), ); - authorize_upgrade(&headers, &policy).expect("matching token should authorize"); + authorize_upgrade(&uri, &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"); + let err = authorize_upgrade(&uri, &headers, &policy).expect_err("wrong token should fail"); assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); } + #[test] + fn generated_query_token_is_required_by_default() { + let policy = policy_from_settings(&AppServerWebsocketAuthSettings::default()) + .expect("generated token policy should build"); + let token = policy + .generated_query_token() + .expect("query-token policy should expose its token"); + + let valid_uri: Uri = format!("/?token={token}").parse().expect("valid token URI"); + authorize_upgrade(&valid_uri, &HeaderMap::new(), &policy) + .expect("matching query token should authorize"); + + for uri in [ + Uri::from_static("/"), + Uri::from_static("/?token=wrong"), + Uri::from_static("/?token=one&token=two"), + ] { + let err = authorize_upgrade(&uri, &HeaderMap::new(), &policy) + .expect_err("missing or invalid query token should fail"); + assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED); + } + } + + #[test] + fn no_token_check_allows_invalid_tokens() { + let settings = AppServerWebsocketAuthArgs { + no_token_check: true, + ..Default::default() + } + .try_into_settings() + .expect("no-token-check args should parse"); + assert_eq!( + settings, + AppServerWebsocketAuthSettings { + config: AppServerWebsocketAuthConfig::QueryToken { enforce: false }, + } + ); + let policy = policy_from_settings(&settings).expect("generated token policy should build"); + + for uri in [ + Uri::from_static("/"), + Uri::from_static("/?token=wrong"), + Uri::from_static("/?token=one&token=two"), + ] { + authorize_upgrade(&uri, &HeaderMap::new(), &policy) + .expect("missing or invalid token should be allowed"); + } + } + #[test] fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() { let err = AppServerWebsocketAuthArgs { @@ -624,13 +760,13 @@ mod tests { assert_eq!( settings, AppServerWebsocketAuthSettings { - config: Some(AppServerWebsocketAuthConfig::SignedBearerToken { + config: 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, - }), + }, } ); } diff --git a/codex-rs/app-server-transport/src/transport/websocket.rs b/codex-rs/app-server-transport/src/transport/websocket.rs index 5677faae672c..8011b9a89292 100644 --- a/codex-rs/app-server-transport/src/transport/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/websocket.rs @@ -3,7 +3,6 @@ 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; @@ -19,6 +18,7 @@ use axum::extract::ws::WebSocketUpgrade; use axum::http::HeaderMap; use axum::http::Request; use axum::http::StatusCode; +use axum::http::Uri; use axum::http::header::ORIGIN; use axum::middleware; use axum::middleware::Next; @@ -54,10 +54,14 @@ fn colorize(text: &str, style: Style) -> String { } #[allow(clippy::print_stderr)] -fn print_websocket_startup_banner(addr: SocketAddr) { +fn print_websocket_startup_banner(addr: SocketAddr, auth_policy: &WebsocketAuthPolicy) { 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 raw_listen_url = match auth_policy.generated_query_token() { + Some(token) => format!("ws://{addr}/?token={token}"), + None => format!("ws://{addr}"), + }; + let listen_url = colorize(&raw_listen_url, 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()); @@ -106,9 +110,10 @@ async fn websocket_upgrade_handler( websocket: WebSocketUpgrade, ConnectInfo(peer_addr): ConnectInfo, State(state): State, + uri: Uri, headers: HeaderMap, ) -> impl IntoResponse { - if let Err(err) = authorize_upgrade(&headers, state.auth_policy.as_ref()) { + if let Err(err) = authorize_upgrade(&uri, &headers, state.auth_policy.as_ref()) { warn!( %peer_addr, message = err.message(), @@ -132,17 +137,9 @@ pub async fn start_websocket_acceptor( 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); + print_websocket_startup_banner(local_addr, &auth_policy); info!("app-server websocket listening on ws://{local_addr}"); let router = Router::new() diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index c9ac2f5fb27c..0c8f8120f569 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -34,6 +34,20 @@ When running with `--listen ws://IP:PORT`, the same listener also serves basic H - `GET /healthz` returns `200 OK` when no `Origin` header is present. - Any request carrying an `Origin` header is rejected with `403 Forbidden`. +By default, app-server uses query-token authentication: it generates a 256-bit connection token +and prints the complete connection URL to stderr, for example: + +```text +ws://127.0.0.1:4500/?token= +``` + +By default, clients must provide that exact `token` query parameter during the websocket upgrade. +Treat the printed URL as a secret. Pass `--no-token-check` to accept clients that omit the token or +provide an incorrect token; app-server still prints the tokenized URL. + +When `--ws-auth capability-token` or `--ws-auth signed-bearer-token` is configured, app-server +prints the token-free listen URL and requires `Authorization: Bearer ` during upgrade. + 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` 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 648d60f3b594..ea7f05963b02 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 @@ -160,6 +160,96 @@ async fn websocket_transport_rejects_browser_origin_without_auth() -> Result<()> Ok(()) } +#[tokio::test] +async fn websocket_transport_requires_generated_query_token() -> 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, token) = + spawn_websocket_server_with_generated_token(codex_home.path(), "ws://127.0.0.1:0", &[]) + .await?; + + assert_websocket_query_connect_rejected(bind_addr, /*query_token*/ None).await?; + assert_websocket_query_connect_rejected(bind_addr, Some("wrong-token")).await?; + + let mut ws = connect_websocket_with_query_token(bind_addr, &token).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_generated_token_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_uses_remote_ws_token() -> 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 expected_token = "configured-remote-ws-token"; + + let (mut process, bind_addr, token) = spawn_websocket_server_and_read_generated_token( + codex_home.path(), + "ws://127.0.0.1:0", + &[], + &[("CODEX_REMOTE_WS_TOKEN", expected_token)], + ) + .await?; + assert_eq!(token.as_deref(), Some(expected_token)); + + let mut ws = connect_websocket_with_query_token(bind_addr, expected_token).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_remote_token_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_no_token_check_accepts_invalid_tokens() -> 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, token) = spawn_websocket_server_with_generated_token( + codex_home.path(), + "ws://127.0.0.1:0", + &["--no-token-check".to_string()], + ) + .await?; + + let mut without_token = connect_websocket(bind_addr).await?; + send_initialize_request(&mut without_token, /*id*/ 1, "ws_missing_token_client").await?; + read_response_for_id(&mut without_token, /*id*/ 1).await?; + + let mut wrong_token = connect_websocket_with_query_token(bind_addr, "wrong-token").await?; + send_initialize_request(&mut wrong_token, /*id*/ 2, "ws_wrong_token_client").await?; + read_response_for_id(&mut wrong_token, /*id*/ 2).await?; + + let mut matching_token = connect_websocket_with_query_token(bind_addr, &token).await?; + send_initialize_request( + &mut matching_token, + /*id*/ 3, + "ws_matching_token_client", + ) + .await?; + read_response_for_id(&mut matching_token, /*id*/ 3).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; @@ -322,24 +412,24 @@ async fn websocket_transport_rejects_short_signed_bearer_secret_configuration() } #[tokio::test] -async fn websocket_transport_rejects_unauthenticated_non_loopback_startup() -> Result<()> { +async fn websocket_transport_authenticates_non_loopback_by_default() -> 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", &[]) + let (mut process, bind_addr, token) = + spawn_websocket_server_with_generated_token(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}" - ); + assert_websocket_query_connect_rejected(bind_addr, /*query_token*/ None).await?; + let mut ws = connect_websocket_with_query_token(bind_addr, &token).await?; + send_initialize_request(&mut ws, /*id*/ 1, "ws_non_loopback_token_client").await?; + read_response_for_id(&mut ws, /*id*/ 1).await?; + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; Ok(()) } @@ -376,7 +466,12 @@ async fn websocket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_tim } 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 + spawn_websocket_server_with_args( + codex_home, + "ws://127.0.0.1:0", + &["--no-token-check".to_string()], + ) + .await } pub(super) async fn spawn_websocket_server_with_args( @@ -384,6 +479,30 @@ pub(super) async fn spawn_websocket_server_with_args( listen_url: &str, extra_args: &[String], ) -> Result<(Child, SocketAddr)> { + let (process, bind_addr, _token) = + spawn_websocket_server_and_read_generated_token(codex_home, listen_url, extra_args, &[]) + .await?; + Ok((process, bind_addr)) +} + +async fn spawn_websocket_server_with_generated_token( + codex_home: &Path, + listen_url: &str, + extra_args: &[String], +) -> Result<(Child, SocketAddr, String)> { + let (process, bind_addr, token) = + spawn_websocket_server_and_read_generated_token(codex_home, listen_url, extra_args, &[]) + .await?; + let token = token.context("websocket app-server did not print a generated query token")?; + Ok((process, bind_addr, token)) +} + +async fn spawn_websocket_server_and_read_generated_token( + codex_home: &Path, + listen_url: &str, + extra_args: &[String], + extra_env: &[(&str, &str)], +) -> Result<(Child, SocketAddr, Option)> { let program = codex_utils_cargo_bin::cargo_bin("codex-app-server") .context("should find app-server binary")?; let mut cmd = Command::new(program); @@ -395,6 +514,7 @@ pub(super) async fn spawn_websocket_server_with_args( .stdout(Stdio::null()) .stderr(Stdio::piped()) .env("CODEX_HOME", codex_home) + .envs(extra_env.iter().copied()) .env("RUST_LOG", "warn"); let mut process = cmd .kill_on_drop(true) @@ -407,7 +527,7 @@ pub(super) async fn spawn_websocket_server_with_args( .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 (bind_addr, generated_token) = loop { let line = timeout( deadline.saturating_duration_since(Instant::now()), stderr_reader.next_line(), @@ -436,13 +556,25 @@ pub(super) async fn spawn_websocket_server_with_args( stripped }; - if let Some(bind_addr) = stripped_line + let Some(raw_url) = stripped_line .split_whitespace() .find_map(|token| token.strip_prefix("ws://")) - .and_then(|addr| addr.parse::().ok()) - { - break bind_addr; - } + else { + continue; + }; + let Ok(url) = url::Url::parse(&format!("ws://{raw_url}")) else { + continue; + }; + let Ok(socket_addrs) = url.socket_addrs(|| None) else { + continue; + }; + let Some(bind_addr) = socket_addrs.first().copied() else { + continue; + }; + let generated_token = url + .query_pairs() + .find_map(|(key, value)| (key == "token").then(|| value.into_owned())); + break (bind_addr, generated_token); }; tokio::spawn(async move { @@ -451,7 +583,7 @@ pub(super) async fn spawn_websocket_server_with_args( } }); - Ok((process, bind_addr)) + Ok((process, bind_addr, generated_token)) } pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result { @@ -478,6 +610,63 @@ pub(super) async fn connect_websocket_with_bearer( } } +async fn connect_websocket_with_query_token( + bind_addr: SocketAddr, + query_token: &str, +) -> Result { + let url = format!( + "ws://{}/?token={query_token}", + connectable_bind_addr(bind_addr) + ); + let request = websocket_request( + url.as_str(), + /*bearer_token*/ None, + /*origin*/ None, + )?; + let deadline = Instant::now() + DEFAULT_READ_TIMEOUT; + loop { + 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; + } + } + } +} + +async fn assert_websocket_query_connect_rejected( + bind_addr: SocketAddr, + query_token: Option<&str>, +) -> Result<()> { + let path_and_query = query_token + .map(|token| format!("/?token={token}")) + .unwrap_or_default(); + let url = format!("ws://{}{path_and_query}", connectable_bind_addr(bind_addr)); + let request = websocket_request( + url.as_str(), + /*bearer_token*/ None, + /*origin*/ None, + )?; + + match connect_async(request).await { + Ok((_stream, response)) => { + bail!( + "expected websocket handshake rejection, got {}", + response.status() + ) + } + // The pinned tungstenite fork can replace an otherwise valid rejected + // upgrade response with a synthetic 400 when it normalizes a response + // that lacks websocket response headers. Unit coverage verifies that + // app-server's authorization result itself is 401. + Err(WsError::Http(_response)) => Ok(()), + Err(err) => bail!("expected http rejection during websocket handshake: {err}"), + } +} + async fn assert_websocket_connect_rejected( bind_addr: SocketAddr, bearer_token: Option<&str>,