gcloud_sdk/token_source/
auth_token_generator.rs

1use hyper::header::HeaderValue;
2use jiff::{SignedDuration, Timestamp};
3use tokio::sync::RwLock;
4
5use crate::token_source::*;
6use tracing::*;
7
8/// A token together with its pre-validated `authorization` header value, so that
9/// serving a cache hit never re-parses or re-allocates the header.
10struct CachedToken {
11    token: Token,
12    authorization: HeaderValue,
13}
14
15impl CachedToken {
16    fn from_token(token: Token) -> crate::error::Result<Self> {
17        let mut authorization = HeaderValue::from_str(&token.header_value())?;
18        authorization.set_sensitive(true);
19        Ok(Self {
20            token,
21            authorization,
22        })
23    }
24}
25
26pub struct GoogleAuthTokenGenerator {
27    token_source: BoxSource,
28    cached_token: RwLock<Option<CachedToken>>,
29}
30
31impl GoogleAuthTokenGenerator {
32    pub async fn new(
33        token_source_type: TokenSourceType,
34        token_scopes: Vec<String>,
35    ) -> crate::error::Result<GoogleAuthTokenGenerator> {
36        let token_source: BoxSource = create_source(token_source_type, token_scopes).await?;
37
38        Ok(GoogleAuthTokenGenerator {
39            token_source,
40            cached_token: RwLock::new(None),
41        })
42    }
43
44    pub async fn clear_cache(&self) {
45        let mut write_state = self.cached_token.write().await;
46        *write_state = None;
47    }
48
49    pub async fn create_token(&self) -> crate::error::Result<Token> {
50        self.with_cached(|cached| cached.token.clone()).await
51    }
52
53    /// The `authorization` header value for the current token, already validated
54    /// and marked sensitive; cloning it is a `Bytes` refcount bump, not an allocation.
55    pub async fn authorization_header(&self) -> crate::error::Result<HeaderValue> {
56        self.with_cached(|cached| cached.authorization.clone())
57            .await
58    }
59
60    /// Runs the double-checked refresh and applies `pick` to the resulting cached
61    /// token by reference, so a cache hit clones only what the caller asks for.
62    async fn with_cached<R>(
63        &self,
64        pick: impl FnOnce(&CachedToken) -> R,
65    ) -> crate::error::Result<R> {
66        let now = Timestamp::now();
67        // Give a bit more time for the network call than the token strictly has left;
68        // both the read-only fast path and the write-lock recheck below must use this
69        // same threshold, or a token can sit inside the margin forever, with every
70        // caller taking the write lock and getting the stale token back without ever
71        // triggering the refresh the margin exists to trigger.
72        let refresh_after = now.add(SignedDuration::from_secs(15));
73
74        {
75            let read_state = self.cached_token.read().await;
76            if let Some(cached) = read_state.as_ref() {
77                if cached.token.expiry.gt(&refresh_after) {
78                    return Ok(pick(cached));
79                }
80            }
81        }
82
83        let mut write_token = self.cached_token.write().await;
84        match write_token.as_ref() {
85            Some(updated_cached) if updated_cached.token.expiry.gt(&refresh_after) => {
86                Ok(pick(updated_cached))
87            }
88            _ => {
89                let new_token = self.token_source.token().await?;
90                debug!(
91                    "Created a new Google OAuth token. Type: {}. Expiring: {}.",
92                    new_token.token_type, new_token.expiry,
93                );
94                let new_cached = CachedToken::from_token(new_token)?;
95                let result = pick(&new_cached);
96                *write_token = Some(new_cached);
97                Ok(result)
98            }
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::token_source::Source;
107    use async_trait::async_trait;
108    use secret_vault_value::SecretValue;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110    use std::sync::Arc;
111
112    struct CountingSource {
113        calls: Arc<AtomicUsize>,
114    }
115
116    #[async_trait]
117    impl Source for CountingSource {
118        async fn token(&self) -> crate::error::Result<Token> {
119            self.calls.fetch_add(1, Ordering::SeqCst);
120            Ok(Token {
121                token_type: "Bearer".to_string(),
122                token: SecretValue::from("cached-token"),
123                expiry: Timestamp::now() + SignedDuration::from_hours(1),
124            })
125        }
126    }
127
128    #[tokio::test]
129    async fn authorization_header_reuses_cached_token() {
130        let calls = Arc::new(AtomicUsize::new(0));
131        let source = CountingSource {
132            calls: calls.clone(),
133        };
134        let generator = GoogleAuthTokenGenerator::new(
135            TokenSourceType::ExternalSource(Box::new(source)),
136            vec![],
137        )
138        .await
139        .unwrap();
140
141        let first = generator.authorization_header().await.unwrap();
142        let second = generator.authorization_header().await.unwrap();
143
144        assert_eq!(first, second);
145        assert!(first.is_sensitive());
146        assert_eq!(calls.load(Ordering::SeqCst), 1);
147    }
148
149    struct FixedExpirySource {
150        calls: Arc<AtomicUsize>,
151        expires_in: SignedDuration,
152    }
153
154    #[async_trait]
155    impl Source for FixedExpirySource {
156        async fn token(&self) -> crate::error::Result<Token> {
157            self.calls.fetch_add(1, Ordering::SeqCst);
158            Ok(Token {
159                token_type: "Bearer".to_string(),
160                token: SecretValue::from("margin-token"),
161                expiry: Timestamp::now() + self.expires_in,
162            })
163        }
164    }
165
166    #[tokio::test]
167    async fn token_inside_refresh_margin_is_refreshed() {
168        let calls = Arc::new(AtomicUsize::new(0));
169        let source = FixedExpirySource {
170            calls: calls.clone(),
171            expires_in: SignedDuration::from_secs(10),
172        };
173        let generator = GoogleAuthTokenGenerator::new(
174            TokenSourceType::ExternalSource(Box::new(source)),
175            vec![],
176        )
177        .await
178        .unwrap();
179
180        generator.authorization_header().await.unwrap();
181        generator.authorization_header().await.unwrap();
182
183        assert_eq!(calls.load(Ordering::SeqCst), 2);
184    }
185}
gcloud_sdk/token_source/auth_token_generator.rs

gcloud_sdk/token_source/
auth_token_generator.rs
