google_ai_rs/

auth.rs

use serde::Serialize;
use thiserror::Error;
use tokio::stream;
use tonic::RawRequestHeaderValue;

#[cfg(feature = "jwt")]
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
#[cfg(feature = "jwt")]
use rsa::{
    pkcs1::{DecodeRsaPrivateKey, Error as Pkcs1Error},
    pkcs1v15::SigningKey,
    pkcs8::{DecodePrivateKey, Error as Pkcs8Error},
    signature::{RandomizedSigner, SignatureEncoding},
    RsaPrivateKey,
};
#[cfg(feature = "jwt")]
use serde::Deserialize;
#[cfg(feature = "jwt")]
use serde_json::json;
#[cfg(feature = "jwt")]
use sha2::Sha256;
#[cfg(feature = "jwt")]
use std::{
    path::Path,
    sync::Arc,
    time::{Duration, SystemTime, SystemTimeError},
};
#[cfg(feature = "jwt")]
use tokio::sync::RwLock;

/// Authentication configuration options
#[derive(Clone, Debug)]
pub enum Auth {
    /// API key authentication (simple and fast but less secure)
    ApiKey(String),
    #[cfg(feature = "jwt")]
    /// JWT-based service account authentication (more secure but with more overhead)
    TokenSource(TokenSource),
}

impl<S: Into<String>> From<S> for Auth {
    fn from(value: S) -> Self {
        Auth::ApiKey(value.into())
    }
}

/// JSON Web Token configuration for service account authentication
#[cfg(feature = "jwt")]
#[derive(Deserialize, Clone, Debug)]
pub struct JWTConfig {
    /// Service account client email (format: name@project.iam.gserviceaccount.com)
    #[serde(rename = "client_email")]
    pub client_email: String,

    /// RSA private key in PEM or DER format (keep secure!)
    #[serde(rename = "private_key")]
    pub private_key: String,

    /// Optional private key identifier from Google Cloud
    #[serde(rename = "private_key_id")]
    pub private_key_id: String,

    /// Token lifetime duration (default: 1 hour)
    #[serde(skip)]
    pub lifetime: Option<Duration>,
}

/// Token generation source types
#[cfg(feature = "jwt")]
#[derive(Clone, Debug)]
pub enum TokenSource {
    #[deprecated(note = "use `TokenSource::JWT`(all-caps) instead")]
    /// JSON Web Token authentication flow
    Jwt { jwt: Box<JwtService> },

    /// JSON Web Token authentication flow
    JWT(JWTConfig),
}

/// Authentication error types
#[derive(Debug, Error)]
pub enum Error {
    #[error("Token generation failed: {0}")]
    TokenGeneration(String),

    #[error("Invalid header value")]
    InvalidHeader,

    #[cfg(feature = "jwt")]
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),

    #[cfg(feature = "jwt")]
    #[error("JSON parsing error: {0}")]
    Json(#[from] serde_json::Error),

    #[error("Private key parsing failed: {0}")]
    #[cfg(feature = "jwt")]
    PrivateKey(#[from] PrivateKeyError),

    #[error("System time error: {0}")]
    #[cfg(feature = "jwt")]
    SystemTime(#[from] SystemTimeError),

    #[error("Invalid token lifetime")]
    #[cfg(feature = "jwt")]
    InvalidLifetime,
}

/// Private key parsing specific errors
#[cfg(feature = "jwt")]
#[derive(Debug, Error)]
pub enum PrivateKeyError {
    #[error("PKCS#1 parsing error: {0}")]
    Pkcs1(#[from] Pkcs1Error),

    #[error("PKCS#8 parsing error: {0}")]
    Pkcs8(#[from] Pkcs8Error),

    #[error("PEM format error: {0}")]
    Pem(#[from] pem::PemError),
}

#[cfg(feature = "jwt")]
const DEFAULT_TOKEN_LIFETIME: Duration = Duration::from_secs(3600);
#[cfg(feature = "jwt")]
const JWT_AUDIENCE: &str = "https://generativelanguage.googleapis.com/";
#[cfg(feature = "jwt")]
const JWT_HEADER: &str = "authorization";

const API_KEY_HEADER: &str = "x-goog-api-key";

impl Auth {
    /// Creates API key authentication
    pub fn new(api_key: &str) -> Self {
        Self::ApiKey(api_key.to_owned())
    }

    /// Creates service account authentication from JSON file
    ///
    ///
    /// # Example
    /// ```
    /// # use google_ai_rs::auth::Auth;
    /// # async fn f() -> Result<(), Box<dyn std::error::Error>> {
    /// let auth = Auth::service_account("path/to/service-account.json")
    ///     .await
    ///     .expect("Valid service account");
    /// # Ok(())
    /// # }
    /// ```
    #[deprecated(note = "use `Auth::service` instead")]
    #[allow(deprecated)]
    #[cfg(feature = "jwt")]
    pub async fn service_account<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        Ok(Self::TokenSource(
            TokenSource::from_service_account(path).await?,
        ))
    }

    /// Creates service account authentication from JSON file
    ///
    /// # Example
    /// ```
    /// # use google_ai_rs::auth::Auth;
    /// # async fn f() -> Result<(), Box<dyn std::error::Error>> {
    /// let auth = Auth::service("path/to/service-account.json")
    ///     .await
    ///     .expect("Valid service account");
    /// # Ok(())
    /// # }
    /// ```
    #[cfg(feature = "jwt")]
    pub async fn service<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        Ok(Self::TokenSource(TokenSource::service(path).await?))
    }

    /// Creates JWT authentication from configuration
    #[deprecated(note = "use `Auth::jwt` instead")]
    #[allow(deprecated)]
    #[cfg(feature = "jwt")]
    pub async fn from_jwt_config(config: JWTConfig) -> Result<Self, Error> {
        Ok(Self::TokenSource(TokenSource::from_jwt(config).await?))
    }

    /// Creates JWT authentication from configuration
    #[cfg(feature = "jwt")]
    pub fn jwt(config: JWTConfig) -> Self {
        Self::TokenSource(TokenSource::jwt(config))
    }
}

#[cfg(feature = "jwt")]
impl TokenSource {
    /// Creates service account authentication from JSON file
    #[deprecated(note = "use `TokenSource::service` instead")]
    #[allow(deprecated)]
    #[cfg(feature = "jwt")]
    pub async fn from_service_account<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let json = tokio::fs::read(path).await?;
        let config: JWTConfig = serde_json::from_slice(&json)?;
        Self::from_jwt(config).await
    }

    /// Creates service account authentication from JSON file
    #[cfg(feature = "jwt")]
    pub async fn service<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let json = tokio::fs::read(path).await?;
        let config: JWTConfig = serde_json::from_slice(&json)?;
        Ok(Self::jwt(config))
    }

    /// Creates JWT authentication from configuration
    #[deprecated(note = "use `TokenSource::jwt` instead")]
    #[allow(deprecated)]
    #[cfg(feature = "jwt")]
    pub async fn from_jwt(config: JWTConfig) -> Result<Self, Error> {
        config.parsed().map(|jwt| Self::Jwt { jwt: Box::new(jwt) })
    }

    /// Creates JWT authentication from configuration
    #[cfg(feature = "jwt")]
    pub fn jwt(config: JWTConfig) -> Self {
        Self::JWT(config)
    }
}

#[derive(Debug)]
pub(crate) enum AuthParsed {
    ApiKey(RawRequestHeaderValue),
    #[cfg(feature = "jwt")]
    JwtKind(JwtService),
}

impl Auth {
    pub(crate) fn parsed(self) -> Result<AuthParsed, Error> {
        match self {
            // FIXME: InvalidHeader is such a poor error variant...
            Auth::ApiKey(api_key) => Ok(AuthParsed::ApiKey(
                api_key.parse().map_err(|_| Error::InvalidHeader)?,
            )),
            #[cfg(feature = "jwt")]
            Auth::TokenSource(token_source) => match token_source {
                #[allow(deprecated)]
                // FIXME: Revalidate the unchecked jwt
                TokenSource::Jwt { jwt } => Ok(AuthParsed::JwtKind(*jwt)),
                TokenSource::JWT(jwtconfig) => jwtconfig.parsed().map(AuthParsed::JwtKind),
            },
        }
    }
}

#[cfg(feature = "jwt")]
impl JWTConfig {
    fn parsed(self) -> Result<JwtService, Error> {
        JwtService::new(self)
    }
}

impl AuthParsed {
    /// Adds authentication headers to gRPC requests
    #[cfg(not(feature = "jwt"))]
    pub(crate) fn _into_request(self, request: &mut tonic::RawRequestHeader) {
        match self {
            Self::ApiKey(metadata_value) => {
                request.insert(API_KEY_HEADER, metadata_value);
            }
        }
    }

    // /// Adds authentication headers to gRPC requests
    // #[cfg(feature = "jwt")]
    // pub(crate) async fn _into_request(
    //     self,
    //     request: &mut tonic::RawRequestHeader,
    // ) {
    //     match self {
    //         Self::ApiKey(metadata_value) => {
    //             request.insert(API_KEY_HEADER, metadata_value);
    //         }
    //         Self::JwtKind(jwt_service) => {
    //             let token = jwt_service.get_token().await;
    //             request.insert(JWT_HEADER, token);
    //         }
    //     }
    // }

    // /// Adds authentication headers to gRPC requests
    // #[cfg(not(feature = "jwt"))]
    // pub(crate) fn to_request(&self, request: &mut tonic::RawRequestHeader) {
    //     match self {
    //         Self::ApiKey(metadata_value) => {
    //             request.insert(API_KEY_HEADER, metadata_value.clone());
    //         }
    //     }
    // }

    /// Adds authentication headers to gRPC requests
    #[cfg(feature = "jwt")]
    pub(crate) async fn to_request(&self, request: &mut tonic::RawRequestHeader) {
        match self {
            Self::ApiKey(metadata_value) => {
                request.insert(API_KEY_HEADER, metadata_value.clone());
            }
            Self::JwtKind(jwt_service) => {
                let token = jwt_service.get_token().await;
                request.insert(JWT_HEADER, token);
            }
        }
    }
}

#[cfg(feature = "jwt")]
use hidden::*;

#[cfg(feature = "jwt")]
mod hidden {
    use std::ops::Deref;

    use super::*;

    /// JWT token service with caching
    #[derive(Clone, Debug)]
    pub struct JwtService {
        pub(super) inner: Arc<JwtServiceInner>,
    }

    #[derive(Debug)]
    pub struct JwtServiceInner {
        pub(super) config: JWTConfig,
        pub(super) signing_key: SigningKey<Sha256>,
        pub(super) cache: RwLock<JwtCache>,
    }

    impl Deref for JwtService {
        type Target = JwtServiceInner;

        fn deref(&self) -> &Self::Target {
            &self.inner
        }
    }
    // TODO: Rename to JwtParsed
}

/// Cached JWT token data
#[derive(Debug)]
#[cfg(feature = "jwt")]
struct JwtCache {
    token: RawRequestHeaderValue,
    expires_at: SystemTime,
}

#[cfg(feature = "jwt")]
impl JwtService {
    fn new(config: JWTConfig) -> Result<Self, Error> {
        let private_key = parse_private_key(config.private_key.as_bytes())?;
        let signing_key = SigningKey::<Sha256>::new(private_key);

        let (token, expires_at) = Self::generate_token_(&config, &signing_key)?;

        Ok(JwtService {
            inner: Arc::new(JwtServiceInner {
                config,
                signing_key,
                cache: RwLock::new(JwtCache { token, expires_at }),
            }),
        })
    }

    #[inline(always)]
    fn generate_token_(
        config: &JWTConfig,
        signing_key: &SigningKey<Sha256>,
    ) -> Result<(RawRequestHeaderValue, SystemTime), Error> {
        #[derive(Serialize)]
        struct JwtHeader<'a> {
            alg: &'static str,
            typ: &'static str,
            kid: &'a str,
        }

        let header = JwtHeader {
            alg: "RS256",
            typ: "JWT",
            kid: &config.private_key_id,
        };

        let now = SystemTime::now();
        // FIXME: Have fallback?
        let iat = now.duration_since(SystemTime::UNIX_EPOCH)?.as_secs();
        let lifetime = config.lifetime.unwrap_or(DEFAULT_TOKEN_LIFETIME);

        let exp = iat + lifetime.as_secs();

        #[derive(Serialize)]
        struct JwtClaims<'a> {
            iss: &'a str,
            sub: &'a str,
            aud: &'static str,
            exp: u64,
            iat: u64,
        }

        let claims = JwtClaims {
            iss: &config.client_email,
            sub: &config.client_email,
            aud: JWT_AUDIENCE,
            exp,
            iat,
        };

        // The serializations can't fail
        let encoded_header = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).unwrap());
        let encoded_claims = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap());
        let message = format!("{encoded_header}.{encoded_claims}");

        let signature = signing_key
            .sign_with_rng(&mut rand::thread_rng(), message.as_bytes())
            .to_bytes();

        let encoded_sig = URL_SAFE_NO_PAD.encode(signature);
        let jwt_token = format!("Bearer {message}.{encoded_sig}")
            .parse()
            // Once the start auth is valid, it'll continue to be valid
            .map_err(|_| Error::InvalidHeader)?;
        let expires_at = now + lifetime;

        Ok((jwt_token, expires_at))
    }

    /// Generates a new signed JWT token
    fn generate_token_infallibly(&self) -> (RawRequestHeaderValue, SystemTime) {
        Self::generate_token_(&self.config, &self.signing_key).unwrap()
    }

    /// Retrieves valid token from cache or generates new one
    async fn get_token(&self) -> RawRequestHeaderValue {
        // Fast path: check cache with read lock
        {
            let cache = self.cache.read().await;
            if SystemTime::now() < cache.expires_at {
                return cache.token.clone();
            }
        }

        // Slow path: regenerate token with write lock

        // Once the start auth is valid (parsed is called and it's called in client builder),
        // it'll continue to be valid
        let (new_token, expires_at) = self.generate_token_infallibly();
        let mut cache = self.cache.write().await;

        *cache = JwtCache {
            token: new_token,
            expires_at,
        };

        cache.token.clone()
    }
}

/// Parses RSA private key from multiple formats
#[cfg(feature = "jwt")]
fn parse_private_key(bytes: &[u8]) -> Result<RsaPrivateKey, PrivateKeyError> {
    // Try PEM format first
    if let Ok(pem) = pem::parse(bytes) {
        return RsaPrivateKey::from_pkcs8_der(pem.contents())
            .or_else(|_| RsaPrivateKey::from_pkcs1_der(pem.contents()))
            .map_err(Into::into);
    }

    // Fallback to DER format
    RsaPrivateKey::from_pkcs8_der(bytes)
        .or_else(|_| RsaPrivateKey::from_pkcs1_der(bytes))
        .map_err(Into::into)
}