tryaudex_core/

sso.rs

use serde::{Deserialize, Serialize};

use crate::error::{AvError, Result};

/// SSO provider type.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SsoProvider {
    Okta,
    Auth0,
    Google,
    /// Generic OIDC provider
    Oidc,
}

/// SSO configuration in `[sso]` config section.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SsoConfig {
    /// SSO provider type
    pub provider: SsoProvider,
    /// OIDC issuer URL (e.g. "https://myorg.okta.com", "https://accounts.google.com")
    pub issuer: String,
    /// OIDC client ID
    pub client_id: String,
    /// OIDC client secret (optional, for confidential clients)
    pub client_secret: Option<String>,
    /// Redirect URI for the device/local auth flow
    #[serde(default = "default_redirect")]
    pub redirect_uri: String,
    /// OIDC scopes to request
    #[serde(default = "default_scopes")]
    pub scopes: Vec<String>,
    /// Claim to use as the identity (default: "email")
    #[serde(default = "default_identity_claim")]
    pub identity_claim: String,
    /// Claim to use for group/team membership (default: "groups")
    #[serde(default = "default_groups_claim")]
    pub groups_claim: String,
    /// Whether SSO is required (if false, falls back to local identity)
    #[serde(default)]
    pub required: bool,
}

fn default_redirect() -> String {
    "http://localhost:8400/callback".to_string()
}

fn default_scopes() -> Vec<String> {
    vec![
        "openid".to_string(),
        "email".to_string(),
        "profile".to_string(),
    ]
}

fn default_identity_claim() -> String {
    "email".to_string()
}

fn default_groups_claim() -> String {
    "groups".to_string()
}

/// Token response from the OIDC provider.
#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
    id_token: Option<String>,
    #[allow(dead_code)]
    token_type: Option<String>,
    #[allow(dead_code)]
    expires_in: Option<u64>,
}

/// Decoded identity from an SSO token.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SsoIdentity {
    /// The user's identity (email, subject, etc.)
    pub identity: String,
    /// Group/team memberships from the SSO provider
    pub groups: Vec<String>,
    /// Raw claims from the ID token
    pub claims: serde_json::Value,
}

/// Cached SSO session stored on disk.
#[derive(Debug, Serialize, Deserialize)]
struct CachedSsoSession {
    identity: SsoIdentity,
    expires_at: i64,
    /// Issuer URL that issued this identity — prevents cross-provider reuse.
    #[serde(default)]
    issuer: String,
    /// Client ID the token was issued for.
    #[serde(default)]
    client_id: String,
    /// Machine fingerprint at time of caching — prevents replay on another host.
    #[serde(default)]
    machine_id: String,
}

/// Read the machine-unique identifier from the OS.
/// Returns an empty string if unavailable (non-fatal; validation will fail on mismatch).
fn read_machine_id() -> String {
    // Linux / systemd: /etc/machine-id
    if let Ok(id) = std::fs::read_to_string("/etc/machine-id") {
        let trimmed = id.trim().to_string();
        if !trimmed.is_empty() {
            return trimmed;
        }
    }
    // macOS: IOPlatformUUID via ioreg
    #[cfg(target_os = "macos")]
    {
        if let Ok(out) = std::process::Command::new("ioreg")
            .args(["-rd1", "-c", "IOPlatformExpertDevice"])
            .output()
        {
            let text = String::from_utf8_lossy(&out.stdout);
            for line in text.lines() {
                if line.contains("IOPlatformUUID") {
                    if let Some(start) = line.rfind('"') {
                        let rest = &line[..start];
                        if let Some(end) = rest.rfind('"') {
                            return rest[end + 1..].to_string();
                        }
                    }
                }
            }
        }
    }
    String::new()
}

fn cache_path() -> std::path::PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(|| std::path::PathBuf::from("."))
        .join("audex")
        .join("sso_session.json")
}

/// Load cached SSO session if still valid and from the same provider and machine (encrypted at rest).
pub fn load_cached_identity(issuer: &str, client_id: &str) -> Option<SsoIdentity> {
    let path = cache_path();
    let cached: CachedSsoSession = crate::keystore::decrypt_from_file(&path).ok()??;
    let now = chrono::Utc::now().timestamp();
    if now >= cached.expires_at {
        return None;
    }
    // Reject cached identity from a different provider/client
    if cached.issuer != issuer || cached.client_id != client_id {
        tracing::info!("SSO cache issuer/client mismatch — re-authenticating");
        return None;
    }
    // Reject cached identity from a different machine (prevents exfiltrated cache replay)
    let current_machine_id = read_machine_id();
    if !current_machine_id.is_empty() && cached.machine_id != current_machine_id {
        tracing::info!("SSO cache machine-id mismatch — re-authenticating");
        return None;
    }
    Some(cached.identity)
}

fn save_cached_identity(identity: &SsoIdentity, ttl_secs: u64, issuer: &str, client_id: &str) {
    let cached = CachedSsoSession {
        identity: identity.clone(),
        expires_at: chrono::Utc::now().timestamp() + ttl_secs as i64,
        issuer: issuer.to_string(),
        client_id: client_id.to_string(),
        machine_id: read_machine_id(),
    };
    let path = cache_path();
    if let Some(parent) = path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    if let Err(e) = crate::keystore::encrypt_to_file(&path, &cached) {
        tracing::warn!(error = %e, "Failed to cache SSO identity — next login will require re-authentication");
    }
}

/// Authenticate via SSO using the device authorization flow.
/// Opens a browser for login and listens on a local callback URL.
pub async fn authenticate(config: &SsoConfig) -> Result<SsoIdentity> {
    // Check cache first — only reuse if same issuer/client
    if let Some(cached) = load_cached_identity(&config.issuer, &config.client_id) {
        return Ok(cached);
    }

    let discovery_url = format!(
        "{}/.well-known/openid-configuration",
        config.issuer.trim_end_matches('/')
    );

    // R6-H23: do NOT follow redirects on the discovery request.  reqwest's
    // default policy transparently follows up to 10 hops, which means an
    // attacker who can inject a 302 on the discovery URL (a compromised
    // CDN, a misconfigured reverse proxy, a rogue intermediate) can
    // redirect the OIDC client to their own IdP's discovery document —
    // returning whatever authorization_endpoint and jwks_uri they like.
    // The subsequent set_issuer check is the *second* line of defence;
    // blocking the redirect at the transport is the first.
    // R6-M37: enforce HTTP timeouts on every SSO network call. Without a
    // timeout, a slow-loris IdP (or a hung reverse proxy) would hang the
    // CLI indefinitely while it waits for discovery/JWKS/userinfo bytes.
    let discovery_client = reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::none())
        .timeout(std::time::Duration::from_secs(30))
        .connect_timeout(std::time::Duration::from_secs(10))
        .build()
        .map_err(|e| AvError::InvalidPolicy(format!("SSO client build failed: {}", e)))?;

    // Fetch OIDC discovery document.
    // R6-M38: cap the discovery document at 1 MiB. A hostile/compromised
    // IdP could otherwise stream an arbitrarily large JSON body into
    // `.json()` and OOM the client.
    let discovery: serde_json::Value = fetch_json_capped(
        &discovery_client,
        &discovery_url,
        None,
        1024 * 1024,
        "SSO discovery",
    )
    .await?;

    // R6-H23: validate `issuer` in the discovery document matches the
    // configured issuer exactly (OIDC Discovery 1.0 §4.3).  Without this,
    // a discovery document returned from a different origin (e.g. via a
    // proxy that rewrites the host) would pass through and mint a JWT
    // validation config pointing at the wrong issuer — later set_issuer
    // on the JWT would then "succeed" because both sides come from the
    // same attacker-controlled document.
    let discovery_issuer = discovery["issuer"].as_str().ok_or_else(|| {
        AvError::InvalidPolicy("OIDC discovery document missing `issuer` field".to_string())
    })?;
    if discovery_issuer.trim_end_matches('/') != config.issuer.trim_end_matches('/') {
        return Err(AvError::InvalidPolicy(format!(
            "OIDC discovery issuer mismatch: configured={}, received={}. \
             This can indicate a redirect on the discovery URL or a misconfigured IdP.",
            config.issuer, discovery_issuer
        )));
    }

    // R6-M37: same timeouts for the token/userinfo client.
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(30))
        .connect_timeout(std::time::Duration::from_secs(10))
        .build()
        .map_err(|e| AvError::InvalidPolicy(format!("SSO client build failed: {}", e)))?;

    let auth_endpoint = discovery["authorization_endpoint"]
        .as_str()
        .ok_or_else(|| {
            AvError::InvalidPolicy("Missing authorization_endpoint in OIDC discovery".to_string())
        })?;
    let token_endpoint = discovery["token_endpoint"].as_str().ok_or_else(|| {
        AvError::InvalidPolicy("Missing token_endpoint in OIDC discovery".to_string())
    })?;

    // Generate PKCE challenge
    let verifier = generate_pkce_verifier();
    let challenge = generate_pkce_challenge(&verifier);
    let state = generate_state();
    let nonce = generate_state(); // same format: random 16-byte hex

    let scopes = config.scopes.join(" ");
    let auth_url = format!(
        "{}?client_id={}&redirect_uri={}&response_type=code&scope={}&state={}&nonce={}&code_challenge={}&code_challenge_method=S256",
        auth_endpoint,
        urlencoding::encode(&config.client_id),
        urlencoding::encode(&config.redirect_uri),
        urlencoding::encode(&scopes),
        urlencoding::encode(&state),
        urlencoding::encode(&nonce),
        urlencoding::encode(&challenge),
    );

    eprintln!("Opening browser for SSO login...");
    eprintln!("If the browser doesn't open, visit:\n  {}\n", auth_url);

    // Try to open browser
    let _ = open_browser(&auth_url);

    // Start local server to receive the callback
    let code = listen_for_callback(&config.redirect_uri, &state).await?;

    // Exchange code for tokens
    let mut token_params = vec![
        ("grant_type", "authorization_code".to_string()),
        ("code", code),
        ("redirect_uri", config.redirect_uri.clone()),
        ("client_id", config.client_id.clone()),
        ("code_verifier", verifier),
    ];
    if let Some(ref secret) = config.client_secret {
        token_params.push(("client_secret", secret.clone()));
    }

    let token_resp: TokenResponse = client
        .post(token_endpoint)
        .form(&token_params)
        .send()
        .await
        .map_err(|e| AvError::InvalidPolicy(format!("SSO token exchange failed: {}", e)))?
        .json()
        .await
        .map_err(|e| AvError::InvalidPolicy(format!("SSO token parse failed: {}", e)))?;

    // Decode and verify the ID token (signature + claims)
    //
    // R6-H24: previously `jwks_uri` defaulted to the empty string when
    // missing from the discovery document, which caused
    // `decode_and_verify_jwt` to skip signature verification and fall
    // back to transport-trust-only decoding with only a tracing::warn.
    // An attacker who could influence the discovery document — via a
    // compromised IdP, a misconfigured reverse proxy, or the redirect
    // vector covered in R6-H23 — could return `{"jwks_uri": ""}` (or
    // simply omit the key) and forge arbitrary ID tokens whose only
    // protection was HTTPS transport.  Treat missing/empty jwks_uri as
    // fatal when we actually have an id_token to verify.
    let claims = if let Some(ref id_token) = token_resp.id_token {
        let jwks_uri = discovery["jwks_uri"].as_str().unwrap_or_default();
        if jwks_uri.is_empty() {
            return Err(AvError::InvalidPolicy(
                "OIDC discovery document is missing `jwks_uri`; refusing to accept an \
                 id_token without signature verification. If the provider genuinely does \
                 not expose JWKS, configure audex to use the userinfo endpoint instead."
                    .to_string(),
            ));
        }
        decode_and_verify_jwt(
            &client,
            jwks_uri,
            id_token,
            &config.issuer,
            &config.client_id,
        )
        .await?
    } else {
        // Use userinfo endpoint as fallback
        let userinfo_url = discovery["userinfo_endpoint"].as_str().ok_or_else(|| {
            AvError::InvalidPolicy("No id_token and no userinfo_endpoint".to_string())
        })?;

        // R6-M38: cap userinfo body at 1 MiB.
        fetch_json_capped(
            &client,
            userinfo_url,
            Some(&token_resp.access_token),
            1024 * 1024,
            "SSO userinfo",
        )
        .await?
    };

    // R6-M36: require nonce presence and equality. Without the explicit
    // presence check, a token missing the nonce claim would silently
    // bypass replay protection mandated by OIDC Core §3.1.3.7. Userinfo
    // responses (the id_token-less fallback above) legitimately do not
    // carry a nonce; only enforce the check when we are verifying an
    // id_token.
    if token_resp.id_token.is_some() {
        let token_nonce = claims
            .get("nonce")
            .and_then(|v| v.as_str())
            .ok_or_else(|| {
                AvError::InvalidPolicy(
                    "SSO id_token is missing required `nonce` claim — \
                     rejecting to prevent replay (OIDC Core §3.1.3.7)"
                        .to_string(),
                )
            })?;
        if token_nonce != nonce {
            return Err(AvError::InvalidPolicy(
                "SSO nonce mismatch — possible token replay".to_string(),
            ));
        }
    }

    // R6-M41: require the identity claim to be present and a string. The
    // previous `.unwrap_or("unknown")` let an IdP whose userinfo returned
    // a JSON scalar/array silently fall through to `"unknown"` as the
    // principal identity — bypassing identity_claim entirely and letting
    // every such caller map to the same "unknown" user in policy.
    let identity_str = claims[&config.identity_claim]
        .as_str()
        .ok_or_else(|| {
            AvError::InvalidPolicy(format!(
                "SSO identity claim `{}` is missing or not a string. \
                 Check the identity_claim configuration against the IdP's \
                 token/userinfo shape.",
                config.identity_claim
            ))
        })?
        .to_string();

    let groups: Vec<String> = claims[&config.groups_claim]
        .as_array()
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect()
        })
        .unwrap_or_default();

    let sso_identity = SsoIdentity {
        identity: identity_str,
        groups,
        claims,
    };

    // Cache for 1 hour
    save_cached_identity(&sso_identity, 3600, &config.issuer, &config.client_id);

    Ok(sso_identity)
}

/// Decode and verify a JWT ID token against the provider's JWKS keys.
///
/// Fetches the JWKS from the provider's `jwks_uri`, verifies the cryptographic
/// signature, and validates `iss`, `aud`, and `exp` claims.
///
/// Falls back to transport-trust-only decoding if `jwks_uri` is unavailable,
/// but logs a prominent warning.
async fn decode_and_verify_jwt(
    client: &reqwest::Client,
    jwks_uri: &str,
    token: &str,
    expected_issuer: &str,
    expected_audience: &str,
) -> Result<serde_json::Value> {
    use jsonwebtoken::{decode, DecodingKey, Validation};

    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(AvError::InvalidPolicy("Invalid JWT format".to_string()));
    }

    // Try JWKS-based signature verification
    if !jwks_uri.is_empty() {
        let header = jsonwebtoken::decode_header(token)
            .map_err(|e| AvError::InvalidPolicy(format!("JWT header decode failed: {}", e)))?;

        let algorithm = header.alg;

        // Enforce an allowlist of expected asymmetric algorithms to prevent
        // algorithm confusion attacks where an attacker pairs a JWK with an
        // unexpected algorithm (e.g. HMAC symmetric key as RSA).
        const ALLOWED_ALGORITHMS: &[jsonwebtoken::Algorithm] = &[
            jsonwebtoken::Algorithm::RS256,
            jsonwebtoken::Algorithm::RS384,
            jsonwebtoken::Algorithm::RS512,
            jsonwebtoken::Algorithm::ES256,
            jsonwebtoken::Algorithm::ES384,
            jsonwebtoken::Algorithm::PS256,
            jsonwebtoken::Algorithm::PS384,
            jsonwebtoken::Algorithm::PS512,
            jsonwebtoken::Algorithm::EdDSA,
        ];
        if !ALLOWED_ALGORITHMS.contains(&algorithm) {
            return Err(AvError::InvalidPolicy(format!(
                "JWT algorithm {:?} is not in the allowed set of asymmetric algorithms",
                algorithm
            )));
        }

        // Fetch JWKS from provider.
        // R6-M38: cap JWKS body at 1 MiB so a hostile IdP can't OOM the
        // verifier with a giant key set.
        let jwks_resp: serde_json::Value =
            fetch_json_capped(client, jwks_uri, None, 1024 * 1024, "JWKS").await?;

        // Parse JWKs into typed structs
        let jwk_set: jsonwebtoken::jwk::JwkSet = serde_json::from_value(jwks_resp)
            .map_err(|e| AvError::InvalidPolicy(format!("JWKS parse into JwkSet failed: {}", e)))?;

        // R6-M40: require the token header to carry a `kid` and match it
        // exactly. The previous `jwk_set.keys.first()` fallback let a
        // token with no kid verify against an arbitrary key from the
        // JWKS — if the IdP rotated keys or an attacker could influence
        // which entry landed first, verification could succeed against
        // the wrong key.
        let kid = header.kid.as_deref().ok_or_else(|| {
            AvError::InvalidPolicy(
                "JWT header is missing `kid` — refusing to pick an \
                 arbitrary key from the JWKS. Reject tokens without kid \
                 to prevent key-confusion attacks."
                    .to_string(),
            )
        })?;
        let matching_key = jwk_set
            .keys
            .iter()
            .find(|k| k.common.key_id.as_deref() == Some(kid));

        if let Some(jwk) = matching_key {
            let decoding_key = DecodingKey::from_jwk(jwk)
                .map_err(|e| AvError::InvalidPolicy(format!("JWK decode failed: {}", e)))?;

            let mut validation = Validation::new(algorithm);
            validation.set_issuer(&[expected_issuer]);
            validation.set_audience(&[expected_audience]);

            let token_data = decode::<serde_json::Value>(token, &decoding_key, &validation)
                .map_err(|e| {
                    AvError::InvalidPolicy(format!("JWT signature verification failed: {}", e))
                })?;

            return Ok(token_data.claims);
        }

        return Err(AvError::InvalidPolicy(format!(
            "No matching JWK found for token kid={:?}",
            kid
        )));
    }

    // R6-H24: defence-in-depth.  The caller already rejects empty
    // jwks_uri before we get here, but leave an explicit refusal so a
    // future caller cannot reintroduce the "silently fall back to
    // unverified" bug.  Any path that reaches this point with an
    // empty jwks_uri is a programmer error, not a runtime condition
    // we should paper over.
    Err(AvError::InvalidPolicy(
        "JWT signature verification cannot be skipped: jwks_uri is empty. \
         This is a programming error — callers must reject missing jwks_uri \
         before invoking decode_and_verify_jwt."
            .to_string(),
    ))
}

/// Decode JWT claims without cryptographic signature verification.
///
/// Validates token format, `exp`, `iss`, and `aud` claims, but does NOT
/// verify the JWT signature. Retained only for unit tests that exercise
/// the claim-level validation logic — the production fallback that used
/// to call this was removed in R6-H24 to prevent silent signature-skip.
#[cfg(test)]
fn decode_jwt_claims_unverified(
    token: &str,
    expected_issuer: &str,
    expected_audience: &str,
) -> Result<serde_json::Value> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(AvError::InvalidPolicy("Invalid JWT format".to_string()));
    }

    use base64::Engine;
    let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(parts[1])
        .map_err(|e| AvError::InvalidPolicy(format!("JWT base64 decode failed: {}", e)))?;

    let claims: serde_json::Value = serde_json::from_slice(&payload)
        .map_err(|e| AvError::InvalidPolicy(format!("JWT claims parse failed: {}", e)))?;

    // Reject expired tokens to prevent replay attacks.
    if let Some(exp) = claims.get("exp").and_then(|v| v.as_i64()) {
        let now = chrono::Utc::now().timestamp();
        if now > exp {
            return Err(AvError::InvalidPolicy(format!(
                "SSO token expired at {} (current time: {})",
                exp, now
            )));
        }
    } else {
        tracing::warn!("SSO token has no 'exp' claim — cannot verify expiry");
    }

    // Validate issuer
    match claims.get("iss").and_then(|v| v.as_str()) {
        Some(iss) if iss == expected_issuer => {}
        Some(iss) => {
            return Err(AvError::InvalidPolicy(format!(
                "SSO token issuer mismatch: got '{}', expected '{}'",
                iss, expected_issuer
            )));
        }
        None => {
            return Err(AvError::InvalidPolicy(
                "SSO token missing 'iss' claim".to_string(),
            ));
        }
    }

    // Validate audience
    let aud_valid = match claims.get("aud") {
        Some(serde_json::Value::String(aud)) => aud == expected_audience,
        Some(serde_json::Value::Array(arr)) => {
            arr.iter().any(|v| v.as_str() == Some(expected_audience))
        }
        _ => false,
    };
    if !aud_valid {
        return Err(AvError::InvalidPolicy(format!(
            "SSO token audience does not include expected client_id '{}'",
            expected_audience
        )));
    }

    Ok(claims)
}

/// Generate a random PKCE code verifier.
fn generate_pkce_verifier() -> String {
    use rand::Rng;
    let mut rng = rand::rng();
    let bytes: Vec<u8> = (0..32).map(|_| rng.random::<u8>()).collect();
    use base64::Engine;
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&bytes)
}

/// Generate PKCE code challenge from verifier (S256).
fn generate_pkce_challenge(verifier: &str) -> String {
    use sha2::Digest;
    let hash = sha2::Sha256::digest(verifier.as_bytes());
    use base64::Engine;
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash)
}

/// Generate a random state parameter.
fn generate_state() -> String {
    use rand::Rng;
    let mut rng = rand::rng();
    let bytes: Vec<u8> = (0..16).map(|_| rng.random::<u8>()).collect();
    hex::encode(bytes)
}

/// Escape HTML special characters to prevent XSS when embedding user input in HTML responses.
fn html_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '&' => out.push_str("&amp;"),
            '<' => out.push_str("&lt;"),
            '>' => out.push_str("&gt;"),
            '"' => out.push_str("&quot;"),
            _ => out.push(ch),
        }
    }
    out
}

/// Try to open a URL in the default browser.
fn open_browser(url: &str) -> std::io::Result<()> {
    #[cfg(target_os = "macos")]
    {
        std::process::Command::new("open").arg(url).spawn()?;
    }
    #[cfg(target_os = "linux")]
    {
        std::process::Command::new("xdg-open").arg(url).spawn()?;
    }
    #[cfg(target_os = "windows")]
    {
        std::process::Command::new("cmd")
            .args(["/c", "start", url])
            .spawn()?;
    }
    Ok(())
}

/// R6-M38: fetch a JSON document with a hard body-size cap.
///
/// `.json()` on a reqwest response eagerly buffers the whole body with no
/// limit, so a hostile or compromised IdP could stream gigabytes into
/// discovery/userinfo/JWKS and OOM the client. Read the body one chunk
/// at a time with a running byte counter and bail out at `max_bytes`.
async fn fetch_json_capped(
    client: &reqwest::Client,
    url: &str,
    bearer: Option<&str>,
    max_bytes: usize,
    label: &str,
) -> Result<serde_json::Value> {
    let mut req = client.get(url);
    if let Some(token) = bearer {
        req = req.bearer_auth(token);
    }
    let mut resp = req
        .send()
        .await
        .map_err(|e| AvError::InvalidPolicy(format!("{} failed: {}", label, e)))?;
    if !resp.status().is_success() {
        return Err(AvError::InvalidPolicy(format!(
            "{} returned HTTP {}",
            label,
            resp.status()
        )));
    }
    let mut body: Vec<u8> = Vec::new();
    loop {
        match resp.chunk().await {
            Ok(Some(chunk)) => {
                if body.len() + chunk.len() > max_bytes {
                    return Err(AvError::InvalidPolicy(format!(
                        "{} response exceeds {} byte cap",
                        label, max_bytes
                    )));
                }
                body.extend_from_slice(&chunk);
            }
            Ok(None) => break,
            Err(e) => {
                return Err(AvError::InvalidPolicy(format!(
                    "{} body read failed: {}",
                    label, e
                )));
            }
        }
    }
    serde_json::from_slice(&body)
        .map_err(|e| AvError::InvalidPolicy(format!("{} parse failed: {}", label, e)))
}

/// Listen on the callback URL for the authorization code.
async fn listen_for_callback(redirect_uri: &str, expected_state: &str) -> Result<String> {
    use tokio::io::AsyncReadExt;
    use tokio::io::AsyncWriteExt;

    // Parse port from redirect URI
    let url: url::Url = redirect_uri
        .parse()
        .map_err(|e| AvError::InvalidPolicy(format!("Invalid redirect URI: {}", e)))?;
    let port = url.port().unwrap_or(8400);

    let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{}", port))
        .await
        .map_err(|e| {
            AvError::InvalidPolicy(format!(
                "Failed to start callback listener on port {}: {}",
                port, e
            ))
        })?;

    eprintln!(
        "Waiting for SSO callback on port {} (5 minute timeout)...",
        port
    );

    let (mut stream, _) =
        tokio::time::timeout(std::time::Duration::from_secs(300), listener.accept())
            .await
            .map_err(|_| AvError::InvalidPolicy("SSO login timed out after 5 minutes".to_string()))?
            .map_err(|e| AvError::InvalidPolicy(format!("Failed to accept callback: {}", e)))?;

    // R6-M39: read the HTTP request until we see the end-of-headers
    // delimiter (`\r\n\r\n`) or hit a bounded total size. The previous
    // fixed 4096-byte single read truncated real-world browser callbacks
    // whose cookies + referer + user-agent blew past 4 KiB, causing the
    // code/state query string to be silently split and the login to
    // fail with "Failed to parse callback URL". Cap the total read at
    // 16 KiB to bound memory use under a malicious client.
    const MAX_CALLBACK_REQUEST_BYTES: usize = 16 * 1024;
    let mut buf: Vec<u8> = Vec::with_capacity(4096);
    let mut scratch = [0u8; 4096];
    loop {
        let n = stream
            .read(&mut scratch)
            .await
            .map_err(|e| AvError::InvalidPolicy(format!("Failed to read callback: {}", e)))?;
        if n == 0 {
            break;
        }
        if buf.len() + n > MAX_CALLBACK_REQUEST_BYTES {
            return Err(AvError::InvalidPolicy(format!(
                "SSO callback request exceeds {} byte cap",
                MAX_CALLBACK_REQUEST_BYTES
            )));
        }
        buf.extend_from_slice(&scratch[..n]);
        // Only the request line + headers are needed to extract the
        // code/state query parameters, so stop once we see the
        // end-of-headers sentinel.
        if buf.windows(4).any(|w| w == b"\r\n\r\n") {
            break;
        }
    }
    let request = String::from_utf8_lossy(&buf);

    // Parse the query string from GET request
    let first_line = request.lines().next().unwrap_or("");
    let path = first_line.split_whitespace().nth(1).unwrap_or("/");
    let query_url = format!("http://localhost{}", path);
    let parsed: url::Url = query_url
        .parse()
        .map_err(|_| AvError::InvalidPolicy("Failed to parse callback URL".to_string()))?;

    let params: std::collections::HashMap<String, String> = parsed
        .query_pairs()
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .collect();

    // Verify state
    let state = params.get("state").cloned().unwrap_or_default();
    // Constant-time comparison for state to be consistent with other
    // security-sensitive comparisons, even though loopback + random state
    // makes timing attacks impractical.
    let state_valid = {
        let a = state.as_bytes();
        let b = expected_state.as_bytes();
        if a.len() != b.len() {
            false
        } else {
            a.iter()
                .zip(b.iter())
                .fold(0u8, |acc, (x, y)| acc | (x ^ y))
                == 0
        }
    };
    if !state_valid {
        let response = "HTTP/1.1 400 Bad Request\r\nContent-Type: text/html\r\n\r\n<h1>SSO Error</h1><p>State mismatch. Please try again.</p>";
        let _ = stream.write_all(response.as_bytes()).await;
        return Err(AvError::InvalidPolicy(
            "SSO state mismatch — possible CSRF".to_string(),
        ));
    }

    // Check for error
    if let Some(error) = params.get("error") {
        let desc = params.get("error_description").cloned().unwrap_or_default();
        let escaped_error = html_escape(error);
        let escaped_desc = html_escape(&desc);
        let response = format!("HTTP/1.1 400 Bad Request\r\nContent-Type: text/html\r\n\r\n<h1>SSO Error</h1><p>{}: {}</p>", escaped_error, escaped_desc);
        let _ = stream.write_all(response.as_bytes()).await;
        return Err(AvError::InvalidPolicy(format!(
            "SSO error: {}: {}",
            error, desc
        )));
    }

    let code = params
        .get("code")
        .cloned()
        .ok_or_else(|| AvError::InvalidPolicy("No authorization code in callback".to_string()))?;

    // Send success response
    let response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n<h1>Audex SSO Login Successful</h1><p>You can close this window and return to the terminal.</p>";
    let _ = stream.write_all(response.as_bytes()).await;

    Ok(code)
}

/// Clear the cached SSO session.
pub fn logout() {
    let path = cache_path();
    let _ = std::fs::remove_file(path);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sso_config_deserialize() {
        let toml_str = r#"
provider = "okta"
issuer = "https://myorg.okta.com"
client_id = "0oa1234567890"
scopes = ["openid", "email", "groups"]
identity_claim = "email"
groups_claim = "groups"
required = true
"#;
        let config: SsoConfig = toml::from_str(toml_str).unwrap();
        assert!(matches!(config.provider, SsoProvider::Okta));
        assert_eq!(config.issuer, "https://myorg.okta.com");
        assert_eq!(config.client_id, "0oa1234567890");
        assert!(config.required);
        assert_eq!(config.scopes.len(), 3);
    }

    #[test]
    fn test_sso_config_google() {
        let toml_str = r#"
provider = "google"
issuer = "https://accounts.google.com"
client_id = "12345.apps.googleusercontent.com"
client_secret = "GOCSPX-secret"
"#;
        let config: SsoConfig = toml::from_str(toml_str).unwrap();
        assert!(matches!(config.provider, SsoProvider::Google));
        assert!(config.client_secret.is_some());
        // Defaults
        assert_eq!(config.identity_claim, "email");
        assert_eq!(config.groups_claim, "groups");
        assert_eq!(config.redirect_uri, "http://localhost:8400/callback");
        assert!(!config.required);
    }

    #[test]
    fn test_sso_identity_serialize() {
        let identity = SsoIdentity {
            identity: "alice@example.com".to_string(),
            groups: vec!["backend-team".to_string(), "devops".to_string()],
            claims: serde_json::json!({"email": "alice@example.com", "sub": "user123"}),
        };
        let json = serde_json::to_string(&identity).unwrap();
        assert!(json.contains("alice@example.com"));
        assert!(json.contains("backend-team"));
        // Roundtrip
        let parsed: SsoIdentity = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.identity, "alice@example.com");
        assert_eq!(parsed.groups.len(), 2);
    }

    #[test]
    fn test_pkce_verifier_and_challenge() {
        let verifier = generate_pkce_verifier();
        assert!(!verifier.is_empty());
        let challenge = generate_pkce_challenge(&verifier);
        assert!(!challenge.is_empty());
        assert_ne!(verifier, challenge);
    }

    #[test]
    fn test_state_generation() {
        let state1 = generate_state();
        let state2 = generate_state();
        assert_ne!(state1, state2);
        assert_eq!(state1.len(), 32); // 16 bytes = 32 hex chars
    }

    #[test]
    fn test_decode_jwt_claims_unverified() {
        // Create a minimal JWT with base64url-encoded payload
        use base64::Engine;
        let future_exp = chrono::Utc::now().timestamp() + 3600;
        let claims = serde_json::json!({
            "email": "test@example.com",
            "groups": ["admin"],
            "exp": future_exp,
            "iss": "https://accounts.example.com",
            "aud": "my-client-id"
        });
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(serde_json::to_vec(&claims).unwrap());
        let fake_jwt = format!("header.{}.signature", payload);

        let decoded =
            decode_jwt_claims_unverified(&fake_jwt, "https://accounts.example.com", "my-client-id")
                .unwrap();
        assert_eq!(decoded["email"], "test@example.com");
        assert_eq!(decoded["groups"][0], "admin");
    }

    #[test]
    fn test_decode_jwt_invalid() {
        assert!(decode_jwt_claims_unverified("not-a-jwt", "iss", "aud").is_err());
        assert!(decode_jwt_claims_unverified("a.b", "iss", "aud").is_err());
    }

    #[test]
    fn test_decode_jwt_expired() {
        use base64::Engine;
        let past_exp = chrono::Utc::now().timestamp() - 3600;
        let claims = serde_json::json!({"email": "test@example.com", "exp": past_exp, "iss": "https://iss", "aud": "aud"});
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(serde_json::to_vec(&claims).unwrap());
        let fake_jwt = format!("header.{}.signature", payload);
        let err = decode_jwt_claims_unverified(&fake_jwt, "https://iss", "aud").unwrap_err();
        assert!(err.to_string().contains("expired"));
    }

    #[test]
    fn test_decode_jwt_wrong_issuer() {
        use base64::Engine;
        let future_exp = chrono::Utc::now().timestamp() + 3600;
        let claims =
            serde_json::json!({"exp": future_exp, "iss": "https://evil.com", "aud": "my-client"});
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(serde_json::to_vec(&claims).unwrap());
        let fake_jwt = format!("header.{}.signature", payload);
        let err =
            decode_jwt_claims_unverified(&fake_jwt, "https://legit.com", "my-client").unwrap_err();
        assert!(err.to_string().contains("issuer mismatch"));
    }

    #[test]
    fn test_decode_jwt_wrong_audience() {
        use base64::Engine;
        let future_exp = chrono::Utc::now().timestamp() + 3600;
        let claims =
            serde_json::json!({"exp": future_exp, "iss": "https://iss", "aud": "other-client"});
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(serde_json::to_vec(&claims).unwrap());
        let fake_jwt = format!("header.{}.signature", payload);
        let err = decode_jwt_claims_unverified(&fake_jwt, "https://iss", "my-client").unwrap_err();
        assert!(err.to_string().contains("audience"));
    }
}