1mod external_account;
2mod impersonated_service_account;
3mod service_account;
4
5use std::fmt::Debug;
6use std::fmt::Formatter;
7use std::sync::Arc;
8use std::sync::Mutex;
9
10use anyhow::Result;
11use async_trait::async_trait;
12use reqwest::Client;
13use serde::Deserialize;
14use serde::Serialize;
15
16use super::credential::Credential;
17use crate::time::now;
18use crate::time::DateTime;
19
20#[derive(Clone, Deserialize, Default)]
24#[serde(default)]
25pub struct Token {
26 access_token: String,
27 scope: String,
28 token_type: String,
29 expires_in: usize,
30}
31
32impl Token {
33 pub fn new(access_token: &str, expires_in: usize, scope: &str) -> Self {
37 Self {
38 access_token: access_token.to_string(),
39 scope: scope.to_string(),
40 expires_in,
41 token_type: "Bearer".to_string(),
42 }
43 }
44
45 pub(crate) fn access_token(&self) -> &str {
47 &self.access_token
48 }
49
50 pub(crate) fn expires_in(&self) -> usize {
52 self.expires_in
53 }
54}
55
56impl Debug for Token {
58 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
59 f.debug_struct("Token")
60 .field("access_token", &"<redacted>")
61 .field("scope", &self.scope)
62 .field("token_type", &self.token_type)
63 .field("expires_in", &self.expires_in)
64 .finish()
65 }
66}
67
68#[derive(Debug, Serialize)]
80pub struct Claims {
81 iss: String,
82 scope: String,
83 aud: String,
84 exp: u64,
85 iat: u64,
86}
87
88impl Claims {
89 pub fn new(client_email: &str, scope: &str) -> Claims {
90 let current = now().timestamp() as u64;
91
92 Claims {
93 iss: client_email.to_string(),
94 scope: scope.to_string(),
95 aud: "https://oauth2.googleapis.com/token".to_string(),
96 exp: current + 3600,
97 iat: current,
98 }
99 }
100}
101
102#[async_trait]
104pub trait TokenLoad: 'static + Send + Sync + Debug {
105 async fn load(&self, client: Client) -> Result<Option<Token>>;
111}
112
113#[cfg_attr(test, derive(Debug))]
115pub struct TokenLoader {
116 scope: String,
117 client: Client,
118
119 credential: Option<Credential>,
120 disable_vm_metadata: bool,
121 service_account: Option<String>,
122 customized_token_loader: Option<Box<dyn TokenLoad>>,
123
124 token: Arc<Mutex<Option<(Token, DateTime)>>>,
125}
126
127impl TokenLoader {
128 pub fn new(scope: &str, client: Client) -> Self {
140 Self {
141 scope: scope.to_string(),
142 client,
143
144 credential: None,
145 disable_vm_metadata: false,
146 service_account: None,
147 customized_token_loader: None,
148
149 token: Arc::default(),
150 }
151 }
152
153 pub fn with_credentials(mut self, credentials: Credential) -> Self {
155 self.credential = Some(credentials);
156 self
157 }
158
159 pub fn with_disable_vm_metadata(mut self, disable_vm_metadata: bool) -> Self {
161 self.disable_vm_metadata = disable_vm_metadata;
162 self
163 }
164
165 pub fn with_service_account(mut self, service_account: &str) -> Self {
167 self.service_account = Some(service_account.to_string());
168 self
169 }
170
171 pub fn with_customized_token_loader(
173 mut self,
174 customized_token_loader: Box<dyn TokenLoad>,
175 ) -> Self {
176 self.customized_token_loader = Some(customized_token_loader);
177 self
178 }
179
180 pub async fn load(&self) -> Result<Option<Token>> {
182 match self.token.lock().expect("lock poisoned").clone() {
183 Some((token, expire_in))
184 if now()
185 < expire_in - chrono::TimeDelta::try_seconds(2 * 60).expect("in bounds") =>
186 {
187 return Ok(Some(token))
188 }
189 _ => (),
190 }
191
192 let token = if let Some(token) = self.load_inner().await? {
193 token
194 } else {
195 return Ok(None);
196 };
197
198 let expire_in =
199 now() + chrono::TimeDelta::try_seconds(token.expires_in() as i64).expect("in bounds");
200
201 let mut lock = self.token.lock().expect("lock poisoned");
202 *lock = Some((token.clone(), expire_in));
203
204 Ok(Some(token))
205 }
206
207 async fn load_inner(&self) -> Result<Option<Token>> {
208 if let Some(token) = self.load_via_customized_token_loader().await? {
209 return Ok(Some(token));
210 }
211
212 if let Some(token) = self.load_via_service_account().await? {
213 return Ok(Some(token));
214 }
215
216 if let Some(token) = self.load_via_impersonated_service_account().await? {
217 return Ok(Some(token));
218 }
219
220 if let Some(token) = self.load_via_external_account().await? {
221 return Ok(Some(token));
222 }
223
224 if let Some(token) = self.load_via_vm_metadata().await? {
225 return Ok(Some(token));
226 }
227
228 Ok(None)
229 }
230
231 async fn load_via_customized_token_loader(&self) -> Result<Option<Token>> {
232 match &self.customized_token_loader {
233 Some(f) => f.load(self.client.clone()).await,
234 None => Ok(None),
235 }
236 }
237
238 async fn load_via_vm_metadata(&self) -> Result<Option<Token>> {
240 if self.disable_vm_metadata {
241 return Ok(None);
242 }
243
244 let service_account = self.service_account.as_deref().unwrap_or("default");
246
247 let url = format!("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/{service_account}/token?scopes={}", self.scope);
248
249 let resp = self
250 .client
251 .get(&url)
252 .header("Metadata-Flavor", "Google")
253 .send()
254 .await?;
255
256 let token: Token = serde_json::from_slice(&resp.bytes().await?)?;
257 Ok(Some(token))
258 }
259}