Skip to content

Commit

Permalink
resolve comments
Browse files Browse the repository at this point in the history
  • Loading branch information
wcy-fdu committed Apr 15, 2024
1 parent 6d2023f commit a803425
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 20 deletions.
16 changes: 7 additions & 9 deletions src/azure/storage/credential.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use crate::time::DateTime;

/// Credential that holds the access_key and secret_key.
#[derive(Clone)]
#[cfg_attr(test, derive(Debug))]
Expand All @@ -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 {
Expand All @@ -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;
}
};

Expand All @@ -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(),
}
}
}
25 changes: 17 additions & 8 deletions src/azure/storage/loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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)
}
Expand All @@ -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),
}
}
Expand Down
5 changes: 2 additions & 3 deletions src/azure/storage/signer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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());
Expand Down

0 comments on commit a803425

Please sign in to comment.