mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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>
This commit is contained in:
viyatb-oai
Eric Traut
@@ -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::<TransportEvent>(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 {
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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<TransportEvent>,
|
||||
connection_counter: Arc<AtomicU64>,
|
||||
auth_policy: Arc<WebsocketAuthPolicy>,
|
||||
}
|
||||
|
||||
async fn health_check_handler() -> StatusCode {
|
||||
@@ -118,12 +125,23 @@ async fn websocket_upgrade_handler(
|
||||
websocket: WebSocketUpgrade,
|
||||
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<WebSocketListenerState>,
|
||||
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<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
auth_policy: WebsocketAuthPolicy,
|
||||
) -> IoResult<JoinHandle<()>> {
|
||||
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,
|
||||
|
||||
@@ -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<WebsocketAuthCliMode>,
|
||||
|
||||
/// Absolute path to the capability-token file.
|
||||
#[arg(long = "ws-token-file", value_name = "PATH")]
|
||||
pub ws_token_file: Option<PathBuf>,
|
||||
|
||||
/// 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<PathBuf>,
|
||||
|
||||
/// Expected issuer for signed JWT bearer tokens.
|
||||
#[arg(long = "ws-issuer", value_name = "ISSUER")]
|
||||
pub ws_issuer: Option<String>,
|
||||
|
||||
/// Expected audience for signed JWT bearer tokens.
|
||||
#[arg(long = "ws-audience", value_name = "AUDIENCE")]
|
||||
pub ws_audience: Option<String>,
|
||||
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
|
||||
pub enum WebsocketAuthCliMode {
|
||||
CapabilityToken,
|
||||
SignedBearerToken,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AppServerWebsocketAuthSettings {
|
||||
pub config: Option<AppServerWebsocketAuthConfig>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AppServerWebsocketAuthConfig {
|
||||
CapabilityToken {
|
||||
token_file: AbsolutePathBuf,
|
||||
},
|
||||
SignedBearerToken {
|
||||
shared_secret_file: AbsolutePathBuf,
|
||||
issuer: Option<String>,
|
||||
audience: Option<String>,
|
||||
max_clock_skew_seconds: u64,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub(crate) struct WebsocketAuthPolicy {
|
||||
pub(crate) mode: Option<WebsocketAuthMode>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum WebsocketAuthMode {
|
||||
CapabilityToken {
|
||||
token_sha256: [u8; 32],
|
||||
},
|
||||
SignedBearerToken {
|
||||
shared_secret: Vec<u8>,
|
||||
issuer: Option<String>,
|
||||
audience: Option<String>,
|
||||
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<i64>,
|
||||
iss: Option<String>,
|
||||
aud: Option<JwtAudienceClaim>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum JwtAudienceClaim {
|
||||
Single(String),
|
||||
Multiple(Vec<String>),
|
||||
}
|
||||
|
||||
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<AppServerWebsocketAuthSettings> {
|
||||
let normalize = |value: Option<String>| {
|
||||
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<WebsocketAuthPolicy> {
|
||||
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<JwtClaims, WebsocketAuthError> {
|
||||
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::<JwtClaims>(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<String> {
|
||||
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> {
|
||||
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<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 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}"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user