[arrow-rs] branch master updated: Don't panic on invalid Azure access key (#4972) (#4974)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 6b4fd2f222 Don't panic on invalid Azure access key (#4972) (#4974)
6b4fd2f222 is described below

commit 6b4fd2f2224f3eb08a3eb55cf54bde9b2f1d2793
Author: Raphael Taylor-Davies <1781103+tustv...@users.noreply.github.com>
AuthorDate: Thu Oct 26 09:17:23 2023 +0100

    Don't panic on invalid Azure access key (#4972) (#4974)
---
 object_store/src/azure/builder.rs    | 14 ++++++++------
 object_store/src/azure/credential.rs | 23 +++++++++++++++++++----
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/object_store/src/azure/builder.rs 
b/object_store/src/azure/builder.rs
index 915e4c59a8..02e0762b6d 100644
--- a/object_store/src/azure/builder.rs
+++ b/object_store/src/azure/builder.rs
@@ -17,7 +17,7 @@
 
 use crate::azure::client::{AzureClient, AzureConfig};
 use crate::azure::credential::{
-    AzureCliCredential, ClientSecretOAuthProvider, ImdsManagedIdentityProvider,
+    AzureAccessKey, AzureCliCredential, ClientSecretOAuthProvider, 
ImdsManagedIdentityProvider,
     WorkloadIdentityOAuthProvider,
 };
 use crate::azure::{AzureCredential, AzureCredentialProvider, MicrosoftAzure, 
STORE};
@@ -800,11 +800,12 @@ impl MicrosoftAzureBuilder {
             // Allow overriding defaults. Values taken from
             // from 
https://docs.rs/azure_storage/0.2.0/src/azure_storage/core/clients/storage_account_client.rs.html#129-141
             let url = url_from_env("AZURITE_BLOB_STORAGE_URL", 
"http://127.0.0.1:10000";)?;
-            let account_key = self
-                .access_key
-                .unwrap_or_else(|| EMULATOR_ACCOUNT_KEY.to_string());
+            let key = match self.access_key {
+                Some(k) => AzureAccessKey::try_new(&k)?,
+                None => AzureAccessKey::try_new(EMULATOR_ACCOUNT_KEY)?,
+            };
 
-            let credential = 
static_creds(AzureCredential::AccessKey(account_key));
+            let credential = static_creds(AzureCredential::AccessKey(key));
 
             self.client_options = self.client_options.with_allow_http(true);
             (true, url, credential, account_name)
@@ -828,7 +829,8 @@ impl MicrosoftAzureBuilder {
             } else if let Some(bearer_token) = self.bearer_token {
                 static_creds(AzureCredential::BearerToken(bearer_token))
             } else if let Some(access_key) = self.access_key {
-                static_creds(AzureCredential::AccessKey(access_key))
+                let key = AzureAccessKey::try_new(&access_key)?;
+                static_creds(AzureCredential::AccessKey(key))
             } else if let (Some(client_id), Some(tenant_id), 
Some(federated_token_file)) =
                 (&self.client_id, &self.tenant_id, self.federated_token_file)
             {
diff --git a/object_store/src/azure/credential.rs 
b/object_store/src/azure/credential.rs
index 283d7ff9d7..2b8788d333 100644
--- a/object_store/src/azure/credential.rs
+++ b/object_store/src/azure/credential.rs
@@ -75,6 +75,9 @@ pub enum Error {
     #[snafu(display("Error reading federated token file "))]
     FederatedTokenFile,
 
+    #[snafu(display("Invalid Access Key: {}", source))]
+    InvalidAccessKey { source: base64::DecodeError },
+
     #[snafu(display("'az account get-access-token' command failed: 
{message}"))]
     AzureCli { message: String },
 
@@ -93,13 +96,25 @@ impl From<Error> for crate::Error {
     }
 }
 
+/// A shared Azure Storage Account Key
+#[derive(Debug, Eq, PartialEq)]
+pub struct AzureAccessKey(Vec<u8>);
+
+impl AzureAccessKey {
+    /// Create a new [`AzureAccessKey`], checking it for validity
+    pub fn try_new(key: &str) -> Result<Self> {
+        let key = BASE64_STANDARD.decode(key).context(InvalidAccessKeySnafu)?;
+        Ok(Self(key))
+    }
+}
+
 /// An Azure storage credential
 #[derive(Debug, Eq, PartialEq)]
 pub enum AzureCredential {
     /// A shared access key
     ///
     /// 
<https://learn.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key>
-    AccessKey(String),
+    AccessKey(AzureAccessKey),
     /// A shared access signature
     ///
     /// 
<https://learn.microsoft.com/en-us/rest/api/storageservices/delegate-access-with-shared-access-signature>
@@ -149,7 +164,7 @@ impl CredentialExt for RequestBuilder {
                     request.url(),
                     request.method(),
                     account,
-                    key.as_str(),
+                    key,
                 );
 
                 // "signature" is a base 64 encoded string so it should never
@@ -174,10 +189,10 @@ fn generate_authorization(
     u: &Url,
     method: &Method,
     account: &str,
-    key: &str,
+    key: &AzureAccessKey,
 ) -> String {
     let str_to_sign = string_to_sign(h, u, method, account);
-    let auth = hmac_sha256(BASE64_STANDARD.decode(key).unwrap(), str_to_sign);
+    let auth = hmac_sha256(&key.0, str_to_sign);
     format!("SharedKey {}:{}", account, BASE64_STANDARD.encode(auth))
 }