add AutoRefreshingTokenCredential support to StorageAccountClient (#757)

Author: yvespp

sdk/storage/src/core/clients/storage_account_client.rs

@@ -6,6 +6,7 @@ use crate::{
         AccountSharedAccessSignatureBuilder, ClientAccountSharedAccessSignature,
     },
 };
+use azure_core::auth::TokenCredential;
 use azure_core::headers::*;
 use azure_core::HttpClient;
 use bytes::Bytes;
@@ -26,15 +27,29 @@ pub const EMULATOR_ACCOUNT: &str = "devstoreaccount1";
 pub const EMULATOR_ACCOUNT_KEY: &str =
     "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==";
 
+pub const STORAGE_TOKEN_SCOPE: &str = "https://storage.azure.com/";
+
 const HEADER_VERSION: &str = "x-ms-version";
 
 const AZURE_VERSION: &str = "2019-12-12";
 
-#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum StorageCredentials {
     Key(String, String),
     SASToken(Vec<(String, String)>),
     BearerToken(String),
+    TokenCredential(Box<dyn TokenCredential>),
+}
+
+impl std::fmt::Debug for StorageCredentials {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self {
+            StorageCredentials::TokenCredential(_) => f
+                .debug_struct("StorageCredentials")
+                .field("credential", &"TokenCredential")
+                .finish(),
+            _ => self.fmt(f),
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -45,7 +60,7 @@ pub enum ServiceType {
     Table,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct StorageAccountClient {
     storage_credentials: StorageCredentials,
     http_client: Arc<dyn HttpClient>,
@@ -236,6 +251,36 @@ impl StorageAccountClient {
         })
     }
 
+    pub fn new_token_credential<A>(
+        http_client: Arc<dyn HttpClient>,
+        account: A,
+        token_credential: Box<dyn TokenCredential>,
+    ) -> Arc<Self>
+    where
+        A: Into<String>,
+    {
+        let account = account.into();
+
+        Arc::new(Self {
+            blob_storage_url: Url::parse(&format!("https://{}.blob.core.windows.net", &account))
+                .unwrap(),
+            table_storage_url: Url::parse(&format!("https://{}.table.core.windows.net", &account))
+                .unwrap(),
+            queue_storage_url: Url::parse(&format!("https://{}.queue.core.windows.net", &account))
+                .unwrap(),
+            queue_storage_secondary_url: Url::parse(&format!(
+                "https://{}-secondary.queue.core.windows.net",
+                &account
+            ))
+            .unwrap(),
+            filesystem_url: Url::parse(&format!("https://{}.dfs.core.windows.net", &account))
+                .unwrap(),
+            storage_credentials: StorageCredentials::TokenCredential(token_credential),
+            http_client,
+            account,
+        })
+    }
+
     pub fn new_connection_string(
         http_client: Arc<dyn HttpClient>,
         connection_string: &str,
@@ -407,6 +452,15 @@ impl StorageAccountClient {
             StorageCredentials::BearerToken(token) => {
                 request.header(AUTHORIZATION, format!("Bearer {}", token))
             }
+            StorageCredentials::TokenCredential(token_credential) => {
+                let bearer_token_future = token_credential.get_token(STORAGE_TOKEN_SCOPE);
+                let bearer_token = futures::executor::block_on(bearer_token_future)?;
+
+                request.header(
+                    AUTHORIZATION,
+                    format!("Bearer {}", bearer_token.token.secret()),
+                )
+            }
         };
 
         let request = if let Some(request_body) = request_body {

sdk/storage_blobs/examples/blob_06_auto_refreshing_credentials.rs

@@ -0,0 +1,43 @@
+#[macro_use]
+extern crate log;
+
+use azure_identity::token_credentials::AutoRefreshingTokenCredential;
+use azure_identity::token_credentials::DefaultAzureCredential;
+use azure_storage::core::prelude::*;
+use azure_storage_blobs::prelude::*;
+use std::error::Error;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
+    env_logger::init();
+    // First we retrieve the account name, container and blob name from command line args
+
+    let account = std::env::args()
+        .nth(1)
+        .expect("please specify the account name as first command line parameter");
+    let container = std::env::args()
+        .nth(2)
+        .expect("please specify the container name as second command line parameter");
+    let blob = std::env::args()
+        .nth(3)
+        .expect("please specify the blob name as third command line parameter");
+
+    let creds = std::sync::Arc::new(DefaultAzureCredential::default());
+    let auto_creds = Box::new(AutoRefreshingTokenCredential::new(creds));
+
+    let http_client = azure_core::new_http_client();
+    let blob_client =
+        StorageAccountClient::new_token_credential(http_client.clone(), &account, auto_creds)
+            .as_container_client(&container)
+            .as_blob_client(&blob);
+
+    trace!("Requesting blob");
+
+    let response = blob_client.get().execute().await?;
+
+    let s_content = String::from_utf8(response.data.to_vec())?;
+    println!("blob == {:?}", blob);
+    println!("s_content == {}", s_content);
+
+    Ok(())
+}