From 75c054da0105a3869e17eb7bdf68b6c40a97eb07 Mon Sep 17 00:00:00 2001 From: Drew Goddyn Date: Thu, 12 Mar 2026 07:41:32 -0700 Subject: [PATCH] feat: support GOOGLE_WORKSPACE_CLI_API_BASE to override Google API base URL Override the base URL for all Google API calls via a single env var. When set, discovery document fetches and all API requests (including helper commands like +triage, +reply, +send) route through the specified URL instead of *.googleapis.com. When unset, behavior is unchanged. --- Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 1 + src/discovery.rs | 141 +++++++++++++++++++++++++++++++++--- src/helpers/gmail/mod.rs | 14 +++- src/helpers/gmail/reply.rs | 2 +- src/helpers/gmail/triage.rs | 6 +- src/helpers/gmail/watch.rs | 9 ++- src/helpers/workflows.rs | 8 +- 9 files changed, 163 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 68931d42..10f40f6f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -845,7 +845,7 @@ dependencies = [ [[package]] name = "gws" -version = "0.11.1" +version = "0.11.5" dependencies = [ "aes-gcm", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index 44bf0235..d9bf414a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ [package] name = "gws" -version = "0.11.1" +version = "0.11.5" edition = "2021" description = "Google Workspace CLI — dynamic command surface from Discovery Service" license = "Apache-2.0" diff --git a/README.md b/README.md index 9130ecbd..2dff8d9a 100644 --- a/README.md +++ b/README.md @@ -314,6 +314,7 @@ All variables are optional. See [`.env.example`](.env.example) for a copy-paste | `GOOGLE_WORKSPACE_CLI_SANITIZE_TEMPLATE` | Default Model Armor template | | `GOOGLE_WORKSPACE_CLI_SANITIZE_MODE` | `warn` (default) or `block` | | `GOOGLE_WORKSPACE_PROJECT_ID` | GCP project ID override for quota/billing and fallback for helper commands | +| `GOOGLE_WORKSPACE_CLI_API_BASE` | Override base URL for all Google API calls (e.g. `https://my-proxy.example.com`). Routes discovery doc fetches and API requests through the specified URL instead of `*.googleapis.com` | Environment variables can also be set in a `.env` file (loaded via [dotenvy](https://crates.io/crates/dotenvy)). diff --git a/src/discovery.rs b/src/discovery.rs index b4fa9ff1..7b6e5ffd 100644 --- a/src/discovery.rs +++ b/src/discovery.rs @@ -183,6 +183,14 @@ pub struct JsonSchemaProperty { pub additional_properties: Option>, } +fn apply_api_base_override(doc: &mut RestDescription, api_base_override: &Option) { + if let Some(ref override_url) = api_base_override { + let base = override_url.trim_end_matches('/'); + doc.root_url = format!("{base}/"); + doc.base_url = Some(format!("{base}/{}", doc.service_path)); + } +} + /// Fetches and caches a Google Discovery Document. pub async fn fetch_discovery_document( service: &str, @@ -200,21 +208,29 @@ pub async fn fetch_discovery_document( let cache_file = cache_dir.join(format!("{service}_{version}.json")); + let api_base_override = std::env::var("GOOGLE_WORKSPACE_CLI_API_BASE").ok(); + // Check cache (24hr TTL) if cache_file.exists() { if let Ok(metadata) = std::fs::metadata(&cache_file) { if let Ok(modified) = metadata.modified() { if modified.elapsed().unwrap_or_default() < std::time::Duration::from_secs(86400) { let data = std::fs::read_to_string(&cache_file)?; - let doc: RestDescription = serde_json::from_str(&data)?; + let mut doc: RestDescription = serde_json::from_str(&data)?; + apply_api_base_override(&mut doc, &api_base_override); return Ok(doc); } } } } + let discovery_base = api_base_override + .as_deref() + .unwrap_or("https://www.googleapis.com"); + let url = format!( - "https://www.googleapis.com/discovery/v1/apis/{}/{}/rest", + "{}/discovery/v1/apis/{}/{}/rest", + discovery_base, crate::validate::encode_path_segment(service), crate::validate::encode_path_segment(version), ); @@ -226,12 +242,18 @@ pub async fn fetch_discovery_document( resp.text().await? } else { // Try the $discovery/rest URL pattern used by newer APIs (Forms, Keep, Meet, etc.) - let alt_url = format!("https://{service}.googleapis.com/$discovery/rest"); - let alt_resp = client - .get(&alt_url) - .query(&[("version", version)]) - .send() - .await?; + // When an override is set, preserve the service name in the path so the proxy + // can route to the correct upstream. + let alt_url = if let Some(ref base) = api_base_override { + format!("{}/$discovery/rest", base.trim_end_matches('/')) + } else { + format!("https://{service}.googleapis.com/$discovery/rest") + }; + let mut alt_req = client.get(&alt_url).query(&[("version", version)]); + if api_base_override.is_some() { + alt_req = alt_req.query(&[("service", service)]); + } + let alt_resp = alt_req.send().await?; if !alt_resp.status().is_success() { anyhow::bail!( "Failed to fetch Discovery Document for {service}/{version}: HTTP {} (tried both standard and $discovery URLs)", @@ -243,11 +265,12 @@ pub async fn fetch_discovery_document( // Write to cache if let Err(e) = std::fs::write(&cache_file, &body) { - // Non-fatal: just warn via stderr-safe approach let _ = e; } - let doc: RestDescription = serde_json::from_str(&body)?; + let mut doc: RestDescription = serde_json::from_str(&body)?; + apply_api_base_override(&mut doc, &api_base_override); + Ok(doc) } @@ -320,4 +343,102 @@ mod tests { assert!(doc.resources.is_empty()); assert!(doc.schemas.is_empty()); } + + struct EnvVarGuard { + name: String, + original: Option, + } + + impl EnvVarGuard { + fn set(name: &str, value: &str) -> Self { + let original = std::env::var_os(name); + std::env::set_var(name, value); + Self { + name: name.to_string(), + original, + } + } + + fn remove(name: &str) -> Self { + let original = std::env::var_os(name); + std::env::remove_var(name); + Self { + name: name.to_string(), + original, + } + } + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + match &self.original { + Some(v) => std::env::set_var(&self.name, v), + None => std::env::remove_var(&self.name), + } + } + } + + #[tokio::test] + #[serial_test::serial] + async fn test_api_base_override_rewrites_doc() { + let _guard = EnvVarGuard::set( + "GOOGLE_WORKSPACE_CLI_API_BASE", + "https://my-proxy.example.com", + ); + + let doc_json = r#"{ + "name": "test", + "version": "v1", + "rootUrl": "https://www.googleapis.com/", + "servicePath": "test/v1/" + }"#; + let mut doc: RestDescription = serde_json::from_str(doc_json).unwrap(); + + let api_base_override = std::env::var("GOOGLE_WORKSPACE_CLI_API_BASE").ok(); + apply_api_base_override(&mut doc, &api_base_override); + + assert_eq!(doc.root_url, "https://my-proxy.example.com/"); + assert_eq!( + doc.base_url.as_deref(), + Some("https://my-proxy.example.com/test/v1/") + ); + } + + #[tokio::test] + #[serial_test::serial] + async fn test_api_base_override_unset_preserves_original() { + let _guard = EnvVarGuard::remove("GOOGLE_WORKSPACE_CLI_API_BASE"); + + let doc_json = r#"{ + "name": "drive", + "version": "v3", + "rootUrl": "https://www.googleapis.com/", + "servicePath": "drive/v3/" + }"#; + + let mut doc: RestDescription = serde_json::from_str(doc_json).unwrap(); + + let api_base_override = std::env::var("GOOGLE_WORKSPACE_CLI_API_BASE").ok(); + apply_api_base_override(&mut doc, &api_base_override); + + assert_eq!(doc.root_url, "https://www.googleapis.com/"); + assert!(doc.base_url.is_none()); + } + + #[test] + fn test_api_base_override_strips_trailing_slash() { + let mut doc: RestDescription = serde_json::from_str( + r#"{"name": "test", "version": "v1", "rootUrl": "https://x.com/", "servicePath": "t/v1/"}"#, + ) + .unwrap(); + + let override_url = Some("https://my-proxy.example.com/".to_string()); + apply_api_base_override(&mut doc, &override_url); + + assert_eq!(doc.root_url, "https://my-proxy.example.com/"); + assert_eq!( + doc.base_url.as_deref(), + Some("https://my-proxy.example.com/t/v1/") + ); + } } diff --git a/src/helpers/gmail/mod.rs b/src/helpers/gmail/mod.rs index bada6160..1f5fadd3 100644 --- a/src/helpers/gmail/mod.rs +++ b/src/helpers/gmail/mod.rs @@ -41,6 +41,17 @@ pub(super) const GMAIL_SCOPE: &str = "https://www.googleapis.com/auth/gmail.modi pub(super) const GMAIL_READONLY_SCOPE: &str = "https://www.googleapis.com/auth/gmail.readonly"; pub(super) const PUBSUB_SCOPE: &str = "https://www.googleapis.com/auth/pubsub"; +static GMAIL_API_BASE: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("GOOGLE_WORKSPACE_CLI_API_BASE") + .unwrap_or_else(|_| "https://gmail.googleapis.com".to_string()) + .trim_end_matches('/') + .to_string() +}); + +pub(super) fn gmail_api_base() -> &'static str { + &GMAIL_API_BASE +} + pub(super) struct OriginalMessage { pub thread_id: String, pub message_id_header: String, @@ -170,7 +181,8 @@ pub(super) async fn fetch_message_metadata( message_id: &str, ) -> Result { let url = format!( - "https://gmail.googleapis.com/gmail/v1/users/me/messages/{}", + "{}/gmail/v1/users/me/messages/{}", + gmail_api_base(), crate::validate::encode_path_segment(message_id) ); diff --git a/src/helpers/gmail/reply.rs b/src/helpers/gmail/reply.rs index 0582b514..fad6383b 100644 --- a/src/helpers/gmail/reply.rs +++ b/src/helpers/gmail/reply.rs @@ -132,7 +132,7 @@ pub(super) struct ReplyConfig { async fn fetch_user_email(client: &reqwest::Client, token: &str) -> Result { let resp = crate::client::send_with_retry(|| { client - .get("https://gmail.googleapis.com/gmail/v1/users/me/profile") + .get(&format!("{}/gmail/v1/users/me/profile", gmail_api_base())) .bearer_auth(token) }) .await diff --git a/src/helpers/gmail/triage.rs b/src/helpers/gmail/triage.rs index d8adffba..f7a5da96 100644 --- a/src/helpers/gmail/triage.rs +++ b/src/helpers/gmail/triage.rs @@ -44,7 +44,8 @@ pub async fn handle_triage(matches: &ArgMatches) -> Result<(), GwsError> { let client = crate::client::build_client()?; // 1. List message IDs - let list_url = "https://gmail.googleapis.com/gmail/v1/users/me/messages"; + let base = gmail_api_base(); + let list_url = format!("{base}/gmail/v1/users/me/messages"); let list_resp = client .get(list_url) @@ -97,7 +98,8 @@ pub async fn handle_triage(matches: &ArgMatches) -> Result<(), GwsError> { let token = &token; async move { let get_url = format!( - "https://gmail.googleapis.com/gmail/v1/users/me/messages/{}?format=metadata&metadataHeaders=From&metadataHeaders=Subject&metadataHeaders=Date", + "{}/gmail/v1/users/me/messages/{}?format=metadata&metadataHeaders=From&metadataHeaders=Subject&metadataHeaders=Date", + gmail_api_base(), msg_id ); diff --git a/src/helpers/gmail/watch.rs b/src/helpers/gmail/watch.rs index 86401799..00408f8d 100644 --- a/src/helpers/gmail/watch.rs +++ b/src/helpers/gmail/watch.rs @@ -138,7 +138,7 @@ pub(super) async fn handle_watch( } let resp = client - .post("https://gmail.googleapis.com/gmail/v1/users/me/watch") + .post(&format!("{}/gmail/v1/users/me/watch", gmail_api_base())) .bearer_auth(&gmail_token) .header("Content-Type", "application/json") .json(&watch_body) @@ -180,7 +180,7 @@ pub(super) async fn handle_watch( // Get initial historyId for tracking let profile_resp = client - .get("https://gmail.googleapis.com/gmail/v1/users/me/profile") + .get(&format!("{}/gmail/v1/users/me/profile", gmail_api_base())) .bearer_auth(&gmail_token) .send() .await @@ -386,7 +386,7 @@ async fn fetch_and_output_messages( sanitize_config: &crate::helpers::modelarmor::SanitizeConfig, ) -> Result<(), GwsError> { let resp = client - .get("https://gmail.googleapis.com/gmail/v1/users/me/history") + .get(&format!("{}/gmail/v1/users/me/history", gmail_api_base())) .query(&[ ("startHistoryId", &start_history_id.to_string()), ("historyTypes", &"messageAdded".to_string()), @@ -403,7 +403,8 @@ async fn fetch_and_output_messages( for msg_id in msg_ids { // Fetch full message let msg_url = format!( - "https://gmail.googleapis.com/gmail/v1/users/me/messages/{}", + "{}/gmail/v1/users/me/messages/{}", + gmail_api_base(), crate::validate::encode_path_segment(&msg_id), ); let msg_resp = client diff --git a/src/helpers/workflows.rs b/src/helpers/workflows.rs index 05355d0d..af57bb84 100644 --- a/src/helpers/workflows.rs +++ b/src/helpers/workflows.rs @@ -459,7 +459,8 @@ async fn handle_email_to_task(matches: &ArgMatches) -> Result<(), GwsError> { // 1. Fetch the email let msg_url = format!( - "https://gmail.googleapis.com/gmail/v1/users/me/messages/{}", + "{}/gmail/v1/users/me/messages/{}", + crate::helpers::gmail::gmail_api_base(), crate::validate::encode_path_segment(message_id), ); let msg_json = get_json( @@ -588,7 +589,10 @@ async fn handle_weekly_digest(matches: &ArgMatches) -> Result<(), GwsError> { // Fetch unread email count let gmail_json = get_json( &client, - "https://gmail.googleapis.com/gmail/v1/users/me/messages", + &format!( + "{}/gmail/v1/users/me/messages", + crate::helpers::gmail::gmail_api_base() + ), &token, &[("q", "is:unread"), ("maxResults", "1")], )