liter_llm 1.0.0.pre.rc.6

data/vendor/liter-llm/src/auth/vertex_oauth.rs

@@ -0,0 +1,353 @@
+//! Google Vertex AI OAuth2 credential provider (service-account JWT flow).
+//!
+//! Creates a self-signed JWT from a Google Cloud service account key, then
+//! exchanges it for an access token via the Google OAuth2 token endpoint.
+//! Tokens are cached and refreshed automatically when within 5 minutes of
+//! expiry.
+//!
+//! # Environment variables
+//!
+//! | Variable | Description |
+//! |----------|-------------|
+//! | `GOOGLE_APPLICATION_CREDENTIALS` | Path to service account JSON key file |
+//! | `VERTEX_AI_SCOPE` | OAuth scope (defaults to `https://www.googleapis.com/auth/cloud-platform`) |
+
+use std::path::Path;
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
+
+use jsonwebtoken::{Algorithm, EncodingKey, Header};
+use secrecy::{ExposeSecret, SecretString};
+use tokio::sync::RwLock;
+
+use super::{Credential, CredentialProvider};
+use crate::client::BoxFuture;
+use crate::error::LiterLlmError;
+
+/// Default OAuth2 scope for Google Cloud Platform / Vertex AI.
+const DEFAULT_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
+
+/// Google OAuth2 token endpoint.
+const TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
+
+/// JWT grant type for the token exchange.
+const GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
+
+/// JWT lifetime in seconds (Google allows up to 3600).
+const JWT_LIFETIME_SECS: u64 = 3600;
+
+/// Minimum remaining lifetime before a cached token is considered expired.
+const EXPIRY_BUFFER_SECS: u64 = 300;
+
+/// Cached token and its acquisition timestamp + lifetime.
+struct CachedToken {
+    token: SecretString,
+    acquired_at: Instant,
+    expires_in_secs: u64,
+}
+
+impl CachedToken {
+    /// Returns `true` if the token is still valid with the safety buffer.
+    fn is_valid(&self) -> bool {
+        let elapsed = self.acquired_at.elapsed().as_secs();
+        elapsed + EXPIRY_BUFFER_SECS < self.expires_in_secs
+    }
+}
+
+/// Google Vertex AI OAuth2 credential provider using the service-account
+/// JWT assertion flow (two-legged OAuth).
+///
+/// Signs a JWT with the service account's private key, then exchanges it
+/// at `https://oauth2.googleapis.com/token` for a short-lived access token.
+pub struct VertexOAuthCredentialProvider {
+    service_account_email: String,
+    private_key_pem: SecretString,
+    scope: String,
+    cached: RwLock<Option<CachedToken>>,
+    http_client: reqwest::Client,
+}
+
+/// Intentionally redacts the private key to prevent accidental exposure.
+impl std::fmt::Debug for VertexOAuthCredentialProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("VertexOAuthCredentialProvider")
+            .field("service_account_email", &self.service_account_email)
+            .field("private_key_pem", &"[redacted]")
+            .field("scope", &self.scope)
+            .finish_non_exhaustive()
+    }
+}
+
+impl VertexOAuthCredentialProvider {
+    /// Create a provider from a parsed service account JSON object.
+    ///
+    /// Expects `client_email` and `private_key` fields in the JSON value.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`LiterLlmError::Authentication`] if required fields are missing.
+    pub fn from_service_account_json(json: &serde_json::Value) -> Result<Self, LiterLlmError> {
+        let email = json
+            .get("client_email")
+            .and_then(serde_json::Value::as_str)
+            .ok_or_else(|| LiterLlmError::Authentication {
+                message: "service account JSON missing 'client_email' field".into(),
+            })?
+            .to_owned();
+
+        let key = json
+            .get("private_key")
+            .and_then(serde_json::Value::as_str)
+            .ok_or_else(|| LiterLlmError::Authentication {
+                message: "service account JSON missing 'private_key' field".into(),
+            })?
+            .to_owned();
+
+        Ok(Self {
+            service_account_email: email,
+            private_key_pem: SecretString::from(key),
+            scope: DEFAULT_SCOPE.to_owned(),
+            cached: RwLock::new(None),
+            http_client: reqwest::Client::new(),
+        })
+    }
+
+    /// Create a provider from a service account JSON key file on disk.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`LiterLlmError::Authentication`] if the file cannot be read or
+    /// parsed, or if required fields are missing.
+    pub fn from_key_file(path: &Path) -> Result<Self, LiterLlmError> {
+        let contents = std::fs::read_to_string(path).map_err(|e| LiterLlmError::Authentication {
+            message: format!("failed to read service account key file {}: {e}", path.display()),
+        })?;
+
+        let json: serde_json::Value = serde_json::from_str(&contents).map_err(|e| LiterLlmError::Authentication {
+            message: format!("failed to parse service account key file: {e}"),
+        })?;
+
+        Self::from_service_account_json(&json)
+    }
+
+    /// Create a provider from the `GOOGLE_APPLICATION_CREDENTIALS` environment
+    /// variable, which should point to a service account JSON key file.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`LiterLlmError::Authentication`] if the environment variable is
+    /// not set, the file cannot be read, or it cannot be parsed.
+    pub fn from_env() -> Result<Self, LiterLlmError> {
+        let path = std::env::var("GOOGLE_APPLICATION_CREDENTIALS").map_err(|_| LiterLlmError::Authentication {
+            message: "missing required environment variable: GOOGLE_APPLICATION_CREDENTIALS".into(),
+        })?;
+
+        let mut provider = Self::from_key_file(Path::new(&path))?;
+
+        if let Ok(scope) = std::env::var("VERTEX_AI_SCOPE") {
+            provider.scope = scope;
+        }
+
+        Ok(provider)
+    }
+
+    /// Override the OAuth2 scope (default: `https://www.googleapis.com/auth/cloud-platform`).
+    #[must_use]
+    pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
+        self.scope = scope.into();
+        self
+    }
+
+    /// Override the HTTP client used for token requests.
+    #[must_use]
+    pub fn with_http_client(mut self, client: reqwest::Client) -> Self {
+        self.http_client = client;
+        self
+    }
+
+    /// Build a signed JWT assertion and exchange it for an access token.
+    async fn fetch_token(&self) -> Result<CachedToken, LiterLlmError> {
+        let now = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .map_err(|e| LiterLlmError::Authentication {
+                message: format!("system clock error: {e}"),
+            })?
+            .as_secs();
+
+        let claims = JwtClaims {
+            iss: &self.service_account_email,
+            scope: &self.scope,
+            aud: TOKEN_ENDPOINT,
+            iat: now,
+            exp: now + JWT_LIFETIME_SECS,
+        };
+
+        let header = Header::new(Algorithm::RS256);
+        let encoding_key = EncodingKey::from_rsa_pem(self.private_key_pem.expose_secret().as_bytes()).map_err(|e| {
+            LiterLlmError::Authentication {
+                message: format!("invalid RSA private key: {e}"),
+            }
+        })?;
+
+        let assertion =
+            jsonwebtoken::encode(&header, &claims, &encoding_key).map_err(|e| LiterLlmError::Authentication {
+                message: format!("JWT signing failed: {e}"),
+            })?;
+
+        let resp = self
+            .http_client
+            .post(TOKEN_ENDPOINT)
+            .form(&[("grant_type", GRANT_TYPE), ("assertion", &assertion)])
+            .send()
+            .await
+            .map_err(|e| LiterLlmError::Authentication {
+                message: format!("Vertex OAuth token request failed: {e}"),
+            })?;
+
+        let status = resp.status();
+        let body = resp.text().await.map_err(|e| LiterLlmError::Authentication {
+            message: format!("Vertex OAuth token response unreadable: {e}"),
+        })?;
+
+        if !status.is_success() {
+            return Err(LiterLlmError::Authentication {
+                message: format!("Vertex OAuth token request returned {status}: {body}"),
+            });
+        }
+
+        let parsed: TokenResponse = serde_json::from_str(&body).map_err(|e| LiterLlmError::Authentication {
+            message: format!("Vertex OAuth token response parse error: {e}"),
+        })?;
+
+        Ok(CachedToken {
+            token: SecretString::from(parsed.access_token),
+            acquired_at: Instant::now(),
+            expires_in_secs: parsed.expires_in,
+        })
+    }
+}
+
+impl CredentialProvider for VertexOAuthCredentialProvider {
+    fn resolve(&self) -> BoxFuture<'_, Credential> {
+        Box::pin(async move {
+            // Fast path: read lock to check cache.
+            {
+                let guard = self.cached.read().await;
+                if let Some(ref cached) = *guard
+                    && cached.is_valid()
+                {
+                    return Ok(Credential::BearerToken(cached.token.clone()));
+                }
+            }
+
+            // Slow path: write lock to refresh.
+            let mut guard = self.cached.write().await;
+
+            // Double-check after acquiring write lock.
+            if let Some(ref cached) = *guard
+                && cached.is_valid()
+            {
+                return Ok(Credential::BearerToken(cached.token.clone()));
+            }
+
+            let fresh = self.fetch_token().await?;
+            let token = fresh.token.clone();
+            *guard = Some(fresh);
+
+            Ok(Credential::BearerToken(token))
+        })
+    }
+}
+
+/// JWT claims for the Google OAuth2 service-account assertion.
+#[derive(serde::Serialize)]
+struct JwtClaims<'a> {
+    iss: &'a str,
+    scope: &'a str,
+    aud: &'a str,
+    iat: u64,
+    exp: u64,
+}
+
+/// Minimal deserialization of the Google OAuth2 token response.
+#[derive(serde::Deserialize)]
+struct TokenResponse {
+    access_token: String,
+    expires_in: u64,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn cached_token_validity() {
+        let cached = CachedToken {
+            token: SecretString::from("test-token".to_owned()),
+            acquired_at: Instant::now(),
+            expires_in_secs: 3600,
+        };
+        assert!(cached.is_valid());
+    }
+
+    #[test]
+    fn cached_token_expired() {
+        let cached = CachedToken {
+            token: SecretString::from("test-token".to_owned()),
+            // Zero lifetime is immediately expired; avoids Duration subtraction panic on Windows.
+            acquired_at: Instant::now(),
+            expires_in_secs: 0,
+        };
+        assert!(!cached.is_valid());
+    }
+
+    #[test]
+    fn from_service_account_json_valid() {
+        let json = serde_json::json!({
+            "client_email": "test@project.iam.gserviceaccount.com",
+            "private_key": "-----BEGIN TEST RSA KEY-----\nMIIE...\n-----END TEST RSA KEY-----\n"
+        });
+        let provider = VertexOAuthCredentialProvider::from_service_account_json(&json);
+        assert!(provider.is_ok());
+        let provider = provider.expect("should succeed");
+        assert_eq!(provider.service_account_email, "test@project.iam.gserviceaccount.com");
+        assert_eq!(provider.scope, DEFAULT_SCOPE);
+    }
+
+    #[test]
+    fn from_service_account_json_missing_email() {
+        let json = serde_json::json!({
+            "private_key": "-----BEGIN TEST RSA KEY-----\nMIIE...\n-----END TEST RSA KEY-----\n"
+        });
+        let err = VertexOAuthCredentialProvider::from_service_account_json(&json).unwrap_err();
+        assert!(err.to_string().contains("client_email"));
+    }
+
+    #[test]
+    fn from_service_account_json_missing_key() {
+        let json = serde_json::json!({
+            "client_email": "test@project.iam.gserviceaccount.com"
+        });
+        let err = VertexOAuthCredentialProvider::from_service_account_json(&json).unwrap_err();
+        assert!(err.to_string().contains("private_key"));
+    }
+
+    #[test]
+    fn with_scope_override() {
+        let json = serde_json::json!({
+            "client_email": "test@project.iam.gserviceaccount.com",
+            "private_key": "-----BEGIN TEST RSA KEY-----\nMIIE...\n-----END TEST RSA KEY-----\n"
+        });
+        let provider = VertexOAuthCredentialProvider::from_service_account_json(&json)
+            .expect("should succeed")
+            .with_scope("https://custom.scope");
+        assert_eq!(provider.scope, "https://custom.scope");
+    }
+
+    #[tokio::test]
+    #[ignore] // Requires network access and a valid service account key file.
+    async fn live_vertex_oauth_token_exchange() {
+        let provider = VertexOAuthCredentialProvider::from_env().expect("GOOGLE_APPLICATION_CREDENTIALS not set");
+        let credential = provider.resolve().await.expect("token exchange failed");
+        assert!(matches!(credential, Credential::BearerToken(_)));
+    }
+}