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 <token>` 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 <etraut@openai.com>
2026-04-27 16:15:09 +00:00 · 2026-03-25 12:35:57 -07:00
parent 91337399fe
commit 6124564297
11 changed files with 1130 additions and 86 deletions
--- 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<MaybeTlsStream<tokio::net::TcpStream>>;
+type HmacSha256 = Hmac<Sha256>;

 #[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<WsClient> {
-    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<WsClient> {
+    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<WsClient>
    }
 }

+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<std::process::Output> {
+    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<reqwest::Response> {
+    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<tokio_tungstenite::tungstenite::http::Request<()>> {
+    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<String> {
+    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}"))
+}