From a8034259290b7cf5f35916fd964fa3d51fa5ecdf Mon Sep 17 00:00:00 2001
From: congyi <15605187270@163.com>
Date: Mon, 15 Apr 2024 17:05:37 +0800
Subject: [PATCH] resolve comments

---
 src/azure/storage/credential.rs | 16 +++++++---------
 src/azure/storage/loader.rs     | 25 +++++++++++++++++--------
 src/azure/storage/signer.rs     |  5 ++---
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/azure/storage/credential.rs b/src/azure/storage/credential.rs
index d610e58..8768f0f 100644
--- a/src/azure/storage/credential.rs
+++ b/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,7 +18,7 @@ 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, String),
+    BearerToken(String, DateTime),
 }
 
 impl Credential {
@@ -26,11 +28,9 @@ impl Credential {
             return false;
         }
         if let Credential::BearerToken(_, expires_on) = self {
-            if let Ok(expires) = chrono::DateTime::parse_from_rfc3339(expires_on) {
-                let buffer = chrono::Duration::try_minutes(2).expect("in bounds");
-                if expires > (chrono::Utc::now() + buffer) {
-                    return false;
-                }
+            let buffer = chrono::TimeDelta::try_minutes(2).expect("in bounds");
+            if expires_on > &(chrono::Utc::now() + buffer) {
+                return false;
             }
         };
 
@@ -43,9 +43,7 @@ impl Credential {
                 account_name.is_empty() || account_key.is_empty()
             }
             Credential::SharedAccessSignature(sas_token) => sas_token.is_empty(),
-            Credential::BearerToken(bearer_token, expire_on) => {
-                bearer_token.is_empty() || expire_on.is_empty()
-            }
+            Credential::BearerToken(bearer_token, _) => bearer_token.is_empty(),
         }
     }
 }
diff --git a/src/azure/storage/loader.rs b/src/azure/storage/loader.rs
index 337f982..067cd66 100644
--- a/src/azure/storage/loader.rs
+++ b/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};
@@ -72,10 +74,11 @@ 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,
-            token.expires_on,
-        ));
+        let expires_on = match token.expires_on.is_empty() {
+            true => now() + chrono::TimeDelta::try_minutes(10).expect("in bounds"),
+            false => parse_rfc3339(&token.expires_on)?,
+        };
+        let cred = Some(Credential::BearerToken(token.access_token, expires_on));
 
         Ok(cred)
     }
@@ -84,10 +87,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,
-                token.expires_on.unwrap_or("".to_string()),
-            ))),
+            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),
         }
     }
diff --git a/src/azure/storage/signer.rs b/src/azure/storage/signer.rs
index 805decd..2019ae5 100644
--- a/src/azure/storage/signer.rs
+++ b/src/azure/storage/signer.rs
@@ -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,8 +307,7 @@ mod tests {
             .uri("https://test.blob.core.windows.net/testbucket/testblob")
             .body(())
             .unwrap();
-        let cred =
-            AzureStorageCredential::BearerToken("token".to_string(), "expires_on".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());