diff --git a/src/azure/storage/credential.rs b/src/azure/storage/credential.rs index 96243f2..b10363e 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,5 +18,32 @@ pub enum Credential { /// associated with the subscription that contains the storage account. /// /// ref: - 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(), + } + } } diff --git a/src/azure/storage/loader.rs b/src/azure/storage/loader.rs index b659b4a..dda61dc 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}; @@ -28,8 +30,9 @@ impl Loader { /// Load credential. pub async fn load(&self) -> Result> { // 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> { 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), } } diff --git a/src/azure/storage/signer.rs b/src/azure/storage/signer.rs index fbc2287..2019ae5 100644 --- a/src/azure/storage/signer.rs +++ b/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());