From 37541625b6b52e1d6e3be8000753faad7a45bd3d Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 24 May 2026 14:05:19 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat(vertex):=20in-process=20SA=20JSON=20?= =?UTF-8?q?=E2=86=92=20JWT=20=E2=86=92=20OAuth=20+=20token=20cache=20(#302?= =?UTF-8?q?=20Phase=20E=20D5.1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the second half of #302 Phase E productionization (Step 2a streaming merged as ai-gateway#386). Before this PR, `ProviderKey .secret` required a pre-minted GCP OAuth2 `access_token` — operator was responsible for refresh (GCP token TTL ~1h), which is awkward for production-scale deployments with many provider keys. This PR extends `VertexSecret` to accept EITHER mode, with mutual- exclusion validation at parse time: - **Pre-minted token** (backward-compat): set `access_token`. The bridge uses it verbatim. Existing deployments are unaffected. - **In-process SA mint** (new): set `service_account_json` to the full GCP SA JSON key. The bridge signs a JWT with the SA's RSA `private_key`, exchanges it for an OAuth2 access token via the SA's `token_uri` (JWT-bearer assertion grant per RFC 7523), and caches it in-process keyed by SA `client_email` with TTL refresh ~60s before the upstream-reported expiry. ## New module: `token_mint` `token_mint::TokenMinter` owns the cache + mint pipeline: - `ServiceAccountKey` — minimum SA JSON fields needed for minting (type, private_key, client_email, token_uri). Other on-disk SA fields (project_id, private_key_id, client_id, auth_uri, *_x509_cert_url) deserialize-skipped; operators can paste the whole SA JSON verbatim. - `validate()` — cheap shape checks (type == "service_account", private_key starts with -----BEGIN, client_email + token_uri non-empty). Run at parse time so the operator gets an actionable error before the first chat. - JWT claims: `{iss=client_email, scope=cloud-platform, aud=token_uri, iat=now, exp=now+3600}` per Google's documented flow. Signed RS256 via `jsonwebtoken` crate. - Token endpoint POST: form-encoded `grant_type=urn:ietf:params :oauth:grant-type:jwt-bearer` + `assertion=`. Response is `{access_token, expires_in}`; `token_type` always "Bearer" so discarded. - Cache: `Arc>>` keyed by `client_email`. Read-lock on the common cache-hit path; write- lock only on mint. Cache stores `expires_at = now + expires_in - 60s` to refresh proactively. - Endpoint error handling: 4xx/5xx from token endpoint becomes `BridgeError::Config` with body truncated to 500 chars (GCP returns operator-actionable OAuth-shape errors like `invalid_grant`; doesn't echo the SA private key). ## Bridge wiring `VertexBridge` now holds `Arc`. Both `chat_gemini` and `chat_gemini_stream` resolve the bearer via `creds.resolve_access_token(&minter).await?` BEFORE entering the deadline-wrapped request future, so token-mint failures surface as direct `Err` returns from `chat()` / `chat_stream()` rather than being yielded mid-stream. The two streams test seams (`with_api_base_override` for Vertex host, `with_token_endpoint_override` for OAuth host) are independent — a wiremock test can override either or both, e.g. exercise the full SA mint + chat round-trip end-to-end against two separate wiremock servers. ## Tests (10 new across token_mint + bridge, all wiremock-based) token_mint::tests: - mint_signs_jwt_with_correct_claims_and_posts_to_token_uri: captures the assertion JWT, decodes against the matching public key, asserts iss/scope/aud/iat/exp claims have the expected values + relationships - get_token_caches_within_ttl_and_only_calls_endpoint_once: 3 calls, wiremock `.expect(1)` enforces single mint - token_endpoint_5xx_surfaces_as_config_error_with_body_truncated: 503 + arbitrary body → Config error with HTTP status + body - invalid_pem_surfaces_clear_error_before_endpoint_call: bad PEM → Config error, `.expect(0)` on endpoint enforces no network call - wrong_type_field_rejected_before_endpoint_call: type="external_account" → Config error before network call bridge::tests: - vertex_secret_accepts_service_account_json_path: SA-only secret parses cleanly + SA validate runs - vertex_secret_rejects_both_credential_modes_set: mutual exclusion - vertex_secret_rejects_neither_credential_mode_set: mutual exclusion - vertex_secret_rejects_empty_access_token_string: edge case where operator pastes "" rather than omitting the field - chat_gemini_sa_path_mints_token_and_forwards_as_bearer: full end-to-end pipeline — SA JSON in secret → token minted via wiremock OAuth endpoint → minted token forwarded as Authorization Bearer on Vertex chat call. Asserts the OAuth endpoint is hit exactly ONCE across two chats (cache hit on the second). Plus the existing `vertex_secret_parses_full_form` is updated for the new `Option` shape. ## Test fixtures Embedded 2048-bit RSA test keypair at `crates/aisix-provider-vertex/test-fixtures/test_sa_{private,public}.pem`. Generated via `openssl genpkey -algorithm RSA -pkeyopt rsa_keygen_bits:2048` + matching pubout. Deterministic so JWT byte sequences reproduce across runs. NOT a real GCP key — safe to commit. `cargo test -p aisix-provider-vertex` → 60/60 PASS. `cargo clippy -p aisix-provider-vertex --all-targets -- -D warnings` clean. `cargo fmt --all` applied. ## Dependency added `jsonwebtoken = "9"` — actively maintained, audited via cargo-audit, no transitive heavy deps. Alternatives considered + rejected: `yup-oauth2` (much larger surface, includes a full OAuth client we don't need) and `gcp_auth` (pulls hyper directly, doesn't compose with the existing reqwest::Client). ## References - GCP SA OAuth flow: https://developers.google.com/identity/protocols/oauth2/service-account#authorizingrequests - JWT Bearer grant (RFC 7523): https://www.rfc-editor.org/rfc/rfc7523 - Standard SA JSON shape: emitted by `gcloud iam service-accounts keys create`. Field set confirmed against the python-aiplatform SDK's `google.oauth2.service_account.Credentials.from_service_account_info()` consumer signature. ## Unblocks AC.11 partial in api7/AISIX-Cloud#302: "Vertex 上 Gemini / Claude / Llama 都通过单一 google-vertex catalog 入口能 stream" — Gemini half is now fully done (chat + stream wired by ai-gateway#386, SA OAuth wired by this PR). Anthropic + Llama publishers remain blocked on D5.3 / D5.4 per the roadmap. ## Out of scope - Live-mode e2e against a real GCP project (the cp-api/dashboard matrix spec is a separate AISIX-Cloud PR — gated on #481 wiring mock-vertex into compose first) - TTL-expiry behavior under simulated time advance (clock-mocking infra not yet established in the crate; manual sanity-checked the expires_at arithmetic; covered indirectly by the cache-hit test) --- Cargo.lock | 28 ++ crates/aisix-provider-vertex/Cargo.toml | 1 + crates/aisix-provider-vertex/src/bridge.rs | 303 +++++++++++- crates/aisix-provider-vertex/src/lib.rs | 15 +- .../aisix-provider-vertex/src/token_mint.rs | 445 ++++++++++++++++++ .../test-fixtures/test_sa_private.pem | 28 ++ .../test-fixtures/test_sa_public.pem | 9 + 7 files changed, 809 insertions(+), 20 deletions(-) create mode 100644 crates/aisix-provider-vertex/src/token_mint.rs create mode 100644 crates/aisix-provider-vertex/test-fixtures/test_sa_private.pem create mode 100644 crates/aisix-provider-vertex/test-fixtures/test_sa_public.pem diff --git a/Cargo.lock b/Cargo.lock index f4fccd41..b9076357 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -288,6 +288,7 @@ dependencies = [ "bytes", "futures", "http 1.4.0", + "jsonwebtoken", "reqwest", "serde", "serde_json", @@ -2625,6 +2626,21 @@ dependencies = [ "uuid-simd", ] +[[package]] +name = "jsonwebtoken" +version = "9.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a87cc7a48537badeae96744432de36f4be2b4a34a05a5ef32e9dd8a1c169dde" +dependencies = [ + "base64 0.22.1", + "js-sys", + "pem", + "ring", + "serde", + "serde_json", + "simple_asn1", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -4142,6 +4158,18 @@ version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" +[[package]] +name = "simple_asn1" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d585997b0ac10be3c5ee635f1bab02d512760d14b7c468801ac8a01d9ae5f1d" +dependencies = [ + "num-bigint", + "num-traits", + "thiserror 2.0.18", + "time", +] + [[package]] name = "sketches-ddsketch" version = "0.3.1" diff --git a/crates/aisix-provider-vertex/Cargo.toml b/crates/aisix-provider-vertex/Cargo.toml index 3f1e788c..993f8ddc 100644 --- a/crates/aisix-provider-vertex/Cargo.toml +++ b/crates/aisix-provider-vertex/Cargo.toml @@ -15,6 +15,7 @@ async-trait.workspace = true async-stream = "0.3" bytes.workspace = true futures.workspace = true +jsonwebtoken = "9" thiserror.workspace = true tracing.workspace = true reqwest.workspace = true diff --git a/crates/aisix-provider-vertex/src/bridge.rs b/crates/aisix-provider-vertex/src/bridge.rs index 60b28d5a..25269f2e 100644 --- a/crates/aisix-provider-vertex/src/bridge.rs +++ b/crates/aisix-provider-vertex/src/bridge.rs @@ -29,8 +29,10 @@ use http::{ }; use reqwest::{header, Client, StatusCode}; use serde::{Deserialize, Serialize}; +use std::sync::Arc; use std::time::{Duration, Instant}; +use crate::token_mint::{ServiceAccountKey, TokenMinter}; use crate::wire; /// Family Bridge for Google Vertex AI. @@ -40,6 +42,11 @@ pub struct VertexBridge { /// metrics dashboards keep their existing `provider="vertex"` /// filters working. name: &'static str, + /// In-process GCP OAuth2 token minter + cache. Used only on the + /// `service_account_json` secret path; the pre-minted-token path + /// bypasses this entirely. Wrapped in `Arc` so the bridge stays + /// `Clone`-friendly for callers that share it across Hub registrations. + token_minter: Arc, /// Test-only Vertex API base override (e.g. wiremock URI). When /// set, replaces the canonical `-aiplatform.googleapis.com` /// host so wiremock can stand in. @@ -57,6 +64,7 @@ impl VertexBridge { /// when downstream callers want to share a connection pool. pub fn with_client(client: Client) -> Self { Self { + token_minter: Arc::new(TokenMinter::new(client.clone())), client, name: "vertex", #[cfg(test)] @@ -74,6 +82,16 @@ impl VertexBridge { self } + /// Test-only seam: replace the SA `token_uri` host on the + /// internal token minter. Used by the SA-flow tests to redirect + /// JWT-bearer assertions to a wiremock endpoint. + #[cfg(test)] + pub(crate) fn with_token_endpoint_override(mut self, url: impl Into) -> Self { + let new_minter = TokenMinter::new(self.client.clone()).with_token_endpoint_override(url); + self.token_minter = Arc::new(new_minter); + self + } + /// Resolve the base host the bridge POSTs to. Production: /// `https://-aiplatform.googleapis.com`. Tests can pin /// the host via [`Self::with_api_base_override`]. @@ -165,15 +183,34 @@ impl VertexPublisher { /// `ProviderKey.secret` schema for a Vertex provider key. /// /// Convention: GCP credentials are JSON-encoded into the `secret` -/// field. `access_token` is a pre-minted OAuth2 bearer (operator -/// manages refresh; ~1-hour GCP TTL). D5.1 follow-up adds in-process -/// JWT signing via `service_account_json`. +/// field. The operator chooses ONE of two credential modes: +/// +/// 1. **Pre-minted token** — set `access_token`; operator manages +/// refresh (GCP token TTL ~1h). Backward-compatible with the +/// original D5.2.a schema; useful for short-lived testing +/// rigs and for operators who already have a token-mint pipeline. +/// +/// 2. **In-process SA mint** (D5.1) — set `service_account_json` +/// to the full GCP service-account JSON key. The bridge signs a +/// JWT with the SA's RSA private key, exchanges it for an OAuth2 +/// access token via the SA's `token_uri`, and caches it +/// in-process with TTL refresh. +/// +/// Exactly one of the two must be set. Setting both — or neither — +/// fails at parse time so the operator gets an actionable error +/// before the first chat. #[derive(Debug, Deserialize)] struct VertexSecret { /// Pre-minted GCP OAuth2 access token (operator manages refresh). - /// D5.1 follow-up will accept a `service_account_json` field in - /// addition and mint tokens in-process. - access_token: String, + /// Mutually exclusive with `service_account_json`. + #[serde(default)] + access_token: Option, + /// GCP service-account JSON key (the on-disk shape `gcloud iam + /// service-accounts keys create` emits). When present, the bridge + /// mints + caches tokens in-process. Mutually exclusive with + /// `access_token`. + #[serde(default)] + service_account_json: Option, /// GCP project id (numeric or named, e.g. `my-org-prod`). project: String, /// GCP region the Vertex AI deployment targets @@ -182,7 +219,8 @@ struct VertexSecret { } impl VertexSecret { - /// Parse the JSON-encoded credential blob. + /// Parse the JSON-encoded credential blob and validate the + /// mutually-exclusive credential modes. /// /// **Audit-aware:** error messages MUST NOT echo raw secret /// bytes (serde error messages can leak partial content via @@ -191,17 +229,69 @@ impl VertexSecret { if secret.trim().is_empty() { return Err(BridgeError::Config( "vertex provider_key.secret is empty — \ - expected JSON {access_token, project, region}" + expected JSON with project, region, and either access_token \ + or service_account_json" .into(), )); } - serde_json::from_str::(secret).map_err(|_e| { + let parsed: VertexSecret = serde_json::from_str(secret).map_err(|_e| { BridgeError::Config( "vertex provider_key.secret must be valid JSON: \ - {access_token, project, region}" + {project, region, and either access_token or service_account_json}" .into(), ) - }) + })?; + // Enforce mutual exclusion. Both-set is suspect (which one + // wins?); neither-set is unusable. Empty-string token is a + // distinct error so the operator gets a clearer message than + // generic "neither set". + if parsed.access_token.as_deref().is_some_and(str::is_empty) { + return Err(BridgeError::Config( + "vertex provider_key.secret.access_token is empty".into(), + )); + } + let has_token = parsed + .access_token + .as_deref() + .is_some_and(|t| !t.is_empty()); + let has_sa = parsed.service_account_json.is_some(); + if has_token && has_sa { + return Err(BridgeError::Config( + "vertex provider_key.secret must set exactly one of access_token \ + or service_account_json (both were provided)" + .into(), + )); + } + if !has_token && !has_sa { + return Err(BridgeError::Config( + "vertex provider_key.secret must set either access_token or \ + service_account_json (neither was provided)" + .into(), + )); + } + // Validate the SA shape eagerly if present so the operator + // hits the actionable error at parse, not at first chat. + if let Some(sa) = &parsed.service_account_json { + sa.validate()?; + } + Ok(parsed) + } + + /// Resolve the bearer token to use on this request. Returns the + /// pre-minted token verbatim if set; otherwise mints (or pulls + /// from cache) via the bridge's [`TokenMinter`]. + async fn resolve_access_token(&self, minter: &TokenMinter) -> Result { + if let Some(token) = &self.access_token { + return Ok(token.clone()); + } + if let Some(sa) = &self.service_account_json { + return minter.get_token(sa).await; + } + // parse() rejects neither-set, so this is unreachable in + // practice — keep the explicit error for defense in depth. + Err(BridgeError::Config( + "internal: VertexSecret has neither token nor SA after parse".into(), + )) } } @@ -417,7 +507,11 @@ impl VertexBridge { .into(), )); } - let headers = build_request_headers(&creds.access_token, &ctx.request_id)?; + // Resolve bearer: pre-minted token verbatim, or mint+cache + // via the in-process token minter from SA JSON. Failure + // surfaces as a Config error (operator-actionable). + let access_token = creds.resolve_access_token(&self.token_minter).await?; + let headers = build_request_headers(&access_token, &ctx.request_id)?; let client = self.client.clone(); let started = Instant::now(); @@ -488,7 +582,11 @@ impl VertexBridge { .into(), )); } - let headers = build_request_headers(&creds.access_token, &ctx.request_id)?; + // Resolve bearer (pre-minted OR minted-from-SA) BEFORE + // entering the stream future so token-mint errors surface + // as a direct Err return rather than being yielded mid-stream. + let access_token = creds.resolve_access_token(&self.token_minter).await?; + let headers = build_request_headers(&access_token, &ctx.request_id)?; let client = self.client.clone(); let started = Instant::now(); @@ -1006,7 +1104,8 @@ mod tests { fn vertex_secret_parses_full_form() { let json = r#"{"access_token":"ya29.test","project":"my-proj","region":"us-central1"}"#; let s = VertexSecret::parse(json).unwrap(); - assert_eq!(s.access_token, "ya29.test"); + assert_eq!(s.access_token.as_deref(), Some("ya29.test")); + assert!(s.service_account_json.is_none()); assert_eq!(s.project, "my-proj"); assert_eq!(s.region, "us-central1"); } @@ -1050,6 +1149,83 @@ mod tests { } } + #[test] + fn vertex_secret_accepts_service_account_json_path() { + let json = serde_json::json!({ + "service_account_json": { + "type": "service_account", + "private_key": "-----BEGIN PRIVATE KEY-----\nFAKE_PEM_BUT_VALID_HEADER\n-----END PRIVATE KEY-----", + "client_email": "tester@my-proj.iam.gserviceaccount.com", + "token_uri": "https://oauth2.googleapis.com/token", + }, + "project": "my-proj", + "region": "us-central1" + }); + let s = VertexSecret::parse(&json.to_string()).unwrap(); + assert!(s.access_token.is_none()); + let sa = s.service_account_json.as_ref().unwrap(); + assert_eq!(sa.client_email, "tester@my-proj.iam.gserviceaccount.com"); + assert_eq!(sa.token_uri, "https://oauth2.googleapis.com/token"); + } + + #[test] + fn vertex_secret_rejects_both_credential_modes_set() { + let json = serde_json::json!({ + "access_token": "ya29.foo", + "service_account_json": { + "type": "service_account", + "private_key": "-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----", + "client_email": "x@y.z", + "token_uri": "https://oauth2.googleapis.com/token", + }, + "project": "my-proj", + "region": "us-central1" + }); + let err = VertexSecret::parse(&json.to_string()).unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!( + msg.contains("exactly one of access_token or service_account_json"), + "got: {msg}" + ); + assert!(msg.contains("both were provided")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[test] + fn vertex_secret_rejects_neither_credential_mode_set() { + let json = r#"{"project":"my-proj","region":"us-central1"}"#; + let err = VertexSecret::parse(json).unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!( + msg.contains("either access_token or service_account_json"), + "got: {msg}" + ); + assert!(msg.contains("neither was provided")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[test] + fn vertex_secret_rejects_empty_access_token_string() { + // Edge case: operator pastes an empty string into access_token + // rather than omitting the field entirely. The pre-mint path + // would build an Authorization header of "Bearer " which would + // 401 upstream — better to fail at parse time. + let json = r#"{"access_token":"","project":"my-proj","region":"us-central1"}"#; + let err = VertexSecret::parse(json).unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("access_token is empty"), "got: {msg}"); + } + other => panic!("expected Config error, got {other:?}"), + } + } + // ─── URL token validation ────────────────────────────────────────── #[test] @@ -2260,4 +2436,103 @@ mod tests { "expected ?alt=sse query, got {url}" ); } + + // ─── Service-account credential path (D5.1) ───────────────────── + + /// Embedded test SA private key (PEM). Generated via: + /// `openssl genpkey -algorithm RSA -pkeyopt rsa_keygen_bits:2048` + /// Deterministic so the JWT byte sequence is reproducible. NOT a + /// real GCP key — safe to commit. + const TEST_SA_PRIVATE_PEM: &str = include_str!("../test-fixtures/test_sa_private.pem"); + + fn sa_credential_secret(token_uri: &str) -> String { + // Operator's secret JSON, SA-mode (no access_token field). + // Serialize via serde_json so the multi-line PEM lands as a + // proper JSON-escaped string ("\n" escapes). + serde_json::json!({ + "service_account_json": { + "type": "service_account", + "private_key": TEST_SA_PRIVATE_PEM, + "client_email": "tester@my-proj.iam.gserviceaccount.com", + "token_uri": token_uri, + }, + "project": "my-proj", + "region": "us-central1" + }) + .to_string() + } + + /// End-to-end: SA JSON in secret → bridge resolves token via + /// in-process minter → minted token used as Authorization Bearer + /// on the upstream Gemini chat call. + /// + /// Pins the full D5.1 pipeline: VertexSecret SA parse + + /// TokenMinter JWT sign + mint endpoint POST + cache insert + + /// chat_gemini header forwarding. A regression that broke any + /// link in this chain surfaces here. + #[tokio::test] + async fn chat_gemini_sa_path_mints_token_and_forwards_as_bearer() { + // Two wiremock servers: one for GCP OAuth token endpoint, + // one for Vertex AI Gemini chat. + let oauth_server = MockServer::start().await; + let vertex_server = MockServer::start().await; + + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "ya29.minted-by-mock-oauth", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) // exactly one mint despite two chats below + .mount(&oauth_server) + .await; + + // Capture the Authorization header on the Vertex chat call so + // we can assert the bridge actually forwarded the minted + // token (not a hardcoded placeholder). + let captured_auth: std::sync::Arc>> = + std::sync::Arc::new(std::sync::Mutex::new(None)); + let captured_for_responder = captured_auth.clone(); + Mock::given(method("POST")) + .and(path( + "/v1/projects/my-proj/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + )) + .respond_with(move |req: &MockRequest| { + let auth = req + .headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + *captured_for_responder.lock().unwrap() = auth; + default_gemini_response_template() + }) + .mount(&vertex_server) + .await; + + let bridge = VertexBridge::new() + .with_api_base_override(vertex_server.uri()) + .with_token_endpoint_override(oauth_server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("gemini-1.5-pro"), + sample_pk_with_secret(&sa_credential_secret(&oauth_server.uri())), + ); + let req = ChatFormat::new("my-gemini", vec![ChatMessage::user("hi")]); + + // Two chats — second one MUST hit the cache (no second mint). + let _r1 = bridge.chat(&req, &ctx).await.unwrap(); + let _r2 = bridge.chat(&req, &ctx).await.unwrap(); + + let auth = captured_auth + .lock() + .unwrap() + .clone() + .expect("authorization captured"); + assert_eq!( + auth, "Bearer ya29.minted-by-mock-oauth", + "bridge must forward the minted token verbatim" + ); + // wiremock's .expect(1) on the OAuth mock fires here at drop — + // proves the cache prevented a second mint on the second chat. + } } diff --git a/crates/aisix-provider-vertex/src/lib.rs b/crates/aisix-provider-vertex/src/lib.rs index 437c0ac2..d12c205a 100644 --- a/crates/aisix-provider-vertex/src/lib.rs +++ b/crates/aisix-provider-vertex/src/lib.rs @@ -4,18 +4,20 @@ //! //! ## Status (issue #302 Phase E) //! +//! - [x] D5.1 — In-process GCP OAuth2 token mint. The bridge now +//! accepts EITHER a pre-minted `access_token` (operator manages +//! refresh, backward-compatible) OR a full `service_account_json` +//! in `ProviderKey.secret`. When the SA path is taken, the bridge +//! signs a JWT with the SA's RSA private key, posts to the SA's +//! `token_uri` to mint an OAuth2 access token, and caches it +//! in-process keyed by SA `client_email` with TTL refresh ~60s +//! before the upstream-reported expiry. //! - [x] D5.2.a — Gemini publisher chat dispatch //! (`publishers/google/models/:generateContent`) //! - [x] D5.2.b — Gemini streaming via //! `:streamGenerateContent?alt=sse` (SSE chunks, no `[DONE]` //! sentinel — Gemini closes the connection cleanly) //! - [x] D5.5 — `BridgeContext.deadline` plumbing on `chat()` -//! - [ ] D5.1 — In-process GCP OAuth2 token mint (`yup-oauth2` / -//! `gcp_auth` service-account JSON → access token with -//! auto-refresh). **Today the bridge expects a pre-minted -//! access token** in `ProviderKey.secret.access_token`; -//! operators are responsible for refresh (GCP tokens TTL -//! ~1 hour). Follow-up will lift that burden into the bridge. //! - [ ] D5.3 — Anthropic-on-Vertex dispatch //! (`publishers/anthropic/models/:rawPredict`, //! `anthropic_version: "vertex-2023-10-16"`) @@ -51,6 +53,7 @@ #![deny(rust_2018_idioms)] mod bridge; +mod token_mint; mod wire; pub use bridge::{VertexBridge, VertexPublisher}; diff --git a/crates/aisix-provider-vertex/src/token_mint.rs b/crates/aisix-provider-vertex/src/token_mint.rs new file mode 100644 index 00000000..03cfae42 --- /dev/null +++ b/crates/aisix-provider-vertex/src/token_mint.rs @@ -0,0 +1,445 @@ +//! GCP service-account → OAuth2 access token minting + in-process cache. +//! +//! Maps a service-account JSON credential into a short-lived bearer +//! token via the standard JWT-bearer assertion grant flow. +//! +//! First, build a JWT with claims `{iss, scope, aud, iat, exp}` and +//! sign with RS256 using the SA's RSA `private_key`. Then POST +//! `grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer` plus +//! `assertion=` to the SA's `token_uri` (typically +//! `https://oauth2.googleapis.com/token`). Parse `{access_token, +//! expires_in}` from the response. Cache keyed by SA `client_email` +//! with TTL refresh ~60s before expiry so an in-flight request never +//! lands on an expired token. +//! +//! # References +//! +//! - GCP OAuth2 SA flow: +//! +//! - JWT Bearer grant (RFC 7523): +//! +//! - Standard SA JSON shape: emitted verbatim by +//! `gcloud iam service-accounts keys create`. + +use aisix_gateway::BridgeError; +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use tokio::sync::RwLock; + +/// Default scope for Vertex AI access — same as gcloud + the official +/// python-aiplatform SDK default. Narrower scopes (e.g. +/// `aiplatform.googleapis.com/cloud-platform`) work too but `cloud- +/// platform` is the SA-key default GCP recommends. +const VERTEX_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +/// JWT validity window per Google's docs: 1 hour. The returned token +/// has its OWN TTL (also typically 1h); cache uses that, not this. +const JWT_EXPIRY_SECS: u64 = 3600; + +/// Refresh cached tokens at least this many seconds before their +/// reported expiry. Prevents a request from picking up a token that +/// expires while the request is mid-flight. +const TOKEN_REFRESH_SAFETY_MARGIN: Duration = Duration::from_secs(60); + +/// Standard GCP service-account JSON key shape — minimum fields +/// needed for minting. Field names match the on-disk JSON +/// `gcloud iam service-accounts keys create` emits; fields we +/// don't consume (`project_id`, `private_key_id`, `client_id`, +/// `auth_uri`, `auth_provider_x509_cert_url`, +/// `client_x509_cert_url`) are omitted from the struct — serde's +/// default deserializer silently ignores them, so the operator +/// can paste the whole SA JSON verbatim without trimming. +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct ServiceAccountKey { + /// Discriminator; must equal `"service_account"` for normal SA + /// keys (vs. `"external_account"` / `"authorized_user"`). We + /// only support SA keys. + #[serde(rename = "type")] + pub typ: String, + /// PEM-encoded RSA private key. Multi-line in source JSON + /// (`\n` escapes); jsonwebtoken's `from_rsa_pem` decodes from + /// the byte slice directly. + pub private_key: String, + pub client_email: String, + pub token_uri: String, +} + +impl ServiceAccountKey { + /// Cheap shape checks at parse time so the operator gets a fast, + /// actionable error rather than waiting for a JWT-sign failure on + /// the first chat. We intentionally do NOT attempt PEM parsing + /// here (that's a heavier operation with its own error class); + /// the first mint will catch a malformed key with a clear message. + pub fn validate(&self) -> Result<(), BridgeError> { + if self.typ != "service_account" { + return Err(BridgeError::Config(format!( + "vertex service_account_json.type = {:?}, want \"service_account\"", + self.typ + ))); + } + if !self.private_key.starts_with("-----BEGIN") { + return Err(BridgeError::Config( + "vertex service_account_json.private_key is not PEM-formatted \ + (expected `-----BEGIN PRIVATE KEY-----` or `-----BEGIN RSA PRIVATE KEY-----`)" + .into(), + )); + } + if self.client_email.is_empty() { + return Err(BridgeError::Config( + "vertex service_account_json.client_email is empty".into(), + )); + } + if self.token_uri.is_empty() { + return Err(BridgeError::Config( + "vertex service_account_json.token_uri is empty".into(), + )); + } + Ok(()) + } +} + +/// JWT claims for the bearer assertion. Field names per RFC 7523 §3. +#[derive(Serialize)] +struct JwtClaims<'a> { + iss: &'a str, + scope: &'a str, + aud: &'a str, + iat: u64, + exp: u64, +} + +/// Token-endpoint response shape per OAuth2 spec. `token_type` field +/// is always `"Bearer"` for SA flows; discarded here. +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + expires_in: u64, +} + +#[derive(Clone)] +struct CachedToken { + access_token: String, + expires_at: Instant, +} + +/// In-process token cache + minter. One instance per VertexBridge. +/// Cache is keyed by SA's `client_email` so multiple ProviderKeys +/// backed by the same SA share a token slot. +pub(crate) struct TokenMinter { + client: Client, + cache: Arc>>, + /// Test-only override for the SA-supplied `token_uri`. In + /// production we POST directly to the SA's own `token_uri` + /// (typically `https://oauth2.googleapis.com/token`). Tests + /// substitute a wiremock URI here so the assertion flow runs + /// without leaving the test process. + #[cfg(test)] + token_endpoint_override: Option, +} + +impl TokenMinter { + pub fn new(client: Client) -> Self { + Self { + client, + cache: Arc::new(RwLock::new(HashMap::new())), + #[cfg(test)] + token_endpoint_override: None, + } + } + + /// Test-only seam: replace the SA's `token_uri` host with this + /// URL. JWT contents are unchanged so the assertion shape is + /// still verifiable end-to-end. + #[cfg(test)] + pub(crate) fn with_token_endpoint_override(mut self, url: impl Into) -> Self { + self.token_endpoint_override = Some(url.into()); + self + } + + /// Resolve an access token for `sa`. Returns the cached token + /// when one exists and is unexpired; otherwise mints a fresh one + /// and caches it. + pub async fn get_token(&self, sa: &ServiceAccountKey) -> Result { + // Read-lock for the common path. + { + let cache = self.cache.read().await; + if let Some(cached) = cache.get(&sa.client_email) { + if cached.expires_at > Instant::now() { + return Ok(cached.access_token.clone()); + } + } + } + // Cache miss or expired — mint fresh under write-lock. + let (access_token, expires_in_secs) = self.mint(sa).await?; + let cached = CachedToken { + access_token: access_token.clone(), + expires_at: Instant::now() + + Duration::from_secs(expires_in_secs).saturating_sub(TOKEN_REFRESH_SAFETY_MARGIN), + }; + self.cache + .write() + .await + .insert(sa.client_email.clone(), cached); + Ok(access_token) + } + + /// Mint a fresh token by signing the JWT and POSTing to the + /// token endpoint. Returns `(access_token, expires_in_seconds)`. + async fn mint(&self, sa: &ServiceAccountKey) -> Result<(String, u64), BridgeError> { + sa.validate()?; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|e| { + BridgeError::Config(format!( + "vertex token mint: system clock before UNIX epoch: {e}" + )) + })? + .as_secs(); + let claims = JwtClaims { + iss: &sa.client_email, + scope: VERTEX_SCOPE, + aud: &sa.token_uri, + iat: now, + exp: now + JWT_EXPIRY_SECS, + }; + let header = Header::new(Algorithm::RS256); + let key = EncodingKey::from_rsa_pem(sa.private_key.as_bytes()).map_err(|e| { + BridgeError::Config(format!( + "vertex service_account_json.private_key invalid PEM: {e}" + )) + })?; + let jwt = encode(&header, &claims, &key) + .map_err(|e| BridgeError::Config(format!("vertex JWT sign failed: {e}")))?; + + let endpoint = self.resolve_token_endpoint(sa); + let resp = self + .client + .post(&endpoint) + .form(&[ + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ("assertion", &jwt), + ]) + .send() + .await + .map_err(|e| { + BridgeError::Transport(format!("vertex token mint POST {endpoint}: {e}")) + })?; + + let status = resp.status(); + if !status.is_success() { + // Cap body to 500 chars to keep error messages bounded. + // GCP returns OAuth-shape errors `{error, error_description}` + // — text is operator-actionable (invalid_grant etc.) and + // does not echo the SA's private key. + let body = resp.text().await.unwrap_or_default(); + let truncated: String = body.chars().take(500).collect(); + return Err(BridgeError::Config(format!( + "vertex token mint upstream returned HTTP {status}: {truncated}" + ))); + } + let parsed: TokenResponse = resp + .json() + .await + .map_err(|e| BridgeError::UpstreamDecode(format!("vertex token mint response: {e}")))?; + Ok((parsed.access_token, parsed.expires_in)) + } + + fn resolve_token_endpoint(&self, sa: &ServiceAccountKey) -> String { + #[cfg(test)] + if let Some(base) = &self.token_endpoint_override { + return base.clone(); + } + sa.token_uri.clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use jsonwebtoken::{decode, DecodingKey, Validation}; + use wiremock::matchers::{body_string_contains, method}; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + + /// Generate a fresh 2048-bit RSA key pair PEM-encoded for tests. + /// Returns `(private_pem, public_pem)`. Uses jsonwebtoken's + /// internal RustCrypto dep transitively — we keep it test-only + /// to avoid pulling another crate into prod. + fn test_key_pair() -> (String, String) { + // Hand-baked deterministic 2048-bit RSA key pair for tests. + // Generated once via `openssl genpkey -algorithm RSA -out test.pem -pkeyopt rsa_keygen_bits:2048` + // + `openssl rsa -in test.pem -pubout -out test.pub.pem`. + // Deterministic so tests reproduce the same JWT signature byte + // sequence across runs. NOT a real GCP key — safe to commit. + let private_pem = include_str!("../test-fixtures/test_sa_private.pem").to_string(); + let public_pem = include_str!("../test-fixtures/test_sa_public.pem").to_string(); + (private_pem, public_pem) + } + + fn sample_sa(private_pem: &str, token_uri: &str) -> ServiceAccountKey { + ServiceAccountKey { + typ: "service_account".into(), + private_key: private_pem.to_string(), + client_email: "tester@my-proj.iam.gserviceaccount.com".into(), + token_uri: token_uri.to_string(), + } + } + + #[tokio::test] + async fn mint_signs_jwt_with_correct_claims_and_posts_to_token_uri() { + let (private_pem, public_pem) = test_key_pair(); + let server = MockServer::start().await; + + // Capture the inbound JWT assertion so we can decode + verify + // its claims independently using the public key. + let captured: Arc>> = Arc::new(std::sync::Mutex::new(None)); + let captured_for_responder = captured.clone(); + Mock::given(method("POST")) + .and(body_string_contains("grant_type=urn")) + .respond_with(move |req: &Request| { + let body = String::from_utf8(req.body.clone()).unwrap(); + // Body is form-encoded: grant_type=...&assertion= + let assertion = body + .split('&') + .find_map(|kv| kv.strip_prefix("assertion=")) + .unwrap_or_default(); + *captured_for_responder.lock().unwrap() = Some(urlencoding_decode(assertion)); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "ya29.minted-by-mock", + "expires_in": 3600, + "token_type": "Bearer" + })) + }) + .mount(&server) + .await; + + let sa = sample_sa(&private_pem, &server.uri()); + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let token = minter.get_token(&sa).await.unwrap(); + assert_eq!(token, "ya29.minted-by-mock"); + + // Decode the JWT against the matching public key and check claims. + let jwt = captured.lock().unwrap().clone().expect("JWT captured"); + let decoding_key = DecodingKey::from_rsa_pem(public_pem.as_bytes()).unwrap(); + let mut validation = Validation::new(Algorithm::RS256); + validation.set_audience(&[&server.uri()]); + validation.validate_exp = false; // we test exp directly below + let token_data = decode::(&jwt, &decoding_key, &validation).unwrap(); + assert_eq!( + token_data.claims["iss"].as_str().unwrap(), + "tester@my-proj.iam.gserviceaccount.com" + ); + assert_eq!(token_data.claims["scope"].as_str().unwrap(), VERTEX_SCOPE); + assert_eq!(token_data.claims["aud"].as_str().unwrap(), server.uri()); + let iat = token_data.claims["iat"].as_u64().unwrap(); + let exp = token_data.claims["exp"].as_u64().unwrap(); + assert_eq!(exp - iat, JWT_EXPIRY_SECS); + } + + #[tokio::test] + async fn get_token_caches_within_ttl_and_only_calls_endpoint_once() { + let (private_pem, _) = test_key_pair(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "ya29.cache-hit", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) // critical: must be called EXACTLY once across 3 get_token calls + .mount(&server) + .await; + + let sa = sample_sa(&private_pem, &server.uri()); + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + for _ in 0..3 { + assert_eq!(minter.get_token(&sa).await.unwrap(), "ya29.cache-hit"); + } + } + + #[tokio::test] + async fn token_endpoint_5xx_surfaces_as_config_error_with_body_truncated() { + let (private_pem, _) = test_key_pair(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(503) + .set_body_string("upstream OAuth backend transient failure"), + ) + .mount(&server) + .await; + let sa = sample_sa(&private_pem, &server.uri()); + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sa).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("HTTP 503")); + assert!(msg.contains("upstream OAuth backend transient failure")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[tokio::test] + async fn invalid_pem_surfaces_clear_error_before_endpoint_call() { + let server = MockServer::start().await; + // Mock with .expect(0) — endpoint must NOT be called for a + // pre-flight validation failure. + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let sa = ServiceAccountKey { + typ: "service_account".into(), + private_key: "not a PEM at all".into(), + client_email: "x@y.z".into(), + token_uri: server.uri(), + }; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sa).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("not PEM-formatted")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[tokio::test] + async fn wrong_type_field_rejected_before_endpoint_call() { + let (private_pem, _) = test_key_pair(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let sa = ServiceAccountKey { + typ: "external_account".into(), + private_key: private_pem, + client_email: "x@y.z".into(), + token_uri: server.uri(), + }; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sa).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("type")); + assert!(msg.contains("external_account")); + assert!(msg.contains("service_account")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + /// Minimal URL-decode for the captured assertion field. Sufficient + /// for the JWT character set (`A-Za-z0-9-_.`), which the standard + /// form-encoder leaves untouched. + fn urlencoding_decode(s: &str) -> String { + s.replace('+', " ").replace("%2B", "+").replace("%2F", "/") + } +} diff --git a/crates/aisix-provider-vertex/test-fixtures/test_sa_private.pem b/crates/aisix-provider-vertex/test-fixtures/test_sa_private.pem new file mode 100644 index 00000000..a7c9a8ab --- /dev/null +++ b/crates/aisix-provider-vertex/test-fixtures/test_sa_private.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCnYSL6c/voSJFY +eF9wrJ1SY9x8KmypY2bWP3gPkiNc3A/3AYWsVTUU3oBTkb+SBGGxRImiZBhRBvRQ +snArvu/shy2MKiOT1tgpByABimKQWZMg9nGETnUzHslaISBJPNYPdpaknJf9Cyt0 +KK5/hlMVQm0PNqW5rj2Z5UYkXMYqUQ4HpOxjld0kqOq/uQLC9l/edB3xUSuyJCeY +L/bfAEloTtFd16rPkEZ5Jh05kggK1x2I/Z8lrEp/wFXDbZxKpT2YMfj0ztVrQmf7 +RlYxcfGicRJV3SKR8aVHHrXVjXBNSW/EZpn1x2MQeZIXOfcC5QgP6Ks4bcYGbuji +KLAlXXtnAgMBAAECggEAbYJSHkbQI7OG1Lk8yD1HWOZZFSu0mEaeu8IezSEx3clk +8JigWpYM+rBwiTysd95CBHbxDbwrZKgGJN36IcT0uG3g0Pmo+UrxdjZhLGDcB9Fz +P3e94XBroZyc5EkUFJam/sr52I1Tq40pSwBq2qiJpzkknXWFjCyBxTSKZbQFOx1q +DE0YMtizUo7n3/wsiHdt8mI0ABQ2YGjG7E660r9vihkCa3G08T23qZiTQNwBGRqP +RCXy4pah4MvKlGnEPLweC/vDN06hEbpx6PcLhkj83l83zNgP9+qq6aRRi1usgRCv +KM5w2pAFYWu15XHGUBVR+Msxhjk659NilGJtqMZh0QKBgQDRUY7bItjcxsJKjBi0 +YAFSrBCgPw9WYnNXM+wGB0XVNJ7ewyf2lG8zEiA3Krbfc0MYZXeREqA9vQDmXSri +AcDHUU7DPfgUj2AvDz7QXVIwP/L9n39WkN6lxK14lWlHXf70xEQQVMdJ8R9wnoCK +t1iUHlmYRGTJ9ejvwDHYOboNQwKBgQDMtS8OfrGrJjszRZjN4x5gWqPu+j9PniWr +1Y/zFCjST4GXFjnGGcxfuHVQT6bPYyXgGx854dd4spHlfkiZQ+k1tkv+4Bt+nvDB +rBSVSLkzUPZ6wjBMITQIkiAR94WJ1JJ4dc1v6gyf2r7O5MdxayiFAaLhy9ZxtiQp +x/XBre2FDQKBgFzl038SIikp1UT8lGJJUYz9bIuSMR5np0UGeDPcunN7XR8EghH/ +orKJ0t5pCKx3HUoQjlZGa/O6lFGo+8U+fe53+XrRX+7QCyIXpAsZv8ZGO3Owe/VR +al8rwMmJliXkY6kCCistVR1N9GQpFGd8I8XpCl53zDuN9gmhxP1v8VC9AoGAe6On +M0sERkoGEZakjx3xN+MnBmzxFkZ/nESV+AwiB7xrmfSbmnH0hY/kk0g4iSPqOWxI +NO6Z9NVt1z2p3aAt1/ot9lgnYxfedCtaFzxgV4U8CbMF9sVLJy4S3qcwaaoReV41 +YbXsQBSfkFiPuYouY/80AMrbz7xiJTYX0g4Z2nUCgYEAxafHCeJm5CxaM2jRAml7 +NCBCEnxJw+ijv07cisJEOxBxxi40OQMzH/Lb1XcPpoVJWwgTIi7glBK2taJFddS5 +coBMnvTVDdSi9/rbUzO2VKicA+rfV/0+vK5fdfNHAkHiiZ3tWQyR25pX3hvx1JMm +kK6tKMn5RTtkhtuFKeiGTqM= +-----END PRIVATE KEY----- diff --git a/crates/aisix-provider-vertex/test-fixtures/test_sa_public.pem b/crates/aisix-provider-vertex/test-fixtures/test_sa_public.pem new file mode 100644 index 00000000..0308e33e --- /dev/null +++ b/crates/aisix-provider-vertex/test-fixtures/test_sa_public.pem @@ -0,0 +1,9 @@ +-----BEGIN PUBLIC KEY----- +MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAp2Ei+nP76EiRWHhfcKyd +UmPcfCpsqWNm1j94D5IjXNwP9wGFrFU1FN6AU5G/kgRhsUSJomQYUQb0ULJwK77v +7IctjCojk9bYKQcgAYpikFmTIPZxhE51Mx7JWiEgSTzWD3aWpJyX/QsrdCiuf4ZT +FUJtDzalua49meVGJFzGKlEOB6TsY5XdJKjqv7kCwvZf3nQd8VErsiQnmC/23wBJ +aE7RXdeqz5BGeSYdOZIICtcdiP2fJaxKf8BVw22cSqU9mDH49M7Va0Jn+0ZWMXHx +onESVd0ikfGlRx611Y1wTUlvxGaZ9cdjEHmSFzn3AuUID+irOG3GBm7o4iiwJV17 +ZwIDAQAB +-----END PUBLIC KEY----- From fbd8d9575e978aa5411f58899feb55e746a8aaaf Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 24 May 2026 14:19:40 +0800 Subject: [PATCH 2/2] fix(vertex): classify token-endpoint 5xx as UpstreamStatus (audit MEDIUM on #387) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Audit-aigw-387-vertex-oauth flagged that all non-2xx responses from the GCP token endpoint were returned as BridgeError::Config (HTTP 500 "operator must fix"). A flapping GCP OAuth backend would surface to the customer as an unconditional 500 with no retry hint — wrong classification for a transient upstream failure. Fix per audit's suggested code: - 5xx → BridgeError::upstream_status_with_retry_after(status, msg, retry_after) Surfaces the upstream status (cooldown layer maps to 502), and forwards GCP's Retry-After header so a token-endpoint backoff hint flows end-to-end through the cooldown layer rather than being silently dropped. - 4xx → BridgeError::Config (unchanged) invalid_grant / clock skew / bad SA ARE operator-actionable; Config (500) remains correct. Reads Retry-After BEFORE consuming the body via resp.text() so the header survives the body read. Test updates: - token_endpoint_5xx_surfaces_as_config_error_with_body_truncated → token_endpoint_5xx_surfaces_as_upstream_status_with_retry_hint (rename + reshape; asserts status=503, retry_after=Some(30s), message contains body) - New: token_endpoint_4xx_surfaces_as_config_error — pins the 4xx-still-Config branch with an invalid_grant fixture so a future regression that broadened the upstream_status mapping surfaces here `cargo test -p aisix-provider-vertex` → 61/61 PASS (was 60; +1 new 4xx test). `cargo clippy -p aisix-provider-vertex --all-targets -- -D warnings` clean. Audit LOWs (non-blocking, deferred per CLAUDE.md §8): - LOW: no e2e for chat_gemini_stream SA path — both code paths call resolve_access_token at the same lifecycle point; a regression would have to diverge both sites - LOW: cache-miss race on first-call burst — acceptable (worst case one extra mint; spec doesn't require single-flight) - LOW: expires_in:0 edge case — saturating_sub handles cleanly --- .../aisix-provider-vertex/src/token_mint.rs | 60 +++++++++++++++++-- 1 file changed, 54 insertions(+), 6 deletions(-) diff --git a/crates/aisix-provider-vertex/src/token_mint.rs b/crates/aisix-provider-vertex/src/token_mint.rs index 03cfae42..dcf36d65 100644 --- a/crates/aisix-provider-vertex/src/token_mint.rs +++ b/crates/aisix-provider-vertex/src/token_mint.rs @@ -231,15 +231,28 @@ impl TokenMinter { let status = resp.status(); if !status.is_success() { + // Read Retry-After BEFORE consuming the body via .text() so a + // 429/503 from GCP flows the upstream backoff hint into the + // cooldown layer. + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); // Cap body to 500 chars to keep error messages bounded. // GCP returns OAuth-shape errors `{error, error_description}` // — text is operator-actionable (invalid_grant etc.) and // does not echo the SA's private key. let body = resp.text().await.unwrap_or_default(); let truncated: String = body.chars().take(500).collect(); - return Err(BridgeError::Config(format!( - "vertex token mint upstream returned HTTP {status}: {truncated}" - ))); + let msg = format!("vertex token mint upstream returned HTTP {status}: {truncated}"); + // Audit MEDIUM (PR #387): classify 5xx as a transient + // upstream failure (502 with cooldown semantics) rather + // than `Config` (500, operator must fix). A flapping GCP + // token endpoint should not look operator-actionable to + // the customer. 4xx (invalid_grant / bad SA / clock skew) + // IS operator-actionable, so it stays Config. + return Err(if status.is_server_error() { + BridgeError::upstream_status_with_retry_after(status.as_u16(), msg, retry_after) + } else { + BridgeError::Config(msg) + }); } let parsed: TokenResponse = resp .json() @@ -361,12 +374,16 @@ mod tests { } #[tokio::test] - async fn token_endpoint_5xx_surfaces_as_config_error_with_body_truncated() { + async fn token_endpoint_5xx_surfaces_as_upstream_status_with_retry_hint() { + // Audit MEDIUM (PR #387): 5xx is transient upstream — must + // be UpstreamStatus (502 with cooldown semantics), not + // Config (500 operator-must-fix). let (private_pem, _) = test_key_pair(); let server = MockServer::start().await; Mock::given(method("POST")) .respond_with( ResponseTemplate::new(503) + .insert_header("Retry-After", "30") .set_body_string("upstream OAuth backend transient failure"), ) .mount(&server) @@ -374,10 +391,41 @@ mod tests { let sa = sample_sa(&private_pem, &server.uri()); let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); let err = minter.get_token(&sa).await.err().unwrap(); + match err { + BridgeError::UpstreamStatus { + status, + message, + retry_after, + .. + } => { + assert_eq!(status, 503); + assert!(message.contains("HTTP 503")); + assert!(message.contains("upstream OAuth backend transient failure")); + assert_eq!(retry_after, Some(std::time::Duration::from_secs(30))); + } + other => panic!("expected UpstreamStatus error, got {other:?}"), + } + } + + #[tokio::test] + async fn token_endpoint_4xx_surfaces_as_config_error() { + // Audit MEDIUM (PR #387): 4xx remains Config — invalid_grant + // / bad SA / clock skew are operator-actionable. + let (private_pem, _) = test_key_pair(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(400).set_body_string( + r#"{"error":"invalid_grant","error_description":"JWT iat/exp out of range"}"#, + )) + .mount(&server) + .await; + let sa = sample_sa(&private_pem, &server.uri()); + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sa).await.err().unwrap(); match err { BridgeError::Config(msg) => { - assert!(msg.contains("HTTP 503")); - assert!(msg.contains("upstream OAuth backend transient failure")); + assert!(msg.contains("HTTP 400")); + assert!(msg.contains("invalid_grant")); } other => panic!("expected Config error, got {other:?}"), }