yapi_ai/auth/
google_adc.rs

1//! Google Application Default Credentials: an access token for Vertex AI
2//! from a credentials file.
3//!
4//! Port of the parts of `google-auth-library` 10.6 that pi `v1.0.0`'s Vertex
5//! provider reaches: the file named by `GOOGLE_APPLICATION_CREDENTIALS` or
6//! gcloud's well-known file, holding an `authorized_user`, a
7//! `service_account`, an `impersonated_service_account`, or an
8//! `external_account` whose subject token comes from a file or a URL.
9
10use std::collections::HashMap;
11use std::sync::{Mutex, OnceLock};
12
13use base64::Engine as _;
14use serde_json::{Map, Value, json};
15use tokio_util::sync::CancellationToken;
16
17use super::{AuthError, now_ms};
18
19/// The scope Vertex AI requests.
20pub const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
21/// Google's OAuth 2.0 token endpoint.
22pub const OAUTH2_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
23const JWT_BEARER: &str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
24const TOKEN_EXCHANGE: &str = "urn:ietf:params:oauth:grant-type:token-exchange";
25const ACCESS_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:access_token";
26/// Tokens this close to expiry are fetched again, as the library's
27/// `eagerRefreshThresholdMillis`.
28const EAGER_REFRESH_MS: u64 = 5 * 60 * 1000;
29const NO_CREDENTIALS: &str = "Could not load the default credentials. Browse to https://cloud.google.com/docs/authentication/getting-started for more information.";
30
31/// A token for a request: the bearer token and the quota project billed for
32/// it, sent as `x-goog-user-project`.
33#[derive(Clone, Debug, PartialEq, Eq)]
34pub struct Token {
35    /// The OAuth access token.
36    pub access_token: String,
37    /// The credentials' `quota_project_id`.
38    pub quota_project: Option<String>,
39}
40
41/// The credentials file: `GOOGLE_APPLICATION_CREDENTIALS`, else gcloud's
42/// `application_default_credentials.json` under `CLOUDSDK_CONFIG` or
43/// `~/.config/gcloud`, if it exists. `env` reads the request's provider
44/// settings, then the process environment.
45pub fn credentials_file(env: impl Fn(&str) -> Option<String>) -> Option<String> {
46    if let Some(path) = env("GOOGLE_APPLICATION_CREDENTIALS") {
47        return Some(path);
48    }
49    let dir = env("CLOUDSDK_CONFIG")
50        .map(std::path::PathBuf::from)
51        .or_else(|| env("HOME").map(|home| std::path::Path::new(&home).join(".config/gcloud")))?;
52    let path = dir.join("application_default_credentials.json");
53    path.exists().then(|| path.to_string_lossy().into_owned())
54}
55
56struct Cached {
57    token: Token,
58    expires_at_ms: u64,
59}
60
61fn cache() -> &'static Mutex<HashMap<String, Cached>> {
62    static CACHE: OnceLock<Mutex<HashMap<String, Cached>>> = OnceLock::new();
63    CACHE.get_or_init(Default::default)
64}
65
66/// An access token from the credentials file at `path`, cached until five
67/// minutes before it expires. `oauth2_token_url` serves `authorized_user`
68/// credentials.
69pub async fn token(
70    path: Option<String>,
71    oauth2_token_url: &str,
72    cancel: &CancellationToken,
73) -> Result<Token, String> {
74    let path = path.ok_or_else(|| NO_CREDENTIALS.to_owned())?;
75    let text = tokio::fs::read_to_string(&path)
76        .await
77        .map_err(|err| file_error(&err, &path))?;
78    let key = format!("{path}\0{text}");
79    if let Ok(cache) = cache().lock()
80        && let Some(cached) = cache.get(&key)
81        && cached.expires_at_ms > now_ms() + EAGER_REFRESH_MS
82    {
83        return Ok(cached.token.clone());
84    }
85    let credentials: Value = serde_json::from_str(&text)
86        .map_err(|err| format!("Unexpected token in JSON at {path}: {err}"))?;
87    let (access_token, expires_at_ms) = fetch(&credentials, oauth2_token_url, cancel).await?;
88    let token = Token {
89        access_token,
90        quota_project: credentials["quota_project_id"]
91            .as_str()
92            .filter(|id| !id.is_empty())
93            .map(str::to_owned),
94    };
95    if let Ok(mut cache) = cache().lock() {
96        cache.insert(
97            key,
98            Cached {
99                token: token.clone(),
100                expires_at_ms,
101            },
102        );
103    }
104    Ok(token)
105}
106
107/// Node's message for a file that cannot be read.
108fn file_error(err: &std::io::Error, path: &str) -> String {
109    match err.kind() {
110        std::io::ErrorKind::NotFound => {
111            format!("ENOENT: no such file or directory, open '{path}'")
112        }
113        std::io::ErrorKind::PermissionDenied => {
114            format!("EACCES: permission denied, open '{path}'")
115        }
116        _ => format!("{err}, open '{path}'"),
117    }
118}
119
120fn field<'a>(credentials: &'a Value, name: &str) -> Result<&'a str, String> {
121    credentials[name]
122        .as_str()
123        .filter(|value| !value.is_empty())
124        .ok_or_else(|| format!("The incoming JSON object does not contain a {name} field"))
125}
126
127/// The token and its expiry for one credentials object.
128fn fetch<'a>(
129    credentials: &'a Value,
130    oauth2_token_url: &'a str,
131    cancel: &'a CancellationToken,
132) -> super::BoxFuture<'a, Result<(String, u64), String>> {
133    Box::pin(async move {
134        match credentials["type"].as_str().unwrap_or_default() {
135            "authorized_user" => {
136                let fields = [
137                    ("refresh_token", field(credentials, "refresh_token")?),
138                    ("client_id", field(credentials, "client_id")?),
139                    ("client_secret", field(credentials, "client_secret")?),
140                    ("grant_type", "refresh_token"),
141                ];
142                oauth_token(oauth2_token_url, &fields, cancel).await
143            }
144            "service_account" => {
145                let token_uri = credentials["token_uri"]
146                    .as_str()
147                    .filter(|uri| !uri.is_empty())
148                    .unwrap_or(OAUTH2_TOKEN_URL);
149                let assertion = signed_jwt(
150                    field(credentials, "client_email")?,
151                    field(credentials, "private_key")?,
152                    token_uri,
153                )?;
154                let fields = [("grant_type", JWT_BEARER), ("assertion", &assertion)];
155                oauth_token(token_uri, &fields, cancel).await
156            }
157            "impersonated_service_account" => {
158                let source = &credentials["source_credentials"];
159                if !source.is_object() {
160                    return Err(
161                        "The incoming JSON object does not contain a source_credentials field"
162                            .into(),
163                    );
164                }
165                let (source_token, _) = fetch(source, oauth2_token_url, cancel).await?;
166                let url = field(credentials, "service_account_impersonation_url")?;
167                let delegates = credentials["delegates"].clone();
168                impersonate(url, &source_token, delegates, cancel).await
169            }
170            "external_account" => external_account(credentials, cancel).await,
171            other => Err(format!(
172                "Unsupported Google credentials type for Vertex AI: {}",
173                if other.is_empty() { "undefined" } else { other }
174            )),
175        }
176    })
177}
178
179/// A self-signed RS256 JWT asking `audience` for a cloud-platform token, as
180/// `gtoken` builds it.
181fn signed_jwt(email: &str, private_key: &str, audience: &str) -> Result<String, String> {
182    let encode = |value: &Value| {
183        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(yapi_types::json::stringify(value))
184    };
185    let issued = now_ms() / 1000;
186    let header = encode(&json!({"alg": "RS256"}));
187    let claims = encode(&json!({
188        "iss": email,
189        "scope": CLOUD_PLATFORM_SCOPE,
190        "aud": audience,
191        "exp": issued + 3600,
192        "iat": issued,
193    }));
194    let message = format!("{header}.{claims}");
195    let signature = rs256(private_key, message.as_bytes())?;
196    Ok(format!(
197        "{message}.{}",
198        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(signature)
199    ))
200}
201
202/// Signs `message` with a PEM RSA key (PKCS#8 or PKCS#1).
203fn rs256(pem: &str, message: &[u8]) -> Result<Vec<u8>, String> {
204    use aws_lc_rs::rand::SystemRandom;
205    use aws_lc_rs::signature::{RSA_PKCS1_SHA256, RsaKeyPair};
206    let body: String = pem
207        .lines()
208        .filter(|line| !line.starts_with("-----"))
209        .flat_map(|line| line.trim().chars())
210        .collect();
211    let der = base64::engine::general_purpose::STANDARD
212        .decode(body)
213        .map_err(|_| "Invalid private key in service account credentials".to_owned())?;
214    let key = RsaKeyPair::from_pkcs8(&der)
215        .or_else(|_| RsaKeyPair::from_der(&der))
216        .map_err(|_| "Invalid private key in service account credentials".to_owned())?;
217    let mut signature = vec![0; key.public_modulus_len()];
218    key.sign(
219        &RSA_PKCS1_SHA256,
220        &SystemRandom::new(),
221        message,
222        &mut signature,
223    )
224    .map_err(|_| "Could not sign the service account assertion".to_owned())?;
225    Ok(signature)
226}
227
228async fn post(
229    request: reqwest::RequestBuilder,
230    cancel: &CancellationToken,
231) -> Result<(u16, Value), String> {
232    let response = super::send(request, cancel)
233        .await
234        .map_err(|err| match err {
235            AuthError::Cancelled => crate::http::ABORTED_BEFORE_RESPONSE.to_owned(),
236            AuthError::Failed(message) => message,
237        })?;
238    let status = response.status().as_u16();
239    Ok((status, super::json_body(response).await))
240}
241
242/// The error a failed token request reports: the OAuth `error` code, else
243/// the status.
244fn token_error(status: u16, body: &Value) -> String {
245    match &body["error"] {
246        Value::String(code) => code.clone(),
247        Value::Object(error) => error.get("message").and_then(Value::as_str).map_or_else(
248            || format!("Request failed with status code {status}"),
249            str::to_owned,
250        ),
251        _ => format!("Request failed with status code {status}"),
252    }
253}
254
255/// An `access_token` and `expires_in` from an OAuth token endpoint.
256async fn oauth_token(
257    url: &str,
258    fields: &[(&str, &str)],
259    cancel: &CancellationToken,
260) -> Result<(String, u64), String> {
261    let request = super::post_form(url, fields);
262    let (status, body) = post(request, cancel).await?;
263    if !(200..300).contains(&status) {
264        return Err(token_error(status, &body));
265    }
266    let token = body["access_token"]
267        .as_str()
268        .filter(|token| !token.is_empty())
269        .ok_or_else(|| "No access token in the token response".to_owned())?;
270    let expires_in = body["expires_in"].as_u64().unwrap_or(3600);
271    Ok((token.to_owned(), now_ms() + expires_in * 1000))
272}
273
274/// IAM Credentials' `generateAccessToken` for a service account, authorized
275/// by `source_token`.
276async fn impersonate(
277    url: &str,
278    source_token: &str,
279    delegates: Value,
280    cancel: &CancellationToken,
281) -> Result<(String, u64), String> {
282    let mut body = Map::new();
283    body.insert(
284        "delegates".into(),
285        if delegates.is_array() {
286            delegates
287        } else {
288            json!([])
289        },
290    );
291    body.insert("scope".into(), json!([CLOUD_PLATFORM_SCOPE]));
292    body.insert("lifetime".into(), json!("3600s"));
293    let request = crate::http::client()
294        .post(url)
295        .header("authorization", format!("Bearer {source_token}"))
296        .header("content-type", "application/json")
297        .body(yapi_types::json::stringify(&Value::Object(body)));
298    let (status, body) = post(request, cancel).await?;
299    if !(200..300).contains(&status) {
300        return Err(format!(
301            "unable to impersonate: {}",
302            token_error(status, &body)
303        ));
304    }
305    let token = body["accessToken"]
306        .as_str()
307        .filter(|token| !token.is_empty())
308        .ok_or_else(|| "No access token in the impersonation response".to_owned())?;
309    let expires_at_ms = body["expireTime"]
310        .as_str()
311        .and_then(yapi_types::time::parse_iso)
312        .unwrap_or_else(|| now_ms() + 3600 * 1000);
313    Ok((token.to_owned(), expires_at_ms))
314}
315
316/// Workload identity federation: an external subject token exchanged at the
317/// Security Token Service, optionally impersonating a service account.
318async fn external_account(
319    credentials: &Value,
320    cancel: &CancellationToken,
321) -> Result<(String, u64), String> {
322    let source = &credentials["credential_source"];
323    let raw = if let Some(path) = source["file"].as_str() {
324        tokio::fs::read_to_string(path)
325            .await
326            .map_err(|err| file_error(&err, path))?
327    } else if let Some(url) = source["url"].as_str() {
328        let mut request = crate::http::client().get(url);
329        if let Some(headers) = source["headers"].as_object() {
330            for (name, value) in headers {
331                if let Some(value) = value.as_str() {
332                    request = request.header(name.as_str(), value);
333                }
334            }
335        }
336        let response = super::send(request, cancel)
337            .await
338            .map_err(|err| match err {
339                AuthError::Cancelled => crate::http::ABORTED_BEFORE_RESPONSE.to_owned(),
340                AuthError::Failed(message) => message,
341            })?;
342        response.text().await.unwrap_or_default()
343    } else {
344        return Err(
345            "Only file and URL credential sources are supported for external accounts".into(),
346        );
347    };
348    let subject = match source["format"]["type"].as_str() {
349        Some("json") => {
350            let name = source["format"]["subject_token_field_name"]
351                .as_str()
352                .unwrap_or_default();
353            let json: Value = serde_json::from_str(&raw)
354                .map_err(|_| "Unable to parse the subject token response as JSON".to_owned())?;
355            json[name].as_str().map(str::to_owned).ok_or_else(|| {
356                format!("Unable to parse the subject_token from the credential_source {name} field")
357            })?
358        }
359        _ => raw.trim().to_owned(),
360    };
361    if subject.is_empty() {
362        return Err("Unable to parse the subject_token from the credential_source".into());
363    }
364    let token_url = credentials["token_url"]
365        .as_str()
366        .unwrap_or("https://sts.googleapis.com/v1/token");
367    let mut pairs = vec![
368        ("grant_type", TOKEN_EXCHANGE.to_owned()),
369        ("audience", field(credentials, "audience")?.to_owned()),
370        ("scope", CLOUD_PLATFORM_SCOPE.to_owned()),
371        ("requested_token_type", ACCESS_TOKEN_TYPE.to_owned()),
372        ("subject_token", subject),
373        (
374            "subject_token_type",
375            field(credentials, "subject_token_type")?.to_owned(),
376        ),
377    ];
378    if let Some(project) = credentials["workforce_pool_user_project"].as_str() {
379        pairs.push((
380            "options",
381            yapi_types::json::stringify(&json!({"userProject": project})),
382        ));
383    }
384    let borrowed: Vec<(&str, &str)> = pairs.iter().map(|(k, v)| (*k, v.as_str())).collect();
385    let (sts_token, expires_at_ms) = oauth_token(token_url, &borrowed, cancel).await?;
386    match credentials["service_account_impersonation_url"].as_str() {
387        Some(url) if !url.is_empty() => impersonate(url, &sts_token, json!([]), cancel).await,
388        _ => Ok((sts_token, expires_at_ms)),
389    }
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    /// A fresh RSA key as a PKCS#8 PEM, with its public key.
397    pub(crate) fn test_key() -> (String, Vec<u8>) {
398        use aws_lc_rs::encoding::AsDer;
399        use aws_lc_rs::rsa::{KeyPair, KeySize};
400        use aws_lc_rs::signature::KeyPair as _;
401        let key = KeyPair::generate(KeySize::Rsa2048).unwrap();
402        let der = AsDer::<aws_lc_rs::encoding::Pkcs8V1Der<'static>>::as_der(&key).unwrap();
403        let body = base64::engine::general_purpose::STANDARD.encode(der.as_ref());
404        let lines: Vec<&str> = body
405            .as_bytes()
406            .chunks(64)
407            .map(|chunk| std::str::from_utf8(chunk).unwrap())
408            .collect();
409        let pem = format!(
410            "-----BEGIN PRIVATE KEY-----\n{}\n-----END PRIVATE KEY-----\n",
411            lines.join("\n")
412        );
413        (pem, key.public_key().as_ref().to_vec())
414    }
415
416    #[test]
417    fn signs_service_account_assertions() {
418        use aws_lc_rs::signature::{RSA_PKCS1_2048_8192_SHA256, UnparsedPublicKey};
419        let (pem, public) = test_key();
420        let jwt = signed_jwt("sa@p.iam.gserviceaccount.com", &pem, OAUTH2_TOKEN_URL).unwrap();
421        let parts: Vec<&str> = jwt.split('.').collect();
422        assert_eq!(parts.len(), 3);
423        assert_eq!(parts[0], "eyJhbGciOiJSUzI1NiJ9");
424        let signature = base64::engine::general_purpose::URL_SAFE_NO_PAD
425            .decode(parts[2])
426            .unwrap();
427        UnparsedPublicKey::new(&RSA_PKCS1_2048_8192_SHA256, &public)
428            .verify(format!("{}.{}", parts[0], parts[1]).as_bytes(), &signature)
429            .unwrap();
430        let claims: Value = serde_json::from_slice(
431            &base64::engine::general_purpose::URL_SAFE_NO_PAD
432                .decode(parts[1])
433                .unwrap(),
434        )
435        .unwrap();
436        assert_eq!(claims["scope"], CLOUD_PLATFORM_SCOPE);
437        assert_eq!(claims["aud"], OAUTH2_TOKEN_URL);
438        assert_eq!(
439            claims["exp"].as_u64().unwrap() - claims["iat"].as_u64().unwrap(),
440            3600
441        );
442        assert_eq!(
443            signed_jwt("x", "not a key", OAUTH2_TOKEN_URL),
444            Err("Invalid private key in service account credentials".into())
445        );
446    }
447}
yapi_ai/auth/google_adc.rs

yapi_ai/auth/
google_adc.rs
