cloud_storage/
token.rs

1use std::fmt::{Display, Formatter};
2
3/// Trait that refreshes a token when it is expired
4#[async_trait::async_trait]
5pub trait TokenCache: Sync {
6    /// Returns the token that is currently held within the instance of `TokenCache`, together with
7    /// the expiry of that token as a u64 in seconds sine the Unix Epoch (1 Jan 1970).
8    async fn token_and_exp(&self) -> Option<(String, u64)>;
9
10    /// Updates the token to the value `token`.
11    async fn set_token(&self, token: String, exp: u64) -> crate::Result<()>;
12
13    /// Returns the intended scope for the current token.
14    async fn scope(&self) -> String;
15
16    /// Returns a valid, unexpired token. If the contained token is expired, it updates and returns
17    /// the token.
18    async fn get(&self, client: &reqwest::Client) -> crate::Result<String> {
19        match self.token_and_exp().await {
20            Some((token, exp)) if now() + 300 < exp => Ok(token),
21            _ => {
22                let (token, exp) = self.fetch_token(client).await?;
23                self.set_token(token, exp).await?;
24
25                self.token_and_exp()
26                    .await
27                    .map(|(t, _)| t)
28                    .ok_or(crate::Error::Other("Token is not set".to_string()))
29            }
30        }
31    }
32
33    /// Fetches and returns the token using the service account
34    async fn fetch_token(&self, client: &reqwest::Client) -> crate::Result<(String, u64)>;
35}
36
37#[derive(serde::Serialize)]
38struct Claims {
39    iss: String,
40    scope: String,
41    aud: String,
42    exp: u64,
43    iat: u64,
44}
45
46#[derive(serde::Deserialize, Debug)]
47// #[allow(dead_code)]
48struct TokenResponse {
49    access_token: String,
50    expires_in: u64,
51    // token_type: String,
52}
53
54/// This struct contains a token, an expiry, and an access scope.
55pub struct Token {
56    // this field contains the JWT and the expiry thereof. They are in the same Option because if
57    // one of them is `Some`, we require that the other be `Some` as well.
58    token: tokio::sync::RwLock<Option<DefaultTokenData>>,
59    // store the access scope for later use if we need to refresh the token
60    access_scope: String,
61}
62
63#[derive(Debug, Clone)]
64pub struct DefaultTokenData(String, u64);
65
66impl Display for DefaultTokenData {
67    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
68        write!(f, "{}", self.0)
69    }
70}
71
72impl Default for Token {
73    fn default() -> Self {
74        Token::new("https://www.googleapis.com/auth/devstorage.full_control")
75    }
76}
77
78impl Token {
79    pub(crate) fn new(scope: &str) -> Self {
80        Self {
81            token: tokio::sync::RwLock::new(None),
82            access_scope: scope.to_string(),
83        }
84    }
85}
86
87#[async_trait::async_trait]
88impl TokenCache for Token {
89    async fn scope(&self) -> String {
90        self.access_scope.clone()
91    }
92
93    async fn token_and_exp(&self) -> Option<(String, u64)> {
94        self.token.read().await.as_ref().map(|d| (d.0.clone(), d.1))
95    }
96
97    async fn set_token(&self, token: String, exp: u64) -> crate::Result<()> {
98        *self.token.write().await = Some(DefaultTokenData(token, exp));
99        Ok(())
100    }
101
102    async fn fetch_token(&self, client: &reqwest::Client) -> crate::Result<(String, u64)> {
103        let now = now();
104        let exp = now + 3600;
105
106        let claims = Claims {
107            iss: crate::SERVICE_ACCOUNT.client_email.clone(),
108            scope: self.scope().await.into(),
109            aud: "https://www.googleapis.com/oauth2/v4/token".to_string(),
110            exp,
111            iat: now,
112        };
113        let header = jsonwebtoken::Header {
114            alg: jsonwebtoken::Algorithm::RS256,
115            ..Default::default()
116        };
117        let private_key_bytes = crate::SERVICE_ACCOUNT.private_key.as_bytes();
118        let private_key = jsonwebtoken::EncodingKey::from_rsa_pem(private_key_bytes)?;
119        let jwt = jsonwebtoken::encode(&header, &claims, &private_key)?;
120        let body = [
121            ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
122            ("assertion", &jwt),
123        ];
124        let response: TokenResponse = client
125            .post("https://www.googleapis.com/oauth2/v4/token")
126            .form(&body)
127            .send()
128            .await?
129            .json()
130            .await?;
131        Ok((response.access_token, now + response.expires_in))
132    }
133}
134
135fn now() -> u64 {
136    std::time::SystemTime::now()
137        .duration_since(std::time::SystemTime::UNIX_EPOCH)
138        .unwrap()
139        .as_secs()
140}
cloud_storage/token.rs

cloud_storage/
token.rs
