fix: regularly check whether the federated token has expired

src/azure/storage/credential.rs

@@ -1,3 +1,5 @@
+use crate::time::DateTime;
+
 /// Credential that holds the access_key and secret_key.
 #[derive(Clone)]
 #[cfg_attr(test, derive(Debug))]
@@ -16,5 +18,32 @@ pub enum Credential {
     /// associated with the subscription that contains the storage account.
     ///
     /// ref: <https://docs.microsoft.com/rest/api/storageservices/authorize-with-azure-active-directory>
-    BearerToken(String),
+    BearerToken(String, DateTime),
+}
+
+impl Credential {
+    /// is current cred is valid?
+    pub fn is_valid(&self) -> bool {
+        if self.is_empty() {
+            return false;
+        }
+        if let Credential::BearerToken(_, expires_on) = self {
+            let buffer = chrono::TimeDelta::try_seconds(20).expect("in bounds");
+            if expires_on < &(chrono::Utc::now() + buffer) {
+                return false;
+            }
+        };
+
+        true
+    }
+
+    fn is_empty(&self) -> bool {
+        match self {
+            Credential::SharedKey(account_name, account_key) => {
+                account_name.is_empty() || account_key.is_empty()
+            }
+            Credential::SharedAccessSignature(sas_token) => sas_token.is_empty(),
+            Credential::BearerToken(bearer_token, _) => bearer_token.is_empty(),
+        }
+    }
 }

src/azure/storage/loader.rs

@@ -3,6 +3,8 @@ use std::sync::Mutex;
 
 use anyhow::Result;
 
+use crate::time::{now, parse_rfc3339};
+
 use super::credential::Credential;
 use super::imds_credential;
 use super::{config::Config, workload_identity_credential};
@@ -28,8 +30,9 @@ impl Loader {
     /// Load credential.
     pub async fn load(&self) -> Result<Option<Credential>> {
         // Return cached credential if it's valid.
-        if let Some(cred) = self.credential.lock().expect("lock poisoned").clone() {
-            return Ok(Some(cred));
+        match self.credential.lock().expect("lock poisoned").clone() {
+            Some(cred) if cred.is_valid() => return Ok(Some(cred)),
+            _ => (),
         }
         let cred = self.load_inner().await?;
 
@@ -71,7 +74,12 @@ impl Loader {
     async fn load_via_imds(&self) -> Result<Option<Credential>> {
         let token =
             imds_credential::get_access_token("https://storage.azure.com/", &self.config).await?;
-        let cred = Some(Credential::BearerToken(token.access_token));
+        let expires_on = if token.expires_on.is_empty() {
+            now() + chrono::TimeDelta::try_minutes(10).expect("in bounds")
+        } else {
+            parse_rfc3339(&token.expires_on)?
+        };
+        let cred = Some(Credential::BearerToken(token.access_token, expires_on));
 
         Ok(cred)
     }
@@ -80,7 +88,16 @@ impl Loader {
         let workload_identity_token =
             workload_identity_credential::get_workload_identity_token(&self.config).await?;
         match workload_identity_token {
-            Some(token) => Ok(Some(Credential::BearerToken(token.access_token))),
+            Some(token) => {
+                let expires_on_duration = match token.expires_on {
+                    None => now() + chrono::TimeDelta::try_minutes(10).expect("in bounds"),
+                    Some(expires_on) => parse_rfc3339(&expires_on)?,
+                };
+                Ok(Some(Credential::BearerToken(
+                    token.access_token,
+                    expires_on_duration,
+                )))
+            }
             None => Ok(None),
         }
     }

src/azure/storage/signer.rs

@@ -61,7 +61,7 @@ impl Signer {
                 ctx.query_append(token);
                 return Ok(ctx);
             }
-            Credential::BearerToken(token) => match method {
+            Credential::BearerToken(token, _) => match method {
                 SigningMethod::Query(_) => {
                     return Err(anyhow!("BearerToken can't be used in query string"));
                 }
@@ -269,9 +269,9 @@ mod tests {
     use http::Request;
 
     use super::super::config::Config;
-    use crate::azure::storage::loader::Loader;
     use crate::AzureStorageCredential;
     use crate::AzureStorageSigner;
+    use crate::{azure::storage::loader::Loader, time::now};
 
     #[tokio::test]
     async fn test_sas_url() {
@@ -307,7 +307,7 @@ mod tests {
             .uri("https://test.blob.core.windows.net/testbucket/testblob")
             .body(())
             .unwrap();
-        let cred = AzureStorageCredential::BearerToken("token".to_string());
+        let cred = AzureStorageCredential::BearerToken("token".to_string(), now());
 
         // Can effectively sign request with SigningMethod::Header
         assert!(signer.sign(&mut req, &cred).is_ok());