Skip to main content

reqsign/google/
token.rs

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/// Token is the authentication methods used by google services.
21///
22/// Most of the time, they will be exchanged via application credentials.
23#[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    /// Create a new token.
34    ///
35    /// scope will looks like: `https://www.googleapis.com/auth/devstorage.read_only`.
36    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    /// Notes: don't allow get token from reqsign.
46    pub(crate) fn access_token(&self) -> &str {
47        &self.access_token
48    }
49
50    /// Notes: don't allow get expires_in from reqsign.
51    pub(crate) fn expires_in(&self) -> usize {
52        self.expires_in
53    }
54}
55
56/// Make sure `access_token` is redacted for Token
57impl 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/// Claims is used to build JWT for google cloud.
69///
70/// ```json
71/// {
72///   "iss": "761326798069-r5mljlln1rd4lrbhg75efgigp36m78j5@developer.gserviceaccount.com",
73///   "scope": "https://www.googleapis.com/auth/devstorage.read_only",
74///   "aud": "https://oauth2.googleapis.com/token",
75///   "exp": 1328554385,
76///   "iat": 1328550785
77/// }
78/// ```
79#[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/// Loader trait will try to load credential from different sources.
103#[async_trait]
104pub trait TokenLoad: 'static + Send + Sync + Debug {
105    /// Load credential from sources.
106    ///
107    /// - If succeed, return `Ok(Some(cred))`
108    /// - If not found, return `Ok(None)`
109    /// - If unexpected errors happened, return `Err(err)`
110    async fn load(&self, client: Client) -> Result<Option<Token>>;
111}
112
113/// TokenLoader will load token from different methods.
114#[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    /// Create a new token loader.
129    ///
130    /// ## Scope
131    ///
132    /// For example, valid scopes for google cloud services should be
133    ///
134    /// - read-only: `https://www.googleapis.com/auth/devstorage.read_only`
135    /// - read-write: `https://www.googleapis.com/auth/devstorage.read_write`
136    /// - full-control: `https://www.googleapis.com/auth/devstorage.full_control`
137    ///
138    /// Reference: [Cloud Storage authentication](https://cloud.google.com/storage/docs/authentication)
139    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    /// Set the credential for token loader.
154    pub fn with_credentials(mut self, credentials: Credential) -> Self {
155        self.credential = Some(credentials);
156        self
157    }
158
159    /// Disable vm metadata.
160    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    /// Set the service account for token loader.
166    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    /// Set the customized token loader for token loader.
172    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    /// Load token from different sources.
181    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    /// Exchange token via vm metadata
239    async fn load_via_vm_metadata(&self) -> Result<Option<Token>> {
240        if self.disable_vm_metadata {
241            return Ok(None);
242        }
243
244        // Use `default` if service account not set by user.
245        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}