GH-51349: [C++][Parquet] Add flag to control pass-through of KMS URLs from key material

cpp/src/parquet/encryption/crypto_factory.cc

@@ -178,6 +178,10 @@ std::shared_ptr<FileDecryptionProperties> CryptoFactory::GetFileDecryptionProper
       key_toolkit_, kms_connection_config, decryption_config.cache_lifetime_seconds,
       file_path, file_system);
 
+  if (decryption_config.read_kms_url) {
+    key_retriever->EnableReadingKmsUrl();
+  }
+
   return FileDecryptionProperties::Builder()
       .key_retriever(std::move(key_retriever))
       ->plaintext_files_allowed()
@@ -188,9 +192,9 @@ void CryptoFactory::RotateMasterKeys(
     const KmsConnectionConfig& kms_connection_config,
     const std::string& parquet_file_path,
     const std::shared_ptr<::arrow::fs::FileSystem>& file_system, bool double_wrapping,
-    double cache_lifetime_seconds) {
+    double cache_lifetime_seconds, bool read_kms_url) {
   key_toolkit_->RotateMasterKeys(kms_connection_config, parquet_file_path, file_system,
-                                 double_wrapping, cache_lifetime_seconds);
+                                 double_wrapping, cache_lifetime_seconds, read_kms_url);
 }
 
 }  // namespace parquet::encryption

cpp/src/parquet/encryption/crypto_factory.h

@@ -91,6 +91,12 @@ struct PARQUET_EXPORT DecryptionConfiguration {
   /// objects).
   /// The default is 600 (10 minutes).
   double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds;
+
+  /// Whether the KMS instance URL should be read from Parquet key material if it is
+  /// not configured in the KmsConnectionConfig.
+  /// This should only be enabled when the KMS implementation validates the URL it
+  /// receives, to ensure a KMS access token isn't sent to a malicious URL.
+  bool read_kms_url = false;
 };
 
 /// This is a core class, that translates the parameters of high level encryption (like
@@ -135,11 +141,17 @@ class PARQUET_EXPORT CryptoFactory {
   /// and then re-encrypted with new master keys.
   /// This relies on the KMS supporting versioning, such that the old master key is
   /// used when unwrapping a key, and the latest version is used when wrapping a key.
+  ///
+  /// If read_kms_url is true, the KMS instance URL is read from the key material being
+  /// rotated if it is not provided in the KmsConnectionConfig. This should only be
+  /// enabled when the KMS implementation validates the URL it receives, to ensure a KMS
+  /// access token isn't sent to a malicious URL.
   void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config,
                         const std::string& parquet_file_path,
                         const std::shared_ptr<::arrow::fs::FileSystem>& file_system,
                         bool double_wrapping = kDefaultDoubleWrapping,
-                        double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds);
+                        double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds,
+                        bool read_kms_url = false);
 
  private:
   ColumnPathToEncryptionPropertiesMap GetColumnEncryptionProperties(

cpp/src/parquet/encryption/file_key_unwrapper.cc

@@ -133,25 +133,31 @@ KeyWithMasterId FileKeyUnwrapper::GetDataEncryptionKey(const KeyMaterial& key_ma
   return KeyWithMasterId(std::move(data_key), master_key_id);
 }
 
+void FileKeyUnwrapper::EnableReadingKmsUrl() { read_kms_url_ = true; }
+
 std::shared_ptr<KmsClient> FileKeyUnwrapper::GetKmsClientFromConfigOrKeyMaterial(
     const KeyMaterial& key_material) {
   std::string& kms_instance_id = kms_connection_config_.kms_instance_id;
   if (kms_instance_id.empty()) {
     kms_instance_id = key_material.kms_instance_id();
     if (kms_instance_id.empty()) {
       throw ParquetException(
-          "KMS instance ID is missing both in both kms connection configuration and file "
+          "KMS instance ID is missing in both the KMS connection configuration and file "
           "key material");
     }
   }
 
   std::string& kms_instance_url = kms_connection_config_.kms_instance_url;
   if (kms_instance_url.empty()) {
-    kms_instance_url = key_material.kms_instance_url();
-    if (kms_instance_url.empty()) {
-      throw ParquetException(
-          "KMS instance ID is missing both in both kms connection configuration and file "
-          "key material");
+    if (read_kms_url_) {
+      kms_instance_url = key_material.kms_instance_url();
+      if (kms_instance_url.empty()) {
+        throw ParquetException(
+            "KMS instance URL is missing in both the KMS connection configuration and "
+            "the file key material");
+      }
+    } else {
+      kms_instance_url = KmsClient::kKmsInstanceUrlDefault;
     }
   }
 

cpp/src/parquet/encryption/file_key_unwrapper.h

@@ -70,6 +70,10 @@ class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever {
   /// Get the data key along with the master key id from key material
   KeyWithMasterId GetDataEncryptionKey(const KeyMaterial& key_material);
 
+  /// Enable reading the KMS instance URL from Parquet key material when it is not
+  /// already set.
+  void EnableReadingKmsUrl();
+
  private:
   FileKeyUnwrapper(std::shared_ptr<KeyToolkit> key_toolkit_owner, KeyToolkit* key_toolkit,
                    const KmsConnectionConfig& kms_connection_config,
@@ -91,6 +95,7 @@ class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever {
   std::shared_ptr<FileKeyMaterialStore> key_material_store_;
   const std::string file_path_;
   std::shared_ptr<::arrow::fs::FileSystem> file_system_;
+  bool read_kms_url_ = false;
 };
 
 }  // namespace parquet::encryption

cpp/src/parquet/encryption/key_management_test.cc

@@ -32,6 +32,8 @@
 #include "arrow/util/logging.h"
 
 #include "parquet/encryption/crypto_factory.h"
+#include "parquet/encryption/file_system_key_material_store.h"
+#include "parquet/encryption/key_material.h"
 #include "parquet/encryption/key_toolkit.h"
 #include "parquet/encryption/test_encryption_util.h"
 #include "parquet/encryption/test_in_memory_kms.h"
@@ -194,6 +196,72 @@ class TestEncryptionKeyManagement : public ::testing::Test {
     crypto_factory_.RemoveCacheEntriesForAllTokens();
   }
 
+  // Write a file that uses external key material and records the KMS
+  // instance ID and instance URL in its key material.
+  std::string WriteExternalMaterialFileWithKmsConfig(const std::string& instance_id,
+                                                     const std::string& instance_url) {
+    kms_connection_config_.kms_instance_id = instance_id;
+    kms_connection_config_.kms_instance_url = instance_url;
+    TestOnlyInServerWrapKms::InitializeMasterKeys(key_list_);
+    constexpr bool double_wrapping = true;
+    constexpr int encryption_no = 0;
+    this->WriteEncryptedParquetFile(double_wrapping, /*internal_key_material=*/false,
+                                    encryption_no);
+    return temp_dir_->path().ToString() + GetFileName(double_wrapping, wrap_locally_,
+                                                      /*internal_key_material=*/false,
+                                                      encryption_no);
+  }
+
+  // Rotate the keys of a file written with the KMS ID and URL configured,
+  // and return the KMS connection configurations used to create clients
+  // during key rotation.
+  std::vector<KmsConnectionConfig> RotateKeysWithKmsConfig(
+      const KmsConnectionConfig& rotation_config, const bool read_kms_url) {
+    const auto file_system = std::make_shared<::arrow::fs::LocalFileSystem>();
+    this->SetupCryptoFactory(false);
+
+    const std::string file_path =
+        this->WriteExternalMaterialFileWithKmsConfig("123", "https://example.com/kms");
+
+    auto kms_client_factory = std::make_shared<TestOnlyInMemoryKmsClientFactory>(
+        /*wrap_locally=*/false, key_list_);
+    auto crypto_factory = std::make_shared<CryptoFactory>();
+    crypto_factory->RegisterKmsClientFactory(kms_client_factory);
+
+    TestOnlyInServerWrapKms::StartKeyRotation(new_key_list_);
+    crypto_factory->RotateMasterKeys(rotation_config, file_path, file_system,
+                                     /*double_wrapping=*/true,
+                                     kDefaultCacheLifetimeSeconds, read_kms_url);
+    TestOnlyInServerWrapKms::FinishKeyRotation();
+
+    std::vector<KmsConnectionConfig> creation_requests =
+        kms_client_factory->CreationRequests();
+
+    // The new key material always uses the KMS connection configuration provided,
+    // not the config from the previous key material.
+    // If it's empty, default values are written.
+    const auto key_material_store =
+        FileSystemKeyMaterialStore::Make(file_path, file_system,
+                                         /*use_tmp_prefix=*/false);
+    const KeyMaterial rotated_key_material = KeyMaterial::Parse(
+        key_material_store->GetKeyMaterial(std::string(KeyMaterial::kFooterKeyIdInFile)));
+    const auto& expected_id = rotation_config.kms_instance_id.empty()
+                                  ? KmsClient::kKmsInstanceIdDefault
+                                  : rotation_config.kms_instance_id;
+    const auto& expected_url = rotation_config.kms_instance_url.empty()
+                                   ? KmsClient::kKmsInstanceUrlDefault
+                                   : rotation_config.kms_instance_url;
+    EXPECT_EQ(rotated_key_material.kms_instance_id(), expected_id);
+    EXPECT_EQ(rotated_key_material.kms_instance_url(), expected_url);
+
+    // Check the rotated file is readable
+    const auto file_decryption_properties = crypto_factory->GetFileDecryptionProperties(
+        rotation_config, GetDecryptionConfiguration(), file_path, file_system);
+    decryptor_.DecryptFile(file_path, file_decryption_properties);
+
+    return creation_requests;
+  }
+
   // Create encryption properties without keeping the creating CryptoFactory alive
   std::shared_ptr<FileEncryptionProperties> GetOrphanedFileEncryptionProperties(
       std::shared_ptr<KmsClientFactory> kms_client_factory,
@@ -441,4 +509,100 @@ TEST_F(TestEncryptionKeyManagement, ReadParquetMRExternalKeyMaterialFile) {
   }
 }
 
+TEST_F(TestEncryptionKeyManagement, ReadKmsUrlFromFile) {
+  this->SetupCryptoFactory(true);
+
+  constexpr bool internal_key_material = true;
+  constexpr bool double_wrapping = true;
+  constexpr int encryption_no = 0;
+
+  std::string file_name = "kms-config-test-file.parquet.encrypted";
+  std::string file_path = temp_dir_->path().ToString() + file_name;
+
+  auto encryption_config =
+      GetEncryptionConfiguration(double_wrapping, internal_key_material, encryption_no);
+
+  KmsConnectionConfig write_config;
+  write_config.kms_instance_id = "123";
+  write_config.kms_instance_url = "https://example.com/kms";
+
+  auto file_encryption_properties =
+      crypto_factory_.GetFileEncryptionProperties(write_config, encryption_config);
+  encryptor_.EncryptFile(file_path, file_encryption_properties);
+
+  for (const auto& enable_kms_url_read : {false, true}) {
+    // Create a fresh crypto factory and client factory for each read
+    // to avoid re-using cached clients.
+    CryptoFactory read_crypto_factory;
+    auto kms_client_factory =
+        std::make_shared<TestOnlyInMemoryKmsClientFactory>(true, key_list_);
+    read_crypto_factory.RegisterKmsClientFactory(kms_client_factory);
+
+    auto decryption_config = DecryptionConfiguration();
+    decryption_config.read_kms_url = enable_kms_url_read;
+
+    KmsConnectionConfig read_config;
+
+    auto file_decryption_properties =
+        read_crypto_factory.GetFileDecryptionProperties(read_config, decryption_config);
+
+    decryptor_.DecryptFile(file_path, file_decryption_properties);
+
+    ASSERT_EQ(kms_client_factory->CreationRequests().size(), 1);
+    const auto& request = kms_client_factory->CreationRequests()[0];
+    EXPECT_EQ(request.kms_instance_id, "123");
+    if (enable_kms_url_read) {
+      EXPECT_EQ(request.kms_instance_url, "https://example.com/kms");
+    } else {
+      EXPECT_EQ(request.kms_instance_url, "DEFAULT");
+    }
+  }
+}
+
+TEST_F(TestEncryptionKeyManagement, ReadKmsUrlFromFileDuringKeyRotation) {
+  // Use an empty config for rotation
+  const KmsConnectionConfig rotation_config;
+  const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/true);
+
+  ASSERT_EQ(requests.size(), 2);
+  // The first KMS creation request is for wrapping new keys.
+  // This uses the empty config provided.
+  EXPECT_EQ(requests[0].kms_instance_id, "");
+  EXPECT_EQ(requests[0].kms_instance_url, "");
+  // The KMS client used to unwrap the previous keys should be configured
+  // with the instance ID and url provided at write time.
+  EXPECT_EQ(requests[1].kms_instance_id, "123");
+  EXPECT_EQ(requests[1].kms_instance_url, "https://example.com/kms");
+}
+
+TEST_F(TestEncryptionKeyManagement, KeyRotationWithoutReadingKmsUrl) {
+  // Use an empty config for rotation
+  const KmsConnectionConfig rotation_config;
+  const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/false);
+
+  ASSERT_EQ(requests.size(), 2);
+  // The first KMS creation request is for wrapping new keys.
+  // This uses the empty config provided.
+  EXPECT_EQ(requests[0].kms_instance_id, "");
+  EXPECT_EQ(requests[0].kms_instance_url, "");
+  // When unwrapping the existing keys, the URL in the key material is
+  // ignored and the default used.
+  EXPECT_EQ(requests[1].kms_instance_id, "123");
+  EXPECT_EQ(requests[1].kms_instance_url, KmsClient::kKmsInstanceUrlDefault);
+}
+
+TEST_F(TestEncryptionKeyManagement, KeyRotationUsesProvidedKmsConfig) {
+  KmsConnectionConfig rotation_config;
+  rotation_config.kms_instance_id = "456";
+  rotation_config.kms_instance_url = "https://example.com/kms2";
+  const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/true);
+
+  ASSERT_EQ(requests.size(), 1);
+  // Wrap and unwrap both use the same configuration.
+  // The instance id and url in the existing key material is ignored even though
+  // read_kms_url is enabled. The provided config takes precedence.
+  EXPECT_EQ(requests[0].kms_instance_id, "456");
+  EXPECT_EQ(requests[0].kms_instance_url, "https://example.com/kms2");
+}
+
 }  // namespace parquet::encryption::test

cpp/src/parquet/encryption/key_toolkit.cc

@@ -47,7 +47,7 @@ void KeyToolkit::RotateMasterKeys(
     const KmsConnectionConfig& kms_connection_config,
     const std::string& parquet_file_path,
     const std::shared_ptr<::arrow::fs::FileSystem>& file_system, bool double_wrapping,
-    double cache_lifetime_seconds) {
+    double cache_lifetime_seconds, bool read_kms_url) {
   // If process wrote files with double-wrapped keys, clean KEK cache (since master keys
   // are changing). Only once for each key rotation cycle; not for every file.
   const auto now = internal::CurrentTimePoint();
@@ -65,6 +65,9 @@ void KeyToolkit::RotateMasterKeys(
   // Unwrapper for decrypting encrypted keys
   FileKeyUnwrapper file_key_unwrapper(this, kms_connection_config, cache_lifetime_seconds,
                                       key_material_store);
+  if (read_kms_url) {
+    file_key_unwrapper.EnableReadingKmsUrl();
+  }
 
   // Create a temporary store to hold new key material during rotation,
   // and wrapper that will write material to this store when getting key metadata.

cpp/src/parquet/encryption/key_toolkit.h

@@ -77,7 +77,8 @@ class PARQUET_EXPORT KeyToolkit {
   void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config,
                         const std::string& parquet_file_path,
                         const std::shared_ptr<::arrow::fs::FileSystem>& file_system,
-                        bool double_wrapping, double cache_lifetime_seconds);
+                        bool double_wrapping, double cache_lifetime_seconds,
+                        bool read_kms_url = false);
 
  private:
   TwoLevelCacheWithExpiration<std::shared_ptr<KmsClient>> kms_client_cache_;

cpp/src/parquet/encryption/test_in_memory_kms.h

@@ -90,12 +90,20 @@ class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory {
 
   std::shared_ptr<KmsClient> CreateKmsClient(
       const KmsConnectionConfig& kms_connection_config) {
+    create_requests_.push_back(kms_connection_config);
     if (wrap_locally_) {
       return std::make_shared<TestOnlyLocalWrapInMemoryKms>(kms_connection_config);
     } else {
       return std::make_shared<TestOnlyInServerWrapKms>();
     }
   }
+
+  /// Get the `KmsConnectionConfig` values that have been used to
+  /// create clients with this factory.
+  const std::vector<KmsConnectionConfig>& CreationRequests() { return create_requests_; }
+
+ private:
+  std::vector<KmsConnectionConfig> create_requests_;
 };
 
 }  // namespace parquet::encryption

docs/source/python/parquet/parquet_encryption.rst

@@ -232,6 +232,11 @@ file decryption properties) is optional and it includes the following options:
 
 * ``cache_lifetime``, the lifetime of cached entities (key encryption keys, local
   wrapping keys, KMS client objects) represented as a ``datetime.timedelta``.
+* ``read_kms_url``, whether the KMS instance URL may be read from the key material
+  of the file being read, when it is not set in the ``KmsConnectionConfig``. This
+  defaults to ``False``, and should only be enabled when the KMS implementation
+  validates the URL it receives, to ensure a KMS access token isn't sent to a
+  malicious URL.
 
 External key material and key rotation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -251,6 +256,11 @@ key material file, without changing the Parquet file itself:
    ...     kms_connection_config, parquet_file_path="table.parquet",
    ... )
 
+``rotate_master_keys`` also accepts ``read_kms_url``, which behaves like the
+``DecryptionConfiguration`` option of the same name when the existing key material is
+read. The key material written by key rotation always uses the connection properties
+from the ``KmsConnectionConfig`` that is passed in.
+
 Direct Key Encryption (without KMS)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~