From 6124564297289ddadd5ab364bb961ea367fff8d0 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 25 Mar 2026 12:35:57 -0700 Subject: [PATCH] feat: add websocket auth for app-server (#14847) ## Summary This change adds websocket authentication at the app-server transport boundary and enforces it before JSON-RPC `initialize`, so authenticated deployments reject unauthenticated clients during the websocket handshake rather than after a connection has already been admitted. During rollout, websocket auth is opt-in for non-loopback listeners so we do not break existing remote clients. If `--ws-auth ...` is configured, the server enforces auth during websocket upgrade. If auth is not configured, non-loopback listeners still start, but app-server logs a warning and the startup banner calls out that auth should be configured before real remote use. The server supports two auth modes: a file-backed capability token, and a standard HMAC-signed JWT/JWS bearer token verified with the `jsonwebtoken` crate, with optional issuer, audience, and clock-skew validation. Capability tokens are normalized, hashed, and compared in constant time. Short shared secrets for signed bearer tokens are rejected at startup. Requests carrying an `Origin` header are rejected with `403` by transport middleware, and authenticated clients present credentials as `Authorization: Bearer ` during websocket upgrade. ## Validation - `cargo test -p codex-app-server transport::auth` - `cargo test -p codex-cli app_server_` - `cargo clippy -p codex-app-server --all-targets -- -D warnings` - `just bazel-lock-check` Note: in the broad `cargo test -p codex-app-server connection_handling_websocket` run, the touched websocket auth cases passed, but unrelated Unix shutdown tests failed with a timeout in this environment. --------- Co-authored-by: Eric Traut --- MODULE.bazel.lock | 2 + codex-rs/Cargo.lock | 31 + codex-rs/Cargo.toml | 3 + codex-rs/app-server/Cargo.toml | 4 + codex-rs/app-server/README.md | 9 + codex-rs/app-server/src/lib.rs | 7 + codex-rs/app-server/src/main.rs | 6 + codex-rs/app-server/src/transport.rs | 34 +- codex-rs/app-server/src/transport/auth.rs | 583 ++++++++++++++++++ .../suite/v2/connection_handling_websocket.rs | 375 +++++++++-- codex-rs/cli/src/main.rs | 162 +++-- 11 files changed, 1130 insertions(+), 86 deletions(-) create mode 100644 codex-rs/app-server/src/transport/auth.rs diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index e0d384767..ee89f6243 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -971,6 +971,7 @@ "jni_0.21.1": "{\"dependencies\":[{\"name\":\"cesu8\",\"req\":\"^1.1.0\"},{\"name\":\"cfg-if\",\"req\":\"^1.0.0\"},{\"name\":\"combine\",\"req\":\"^4.1.0\"},{\"name\":\"java-locator\",\"optional\":true,\"req\":\"^0.1\"},{\"name\":\"jni-sys\",\"req\":\"^0.3.0\"},{\"name\":\"libloading\",\"optional\":true,\"req\":\"^0.7\"},{\"name\":\"log\",\"req\":\"^0.4.4\"},{\"name\":\"thiserror\",\"req\":\"^1.0.20\"},{\"kind\":\"dev\",\"name\":\"assert_matches\",\"req\":\"^1.5.0\"},{\"kind\":\"dev\",\"name\":\"lazy_static\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"rusty-fork\",\"req\":\"^0.3.0\"},{\"kind\":\"build\",\"name\":\"walkdir\",\"req\":\"^2\"},{\"features\":[\"Win32_Globalization\"],\"name\":\"windows-sys\",\"req\":\"^0.45.0\",\"target\":\"cfg(windows)\"},{\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.13.0\",\"target\":\"cfg(windows)\"}],\"features\":{\"default\":[],\"invocation\":[\"java-locator\",\"libloading\"]}}", "jobserver_0.1.34": "{\"dependencies\":[{\"features\":[\"std\"],\"name\":\"getrandom\",\"req\":\"^0.3.2\",\"target\":\"cfg(windows)\"},{\"name\":\"libc\",\"req\":\"^0.2.171\",\"target\":\"cfg(unix)\"},{\"features\":[\"fs\"],\"kind\":\"dev\",\"name\":\"nix\",\"req\":\"^0.28.0\",\"target\":\"cfg(unix)\"},{\"kind\":\"dev\",\"name\":\"tempfile\",\"req\":\"^3.10.1\"}],\"features\":{}}", "js-sys_0.3.85": "{\"dependencies\":[{\"default_features\":false,\"name\":\"once_cell\",\"req\":\"^1.12\"},{\"default_features\":false,\"name\":\"wasm-bindgen\",\"req\":\"=0.2.108\"}],\"features\":{\"default\":[\"std\"],\"std\":[\"wasm-bindgen/std\"]}}", + "jsonwebtoken_9.3.1": "{\"dependencies\":[{\"name\":\"base64\",\"req\":\"^0.22\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.4\",\"target\":\"cfg(all(target_arch = \\\"wasm32\\\", not(any(target_os = \\\"emscripten\\\", target_os = \\\"wasi\\\"))))\"},{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.4\",\"target\":\"cfg(not(all(target_arch = \\\"wasm32\\\", not(any(target_os = \\\"emscripten\\\", target_os = \\\"wasi\\\")))))\"},{\"name\":\"js-sys\",\"req\":\"^0.3\",\"target\":\"cfg(target_arch = \\\"wasm32\\\")\"},{\"name\":\"pem\",\"optional\":true,\"req\":\"^3\"},{\"features\":[\"std\"],\"name\":\"ring\",\"req\":\"^0.17.4\",\"target\":\"cfg(not(target_arch = \\\"wasm32\\\"))\"},{\"features\":[\"std\",\"wasm32_unknown_unknown_js\"],\"name\":\"ring\",\"req\":\"^0.17.4\",\"target\":\"cfg(target_arch = \\\"wasm32\\\")\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"req\":\"^1.0\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"simple_asn1\",\"optional\":true,\"req\":\"^0.6\"},{\"features\":[\"wasm-bindgen\"],\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3\",\"target\":\"cfg(all(target_arch = \\\"wasm32\\\", not(any(target_os = \\\"emscripten\\\", target_os = \\\"wasi\\\"))))\"},{\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3\",\"target\":\"cfg(not(all(target_arch = \\\"wasm32\\\", not(any(target_os = \\\"emscripten\\\", target_os = \\\"wasi\\\")))))\"},{\"kind\":\"dev\",\"name\":\"wasm-bindgen-test\",\"req\":\"^0.3.1\"}],\"features\":{\"default\":[\"use_pem\"],\"use_pem\":[\"pem\",\"simple_asn1\"]}}", "keyring_3.6.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"base64\",\"req\":\"^0.22\"},{\"name\":\"byteorder\",\"optional\":true,\"req\":\"^1.2\",\"target\":\"cfg(target_os = \\\"windows\\\")\"},{\"features\":[\"derive\",\"wrap_help\"],\"kind\":\"dev\",\"name\":\"clap\",\"req\":\"^4\"},{\"name\":\"dbus-secret-service\",\"optional\":true,\"req\":\"^4.0.0-rc.1\",\"target\":\"cfg(target_os = \\\"openbsd\\\")\"},{\"name\":\"dbus-secret-service\",\"optional\":true,\"req\":\"^4.0.0-rc.2\",\"target\":\"cfg(target_os = \\\"linux\\\")\"},{\"name\":\"dbus-secret-service\",\"optional\":true,\"req\":\"^4.0.1\",\"target\":\"cfg(target_os = \\\"freebsd\\\")\"},{\"kind\":\"dev\",\"name\":\"doc-comment\",\"req\":\"^0.3\"},{\"kind\":\"dev\",\"name\":\"env_logger\",\"req\":\"^0.11.5\"},{\"kind\":\"dev\",\"name\":\"fastrand\",\"req\":\"^2\"},{\"features\":[\"std\"],\"name\":\"linux-keyutils\",\"optional\":true,\"req\":\"^0.2\",\"target\":\"cfg(target_os = \\\"linux\\\")\"},{\"name\":\"log\",\"req\":\"^0.4.22\"},{\"name\":\"openssl\",\"optional\":true,\"req\":\"^0.10.66\"},{\"kind\":\"dev\",\"name\":\"rpassword\",\"req\":\"^7\"},{\"kind\":\"dev\",\"name\":\"rprompt\",\"req\":\"^2\"},{\"name\":\"secret-service\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"freebsd\\\")\"},{\"name\":\"secret-service\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"linux\\\")\"},{\"name\":\"secret-service\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"openbsd\\\")\"},{\"name\":\"security-framework\",\"optional\":true,\"req\":\"^2\",\"target\":\"cfg(target_os = \\\"ios\\\")\"},{\"name\":\"security-framework\",\"optional\":true,\"req\":\"^3\",\"target\":\"cfg(target_os = \\\"macos\\\")\"},{\"kind\":\"dev\",\"name\":\"whoami\",\"req\":\"^1.5\"},{\"features\":[\"Win32_Foundation\",\"Win32_Security_Credentials\"],\"name\":\"windows-sys\",\"optional\":true,\"req\":\"^0.60\",\"target\":\"cfg(target_os = \\\"windows\\\")\"},{\"name\":\"zbus\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"freebsd\\\")\"},{\"name\":\"zbus\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"linux\\\")\"},{\"name\":\"zbus\",\"optional\":true,\"req\":\"^4\",\"target\":\"cfg(target_os = \\\"openbsd\\\")\"},{\"name\":\"zeroize\",\"req\":\"^1.8.1\",\"target\":\"cfg(target_os = \\\"windows\\\")\"}],\"features\":{\"apple-native\":[\"dep:security-framework\"],\"async-io\":[\"zbus?/async-io\"],\"async-secret-service\":[\"dep:secret-service\",\"dep:zbus\"],\"crypto-openssl\":[\"dbus-secret-service?/crypto-openssl\",\"secret-service?/crypto-openssl\"],\"crypto-rust\":[\"dbus-secret-service?/crypto-rust\",\"secret-service?/crypto-rust\"],\"linux-native\":[\"dep:linux-keyutils\"],\"linux-native-async-persistent\":[\"linux-native\",\"async-secret-service\"],\"linux-native-sync-persistent\":[\"linux-native\",\"sync-secret-service\"],\"sync-secret-service\":[\"dep:dbus-secret-service\"],\"tokio\":[\"zbus?/tokio\"],\"vendored\":[\"dbus-secret-service?/vendored\",\"openssl?/vendored\"],\"windows-native\":[\"dep:windows-sys\",\"dep:byteorder\"]}}", "kqueue-sys_1.0.4": "{\"dependencies\":[{\"name\":\"bitflags\",\"req\":\"^1.2.1\"},{\"name\":\"libc\",\"req\":\"^0.2.74\"}],\"features\":{}}", "kqueue_1.1.1": "{\"dependencies\":[{\"features\":[\"html_reports\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.5\"},{\"kind\":\"dev\",\"name\":\"dhat\",\"req\":\"^0.3.2\"},{\"name\":\"kqueue-sys\",\"req\":\"^1.0.4\"},{\"name\":\"libc\",\"req\":\"^0.2.17\"},{\"kind\":\"dev\",\"name\":\"tempfile\",\"req\":\"^3.1.0\"}],\"features\":{}}", @@ -1281,6 +1282,7 @@ "simd-adler32_0.3.8": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"adler\",\"req\":\"^1.0.2\"},{\"kind\":\"dev\",\"name\":\"adler32\",\"req\":\"^1.2.0\"},{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"}],\"features\":{\"const-generics\":[],\"default\":[\"std\",\"const-generics\"],\"nightly\":[],\"std\":[]}}", "simdutf8_0.1.5": "{\"dependencies\":[],\"features\":{\"aarch64_neon\":[],\"aarch64_neon_prefetch\":[],\"default\":[\"std\"],\"hints\":[],\"public_imp\":[],\"std\":[]}}", "similar_2.7.0": "{\"dependencies\":[{\"default_features\":false,\"name\":\"bstr\",\"optional\":true,\"req\":\"^1.5.0\"},{\"kind\":\"dev\",\"name\":\"console\",\"req\":\"^0.15.0\"},{\"kind\":\"dev\",\"name\":\"insta\",\"req\":\"^1.10.0\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.130\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.68\"},{\"name\":\"unicode-segmentation\",\"optional\":true,\"req\":\"^1.7.1\"},{\"name\":\"web-time\",\"optional\":true,\"req\":\"^1.1\"}],\"features\":{\"bytes\":[\"bstr\",\"text\"],\"default\":[\"text\"],\"inline\":[\"text\"],\"text\":[],\"unicode\":[\"text\",\"unicode-segmentation\",\"bstr?/unicode\",\"bstr?/std\"],\"wasm32_web_time\":[\"web-time\"]}}", + "simple_asn1_0.6.4": "{\"dependencies\":[{\"default_features\":false,\"name\":\"num-bigint\",\"req\":\"^0.4\"},{\"default_features\":false,\"name\":\"num-traits\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"quickcheck\",\"req\":\"^1.0.3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8.4\"},{\"default_features\":false,\"name\":\"thiserror\",\"req\":\"^2\"},{\"default_features\":false,\"features\":[\"formatting\",\"macros\",\"parsing\"],\"name\":\"time\",\"req\":\"^0.3.47\"},{\"default_features\":false,\"features\":[\"formatting\",\"macros\",\"parsing\",\"quickcheck\"],\"kind\":\"dev\",\"name\":\"time\",\"req\":\"^0.3\"}],\"features\":{}}", "siphasher_1.0.2": "{\"dependencies\":[{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"name\":\"serde_json\",\"optional\":true,\"req\":\"^1.0\"}],\"features\":{\"default\":[\"std\"],\"serde_no_std\":[\"serde/alloc\"],\"serde_std\":[\"std\",\"serde/std\"],\"std\":[]}}", "slab_0.4.12": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.95\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1\"}],\"features\":{\"default\":[\"std\"],\"std\":[]}}", "smallvec_1.15.1": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1\"},{\"default_features\":false,\"name\":\"bincode\",\"optional\":true,\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"bincode1\",\"package\":\"bincode\",\"req\":\"^1.0.1\"},{\"kind\":\"dev\",\"name\":\"debugger_test\",\"req\":\"^0.1.0\"},{\"kind\":\"dev\",\"name\":\"debugger_test_parser\",\"req\":\"^0.1.0\"},{\"default_features\":false,\"name\":\"malloc_size_of\",\"optional\":true,\"req\":\"^0.1\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1\"},{\"default_features\":false,\"name\":\"unty\",\"optional\":true,\"req\":\"^0.0.4\"}],\"features\":{\"const_generics\":[],\"const_new\":[\"const_generics\"],\"debugger_visualizer\":[],\"drain_filter\":[],\"drain_keep_rest\":[\"drain_filter\"],\"impl_bincode\":[\"bincode\",\"unty\"],\"may_dangle\":[],\"specialization\":[],\"union\":[],\"write\":[]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index df44a7f83..187fcca79 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1473,8 +1473,11 @@ dependencies = [ "codex-utils-cli", "codex-utils-json-to-toml", "codex-utils-pty", + "constant_time_eq", "core_test_support", "futures", + "hmac", + "jsonwebtoken", "opentelemetry", "opentelemetry_sdk", "owo-colors", @@ -1484,6 +1487,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2", "shlex", "tempfile", "time", @@ -5759,6 +5763,21 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "9.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a87cc7a48537badeae96744432de36f4be2b4a34a05a5ef32e9dd8a1c169dde" +dependencies = [ + "base64 0.22.1", + "js-sys", + "pem", + "ring", + "serde", + "serde_json", + "simple_asn1", +] + [[package]] name = "keyring" version = "3.6.3" @@ -9249,6 +9268,18 @@ version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" +[[package]] +name = "simple_asn1" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d585997b0ac10be3c5ee635f1bab02d512760d14b7c468801ac8a01d9ae5f1d" +dependencies = [ + "num-bigint", + "num-traits", + "thiserror 2.0.18", + "time", +] + [[package]] name = "siphasher" version = "1.0.2" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 10b9485d2..350a98e45 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -202,6 +202,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" crossterm = "0.28.1" csv = "1.3.1" @@ -220,6 +221,7 @@ flate2 = "1.1.4" futures = { version = "0.3", default-features = false } gethostname = "1.1.0" globset = "0.4" +hmac = "0.12.1" http = "1.3.1" icu_decimal = "2.1" icu_locale_core = "2.1" @@ -232,6 +234,7 @@ indexmap = "2.12.0" insta = "1.46.3" inventory = "0.3.19" itertools = "0.14.0" +jsonwebtoken = "9.3.1" keyring = { version = "3.6", default-features = false } landlock = "0.4.4" lazy_static = "1" diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 1eb09e260..f54bb7ad0 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -53,10 +53,14 @@ codex-utils-absolute-path = { workspace = true } codex-utils-json-to-toml = { workspace = true } chrono = { workspace = true } clap = { workspace = true, features = ["derive"] } +constant_time_eq = { workspace = true } futures = { 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 } tempfile = { workspace = true } time = { workspace = true } toml = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 1e186c0d5..a1856af46 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -34,6 +34,15 @@ When running with `--listen ws://IP:PORT`, the same listener also serves basic H Websocket transport is currently experimental and unsupported. Do not rely on it for production workloads. +Security note: + +- Loopback websocket listeners (`ws://127.0.0.1:PORT`) remain appropriate for localhost and SSH port-forwarding workflows. +- Non-loopback websocket listeners currently allow unauthenticated connections by default during rollout. If you expose one remotely, configure websocket auth explicitly now. +- Supported auth modes are app-server flags: + - `--ws-auth capability-token --ws-token-file /absolute/path` + - `--ws-auth signed-bearer-token --ws-shared-secret-file /absolute/path` for HMAC-signed JWT/JWS bearer tokens, with optional `--ws-issuer`, `--ws-audience`, `--ws-max-clock-skew-seconds` +- Clients present the credential as `Authorization: Bearer ` during the websocket handshake. Auth is enforced before JSON-RPC `initialize`. + Tracing/log output: - `RUST_LOG` controls log filtering/verbosity. diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index dc2512aae..247270a26 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -27,6 +27,7 @@ use crate::transport::CHANNEL_CAPACITY; use crate::transport::ConnectionState; use crate::transport::OutboundConnectionState; use crate::transport::TransportEvent; +use crate::transport::auth::policy_from_settings; use crate::transport::route_outgoing_envelope; use crate::transport::start_stdio_connection; use crate::transport::start_websocket_acceptor; @@ -81,6 +82,9 @@ mod transport; 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::auth::AppServerWebsocketAuthArgs; +pub use crate::transport::auth::AppServerWebsocketAuthSettings; +pub use crate::transport::auth::WebsocketAuthCliMode; const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT"; @@ -337,6 +341,7 @@ pub async fn run_main( default_analytics_enabled, AppServerTransport::Stdio, SessionSource::VSCode, + AppServerWebsocketAuthSettings::default(), ) .await } @@ -348,6 +353,7 @@ pub async fn run_main_with_transport( default_analytics_enabled: bool, transport: AppServerTransport, session_source: SessionSource, + auth: AppServerWebsocketAuthSettings, ) -> IoResult<()> { let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); @@ -375,6 +381,7 @@ pub async fn run_main_with_transport( bind_address, transport_event_tx.clone(), shutdown_token.clone(), + policy_from_settings(&auth)?, ) .await?; TransportRuntime::WebSocket { diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index 60fa0a777..fa95f973e 100644 --- a/codex-rs/app-server/src/main.rs +++ b/codex-rs/app-server/src/main.rs @@ -1,5 +1,6 @@ use clap::Parser; use codex_app_server::AppServerTransport; +use codex_app_server::AppServerWebsocketAuthArgs; use codex_app_server::run_main_with_transport; use codex_arg0::Arg0DispatchPaths; use codex_arg0::arg0_dispatch_or_else; @@ -31,6 +32,9 @@ struct AppServerArgs { value_parser = SessionSource::from_startup_arg )] session_source: SessionSource, + + #[command(flatten)] + auth: AppServerWebsocketAuthArgs, } fn main() -> anyhow::Result<()> { @@ -43,6 +47,7 @@ fn main() -> anyhow::Result<()> { }; let transport = args.listen; let session_source = args.session_source; + let auth = args.auth.try_into_settings()?; run_main_with_transport( arg0_paths, @@ -51,6 +56,7 @@ fn main() -> anyhow::Result<()> { /*default_analytics_enabled*/ false, transport, session_source, + auth, ) .await?; Ok(()) diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 21fbc9965..020e441da 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -1,3 +1,8 @@ +pub(crate) mod auth; + +use self::auth::WebsocketAuthPolicy; +use self::auth::authorize_upgrade; +use self::auth::should_warn_about_unauthenticated_non_loopback_listener; use crate::error_code::OVERLOADED_ERROR_CODE; use crate::message_processor::ConnectionSessionState; use crate::outgoing_message::ConnectionId; @@ -12,6 +17,7 @@ use axum::extract::State; use axum::extract::ws::Message as WebSocketMessage; use axum::extract::ws::WebSocket; use axum::extract::ws::WebSocketUpgrade; +use axum::http::HeaderMap; use axum::http::Request; use axum::http::StatusCode; use axum::http::header::ORIGIN; @@ -83,7 +89,7 @@ fn print_websocket_startup_banner(addr: SocketAddr) { ); } else { eprintln!( - " {note_label} this is a raw WS server; consider running behind TLS/auth for real remote use" + " {note_label} websocket auth is opt-in in this build; configure `--ws-auth ...` before real remote use" ); } } @@ -92,6 +98,7 @@ fn print_websocket_startup_banner(addr: SocketAddr) { struct WebSocketListenerState { transport_event_tx: mpsc::Sender, connection_counter: Arc, + auth_policy: Arc, } async fn health_check_handler() -> StatusCode { @@ -118,12 +125,23 @@ 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(); + } let connection_id = ConnectionId(state.connection_counter.fetch_add(1, Ordering::Relaxed)); info!(%peer_addr, "websocket client connected"); - websocket.on_upgrade(move |stream| async move { - run_websocket_connection(connection_id, stream, state.transport_event_tx).await; - }) + websocket + .on_upgrade(move |stream| async move { + run_websocket_connection(connection_id, stream, state.transport_event_tx).await; + }) + .into_response() } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -333,7 +351,14 @@ pub(crate) async fn start_websocket_acceptor( bind_address: SocketAddr, transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, + auth_policy: WebsocketAuthPolicy, ) -> IoResult> { + if should_warn_about_unauthenticated_non_loopback_listener(bind_address, &auth_policy) { + warn!( + %bind_address, + "starting non-loopback websocket listener without auth; websocket auth is opt-in for now and will become the default in a future release" + ); + } let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; print_websocket_startup_banner(local_addr); @@ -347,6 +372,7 @@ pub(crate) async fn start_websocket_acceptor( .with_state(WebSocketListenerState { transport_event_tx, connection_counter: Arc::new(AtomicU64::new(1)), + auth_policy: Arc::new(auth_policy), }); let server = axum::serve( listener, diff --git a/codex-rs/app-server/src/transport/auth.rs b/codex-rs/app-server/src/transport/auth.rs new file mode 100644 index 000000000..1c6080e74 --- /dev/null +++ b/codex-rs/app-server/src/transport/auth.rs @@ -0,0 +1,583 @@ +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, + + /// 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 { + token_file: AbsolutePathBuf, + }, + SignedBearerToken { + shared_secret_file: AbsolutePathBuf, + issuer: Option, + audience: Option, + max_clock_skew_seconds: u64, + }, +} + +#[derive(Clone, Debug, Default)] +pub(crate) 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 token_file = self.ws_token_file.context( + "`--ws-token-file` is required when `--ws-auth capability-token` is set", + )?; + Some(AppServerWebsocketAuthConfig::CapabilityToken { + token_file: absolute_path_arg("--ws-token-file", token_file)?, + }) + } + Some(WebsocketAuthCliMode::SignedBearerToken) => { + if self.ws_token_file.is_some() { + anyhow::bail!( + "`--ws-token-file` requires `--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_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(crate) fn policy_from_settings( + settings: &AppServerWebsocketAuthSettings, +) -> io::Result { + let mode = match settings.config.as_ref() { + Some(AppServerWebsocketAuthConfig::CapabilityToken { token_file }) => { + let token = read_trimmed_secret(token_file.as_ref())?; + Some(WebsocketAuthMode::CapabilityToken { + token_sha256: sha256_digest(token.as_bytes()), + }) + } + 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 should_warn_about_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(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 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 warns_about_unauthenticated_non_loopback_listener() { + let policy = WebsocketAuthPolicy::default(); + assert!(should_warn_about_unauthenticated_non_loopback_listener( + "0.0.0.0:8765".parse().unwrap(), + &policy, + )); + assert!(!should_warn_about_unauthenticated_non_loopback_listener( + "127.0.0.1:8765".parse().unwrap(), + &policy, + )); + assert!(!should_warn_about_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() { + let err = AppServerWebsocketAuthArgs { + ws_auth: Some(WebsocketAuthCliMode::CapabilityToken), + ..Default::default() + } + .try_into_settings() + .expect_err("capability-token mode should require a token file"); + assert!( + err.to_string().contains("--ws-token-file"), + "unexpected error: {err}" + ); + } + + #[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, None, None, 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"), 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, None, Some("audience"), 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", None, None, 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, None, None, 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/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index c5c76bf8e..19158c68c 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 @@ -2,6 +2,8 @@ use anyhow::Context; use anyhow::Result; use anyhow::bail; use app_test_support::create_mock_responses_server_sequence_unchecked; +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; @@ -12,12 +14,16 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; 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::process::Stdio; use tempfile::TempDir; +use time::OffsetDateTime; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Child; @@ -29,15 +35,17 @@ use tokio::time::timeout; use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async; -use tokio_tungstenite::tungstenite::Error as WebSocketError; +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; pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); pub(super) type WsClient = WebSocketStream>; +type HmacSha256 = Hmac; #[tokio::test] async fn websocket_transport_routes_per_connection_handshake_and_responses() -> Result<()> { @@ -112,46 +120,26 @@ async fn websocket_transport_serves_health_endpoints_on_same_listener() -> Resul } #[tokio::test] -async fn websocket_transport_rejects_requests_with_origin_header() -> Result<()> { +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 client = reqwest::Client::new(); - let deadline = Instant::now() + Duration::from_secs(10); - let healthz = loop { - match client - .get(format!("http://{bind_addr}/healthz")) - .header(ORIGIN.as_str(), "https://example.com") - .send() - .await - .with_context(|| format!("failed to GET http://{bind_addr}/healthz with Origin header")) - { - Ok(response) => break response, - Err(err) => { - if Instant::now() >= deadline { - bail!("failed to GET http://{bind_addr}/healthz with Origin header: {err}"); - } - sleep(Duration::from_millis(50)).await; - } - } - }; - assert_eq!(healthz.status(), StatusCode::FORBIDDEN); + let mut ws = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws, 1, "ws_loopback_client").await?; + let init = read_response_for_id(&mut ws, 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + drop(ws); - let url = format!("ws://{bind_addr}"); - let mut request = url.into_client_request()?; - request - .headers_mut() - .insert(ORIGIN, HeaderValue::from_static("https://example.com")); - match connect_async(request).await { - Err(WebSocketError::Http(response)) => { - assert_eq!(response.status(), StatusCode::FORBIDDEN); - } - Ok(_) => bail!("expected websocket handshake with Origin header to be rejected"), - Err(err) => bail!("expected HTTP rejection for Origin header, got {err}"), - } + assert_websocket_connect_rejected_with_headers( + bind_addr, + None, + Some("https://evil.example"), + StatusCode::FORBIDDEN, + ) + .await?; process .kill() @@ -160,12 +148,205 @@ async fn websocket_transport_rejects_requests_with_origin_header() -> Result<()> 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://127.0.0.1:0", &auth_args).await?; + + assert_websocket_connect_rejected(bind_addr, 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, 1, "ws_auth_client").await?; + let init = read_response_for_id(&mut ws, 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, 1, "ws_signed_auth_client").await?; + let init = read_response_for_id(&mut ws, 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_allows_unauthenticated_non_loopback_startup_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 (mut process, bind_addr) = + spawn_websocket_server_with_args(codex_home.path(), "ws://0.0.0.0:0", &[]).await?; + + let mut ws = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws, 1, "ws_non_loopback_default_client").await?; + let init = read_response_for_id(&mut ws, 1).await?; + assert_eq!(init.id, RequestId::Integer(1)); + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + + Ok(()) +} + 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")?; let mut cmd = Command::new(program); cmd.arg("--listen") - .arg("ws://127.0.0.1:0") + .arg(listen_url) + .args(extra_args) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::piped()) @@ -230,10 +411,18 @@ pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, } pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result { - let url = format!("ws://{bind_addr}"); + connect_websocket_with_bearer(bind_addr, 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, None)?; let deadline = Instant::now() + Duration::from_secs(10); loop { - match connect_async(&url).await { + match connect_async(request.clone()).await { Ok((stream, _response)) => return Ok(stream), Err(err) => { if Instant::now() >= deadline { @@ -245,23 +434,83 @@ pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> 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, + 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) + .args(extra_args) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .env("CODEX_HOME", codex_home) + .env("RUST_LOG", "debug"); + timeout(Duration::from_secs(10), 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() + Duration::from_secs(10); loop { match client - .get(format!("http://{bind_addr}{path}")) + .get(format!("http://{connectable_bind_addr}{path}")) .send() .await - .with_context(|| format!("failed to GET http://{bind_addr}{path}")) + .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 GET http://{bind_addr}{path}: {err}"); + bail!("failed to GET http://{connectable_bind_addr}{path}: {err}"); } sleep(Duration::from_millis(50)).await; } @@ -269,6 +518,30 @@ async fn http_get( } } +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, @@ -459,3 +732,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/cli/src/main.rs b/codex-rs/cli/src/main.rs index 82ede6932..5e0d405eb 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -353,6 +353,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, clap::Subcommand)] @@ -643,49 +646,59 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { prepend_config_flags(&mut mcp_cli.config_overrides, root_config_overrides.clone()); mcp_cli.run().await?; } - Some(Subcommand::AppServer(app_server_cli)) => match app_server_cli.subcommand { - None => { - reject_remote_mode_for_subcommand(root_remote.as_deref(), "app-server")?; - let transport = app_server_cli.listen; - codex_app_server::run_main_with_transport( - arg0_paths.clone(), - root_config_overrides, - codex_core::config_loader::LoaderOverrides::default(), - app_server_cli.analytics_default_enabled, - transport, - codex_protocol::protocol::SessionSource::VSCode, - ) - .await?; + Some(Subcommand::AppServer(app_server_cli)) => { + let AppServerCommand { + subcommand, + listen, + analytics_default_enabled, + auth, + } = app_server_cli; + match subcommand { + None => { + reject_remote_mode_for_subcommand(root_remote.as_deref(), "app-server")?; + let transport = listen; + let auth = auth.try_into_settings()?; + codex_app_server::run_main_with_transport( + arg0_paths.clone(), + root_config_overrides, + codex_core::config_loader::LoaderOverrides::default(), + analytics_default_enabled, + transport, + codex_protocol::protocol::SessionSource::VSCode, + auth, + ) + .await?; + } + Some(AppServerSubcommand::GenerateTs(gen_cli)) => { + reject_remote_mode_for_subcommand( + root_remote.as_deref(), + "app-server generate-ts", + )?; + let options = codex_app_server_protocol::GenerateTsOptions { + experimental_api: gen_cli.experimental, + ..Default::default() + }; + codex_app_server_protocol::generate_ts_with_options( + &gen_cli.out_dir, + gen_cli.prettier.as_deref(), + options, + )?; + } + Some(AppServerSubcommand::GenerateJsonSchema(gen_cli)) => { + reject_remote_mode_for_subcommand( + root_remote.as_deref(), + "app-server generate-json-schema", + )?; + codex_app_server_protocol::generate_json_with_experimental( + &gen_cli.out_dir, + gen_cli.experimental, + )?; + } + Some(AppServerSubcommand::GenerateInternalJsonSchema(gen_cli)) => { + codex_app_server_protocol::generate_internal_json_schema(&gen_cli.out_dir)?; + } } - Some(AppServerSubcommand::GenerateTs(gen_cli)) => { - reject_remote_mode_for_subcommand( - root_remote.as_deref(), - "app-server generate-ts", - )?; - let options = codex_app_server_protocol::GenerateTsOptions { - experimental_api: gen_cli.experimental, - ..Default::default() - }; - codex_app_server_protocol::generate_ts_with_options( - &gen_cli.out_dir, - gen_cli.prettier.as_deref(), - options, - )?; - } - Some(AppServerSubcommand::GenerateJsonSchema(gen_cli)) => { - reject_remote_mode_for_subcommand( - root_remote.as_deref(), - "app-server generate-json-schema", - )?; - codex_app_server_protocol::generate_json_with_experimental( - &gen_cli.out_dir, - gen_cli.experimental, - )?; - } - Some(AppServerSubcommand::GenerateInternalJsonSchema(gen_cli)) => { - codex_app_server_protocol::generate_internal_json_schema(&gen_cli.out_dir)?; - } - }, + } #[cfg(target_os = "macos")] Some(Subcommand::App(app_cli)) => { reject_remote_mode_for_subcommand(root_remote.as_deref(), "app")?; @@ -1701,6 +1714,71 @@ mod tests { assert!(parse_result.is_err()); } + #[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([ + "codex", + "app-server", + "--allow-unauthenticated-non-loopback-ws", + ]); + assert!(parse_result.is_err()); + } + #[test] fn features_enable_parses_feature_name() { let cli = MultitoolCli::try_parse_from(["codex", "features", "enable", "unified_exec"])