fix: regularly check whether the federated token has expired

src/azure/storage/credential.rs

@@ -16,5 +16,36 @@ 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, String),
+}
+
+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 {
+ 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;
+ }
+ }
+ };
+
+ 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, expire_on) => {
+ bearer_token.is_empty() || expire_on.is_empty()
+ }
+ }
+ }
 }

src/azure/storage/loader.rs

@@ -28,8 +28,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 +72,10 @@ 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 cred = Some(Credential::BearerToken(
+ token.access_token,
+ token.expires_on,
+ ));
 
  Ok(cred)
  }
@@ -80,7 +84,10 @@ 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) => Ok(Some(Credential::BearerToken(
+ token.access_token,
+ token.expires_on.unwrap_or("".to_string()),
+ ))),
  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"));
  }
@@ -307,7 +307,8 @@ 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(), "expires_on".to_string());
 
  // Can effectively sign request with SigningMethod::Header
  assert!(signer.sign(&mut req, &cred).is_ok());