gemini_genai_rs/transport/auth/
google_token.rs

1//! OAuth2 access tokens for Google Cloud APIs, from wherever the process runs.
2//!
3//! [`GoogleAccessToken::from_env`] finds a token the way Google's client
4//! libraries do for a service: an explicit `GOOGLE_ACCESS_TOKEN`, else the
5//! metadata server that Cloud Run, GKE and Compute Engine provide, else the
6//! `gcloud` CLI on a developer machine. Tokens are cached until shortly
7//! before they expire, so calling [`token`](GoogleAccessToken::token) per
8//! request is cheap.
9
10use std::time::{Duration, Instant};
11
12use crate::protocol::types::AccessToken;
13use crate::session::AuthError;
14
15/// Refresh this long before a token's reported expiry.
16const EARLY: Duration = Duration::from_secs(300);
17/// How often [`GoogleAccessToken::into_access_token`] checks its token.
18const REFRESH_CHECK: Duration = Duration::from_secs(60);
19/// How long a `gcloud` token is trusted (it does not report its expiry;
20/// user tokens last an hour).
21const GCLOUD_LIFETIME: Duration = Duration::from_secs(45 * 60);
22
23#[derive(Debug, Clone)]
24enum Source {
25    Fixed(String),
26    Metadata(String),
27    Gcloud,
28    /// Metadata server if one answers, else `gcloud`.
29    Auto(String),
30}
31
32/// A cached source of Google OAuth2 access tokens. See the module docs.
33#[derive(Debug)]
34pub struct GoogleAccessToken {
35    source: Source,
36    client: reqwest::Client,
37    cached: tokio::sync::Mutex<Option<(String, Instant)>>,
38}
39
40impl GoogleAccessToken {
41    /// `GOOGLE_ACCESS_TOKEN` when set; otherwise the metadata server when one
42    /// answers (Cloud Run, GKE, Compute Engine), else `gcloud auth
43    /// print-access-token`. `GCE_METADATA_HOST` overrides the metadata host.
44    pub fn from_env() -> Self {
45        match std::env::var("GOOGLE_ACCESS_TOKEN") {
46            Ok(token) if !token.trim().is_empty() => Self::fixed(token.trim()),
47            _ => Self::with_source(Source::Auto(metadata_host())),
48        }
49    }
50
51    /// Always this token (it is not refreshed).
52    pub fn fixed(token: impl Into<String>) -> Self {
53        Self::with_source(Source::Fixed(token.into()))
54    }
55
56    /// Tokens for the attached service account, from the metadata server at
57    /// `host` (normally `metadata.google.internal`).
58    pub fn metadata_server(host: impl Into<String>) -> Self {
59        Self::with_source(Source::Metadata(host.into()))
60    }
61
62    /// Tokens from `gcloud auth print-access-token`.
63    pub fn gcloud() -> Self {
64        Self::with_source(Source::Gcloud)
65    }
66
67    fn with_source(source: Source) -> Self {
68        Self {
69            source,
70            client: reqwest::Client::new(),
71            cached: tokio::sync::Mutex::new(None),
72        }
73    }
74
75    /// A valid access token, fetching a new one when the cached one is
76    /// about to expire.
77    pub async fn token(&self) -> Result<String, AuthError> {
78        let mut cached = self.cached.lock().await;
79        if let Some((token, valid_until)) = cached.as_ref()
80            && Instant::now() < *valid_until
81        {
82            return Ok(token.clone());
83        }
84        let (token, lifetime) = match &self.source {
85            Source::Fixed(token) => return Ok(token.clone()),
86            Source::Metadata(host) => self.fetch_metadata(host).await?,
87            Source::Gcloud => (gcloud_token().await?, GCLOUD_LIFETIME),
88            Source::Auto(host) => match self.fetch_metadata(host).await {
89                Ok(fetched) => fetched,
90                Err(metadata) => (
91                    gcloud_token().await.map_err(|gcloud| {
92                        AuthError::TokenFetchFailed(format!(
93                            "no Google credentials: set GOOGLE_ACCESS_TOKEN, run on Google \
94                             Cloud, or install the gcloud CLI ({metadata}; {gcloud})"
95                        ))
96                    })?,
97                    GCLOUD_LIFETIME,
98                ),
99            },
100        };
101        *cached = Some((
102            token.clone(),
103            Instant::now() + lifetime.saturating_sub(EARLY),
104        ));
105        Ok(token)
106    }
107
108    /// Forget the cached token, e.g. after a request was rejected with 401.
109    pub async fn invalidate(&self) {
110        *self.cached.lock().await = None;
111    }
112
113    /// An [`AccessToken`] that stays valid for as long as it is held.
114    ///
115    /// [`AccessToken`] is read synchronously on every connection attempt,
116    /// and fetching a token is asynchronous, so the token is fetched here
117    /// once (an error means there are no usable credentials) and then kept
118    /// fresh by a background task. The task checks every minute and fetches
119    /// a new token shortly before the current one expires. It stops once
120    /// every clone of the returned `AccessToken` is dropped. If a refresh
121    /// fails, the previous token is kept and the next check retries.
122    ///
123    /// Must be called inside a Tokio runtime.
124    pub async fn into_access_token(self) -> Result<AccessToken, AuthError> {
125        self.into_access_token_every(REFRESH_CHECK).await
126    }
127
128    async fn into_access_token_every(self, every: Duration) -> Result<AccessToken, AuthError> {
129        let first = self.token().await?;
130        let current = std::sync::Arc::new(parking_lot::RwLock::new(first));
131        let weak = std::sync::Arc::downgrade(&current);
132        tokio::spawn(async move {
133            loop {
134                tokio::time::sleep(every).await;
135                let Some(current) = weak.upgrade() else {
136                    break;
137                };
138                match self.token().await {
139                    Ok(token) => *current.write() = token,
140                    Err(e) => tracing::warn!("access token refresh failed: {e}"),
141                }
142            }
143        });
144        Ok(AccessToken::from_fn(move || current.read().clone()))
145    }
146
147    async fn fetch_metadata(&self, host: &str) -> Result<(String, Duration), AuthError> {
148        #[derive(serde::Deserialize)]
149        struct Token {
150            access_token: String,
151            expires_in: u64,
152        }
153        let url =
154            format!("http://{host}/computeMetadata/v1/instance/service-accounts/default/token");
155        let response = self
156            .client
157            .get(&url)
158            .header("Metadata-Flavor", "Google")
159            .timeout(Duration::from_secs(2))
160            .send()
161            .await
162            .map_err(|e| AuthError::TokenFetchFailed(format!("metadata server: {e}")))?;
163        if !response.status().is_success() {
164            return Err(AuthError::TokenFetchFailed(format!(
165                "metadata server: HTTP {}",
166                response.status()
167            )));
168        }
169        let token: Token = response
170            .json()
171            .await
172            .map_err(|e| AuthError::TokenFetchFailed(format!("metadata server: {e}")))?;
173        Ok((token.access_token, Duration::from_secs(token.expires_in)))
174    }
175}
176
177fn metadata_host() -> String {
178    std::env::var("GCE_METADATA_HOST")
179        .ok()
180        .filter(|h| !h.trim().is_empty())
181        .unwrap_or_else(|| "metadata.google.internal".to_string())
182}
183
184async fn gcloud_token() -> Result<String, AuthError> {
185    let output = tokio::task::spawn_blocking(|| {
186        std::process::Command::new("gcloud")
187            .args(["auth", "print-access-token"])
188            .output()
189    })
190    .await
191    .map_err(|e| AuthError::TokenFetchFailed(format!("gcloud: {e}")))?
192    .map_err(|e| AuthError::TokenFetchFailed(format!("gcloud: {e}")))?;
193    let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
194    if !output.status.success() || token.is_empty() {
195        return Err(AuthError::TokenFetchFailed(format!(
196            "gcloud: {}",
197            String::from_utf8_lossy(&output.stderr).trim()
198        )));
199    }
200    Ok(token)
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use tokio::io::{AsyncReadExt, AsyncWriteExt};
207
208    /// A metadata server that counts its requests and checks the header.
209    async fn fake_metadata(
210        expires_in: u64,
211    ) -> (String, std::sync::Arc<std::sync::atomic::AtomicUsize>) {
212        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
213        let host = listener.local_addr().unwrap().to_string();
214        let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
215        let counter = hits.clone();
216        tokio::spawn(async move {
217            while let Ok((mut socket, _)) = listener.accept().await {
218                let mut buf = vec![0u8; 4096];
219                let n = socket.read(&mut buf).await.unwrap_or(0);
220                let request = String::from_utf8_lossy(&buf[..n]).to_lowercase();
221                let n = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
222                let response = if request.contains("metadata-flavor: google") {
223                    let body = format!(
224                        "{{\"access_token\":\"tok-{n}\",\"expires_in\":{expires_in},\"token_type\":\"Bearer\"}}"
225                    );
226                    format!(
227                        "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
228                        body.len()
229                    )
230                } else {
231                    "HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
232                        .into()
233                };
234                let _ = socket.write_all(response.as_bytes()).await;
235            }
236        });
237        (host, hits)
238    }
239
240    #[tokio::test]
241    async fn a_metadata_token_is_cached_until_it_nears_expiry() {
242        let (host, hits) = fake_metadata(3600).await;
243        let source = GoogleAccessToken::metadata_server(host);
244        assert_eq!(source.token().await.unwrap(), "tok-0");
245        assert_eq!(source.token().await.unwrap(), "tok-0");
246        assert_eq!(hits.load(std::sync::atomic::Ordering::SeqCst), 1);
247        source.invalidate().await;
248        assert_eq!(source.token().await.unwrap(), "tok-1");
249    }
250
251    #[tokio::test]
252    async fn a_token_about_to_expire_is_refetched() {
253        // Expires inside the early-refresh window, so it is never reused.
254        let (host, _) = fake_metadata(60).await;
255        let source = GoogleAccessToken::metadata_server(host);
256        assert_eq!(source.token().await.unwrap(), "tok-0");
257        assert_eq!(source.token().await.unwrap(), "tok-1");
258    }
259
260    #[tokio::test]
261    async fn an_access_token_is_refreshed_in_the_background_until_dropped() {
262        // Every fetch is inside the early-refresh window, so each check
263        // fetches a new token.
264        let (host, hits) = fake_metadata(60).await;
265        let token = GoogleAccessToken::metadata_server(host)
266            .into_access_token_every(Duration::from_millis(20))
267            .await
268            .unwrap();
269        assert_eq!(token.get(), "tok-0");
270        tokio::time::sleep(Duration::from_millis(200)).await;
271        assert_ne!(token.get(), "tok-0", "the background task refreshed it");
272
273        drop(token);
274        tokio::time::sleep(Duration::from_millis(60)).await;
275        let after_drop = hits.load(std::sync::atomic::Ordering::SeqCst);
276        tokio::time::sleep(Duration::from_millis(200)).await;
277        assert_eq!(
278            hits.load(std::sync::atomic::Ordering::SeqCst),
279            after_drop,
280            "the refresher stops once the token is dropped"
281        );
282    }
283
284    #[tokio::test]
285    async fn an_access_token_needs_credentials_up_front() {
286        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
287        let host = listener.local_addr().unwrap().to_string();
288        drop(listener);
289        assert!(
290            GoogleAccessToken::metadata_server(host)
291                .into_access_token()
292                .await
293                .is_err()
294        );
295    }
296
297    #[tokio::test]
298    async fn a_fixed_token_is_returned_as_is() {
299        assert_eq!(
300            GoogleAccessToken::fixed("abc").token().await.unwrap(),
301            "abc"
302        );
303    }
304}
gemini_genai_rs/transport/auth/google_token.rs

gemini_genai_rs/transport/auth/
google_token.rs
