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>,    /// Hex-encoded SHA-256 digest of the capability token.    #[arg(long = "ws-token-sha256", value_name = "HEX")]    pub ws_token_sha256: Option<String>,    /// 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 {        source: AppServerWebsocketCapabilityTokenSource,    },    SignedBearerToken {        shared_secret_file: AbsolutePathBuf,        issuer: Option<String>,        audience: Option<String>,        max_clock_skew_seconds: u64,    },}#[derive(Debug, Clone, PartialEq, Eq)]pub enum AppServerWebsocketCapabilityTokenSource {    TokenFile { token_file: AbsolutePathBuf },    TokenSha256 { token_sha256: [u8; 32] },}#[derive(Clone, Debug, Default)]pub 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 source = match (self.ws_token_file, self.ws_token_sha256) {                    (Some(_), Some(_)) => {                        anyhow::bail!(                            "`--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)?,                        }                    }                    (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(WebsocketAuthCliMode::SignedBearerToken) => {                if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() {                    anyhow::bail!(                        "`--ws-token-file` and `--ws-token-sha256` require `--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_token_sha256.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 fn policy_from_settings(    settings: &AppServerWebsocketAuthSettings,) -> io::Result<WebsocketAuthPolicy> {    let mode = match settings.config.as_ref() {        Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source {            AppServerWebsocketCapabilityTokenSource::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 } => {                Some(WebsocketAuthMode::CapabilityToken {                    token_sha256: *token_sha256,                })            }        },        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 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(    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_arg(flag_name: &str, value: &str) -> anyhow::Result<[u8; 32]> {    let trimmed = value.trim();    if trimmed.len() != 64 {        anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest");    }    let mut digest = [0u8; 32];    for (index, pair) in trimmed.as_bytes().chunks_exact(2).enumerate() {        let high = hex_nibble(flag_name, pair[0])?;        let low = hex_nibble(flag_name, pair[1])?;        digest[index] = (high << 4) | low;    }    Ok(digest)}fn hex_nibble(flag_name: &str, byte: u8) -> anyhow::Result<u8> {    match byte {        b'0'..=b'9' => Ok(byte - b'0'),        b'a'..=b'f' => Ok(byte - b'a' + 10),        b'A'..=b'F' => Ok(byte - b'A' + 10),        _ => anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest"),    }}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 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}"        );    }}