openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Extract WebSocket authentication into codex-websocket-auth (#47447)

Browse Source 
## What changed

Move WebSocket authentication arguments, settings, credential loading, and upgrade authorization into a shared crate using `http` types. Update the app server, transport, and CLI to consume it directly, and remove the `AppServer` prefix from the shared types.

Hide token digests and JWT verification secrets from authentication policy `Debug` output.

## Testing

Move the existing argument validation, capability-token authorization, and signed bearer-token verification tests into the new crate.

GitOrigin-RevId: 1dd20043a3f8efbca55af7eda338b62a62e8fd68
This commit is contained in:
Ruslan Nigmatullin
2026-09-23 04:25:28 +00:00
committed by copyberry
parent 4981037e99
commit c44deff7b1
18 changed files with 422 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+22 -5
View File
@@ -2116,6 +2116,7 @@ dependencies = [
"codex-utils-path-uri",
"codex-utils-pty",
"codex-web-search-extension",
"codex-websocket-auth",
"codex-windows-sandbox",
"core_test_support",
"flate2",
@@ -2267,11 +2268,9 @@ dependencies = [
name = "codex-app-server-transport"
version = "0.0.0"
dependencies = [
"anyhow",
"axum",
"base64 0.22.1",
"chrono",
"clap",
"codex-api",
"codex-app-server-protocol",
"codex-config",
@@ -2284,13 +2283,11 @@ dependencies = [
"codex-uds",
"codex-utils-absolute-path",
"codex-utils-rustls-provider",
"codex-websocket-auth",
"codex-websocket-client",
"constant_time_eq",
"futures",
"gethostname",
"hmac 0.12.1",
"httpdate",
"jsonwebtoken",
"libc",
"owo-colors",
"pretty_assertions",
@@ -2531,6 +2528,7 @@ dependencies = [
"codex-utils-path",
"codex-utils-path-uri",
"codex-utils-pty",
"codex-websocket-auth",
"codex-windows-sandbox",
"crossterm",
"flate2",
@@ -5149,6 +5147,25 @@ dependencies = [
"url",
]
[[package]]
name = "codex-websocket-auth"
version = "0.0.0"
dependencies = [
"anyhow",
"base64 0.22.1",
"clap",
"codex-utils-absolute-path",
"constant_time_eq",
"hmac 0.12.1",
"http 1.4.0",
"jsonwebtoken",
"pretty_assertions",
"serde",
"serde_json",
"sha2 0.10.9",
"time",
]
[[package]]
name = "codex-websocket-client"
version = "0.0.0"
codex-rs/Cargo.toml
+2
View File
@@ -108,6 +108,7 @@ members = [
"user-verification",
"tools",
"v8-poc",
"websocket-auth",
"websocket-client",
"windows-sandbox-service",
"worktree",
@@ -197,6 +198,7 @@ codex-code-mode-protocol = { path = "code-mode-protocol" }
codex-code-mode-runtime = { path = "code-mode-runtime" }
codex-home = { path = "codex-home" }
codex-http-client = { path = "http-client" }
codex-websocket-auth = { path = "websocket-auth" }
codex-websocket-client = { path = "websocket-client" }
codex-config = { path = "config" }
codex-connectors = { path = "connectors" }
codex-rs/app-server-transport/Cargo.toml
+1 -5
View File
@@ -15,7 +15,6 @@ workspace = true
[dependencies]
codex-http-client = { workspace = true }
codex-websocket-client = { workspace = true }
anyhow = { workspace = true }
axum = { workspace = true, default-features = false, features = [
"http1",
"json",
@@ -23,7 +22,6 @@ axum = { workspace = true, default-features = false, features = [
"ws",
] }
base64 = { workspace = true }
clap = { workspace = true, features = ["derive"] }
codex-api = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-core = { workspace = true }
@@ -34,12 +32,10 @@ codex-state = { workspace = true }
codex-uds = { workspace = true }
codex-utils-absolute-path = { workspace = true }
codex-utils-rustls-provider = { workspace = true }
constant_time_eq = { workspace = true }
codex-websocket-auth = { workspace = true }
futures = { workspace = true }
gethostname = { workspace = true }
hmac = { workspace = true }
httpdate = { workspace = true }
jsonwebtoken = { workspace = true }
owo-colors = { workspace = true, features = ["supports-colors"] }
rand = { workspace = true }
serde = { workspace = true, features = ["derive"] }
codex-rs/app-server-transport/src/lib.rs
-1
View File
@@ -35,7 +35,6 @@ pub use transport::TransportEvent;
pub use transport::acquire_app_server_startup_lock;
pub use transport::app_server_control_socket_path;
pub use transport::app_server_startup_lock_path;
pub use transport::auth;
pub use transport::daemon_recovery_file_path;
pub use transport::start_control_socket_acceptor;
pub use transport::start_remote_control;
codex-rs/app-server-transport/src/transport/mod.rs
-2
View File
@@ -1,5 +1,3 @@
pub mod auth;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::OutgoingError;
use crate::outgoing_message::OutgoingMessage;
codex-rs/app-server-transport/src/transport/websocket.rs
+3 -3
View File
@@ -1,9 +1,6 @@
use super::CHANNEL_CAPACITY;
use super::ConnectionOrigin;
use super::TransportEvent;
use super::auth::WebsocketAuthPolicy;
use super::auth::authorize_upgrade;
use super::auth::is_unauthenticated_non_loopback_listener;
use super::forward_incoming_message;
use super::next_connection_id;
use super::serialize_outgoing_message;
@@ -26,6 +23,9 @@ use axum::response::IntoResponse;
use axum::response::Response;
use axum::routing::any;
use axum::routing::get;
use codex_websocket_auth::WebsocketAuthPolicy;
use codex_websocket_auth::authorize_upgrade;
use codex_websocket_auth::is_unauthenticated_non_loopback_listener;
use futures::SinkExt;
use futures::StreamExt;
use owo_colors::OwoColorize;
codex-rs/app-server/Cargo.toml
+1
View File
@@ -73,6 +73,7 @@ codex-login = { workspace = true }
codex-image-generation-extension = { workspace = true }
codex-memories-extension = { workspace = true }
codex-web-search-extension = { workspace = true }
codex-websocket-auth = { workspace = true }
codex-memories-write = { workspace = true }
codex-mcp = { workspace = true }
codex-mcp-extension = { workspace = true }
codex-rs/app-server/src/lib.rs
+4 -6
View File
@@ -42,7 +42,6 @@ use crate::transport::RemoteControlStartConfig;
use crate::transport::TransportEvent;
use crate::transport::acquire_app_server_startup_lock;
use crate::transport::app_server_startup_lock_path;
use crate::transport::auth::policy_from_settings;
use crate::transport::route_outgoing_envelope;
use crate::transport::start_control_socket_acceptor;
use crate::transport::start_remote_control;
@@ -68,6 +67,8 @@ use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use codex_rollout::state_db as rollout_state_db;
use codex_state::log_db;
use codex_websocket_auth::WebsocketAuthSettings;
use codex_websocket_auth::policy_from_settings;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
@@ -147,9 +148,6 @@ pub use crate::error_code::INVALID_PARAMS_ERROR_CODE;
pub use crate::transport::AppServerTransport;
pub use crate::transport::RemoteControlStartupMode;
pub use crate::transport::app_server_control_socket_path;
pub use crate::transport::auth::AppServerWebsocketAuthArgs;
pub use crate::transport::auth::AppServerWebsocketAuthSettings;
pub use crate::transport::auth::WebsocketAuthCliMode;
pub use crate::transport::take_remote_control_disabled_env;
const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT";
@@ -444,7 +442,7 @@ pub async fn run_main(
default_analytics_enabled,
AppServerTransport::Stdio,
SessionSource::VSCode,
AppServerWebsocketAuthSettings::default(),
WebsocketAuthSettings::default(),
AppServerRuntimeOptions::default(),
)
.await
@@ -495,7 +493,7 @@ pub async fn run_main_with_transport_options(
default_analytics_enabled: bool,
transport: AppServerTransport,
session_source: SessionSource,
auth: AppServerWebsocketAuthSettings,
auth: WebsocketAuthSettings,
runtime_options: AppServerRuntimeOptions,
) -> IoResult<AppServerExit> {
#[cfg(target_os = "windows")]
codex-rs/app-server/src/main.rs
+2 -2
View File
@@ -4,7 +4,6 @@ use clap::Parser;
use codex_app_server::AppServerCodeModeHostArgs;
use codex_app_server::AppServerRuntimeOptions;
use codex_app_server::AppServerTransport;
use codex_app_server::AppServerWebsocketAuthArgs;
use codex_app_server::PluginStartupTasks;
use codex_app_server::run_main_with_transport_options;
use codex_arg0::Arg0DispatchPaths;
@@ -12,6 +11,7 @@ use codex_arg0::arg0_dispatch_or_else;
use codex_config::LoaderOverrides;
use codex_protocol::protocol::SessionSource;
use codex_utils_cli::CliConfigOverrides;
use codex_websocket_auth::WebsocketAuthArgs;
use std::path::PathBuf;
#[cfg(all(
@@ -55,7 +55,7 @@ struct AppServerArgs {
session_source: SessionSource,
#[command(flatten)]
auth: AppServerWebsocketAuthArgs,
auth: WebsocketAuthArgs,
/// Fail if config.toml contains unknown configuration fields.
#[arg(long = "strict-config", default_value_t = false)]
codex-rs/app-server/src/transport.rs
-1
View File
@@ -29,7 +29,6 @@ pub(crate) use codex_app_server_transport::TransportEvent;
pub(crate) use codex_app_server_transport::acquire_app_server_startup_lock;
pub use codex_app_server_transport::app_server_control_socket_path;
pub(crate) use codex_app_server_transport::app_server_startup_lock_path;
pub use codex_app_server_transport::auth;
pub(crate) use codex_app_server_transport::start_control_socket_acceptor;
pub(crate) use codex_app_server_transport::start_remote_control;
pub(crate) use codex_app_server_transport::start_stdio_connection;
codex-rs/app-server/tests/suite/v2/remote_control.rs
+2 -2
View File
@@ -15,7 +15,6 @@ use app_test_support::to_response;
use app_test_support::write_chatgpt_auth;
use codex_app_server::AppServerRuntimeOptions;
use codex_app_server::AppServerTransport;
use codex_app_server::AppServerWebsocketAuthSettings;
use codex_app_server::PluginStartupTasks;
use codex_app_server::RemoteControlStartupMode;
use codex_app_server::run_main_with_transport_options;
@@ -47,6 +46,7 @@ use codex_protocol::protocol::SessionSource;
use codex_state::RemoteControlEnrollmentRecord;
use codex_state::StateRuntime;
use codex_utils_cli::CliConfigOverrides;
use codex_websocket_auth::WebsocketAuthSettings;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
@@ -241,7 +241,7 @@ async fn explicit_remote_control_startup_fails_when_disabled_by_requirements() -
/*default_analytics_enabled*/ false,
transport,
SessionSource::VSCode,
AppServerWebsocketAuthSettings::default(),
WebsocketAuthSettings::default(),
AppServerRuntimeOptions {
plugin_startup_tasks: PluginStartupTasks::Skip,
remote_control_startup_mode: RemoteControlStartupMode::EnabledEphemeral,
codex-rs/cli/Cargo.toml
+1
View File
@@ -75,6 +75,7 @@ codex-tui = { workspace = true }
codex-utils-absolute-path = { workspace = true }
codex-utils-path = { workspace = true }
codex-utils-path-uri = { workspace = true }
codex-websocket-auth = { workspace = true }
crossterm = { workspace = true, features = ["event-stream"] }
futures = { workspace = true }
http = { workspace = true }
codex-rs/cli/src/main.rs
+3 -3
View File
@@ -602,7 +602,7 @@ struct AppServerCommand {
analytics_default_enabled: bool,
#[command(flatten)]
auth: codex_app_server::AppServerWebsocketAuthArgs,
auth: codex_websocket_auth::WebsocketAuthArgs,
}
#[derive(Debug, clap::Subcommand)]
@@ -4718,7 +4718,7 @@ mod tests {
);
assert_eq!(
app_server.auth.ws_auth,
Some(codex_app_server::WebsocketAuthCliMode::CapabilityToken)
Some(codex_websocket_auth::WebsocketAuthCliMode::CapabilityToken)
);
assert_eq!(
app_server.auth.ws_token_file,
@@ -4747,7 +4747,7 @@ mod tests {
);
assert_eq!(
app_server.auth.ws_auth,
Some(codex_app_server::WebsocketAuthCliMode::SignedBearerToken)
Some(codex_websocket_auth::WebsocketAuthCliMode::SignedBearerToken)
);
assert_eq!(
app_server.auth.ws_shared_secret_file,
codex-rs/cli/src/remote_control_cmd.rs
+2 -2
View File
@@ -5,7 +5,6 @@ use anyhow::Context;
use clap::Args;
use codex_app_server::AppServerRuntimeOptions;
use codex_app_server::AppServerTransport;
use codex_app_server::AppServerWebsocketAuthSettings;
use codex_app_server_daemon::LifecycleCommand as AppServerLifecycleCommand;
use codex_app_server_daemon::LifecycleOutput as AppServerLifecycleOutput;
use codex_app_server_daemon::LifecycleStatus as AppServerLifecycleStatus;
@@ -19,6 +18,7 @@ use codex_config::LoaderOverrides;
use codex_protocol::protocol::SessionSource;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_cli::CliConfigOverrides;
use codex_websocket_auth::WebsocketAuthSettings;
use serde::Serialize;
use tokio::sync::watch;
use tokio::task::JoinHandle;
@@ -137,7 +137,7 @@ async fn run_foreground_remote_control(
/*default_analytics_enabled*/ false,
transport,
SessionSource::VSCode,
AppServerWebsocketAuthSettings::default(),
WebsocketAuthSettings::default(),
runtime_options,
);
let mut app_server_task = tokio::spawn(async move { app_server.await.map(|_| ()) });
codex-rs/websocket-auth/BUILD.bazel
+6
View File
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "websocket-auth",
crate_name = "codex_websocket_auth",
)
codex-rs/websocket-auth/Cargo.toml
+28
View File
@@ -0,0 +1,28 @@
[package]
name = "codex-websocket-auth"
version.workspace = true
edition.workspace = true
license.workspace = true
[lib]
doctest = false
[lints]
workspace = true
[dependencies]
anyhow = { workspace = true }
clap = { workspace = true, features = ["derive"] }
codex-utils-absolute-path = { workspace = true }
constant_time_eq = { workspace = true }
http = { workspace = true }
jsonwebtoken = { workspace = true }
serde = { workspace = true, features = ["derive"] }
sha2 = { workspace = true }
time = { workspace = true }
[dev-dependencies]
base64 = { workspace = true }
hmac = { workspace = true }
pretty_assertions = { workspace = true }
serde_json = { workspace = true }
codex-rs/websocket-auth/src/auth_tests.rs
+288
View File
@@ -0,0 +1,288 @@
use super::*;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use hmac::Hmac;
use hmac::Mac;
use http::HeaderValue;
use pretty_assertions::assert_eq;
use serde_json::json;
type HmacSha256 = Hmac<Sha256>;
fn signed_token(shared_secret: &[u8], claims: serde_json::Value) -> String {
let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#);
let claims_segment = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap());
let payload = format!("{header}.{claims_segment}");
let mut mac = HmacSha256::new_from_slice(shared_secret).unwrap();
mac.update(payload.as_bytes());
let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
format!("{payload}.{signature}")
}
#[test]
fn detects_unauthenticated_non_loopback_listener() {
let policy = WebsocketAuthPolicy::default();
assert!(is_unauthenticated_non_loopback_listener(
"0.0.0.0:8765".parse().unwrap(),
&policy,
));
assert!(!is_unauthenticated_non_loopback_listener(
"127.0.0.1:8765".parse().unwrap(),
&policy,
));
assert!(!is_unauthenticated_non_loopback_listener(
"0.0.0.0:8765".parse().unwrap(),
&WebsocketAuthPolicy {
mode: Some(WebsocketAuthMode::CapabilityToken {
token_sha256: [0u8; 32],
}),
},
));
}
#[test]
fn capability_token_args_require_token_file_or_hash() {
let err = WebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should require a token source");
assert!(
err.to_string().contains("--ws-token-file")
&& err.to_string().contains("--ws-token-sha256"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_args_accept_token_hash() {
let settings = WebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_sha256: Some("ab".repeat(32)),
..Default::default()
}
.try_into_settings()
.expect("capability-token hash args should parse");
assert_eq!(
settings,
WebsocketAuthSettings {
config: Some(WebsocketAuthConfig::CapabilityToken {
source: WebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: [0xab; 32],
},
}),
}
);
}
#[test]
fn capability_token_args_reject_multiple_token_sources() {
let err = WebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_file: Some(PathBuf::from("/tmp/token")),
ws_token_sha256: Some("ab".repeat(32)),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should reject multiple token sources");
assert!(
err.to_string().contains("mutually exclusive"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_args_reject_malformed_token_hash() {
let err = WebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_sha256: Some("not-a-sha256".to_string()),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should reject malformed token hashes");
assert!(
err.to_string().contains("64-character hex"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_hash_policy_authorizes_matching_bearer_token() {
let settings = WebsocketAuthSettings {
config: Some(WebsocketAuthConfig::CapabilityToken {
source: WebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: sha256_digest(b"super-secret-token"),
},
}),
};
let policy = policy_from_settings(&settings).expect("hash policy should build");
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer super-secret-token"),
);
authorize_upgrade(&headers, &policy).expect("matching token should authorize");
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer wrong-token"),
);
let err = authorize_upgrade(&headers, &policy).expect_err("wrong token should fail");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() {
let err = WebsocketAuthArgs {
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 = WebsocketAuthArgs {
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,
WebsocketAuthSettings {
config: Some(WebsocketAuthConfig::SignedBearerToken {
shared_secret_file: AbsolutePathBuf::from_absolute_path("/tmp/secret")
.expect("absolute path"),
issuer: Some("issuer".to_string()),
audience: None,
max_clock_skew_seconds: DEFAULT_MAX_CLOCK_SKEW_SECONDS,
}),
}
);
}
#[test]
fn signed_bearer_token_verification_rejects_tampering() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
}),
);
let tampered = token.replace(".eyJleHAi", ".eyJleHBi");
let err = verify_signed_bearer_token(
&tampered,
shared_secret,
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("tampered jwt should fail");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_token_verification_accepts_valid_token() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
"iss": "issuer",
"aud": "audience",
}),
);
verify_signed_bearer_token(
&token,
shared_secret,
Some("issuer"),
Some("audience"),
/*max_clock_skew_seconds*/ 30,
)
.expect("valid signed token should verify");
}
#[test]
fn signed_bearer_token_verification_accepts_multiple_audiences() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
"aud": ["other-audience", "audience"],
}),
);
verify_signed_bearer_token(
&token,
shared_secret,
/*issuer*/ None,
Some("audience"),
/*max_clock_skew_seconds*/ 30,
)
.expect("jwt audience arrays should verify");
}
#[test]
fn signed_bearer_token_verification_rejects_alg_none_tokens() {
let claims_segment = URL_SAFE_NO_PAD.encode(
serde_json::to_vec(&json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
}))
.unwrap(),
);
let header_segment = URL_SAFE_NO_PAD.encode(br#"{"alg":"none","typ":"JWT"}"#);
let token = format!("{header_segment}.{claims_segment}.");
let err = verify_signed_bearer_token(
&token,
b"0123456789abcdef0123456789abcdef",
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("alg=none jwt should be rejected");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_token_verification_rejects_missing_exp() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"iss": "issuer",
}),
);
let err = verify_signed_bearer_token(
&token,
shared_secret,
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("jwt without exp should be rejected");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn validate_signed_bearer_secret_rejects_short_secret() {
let err = validate_signed_bearer_secret(Path::new("/tmp/secret"), b"too-short")
.expect_err("short shared secret should be rejected");
assert_eq!(err.kind(), ErrorKind::InvalidInput);
assert!(
err.to_string().contains("must be at least 32 bytes"),
"unexpected error: {err}"
);
}
codex-rs/app-server-transport/src/transport/auth.rs → codex-rs/websocket-auth/src/lib.rs
+57 -329
View File
@@ -1,11 +1,14 @@
//! Shared authentication for incoming app-server and exec-server WebSocket upgrades.
//! Policies retain token digests or JWT verification secrets; callers choose when auth is required.
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 http::HeaderMap;
use http::StatusCode;
use http::header::AUTHORIZATION;
use jsonwebtoken::Algorithm;
use jsonwebtoken::DecodingKey;
use jsonwebtoken::Validation;
@@ -25,8 +28,8 @@ 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.
pub struct WebsocketAuthArgs {
/// Authentication mode for incoming WebSocket connections.
#[arg(long = "ws-auth", value_name = "MODE", value_enum)]
pub ws_auth: Option<WebsocketAuthCliMode>,
@@ -62,14 +65,14 @@ pub enum WebsocketAuthCliMode {
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AppServerWebsocketAuthSettings {
pub config: Option<AppServerWebsocketAuthConfig>,
pub struct WebsocketAuthSettings {
pub config: Option<WebsocketAuthConfig>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AppServerWebsocketAuthConfig {
pub enum WebsocketAuthConfig {
CapabilityToken {
source: AppServerWebsocketCapabilityTokenSource,
source: WebsocketCapabilityTokenSource,
},
SignedBearerToken {
shared_secret_file: AbsolutePathBuf,
@@ -80,18 +83,20 @@ pub enum AppServerWebsocketAuthConfig {
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AppServerWebsocketCapabilityTokenSource {
pub enum WebsocketCapabilityTokenSource {
TokenFile { token_file: AbsolutePathBuf },
TokenSha256 { token_sha256: [u8; 32] },
}
/// Loaded upgrade policy. The default permits unauthenticated connections; each
/// listener decides whether that policy is allowed for its bind address.
#[derive(Clone, Debug, Default)]
pub struct WebsocketAuthPolicy {
pub(crate) mode: Option<WebsocketAuthMode>,
mode: Option<WebsocketAuthMode>,
}
#[derive(Clone, Debug)]
pub(crate) enum WebsocketAuthMode {
#[derive(Clone)]
enum WebsocketAuthMode {
CapabilityToken {
token_sha256: [u8; 32],
},
@@ -103,8 +108,21 @@ pub(crate) enum WebsocketAuthMode {
},
}
impl std::fmt::Debug for WebsocketAuthMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::CapabilityToken { .. } => {
f.debug_struct("CapabilityToken").finish_non_exhaustive()
}
Self::SignedBearerToken { .. } => {
f.debug_struct("SignedBearerToken").finish_non_exhaustive()
}
}
}
}
#[derive(Debug)]
pub(crate) struct WebsocketAuthError {
pub struct WebsocketAuthError {
status_code: StatusCode,
message: &'static str,
}
@@ -125,17 +143,17 @@ enum JwtAudienceClaim {
}
impl WebsocketAuthError {
pub(crate) fn status_code(&self) -> StatusCode {
pub fn status_code(&self) -> StatusCode {
self.status_code
}
pub(crate) fn message(&self) -> &'static str {
pub fn message(&self) -> &'static str {
self.message
}
}
impl AppServerWebsocketAuthArgs {
pub fn try_into_settings(self) -> anyhow::Result<AppServerWebsocketAuthSettings> {
impl WebsocketAuthArgs {
pub fn try_into_settings(self) -> anyhow::Result<WebsocketAuthSettings> {
let normalize = |value: Option<String>| {
value.and_then(|value| {
let trimmed = value.trim();
@@ -160,23 +178,19 @@ impl AppServerWebsocketAuthArgs {
"`--ws-token-file` and `--ws-token-sha256` are mutually exclusive"
);
}
(Some(token_file), None) => {
AppServerWebsocketCapabilityTokenSource::TokenFile {
token_file: absolute_path_arg("--ws-token-file", token_file)?,
}
}
(None, Some(token_sha256)) => {
AppServerWebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: sha256_digest_arg("--ws-token-sha256", &token_sha256)?,
}
}
(Some(token_file), None) => WebsocketCapabilityTokenSource::TokenFile {
token_file: absolute_path_arg("--ws-token-file", token_file)?,
},
(None, Some(token_sha256)) => WebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: sha256_digest_arg("--ws-token-sha256", &token_sha256)?,
},
(None, None) => {
anyhow::bail!(
"`--ws-token-file` or `--ws-token-sha256` is required when `--ws-auth capability-token` is set"
);
}
};
Some(AppServerWebsocketAuthConfig::CapabilityToken { source })
Some(WebsocketAuthConfig::CapabilityToken { source })
}
Some(WebsocketAuthCliMode::SignedBearerToken) => {
if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() {
@@ -187,7 +201,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 {
Some(WebsocketAuthConfig::SignedBearerToken {
shared_secret_file: absolute_path_arg(
"--ws-shared-secret-file",
shared_secret_file,
@@ -215,28 +229,27 @@ impl AppServerWebsocketAuthArgs {
}
};
Ok(AppServerWebsocketAuthSettings { config })
Ok(WebsocketAuthSettings { config })
}
}
pub fn policy_from_settings(
settings: &AppServerWebsocketAuthSettings,
) -> io::Result<WebsocketAuthPolicy> {
/// Loads credentials once; changes to secret files do not update an existing policy.
pub fn policy_from_settings(settings: &WebsocketAuthSettings) -> io::Result<WebsocketAuthPolicy> {
let mode = match settings.config.as_ref() {
Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source {
AppServerWebsocketCapabilityTokenSource::TokenFile { token_file } => {
Some(WebsocketAuthConfig::CapabilityToken { source }) => match source {
WebsocketCapabilityTokenSource::TokenFile { token_file } => {
let token = read_trimmed_secret(token_file.as_ref())?;
Some(WebsocketAuthMode::CapabilityToken {
token_sha256: sha256_digest(token.as_bytes()),
})
}
AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => {
WebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => {
Some(WebsocketAuthMode::CapabilityToken {
token_sha256: *token_sha256,
})
}
},
Some(AppServerWebsocketAuthConfig::SignedBearerToken {
Some(WebsocketAuthConfig::SignedBearerToken {
shared_secret_file,
issuer,
audience,
@@ -263,14 +276,16 @@ pub fn policy_from_settings(
Ok(WebsocketAuthPolicy { mode })
}
pub(crate) fn is_unauthenticated_non_loopback_listener(
pub 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(
/// Checks the bearer credential before upgrading. This does not enforce expiry
/// or revocation after a WebSocket has been established.
pub fn authorize_upgrade(
headers: &HeaderMap,
policy: &WebsocketAuthPolicy,
) -> Result<(), WebsocketAuthError> {
@@ -460,292 +475,5 @@ fn unauthorized(message: &'static str) -> WebsocketAuthError {
}
#[cfg(test)]
mod tests {
use super::*;
use axum::http::HeaderValue;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use hmac::Hmac;
use hmac::Mac;
use serde_json::json;
type HmacSha256 = Hmac<Sha256>;
fn signed_token(shared_secret: &[u8], claims: serde_json::Value) -> String {
let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#);
let claims_segment = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap());
let payload = format!("{header}.{claims_segment}");
let mut mac = HmacSha256::new_from_slice(shared_secret).unwrap();
mac.update(payload.as_bytes());
let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
format!("{payload}.{signature}")
}
#[test]
fn detects_unauthenticated_non_loopback_listener() {
let policy = WebsocketAuthPolicy::default();
assert!(is_unauthenticated_non_loopback_listener(
"0.0.0.0:8765".parse().unwrap(),
&policy,
));
assert!(!is_unauthenticated_non_loopback_listener(
"127.0.0.1:8765".parse().unwrap(),
&policy,
));
assert!(!is_unauthenticated_non_loopback_listener(
"0.0.0.0:8765".parse().unwrap(),
&WebsocketAuthPolicy {
mode: Some(WebsocketAuthMode::CapabilityToken {
token_sha256: [0u8; 32],
}),
},
));
}
#[test]
fn capability_token_args_require_token_file_or_hash() {
let err = AppServerWebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should require a token source");
assert!(
err.to_string().contains("--ws-token-file")
&& err.to_string().contains("--ws-token-sha256"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_args_accept_token_hash() {
let settings = AppServerWebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_sha256: Some("ab".repeat(32)),
..Default::default()
}
.try_into_settings()
.expect("capability-token hash args should parse");
assert_eq!(
settings,
AppServerWebsocketAuthSettings {
config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: [0xab; 32],
},
}),
}
);
}
#[test]
fn capability_token_args_reject_multiple_token_sources() {
let err = AppServerWebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_file: Some(PathBuf::from("/tmp/token")),
ws_token_sha256: Some("ab".repeat(32)),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should reject multiple token sources");
assert!(
err.to_string().contains("mutually exclusive"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_args_reject_malformed_token_hash() {
let err = AppServerWebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
ws_token_sha256: Some("not-a-sha256".to_string()),
..Default::default()
}
.try_into_settings()
.expect_err("capability-token mode should reject malformed token hashes");
assert!(
err.to_string().contains("64-character hex"),
"unexpected error: {err}"
);
}
#[test]
fn capability_token_hash_policy_authorizes_matching_bearer_token() {
let settings = AppServerWebsocketAuthSettings {
config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
token_sha256: sha256_digest(b"super-secret-token"),
},
}),
};
let policy = policy_from_settings(&settings).expect("hash policy should build");
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer super-secret-token"),
);
authorize_upgrade(&headers, &policy).expect("matching token should authorize");
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer wrong-token"),
);
let err = authorize_upgrade(&headers, &policy).expect_err("wrong token should fail");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() {
let err = AppServerWebsocketAuthArgs {
ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")),
..Default::default()
}
.try_into_settings()
.expect_err("mode-specific flags should require --ws-auth");
assert!(
err.to_string().contains("websocket auth flags require"),
"unexpected error: {err}"
);
}
#[test]
fn signed_bearer_args_default_clock_skew_and_trim_optional_claims() {
let settings = AppServerWebsocketAuthArgs {
ws_auth: Some(WebsocketAuthCliMode::SignedBearerToken),
ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")),
ws_issuer: Some(" issuer ".to_string()),
ws_audience: Some(" ".to_string()),
..Default::default()
}
.try_into_settings()
.expect("signed bearer args should parse");
assert_eq!(
settings,
AppServerWebsocketAuthSettings {
config: Some(AppServerWebsocketAuthConfig::SignedBearerToken {
shared_secret_file: AbsolutePathBuf::from_absolute_path("/tmp/secret")
.expect("absolute path"),
issuer: Some("issuer".to_string()),
audience: None,
max_clock_skew_seconds: DEFAULT_MAX_CLOCK_SKEW_SECONDS,
}),
}
);
}
#[test]
fn signed_bearer_token_verification_rejects_tampering() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
}),
);
let tampered = token.replace(".eyJleHAi", ".eyJleHBi");
let err = verify_signed_bearer_token(
&tampered,
shared_secret,
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("tampered jwt should fail");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_token_verification_accepts_valid_token() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
"iss": "issuer",
"aud": "audience",
}),
);
verify_signed_bearer_token(
&token,
shared_secret,
Some("issuer"),
Some("audience"),
/*max_clock_skew_seconds*/ 30,
)
.expect("valid signed token should verify");
}
#[test]
fn signed_bearer_token_verification_accepts_multiple_audiences() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
"aud": ["other-audience", "audience"],
}),
);
verify_signed_bearer_token(
&token,
shared_secret,
/*issuer*/ None,
Some("audience"),
/*max_clock_skew_seconds*/ 30,
)
.expect("jwt audience arrays should verify");
}
#[test]
fn signed_bearer_token_verification_rejects_alg_none_tokens() {
let claims_segment = URL_SAFE_NO_PAD.encode(
serde_json::to_vec(&json!({
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
}))
.unwrap(),
);
let header_segment = URL_SAFE_NO_PAD.encode(br#"{"alg":"none","typ":"JWT"}"#);
let token = format!("{header_segment}.{claims_segment}.");
let err = verify_signed_bearer_token(
&token,
b"0123456789abcdef0123456789abcdef",
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("alg=none jwt should be rejected");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn signed_bearer_token_verification_rejects_missing_exp() {
let shared_secret = b"0123456789abcdef0123456789abcdef";
let token = signed_token(
shared_secret,
json!({
"iss": "issuer",
}),
);
let err = verify_signed_bearer_token(
&token,
shared_secret,
/*issuer*/ None,
/*audience*/ None,
/*max_clock_skew_seconds*/ 30,
)
.expect_err("jwt without exp should be rejected");
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
}
#[test]
fn validate_signed_bearer_secret_rejects_short_secret() {
let err = validate_signed_bearer_secret(Path::new("/tmp/secret"), b"too-short")
.expect_err("short shared secret should be rejected");
assert_eq!(err.kind(), ErrorKind::InvalidInput);
assert!(
err.to_string().contains("must be at least 32 bytes"),
"unexpected error: {err}"
);
}
}
#[path = "auth_tests.rs"]
mod tests;