From a55115cbe9f06bfd40413562ef8aa34418e9728a Mon Sep 17 00:00:00 2001
From: "R. Tyler Croy" <rtyler@brokenco.de>
Date: Tue, 8 Oct 2024 14:17:32 +0000
Subject: [PATCH] fix: properly handle the different flavors of storage options
 keys

This change addresses some of the problems tacked onto #2893 but does
not address the concern with ECS specifically. It does however improve
the handling of `storage_options` since the improved credential
provider code was introduced
---
 crates/aws/Cargo.toml         |  2 +-
 crates/aws/src/credentials.rs | 67 +++++++++++++++++++++++++++++++++--
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/crates/aws/Cargo.toml b/crates/aws/Cargo.toml
index e653d2b242..60179c7f82 100644
--- a/crates/aws/Cargo.toml
+++ b/crates/aws/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "deltalake-aws"
-version = "0.4.0"
+version = "0.4.1"
 authors.workspace = true
 keywords.workspace = true
 readme.workspace = true
diff --git a/crates/aws/src/credentials.rs b/crates/aws/src/credentials.rs
index 71441bf05e..43c9d88a31 100644
--- a/crates/aws/src/credentials.rs
+++ b/crates/aws/src/credentials.rs
@@ -1,6 +1,8 @@
 //! Custom AWS credential providers used by delta-rs
 //!
 
+use std::collections::HashMap;
+use std::str::FromStr;
 use std::sync::Arc;
 
 use aws_config::default_provider::credentials::DefaultCredentialsChain;
@@ -83,13 +85,25 @@ impl OptionsCredentialsProvider {
     /// [Credentials] instance for AWS SDK credential resolution
     fn credentials(&self) -> aws_credential_types::provider::Result {
         debug!("Attempting to pull credentials from `StorageOptions`");
-        let access_key = self.options.0.get(constants::AWS_ACCESS_KEY_ID).ok_or(
+
+        // StorageOptions can have a variety of flavors because
+        // [object_store::aws::AmazonS3ConfigKey] supports a couple different variants for key
+        // names.
+        let config_keys: HashMap<AmazonS3ConfigKey, String> =
+            HashMap::from_iter(self.options.0.iter().filter_map(|(k, v)| {
+                match AmazonS3ConfigKey::from_str(&k.to_lowercase()) {
+                    Ok(k) => Some((k, v.into())),
+                    Err(_) => None,
+                }
+            }));
+
+        let access_key = config_keys.get(&AmazonS3ConfigKey::AccessKeyId).ok_or(
             CredentialsError::not_loaded("access key not in StorageOptions"),
         )?;
-        let secret_key = self.options.0.get(constants::AWS_SECRET_ACCESS_KEY).ok_or(
+        let secret_key = config_keys.get(&AmazonS3ConfigKey::SecretAccessKey).ok_or(
             CredentialsError::not_loaded("secret key not in StorageOptions"),
         )?;
-        let session_token = self.options.0.get(constants::AWS_SESSION_TOKEN).cloned();
+        let session_token = config_keys.get(&AmazonS3ConfigKey::Token).cloned();
 
         Ok(Credentials::new(
             access_key,
@@ -110,6 +124,53 @@ impl ProvideCredentials for OptionsCredentialsProvider {
     }
 }
 
+#[cfg(test)]
+mod options_tests {
+    use super::*;
+    use maplit::hashmap;
+
+    #[test]
+    fn test_empty_options_error() {
+        let options = StorageOptions::default();
+        let provider = OptionsCredentialsProvider { options };
+        let result = provider.credentials();
+        assert!(
+            result.is_err(),
+            "The default StorageOptions don't have credentials!"
+        );
+    }
+
+    #[test]
+    fn test_uppercase_options_resolve() {
+        let mash = hashmap! {
+            "AWS_ACCESS_KEY_ID".into() => "key".into(),
+            "AWS_SECRET_ACCESS_KEY".into() => "secret".into(),
+        };
+        let options = StorageOptions(mash);
+        let provider = OptionsCredentialsProvider { options };
+        let result = provider.credentials();
+        assert!(result.is_ok(), "StorageOptions with at least AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY should resolve");
+        let result = result.unwrap();
+        assert_eq!(result.access_key_id(), "key");
+        assert_eq!(result.secret_access_key(), "secret");
+    }
+
+    #[test]
+    fn test_lowercase_options_resolve() {
+        let mash = hashmap! {
+            "aws_access_key_id".into() => "key".into(),
+            "aws_secret_access_key".into() => "secret".into(),
+        };
+        let options = StorageOptions(mash);
+        let provider = OptionsCredentialsProvider { options };
+        let result = provider.credentials();
+        assert!(result.is_ok(), "StorageOptions with at least AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY should resolve");
+        let result = result.unwrap();
+        assert_eq!(result.access_key_id(), "key");
+        assert_eq!(result.secret_access_key(), "secret");
+    }
+}
+
 /// Generate a random session name for assuming IAM roles
 fn assume_role_sessio_name() -> String {
     let now = chrono::Utc::now();