enable updating current credentials (#1424)

This PR allows replacing credential details for existing
StorageCredential instances.  This allows clients to update the
credentials as runtime, as seen in other SDKs.

Ref: azure-sdk-for-net's [AzureSasCredential.Update](https://learn.microsoft.com/en-us/dotnet/api/azure.azuresascredential.update?view=azure-dotnet#azure-azuresascredential-update(system-string))

Author: demoray

sdk/storage/src/authorization/authorization_policy.rs

@@ -1,10 +1,10 @@
-use crate::clients::ServiceType;
-use crate::StorageCredentials;
-use azure_core::error::{ErrorKind, ResultExt};
-use azure_core::Method;
-use azure_core::{headers::*, Context, Policy, PolicyResult, Request};
-use std::borrow::Cow;
-use std::sync::Arc;
+use crate::{clients::ServiceType, StorageCredentials, StorageCredentialsInner};
+use azure_core::{
+    error::{ErrorKind, ResultExt},
+    headers::*,
+    Context, Method, Policy, PolicyResult, Request,
+};
+use std::{borrow::Cow, ops::Deref, sync::Arc};
 use url::Url;
 
 const STORAGE_TOKEN_SCOPE: &str = "https://storage.azure.com/";
@@ -35,49 +35,51 @@ impl Policy for AuthorizationPolicy {
             !next.is_empty(),
             "Authorization policies cannot be the last policy of a pipeline"
         );
-        let request = match &self.credentials {
-            StorageCredentials::Key(account, key) => {
-                if !request.url().query_pairs().any(|(k, _)| &*k == "sig") {
-                    let auth = generate_authorization(
-                        request.headers(),
-                        request.url(),
-                        *request.method(),
-                        account,
-                        key,
-                        *ctx.get()
-                            .expect("ServiceType must be in the Context at this point"),
-                    )?;
-                    request.insert_header(AUTHORIZATION, auth);
+
+        // lock the credentials within a scope so that it is released as soon as possible
+        {
+            let creds = self.credentials.0.lock().await;
+
+            match creds.deref() {
+                StorageCredentialsInner::Key(account, key) => {
+                    if !request.url().query_pairs().any(|(k, _)| &*k == "sig") {
+                        let auth = generate_authorization(
+                            request.headers(),
+                            request.url(),
+                            *request.method(),
+                            account,
+                            key,
+                            *ctx.get()
+                                .expect("ServiceType must be in the Context at this point"),
+                        )?;
+                        request.insert_header(AUTHORIZATION, auth);
+                    }
                 }
-                request
-            }
-            StorageCredentials::SASToken(query_pairs) => {
-                // Ensure the signature param is not already present
-                if !request.url().query_pairs().any(|(k, _)| &*k == "sig") {
-                    request
-                        .url_mut()
-                        .query_pairs_mut()
-                        .extend_pairs(query_pairs);
+                StorageCredentialsInner::SASToken(query_pairs) => {
+                    // Ensure the signature param is not already present
+                    if !request.url().query_pairs().any(|(k, _)| &*k == "sig") {
+                        request
+                            .url_mut()
+                            .query_pairs_mut()
+                            .extend_pairs(query_pairs);
+                    }
                 }
-                request
-            }
-            StorageCredentials::BearerToken(token) => {
-                request.insert_header(AUTHORIZATION, format!("Bearer {token}"));
-                request
-            }
-            StorageCredentials::TokenCredential(token_credential) => {
-                let bearer_token = token_credential
-                    .get_token(STORAGE_TOKEN_SCOPE)
-                    .await
-                    .context(ErrorKind::Credential, "failed to get bearer token")?;
-
-                request.insert_header(
-                    AUTHORIZATION,
-                    format!("Bearer {}", bearer_token.token.secret()),
-                );
-                request
+                StorageCredentialsInner::BearerToken(token) => {
+                    request.insert_header(AUTHORIZATION, format!("Bearer {token}"));
+                }
+                StorageCredentialsInner::TokenCredential(token_credential) => {
+                    let bearer_token = token_credential
+                        .get_token(STORAGE_TOKEN_SCOPE)
+                        .await
+                        .context(ErrorKind::Credential, "failed to get bearer token")?;
+
+                    request.insert_header(
+                        AUTHORIZATION,
+                        format!("Bearer {}", bearer_token.token.secret()),
+                    );
+                }
+                StorageCredentialsInner::Anonymous => {}
             }
-            StorageCredentials::Anonymous => request,
         };
 
         next[0].send(ctx, request, &next[1..]).await

sdk/storage/src/authorization/mod.rs

@@ -1,16 +1,19 @@
 mod authorization_policy;
 
+pub(crate) use self::authorization_policy::AuthorizationPolicy;
+use crate::clients::{EMULATOR_ACCOUNT, EMULATOR_ACCOUNT_KEY};
 use azure_core::{
     auth::TokenCredential,
     error::{ErrorKind, ResultExt},
 };
-use std::sync::Arc;
+use futures::lock::Mutex;
+use std::{
+    mem::replace,
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
 use url::Url;
 
-pub(crate) use authorization_policy::AuthorizationPolicy;
-
-use crate::clients::{EMULATOR_ACCOUNT, EMULATOR_ACCOUNT_KEY};
-
 /// Credentials for accessing a storage account.
 ///
 /// # Example
@@ -22,7 +25,10 @@ use crate::clients::{EMULATOR_ACCOUNT, EMULATOR_ACCOUNT_KEY};
 /// azure_storage::StorageCredentials::access_key("my_account", "SOMEACCESSKEY");
 /// ```
 #[derive(Clone)]
-pub enum StorageCredentials {
+pub struct StorageCredentials(pub Arc<Mutex<StorageCredentialsInner>>);
+
+#[derive(Clone)]
+pub enum StorageCredentialsInner {
     Key(String, String),
     SASToken(Vec<(String, String)>),
     BearerToken(String),
@@ -31,6 +37,11 @@ pub enum StorageCredentials {
 }
 
 impl StorageCredentials {
+    /// Create a new `StorageCredentials` from a `StorageCredentialsInner`
+    fn wrap(inner: StorageCredentialsInner) -> Self {
+        Self(Arc::new(Mutex::new(inner)))
+    }
+
     /// Create an Access Key based credential
     ///
     /// When you create a storage account, Azure generates two 512-bit storage
@@ -44,7 +55,7 @@ impl StorageCredentials {
         A: Into<String>,
         K: Into<String>,
     {
-        Self::Key(account.into(), key.into())
+        Self::wrap(StorageCredentialsInner::Key(account.into(), key.into()))
     }
 
     /// Create a Shared Access Signature (SAS) token based credential
@@ -60,7 +71,7 @@ impl StorageCredentials {
         S: AsRef<str>,
     {
         let params = get_sas_token_parms(token.as_ref())?;
-        Ok(Self::SASToken(params))
+        Ok(Self::wrap(StorageCredentialsInner::SASToken(params)))
     }
 
     /// Create an Bearer Token based credential
@@ -77,7 +88,7 @@ impl StorageCredentials {
     where
         T: Into<String>,
     {
-        Self::BearerToken(token.into())
+        Self::wrap(StorageCredentialsInner::BearerToken(token.into()))
     }
 
     /// Create a `TokenCredential` based credential
@@ -98,7 +109,7 @@ impl StorageCredentials {
     ///
     /// ref: <https://docs.microsoft.com/rest/api/storageservices/authorize-with-azure-active-directory>
     pub fn token_credential(credential: Arc<dyn TokenCredential>) -> Self {
-        Self::TokenCredential(credential)
+        Self::wrap(StorageCredentialsInner::TokenCredential(credential))
     }
 
     /// Create an anonymous credential
@@ -113,45 +124,71 @@ impl StorageCredentials {
     ///
     /// ref: <https://docs.microsoft.com/azure/storage/blobs/anonymous-read-access-configure>
     pub fn anonymous() -> Self {
-        Self::Anonymous
+        Self::wrap(StorageCredentialsInner::Anonymous)
     }
 
     /// Create an Access Key credential for use with the Azure Storage emulator
     pub fn emulator() -> Self {
         Self::access_key(EMULATOR_ACCOUNT, EMULATOR_ACCOUNT_KEY)
     }
+
+    /// Replace the current credentials with new credentials
+    ///
+    /// This method is useful for updating credentials that are used by multiple
+    /// clients at once.
+    pub async fn replace(&self, other: Self) -> azure_core::Result<()> {
+        if Arc::ptr_eq(&self.0, &other.0) {
+            return Ok(());
+        }
+
+        let mut creds = self.0.lock().await;
+        let other = other.0.lock().await;
+        let creds = creds.deref_mut();
+        let other = other.deref().clone();
+        let _old_creds = replace(creds, other);
+
+        Ok(())
+    }
 }
 
 impl std::fmt::Debug for StorageCredentials {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match &self {
-            StorageCredentials::Key(_, _) => f
-                .debug_struct("StorageCredentials")
-                .field("credential", &"Key")
-                .finish(),
-            StorageCredentials::SASToken(_) => f
-                .debug_struct("StorageCredentials")
-                .field("credential", &"SASToken")
-                .finish(),
-            StorageCredentials::BearerToken(_) => f
-                .debug_struct("StorageCredentials")
-                .field("credential", &"BearerToken")
-                .finish(),
-            StorageCredentials::TokenCredential(_) => f
-                .debug_struct("StorageCredentials")
-                .field("credential", &"TokenCredential")
-                .finish(),
-            StorageCredentials::Anonymous => f
+        let creds = self.0.try_lock();
+
+        match creds.as_deref() {
+            None => f
                 .debug_struct("StorageCredentials")
-                .field("credential", &"Anonymous")
+                .field("credential", &"locked")
                 .finish(),
+            Some(inner) => match &inner {
+                StorageCredentialsInner::Key(_, _) => f
+                    .debug_struct("StorageCredentials")
+                    .field("credential", &"Key")
+                    .finish(),
+                StorageCredentialsInner::SASToken(_) => f
+                    .debug_struct("StorageCredentials")
+                    .field("credential", &"SASToken")
+                    .finish(),
+                StorageCredentialsInner::BearerToken(_) => f
+                    .debug_struct("StorageCredentials")
+                    .field("credential", &"BearerToken")
+                    .finish(),
+                StorageCredentialsInner::TokenCredential(_) => f
+                    .debug_struct("StorageCredentials")
+                    .field("credential", &"TokenCredential")
+                    .finish(),
+                StorageCredentialsInner::Anonymous => f
+                    .debug_struct("StorageCredentials")
+                    .field("credential", &"Anonymous")
+                    .finish(),
+            },
         }
     }
 }
 
 impl From<Arc<dyn TokenCredential>> for StorageCredentials {
     fn from(cred: Arc<dyn TokenCredential>) -> Self {
-        Self::TokenCredential(cred)
+        Self::token_credential(cred)
     }
 }
 
@@ -160,7 +197,7 @@ impl TryFrom<&Url> for StorageCredentials {
     fn try_from(value: &Url) -> Result<Self, Self::Error> {
         match value.query() {
             Some(query) => Self::sas_token(query),
-            None => Ok(Self::Anonymous),
+            None => Ok(Self::anonymous()),
         }
     }
 }
@@ -188,3 +225,30 @@ fn get_sas_token_parms(sas_token: &str) -> azure_core::Result<Vec<(String, Strin
         .map(|p| (String::from(p.0), String::from(p.1)))
         .collect())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_replacement() -> azure_core::Result<()> {
+        let base = StorageCredentials::anonymous();
+        let other = StorageCredentials::bearer_token("foo");
+
+        base.replace(other).await?;
+
+        // check that the value was updated
+        {
+            let inner = base.0.lock().await;
+            let inner_locked = inner.deref();
+            assert!(
+                matches!(&inner_locked, &StorageCredentialsInner::BearerToken(value) if value == "foo")
+            );
+        }
+
+        // updating with the same StorageCredentials shouldn't deadlock
+        base.replace(base.clone()).await?;
+
+        Ok(())
+    }
+}

sdk/storage/src/clients.rs

@@ -1,15 +1,18 @@
-use crate::authorization::AuthorizationPolicy;
-use crate::shared_access_signature::account_sas::{
-    AccountSasPermissions, AccountSasResource, AccountSasResourceType, AccountSharedAccessSignature,
+use crate::{
+    authorization::{AuthorizationPolicy, StorageCredentialsInner},
+    shared_access_signature::account_sas::{
+        AccountSasPermissions, AccountSasResource, AccountSasResourceType,
+        AccountSharedAccessSignature,
+    },
+    StorageCredentials,
 };
-use crate::StorageCredentials;
-use azure_core::date;
 use azure_core::{
+    date,
     error::{Error, ErrorKind},
     headers::*,
     Body, ClientOptions, Method, Pipeline, Request,
 };
-use std::sync::Arc;
+use std::{ops::Deref, sync::Arc};
 use time::OffsetDateTime;
 use url::Url;
 
@@ -44,19 +47,29 @@ impl ServiceType {
     }
 }
 
-pub fn shared_access_signature(
+pub async fn shared_access_signature(
     storage_credentials: &StorageCredentials,
     resource: AccountSasResource,
     resource_type: AccountSasResourceType,
     expiry: OffsetDateTime,
     permissions: AccountSasPermissions,
 ) -> Result<AccountSharedAccessSignature, Error> {
-    match storage_credentials {
-            StorageCredentials::Key(account, key) => {
-                Ok(AccountSharedAccessSignature::new(account.clone(), key.clone(), resource, resource_type, expiry, permissions))
-            }
-            _ => Err(Error::message(ErrorKind::Credential, "failed shared access signature generation. SAS can be generated only from key and account clients")),
-        }
+    let creds = storage_credentials.0.lock().await;
+    let StorageCredentialsInner::Key(account, key) = creds.deref() else {
+        return Err(Error::message(
+            ErrorKind::Credential,
+            "Shared access signature generation - SAS can be generated with access_key clients",
+        ));
+    };
+
+    Ok(AccountSharedAccessSignature::new(
+        account.clone(),
+        key.clone(),
+        resource,
+        resource_type,
+        expiry,
+        permissions,
+    ))
 }
 
 pub fn finalize_request(

sdk/storage/src/lib.rs

@@ -39,7 +39,7 @@ pub mod shared_access_signature;
 
 pub use self::connection_string::{ConnectionString, EndpointProtocol};
 pub use self::connection_string_builder::ConnectionStringBuilder;
-pub use authorization::StorageCredentials;
+pub use authorization::{StorageCredentials, StorageCredentialsInner};
 pub use cloud_location::*;
 pub mod headers;
 pub use copy_id::{copy_id_from_headers, CopyId};