zai-rs 0.5.1

//! JWT generation for GLM-Realtime authentication.
//!
//! The realtime WebSocket supports two auth modes (see
//! <https://github.com/MetaGLM/glm-realtime-sdk/blob/main/GLM-Realtime-doc-for-llm.md>):
//!
//! - **Bearer** (server-side): `Authorization: Bearer {API_KEY}`.
//! - **JWT**: a short-lived JWT signed locally with the API key's secret. Only
//!   the derived token is sent in the WebSocket handshake. Applications that
//!   use an untrusted browser/device must generate the token on a trusted server.
//!
//! The GLM JWT is **non-standard**: header
//! `{"alg":"HS256","sign_type":"SIGN"}` and payload
//! `{"api_key":id,"exp":secs,"timestamp":ms}`, signed with HMAC-SHA256 keyed by
//! the secret half of the `{id}.{secret}` API key.

use base64::Engine;
use chrono::Utc;
use hmac::{Hmac, KeyInit, Mac};
use serde_json::json;
use sha2::Sha256;
use tracing::warn;

use crate::{ZaiResult, client::error::ZaiError};

type HmacSha256 = Hmac<Sha256>;

/// Build the `Authorization` header value for the realtime handshake.
///
/// Always returns a `Bearer <token>` string. With `jwt_seconds = Some(ttl)` the
/// token is a freshly-signed JWT derived from the API key; with `None` it is
/// the raw API key (server-side Bearer auth).
pub fn authorization_header(api_key: &str, jwt_seconds: Option<i64>) -> ZaiResult<String> {
    match jwt_seconds {
        Some(ttl) => Ok(format!("Bearer {}", generate(api_key, ttl)?)),
        None => Ok(format!("Bearer {api_key}")),
    }
}

/// Sign a GLM-Realtime JWT for the given `{id}.{secret}` API key.
///
/// `ttl_seconds` must be in `1..=604_800` (seven days).
pub fn generate(api_key: &str, ttl_seconds: i64) -> ZaiResult<String> {
    let (id, secret) = match api_key.split_once('.') {
        Some(parts) => parts,
        None => {
            warn!("Realtime JWT auth error: API key must be '<id>.<secret>'");
            return Err(ZaiError::RealtimeAuthError(
                "API key must be '<id>.<secret>'".into(),
            ));
        },
    };

    if id.is_empty() || secret.is_empty() {
        warn!("Realtime JWT auth error: API key id and secret must be non-empty");
        return Err(ZaiError::RealtimeAuthError(
            "API key id and secret must be non-empty".into(),
        ));
    }

    // Validate ttl: a non-positive or absurdly large value would otherwise wrap
    // `exp` silently (unguarded `i64` add, inconsistent with the crate's
    // saturating-math pass) and surface as an opaque auth rejection. Seven days
    // is a generous ceiling for any realistic JWT lifetime.
    const MAX_TTL_SECONDS: i64 = 7 * 24 * 3600;
    if ttl_seconds <= 0 || ttl_seconds > MAX_TTL_SECONDS {
        warn!(
            ttl_seconds,
            "Realtime JWT auth error: ttl_seconds out of range"
        );
        return Err(ZaiError::RealtimeAuthError(format!(
            "jwt ttl_seconds must be in 1..={MAX_TTL_SECONDS}, got {ttl_seconds}"
        )));
    }

    let now = Utc::now();
    // Defense-in-depth: `now + ttl` can't overflow once ttl is bounded above, but
    // use `checked_add` so a future change to the bound can't reintroduce a
    // silent wrap.
    let exp = now
        .timestamp()
        .checked_add(ttl_seconds)
        .ok_or_else(|| ZaiError::RealtimeAuthError("jwt exp overflow".into()))?;
    let timestamp_ms = now.timestamp_millis();

    let header = json!({ "alg": "HS256", "sign_type": "SIGN" });
    let payload = json!({
        "api_key": id,
        "exp": exp,
        "timestamp": timestamp_ms,
    });

    let header_b64 = base64url(serde_json::to_string(&header)?.as_bytes());
    let payload_b64 = base64url(serde_json::to_string(&payload)?.as_bytes());
    let signing_input = format!("{header_b64}.{payload_b64}");

    let mut mac = HmacSha256::new_from_slice(secret.as_bytes()).map_err(|e| {
        warn!("Realtime JWT auth error: HMAC key error");
        ZaiError::RealtimeAuthError(format!("HMAC key error: {e}"))
    })?;
    mac.update(signing_input.as_bytes());
    let sig = mac.finalize().into_bytes();
    let sig_b64 = base64url(sig.as_slice());

    Ok(format!("{signing_input}.{sig_b64}"))
}

fn base64url(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

#[cfg(test)]
mod tests {
    use base64::Engine;

    use super::*;

    #[test]
    fn jwt_has_three_urlsafe_segments() {
        let token = generate("abcdefghij.0123456789abcdef", 600).unwrap();
        let parts: Vec<&str> = token.split('.').collect();
        assert_eq!(parts.len(), 3, "JWT must have header.payload.signature");
        // URL-safe, no padding.
        assert!(!token.contains('='));
        assert!(!token.contains('+'));
        assert!(!token.contains('/'));
    }

    #[test]
    fn jwt_header_and_payload_decode_to_expected_fields() {
        let token = generate("abcdefghij.0123456789abcdef", 600).unwrap();
        let mut parts = token.split('.');
        let header_json = String::from_utf8(
            base64::engine::general_purpose::URL_SAFE_NO_PAD
                .decode(parts.next().unwrap())
                .unwrap(),
        )
        .unwrap();
        let payload_json = String::from_utf8(
            base64::engine::general_purpose::URL_SAFE_NO_PAD
                .decode(parts.next().unwrap())
                .unwrap(),
        )
        .unwrap();
        assert!(header_json.contains("\"alg\":\"HS256\""));
        assert!(header_json.contains("\"sign_type\":\"SIGN\""));
        assert!(payload_json.contains("\"api_key\":\"abcdefghij\""));
        assert!(payload_json.contains("\"exp\""));
        assert!(payload_json.contains("\"timestamp\""));
    }

    #[test]
    fn rejects_malformed_keys() {
        assert!(generate("no-dot-here", 600).is_err());
        assert!(generate(".secret", 600).is_err());
        assert!(generate("id.", 600).is_err());
    }

    #[test]
    fn rejects_out_of_range_ttl() {
        // Non-positive and absurdly large values are rejected up front rather
        // than producing a nonsense/overflowing `exp`.
        assert!(generate("abcdefghij.0123456789abcdef", 0).is_err());
        assert!(generate("abcdefghij.0123456789abcdef", -1).is_err());
        assert!(generate("abcdefghij.0123456789abcdef", i64::MAX).is_err());
        // The 7-day ceiling is inclusive; one second over is rejected.
        assert!(generate("abcdefghij.0123456789abcdef", 7 * 24 * 3600).is_ok());
        assert!(generate("abcdefghij.0123456789abcdef", 7 * 24 * 3600 + 1).is_err());
    }

    #[test]
    fn authorization_header_bearer_jwt_and_raw() {
        let h = authorization_header("abcdefghij.0123456789abcdef", Some(3600)).unwrap();
        assert!(h.starts_with("Bearer "));
        assert!(h[7..].split('.').count() == 3);

        let h2 = authorization_header("abcdefghij.0123456789abcdef", None).unwrap();
        assert_eq!(h2, "Bearer abcdefghij.0123456789abcdef");
    }
}