diff --git a/cpp/src/parquet/encryption/crypto_factory.cc b/cpp/src/parquet/encryption/crypto_factory.cc index dbfd20a1b7e0..09f15950cc6b 100644 --- a/cpp/src/parquet/encryption/crypto_factory.cc +++ b/cpp/src/parquet/encryption/crypto_factory.cc @@ -178,6 +178,10 @@ std::shared_ptr 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 diff --git a/cpp/src/parquet/encryption/crypto_factory.h b/cpp/src/parquet/encryption/crypto_factory.h index 3c6a5f29ea88..2d2d302321c9 100644 --- a/cpp/src/parquet/encryption/crypto_factory.h +++ b/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( diff --git a/cpp/src/parquet/encryption/file_key_unwrapper.cc b/cpp/src/parquet/encryption/file_key_unwrapper.cc index 1cc032013779..58963cc8893f 100644 --- a/cpp/src/parquet/encryption/file_key_unwrapper.cc +++ b/cpp/src/parquet/encryption/file_key_unwrapper.cc @@ -133,6 +133,8 @@ 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 FileKeyUnwrapper::GetKmsClientFromConfigOrKeyMaterial( const KeyMaterial& key_material) { std::string& kms_instance_id = kms_connection_config_.kms_instance_id; @@ -140,18 +142,22 @@ std::shared_ptr FileKeyUnwrapper::GetKmsClientFromConfigOrKeyMaterial 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; } } diff --git a/cpp/src/parquet/encryption/file_key_unwrapper.h b/cpp/src/parquet/encryption/file_key_unwrapper.h index c86f68121c87..3123c1f611c5 100644 --- a/cpp/src/parquet/encryption/file_key_unwrapper.h +++ b/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 key_toolkit_owner, KeyToolkit* key_toolkit, const KmsConnectionConfig& kms_connection_config, @@ -91,6 +95,7 @@ class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever { std::shared_ptr key_material_store_; const std::string file_path_; std::shared_ptr<::arrow::fs::FileSystem> file_system_; + bool read_kms_url_ = false; }; } // namespace parquet::encryption diff --git a/cpp/src/parquet/encryption/key_management_test.cc b/cpp/src/parquet/encryption/key_management_test.cc index ed6d15dbb6ac..80ef28b766b3 100644 --- a/cpp/src/parquet/encryption/key_management_test.cc +++ b/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 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( + /*wrap_locally=*/false, key_list_); + auto crypto_factory = std::make_shared(); + 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 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 GetOrphanedFileEncryptionProperties( std::shared_ptr 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(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 diff --git a/cpp/src/parquet/encryption/key_toolkit.cc b/cpp/src/parquet/encryption/key_toolkit.cc index cc56a8af620f..2d0e964e08f4 100644 --- a/cpp/src/parquet/encryption/key_toolkit.cc +++ b/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. diff --git a/cpp/src/parquet/encryption/key_toolkit.h b/cpp/src/parquet/encryption/key_toolkit.h index a0b929877eeb..8ca5e9fac203 100644 --- a/cpp/src/parquet/encryption/key_toolkit.h +++ b/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> kms_client_cache_; diff --git a/cpp/src/parquet/encryption/test_in_memory_kms.h b/cpp/src/parquet/encryption/test_in_memory_kms.h index b9d4169c6345..35d2531be84e 100644 --- a/cpp/src/parquet/encryption/test_in_memory_kms.h +++ b/cpp/src/parquet/encryption/test_in_memory_kms.h @@ -90,12 +90,20 @@ class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory { std::shared_ptr CreateKmsClient( const KmsConnectionConfig& kms_connection_config) { + create_requests_.push_back(kms_connection_config); if (wrap_locally_) { return std::make_shared(kms_connection_config); } else { return std::make_shared(); } } + + /// Get the `KmsConnectionConfig` values that have been used to + /// create clients with this factory. + const std::vector& CreationRequests() { return create_requests_; } + + private: + std::vector create_requests_; }; } // namespace parquet::encryption diff --git a/docs/source/python/parquet/parquet_encryption.rst b/docs/source/python/parquet/parquet_encryption.rst index 98d72b4fe71c..be46b4850463 100644 --- a/docs/source/python/parquet/parquet_encryption.rst +++ b/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) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/python/pyarrow/_parquet_encryption.pyx b/python/pyarrow/_parquet_encryption.pyx index 7fe7fa7491dc..ed88a98b0070 100644 --- a/python/pyarrow/_parquet_encryption.pyx +++ b/python/pyarrow/_parquet_encryption.pyx @@ -198,8 +198,13 @@ cdef class DecryptionConfiguration(_Weakrefable): # Avoid mistakingly creating attributes __slots__ = () - def __init__(self, *, cache_lifetime=None): + def __init__(self, *, cache_lifetime=None, + read_kms_url=None): self.configuration.reset(new CDecryptionConfiguration()) + if cache_lifetime is not None: + self.cache_lifetime = cache_lifetime + if read_kms_url is not None: + self.read_kms_url = read_kms_url @property def cache_lifetime(self): @@ -210,7 +215,27 @@ cdef class DecryptionConfiguration(_Weakrefable): @cache_lifetime.setter def cache_lifetime(self, value): - self.configuration.get().cache_lifetime_seconds = value.total_seconds() + try: + # Expect a timedelta value + seconds = value.total_seconds() + except AttributeError: + # Also accept a number of seconds + seconds = float(value) + self.configuration.get().cache_lifetime_seconds = seconds + + @property + def read_kms_url(self): + """Whether the KMS instance URL may be read from Parquet key material + when 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.""" + return self.configuration.get().read_kms_url + + @read_kms_url.setter + def read_kms_url(self, value): + self.configuration.get().read_kms_url = value cdef inline shared_ptr[CDecryptionConfiguration] unwrap(self) nogil: return self.configuration @@ -430,13 +455,13 @@ cdef class CryptoFactory(_Weakrefable): parquet_file_path : str, pathlib.Path, or None, default None Path to the parquet file to be encrypted. Only required when the internal_key_material attribute of EncryptionConfiguration is set - to False. Used to derive the path for storing key material + to False. Used to derive the path for storing key material specific to this parquet file. filesystem : FileSystem or None, default None - Used only when internal_key_material is set to False on + Used only when internal_key_material is set to False on EncryptionConfiguration. If None, the file system will be inferred - based on parquet_file_path. + based on parquet_file_path. Returns ------- @@ -491,7 +516,7 @@ cdef class CryptoFactory(_Weakrefable): filesystem : FileSystem or None, default None Used only when the parquet file uses external key material. If - None, the file system will be inferred based on parquet_file_path. + None, the file system will be inferred based on parquet_file_path. Returns ------- @@ -538,7 +563,8 @@ cdef class CryptoFactory(_Weakrefable): parquet_file_path, FileSystem filesystem=None, double_wrapping=True, - cache_lifetime_seconds=600): + cache_lifetime_seconds=600, + read_kms_url=False): """ Rotates master encryption keys for a Parquet file that uses external key material. @@ -552,7 +578,7 @@ cdef class CryptoFactory(_Weakrefable): filesystem : FileSystem or None, default None Used only when the parquet file uses external key material. If - None, the file system will be inferred based on parquet_file_path. + None, the file system will be inferred based on parquet_file_path. double_wrapping : bool, default True In the single wrapping mode, encrypts data encryption keys with @@ -563,6 +589,15 @@ cdef class CryptoFactory(_Weakrefable): cache_lifetime_seconds : int or float, default 600 During key rotation, KMS Client and Key Encryption Keys will be cached for this duration. + + read_kms_url : bool, default False + Whether the KMS instance URL may be read from the key material of + the file being rotated, when it is not configured in + kms_connection_config. 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. This only affects + reading the existing key material, the key material written by key + rotation always uses the properties from kms_connection_config. """ cdef: c_string c_parquet_file_path @@ -580,7 +615,8 @@ cdef class CryptoFactory(_Weakrefable): c_parquet_file_path, c_filesystem, double_wrapping, - cache_lifetime_seconds) + cache_lifetime_seconds, + read_kms_url) check_status(status) @@ -665,7 +701,7 @@ cdef class FileSystemKeyMaterialStore(_Weakrefable): filesystem : FileSystem, default None FileSystem where the parquet file is located. If None, - will be inferred based on parquet_file_path. + will be inferred based on parquet_file_path. Returns ------- diff --git a/python/pyarrow/includes/libparquet_encryption.pxd b/python/pyarrow/includes/libparquet_encryption.pxd index 7024f14ac27e..79953433a180 100644 --- a/python/pyarrow/includes/libparquet_encryption.pxd +++ b/python/pyarrow/includes/libparquet_encryption.pxd @@ -86,6 +86,7 @@ cdef extern from "parquet/encryption/crypto_factory.h" \ " parquet::encryption::DecryptionConfiguration": CDecryptionConfiguration() except + double cache_lifetime_seconds + c_bool read_kms_url cdef cppclass CCryptoFactory" parquet::encryption::CryptoFactory": void RegisterKmsClientFactory( @@ -106,7 +107,8 @@ cdef extern from "parquet/encryption/crypto_factory.h" \ const c_string parquet_file_path, const shared_ptr[CFileSystem] file_system, c_bool double_wrapping, - double cache_lifetime_seconds) + double cache_lifetime_seconds, + c_bool read_kms_url) cdef extern from "parquet/encryption/file_key_material_store.h" \ namespace "parquet::encryption" nogil: @@ -184,4 +186,5 @@ cdef extern from "arrow/python/parquet_encryption.h" \ const c_string parquet_file_path, const shared_ptr[CFileSystem] filesystem, c_bool double_wrapping, - double cache_lifetime_seconds) + double cache_lifetime_seconds, + c_bool read_kms_url) diff --git a/python/pyarrow/src/arrow/python/parquet_encryption.cc b/python/pyarrow/src/arrow/python/parquet_encryption.cc index 4fcce64cdbe8..815ee976ae46 100644 --- a/python/pyarrow/src/arrow/python/parquet_encryption.cc +++ b/python/pyarrow/src/arrow/python/parquet_encryption.cc @@ -100,10 +100,10 @@ arrow::Status PyCryptoFactory::SafeRotateMasterKeys( const ::parquet::encryption::KmsConnectionConfig& kms_connection_config, const std::string& parquet_file_path, const std::shared_ptr<::arrow::fs::FileSystem>& filesystem, bool double_wrapping, - double cache_lifetime_seconds) { + double cache_lifetime_seconds, bool read_kms_url) { PARQUET_CATCH_NOT_OK(this->RotateMasterKeys(kms_connection_config, parquet_file_path, filesystem, double_wrapping, - cache_lifetime_seconds)); + cache_lifetime_seconds, read_kms_url)); return arrow::Status::OK(); } diff --git a/python/pyarrow/src/arrow/python/parquet_encryption.h b/python/pyarrow/src/arrow/python/parquet_encryption.h index b485b8b11537..cfff2653cbe8 100644 --- a/python/pyarrow/src/arrow/python/parquet_encryption.h +++ b/python/pyarrow/src/arrow/python/parquet_encryption.h @@ -137,7 +137,7 @@ class ARROW_PYTHON_PARQUET_ENCRYPTION_EXPORT PyCryptoFactory const ::parquet::encryption::KmsConnectionConfig& kms_connection_config, const std::string& parquet_file_path, const std::shared_ptr<::arrow::fs::FileSystem>& filesystem, bool double_wrapping, - double cache_lifetime_seconds); + double cache_lifetime_seconds, bool read_kms_url); }; } // namespace encryption diff --git a/python/pyarrow/tests/parquet/test_encryption.py b/python/pyarrow/tests/parquet/test_encryption.py index 6a3842f3edf8..652cfc292e99 100644 --- a/python/pyarrow/tests/parquet/test_encryption.py +++ b/python/pyarrow/tests/parquet/test_encryption.py @@ -42,6 +42,8 @@ DIRECT_KEY_256 = b"0123456789abcdef0123456789abcdef" DIRECT_AAD_PREFIX = b"test_aad_prefix" +KMS_INSTANCE_ID = "kms-123" +KMS_INSTANCE_URL = "https://example.com/kms" # Marks all of the tests in this module # Ignore these with pytest ... -m 'not parquet_encryption' @@ -463,13 +465,28 @@ def validate_encryption_configuration(encryption_config): def test_encrypted_parquet_decryption_configuration(): + # Test defaults + default_config = pe.DecryptionConfiguration() + assert timedelta(minutes=10.0) == default_config.cache_lifetime + assert default_config.read_kms_url is False + + # Test init parameters decryption_config = pe.DecryptionConfiguration( - cache_lifetime=timedelta(minutes=10.0)) - assert timedelta(minutes=10.0) == decryption_config.cache_lifetime + cache_lifetime=timedelta(minutes=5.0), + read_kms_url=True) + assert timedelta(minutes=5.0) == decryption_config.cache_lifetime + assert decryption_config.read_kms_url is True + # Test setters decryption_config_1 = pe.DecryptionConfiguration() - decryption_config_1.cache_lifetime = timedelta(minutes=10.0) - assert timedelta(minutes=10.0) == decryption_config_1.cache_lifetime + decryption_config_1.cache_lifetime = timedelta(minutes=5.0) + decryption_config_1.read_kms_url = True + assert timedelta(minutes=5.0) == decryption_config_1.cache_lifetime + assert decryption_config_1.read_kms_url is True + + # Can pass integer number of seconds as cache lifetime + decryption_config_2 = pe.DecryptionConfiguration(cache_lifetime=300) + assert timedelta(minutes=5.0) == decryption_config_2.cache_lifetime def test_encrypted_parquet_kms_configuration(): @@ -650,6 +667,108 @@ def check_rotated_external_keys(master_key_id: str) -> None: assert data_table.equals(table_read_after_rotation) +def recording_kms_factory(created_configs, client_class=InMemoryKmsClient): + """Create a KMS client factory that appends the KMS instance ID and URL of + each connection configuration it is given to created_configs""" + def kms_factory(kms_connection_configuration): + created_configs.append( + (kms_connection_configuration.kms_instance_id, + kms_connection_configuration.kms_instance_url)) + return client_class(kms_connection_configuration) + return kms_factory + + +@pytest.mark.parametrize("read_kms_url", [False, True]) +def test_read_kms_url_from_file( + tempdir, data_table, basic_encryption_config, read_kms_url): + """Read a file written with KMS connection properties configured, using a + KmsConnectionConfig that doesn't specify them""" + path = tempdir / PARQUET_NAME + custom_kms_conf = { + FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), + COL_KEY_NAME: COL_KEY.decode("UTF-8"), + } + + write_config = pe.KmsConnectionConfig( + kms_instance_id=KMS_INSTANCE_ID, + kms_instance_url=KMS_INSTANCE_URL, + custom_kms_conf=custom_kms_conf) + write_crypto_factory = pe.CryptoFactory(InMemoryKmsClient) + write_encrypted_parquet(path, data_table, basic_encryption_config, + write_config, write_crypto_factory) + verify_file_encrypted(path) + + # Leave the KMS instance ID and URL unset when reading + read_config = pe.KmsConnectionConfig(custom_kms_conf=custom_kms_conf) + created_configs = [] + read_crypto_factory = pe.CryptoFactory( + recording_kms_factory(created_configs)) + decryption_config = pe.DecryptionConfiguration( + read_kms_url=read_kms_url) + result_table = read_encrypted_parquet( + path, decryption_config, read_config, read_crypto_factory) + assert data_table.equals(result_table) + + if read_kms_url: + # The URL is read from the file key material + assert created_configs == [(KMS_INSTANCE_ID, KMS_INSTANCE_URL)] + else: + # The URL in the key material is ignored and the default provided + # instead. + assert created_configs == [(KMS_INSTANCE_ID, "DEFAULT")] + + +@pytest.mark.parametrize("read_kms_url", [False, True]) +def test_key_rotation_reads_kms_url_from_file(reusable_tempdir, data_table, + read_kms_url): + """Rotate the keys of a file written with KMS connection properties + configured, using a KmsConnectionConfig that doesn't specify them""" + path = reusable_tempdir / PARQUET_NAME + encryption_config = pe.EncryptionConfiguration( + footer_key=FOOTER_KEY_NAME, + column_keys={COL_KEY_NAME: ["a", "b"]}, + internal_key_material=False) + + # Write initial encrypted file with external key material + write_config = pe.KmsConnectionConfig( + kms_instance_id=KMS_INSTANCE_ID, + kms_instance_url=KMS_INSTANCE_URL, + key_access_token="1") + write_crypto_factory = pe.CryptoFactory(MockVersioningKmsClient) + write_encrypted_parquet(path, data_table, encryption_config, write_config, + write_crypto_factory) + + # Rotate keys without specifying the KMS instance ID and URL + rotation_config = pe.KmsConnectionConfig(key_access_token="2") + created_configs = [] + rotation_crypto_factory = pe.CryptoFactory( + recording_kms_factory(created_configs, MockVersioningKmsClient)) + rotation_crypto_factory.rotate_master_keys( + rotation_config, path, read_kms_url=read_kms_url) + + if read_kms_url: + # The empty config provided is used to wrap new keys, + # and the config from the file was used to unwrap the original keys. + assert created_configs == [("", ""), (KMS_INSTANCE_ID, KMS_INSTANCE_URL)] + else: + # The default URL is used when unwrapping the original keys. + assert created_configs == [("", ""), (KMS_INSTANCE_ID, "DEFAULT")] + + # New key material always uses the provided KMS connection configuration + rotated_keys = read_external_keys_to_dict(path) + footer_key_material = rotated_keys[FOOTER_KEY_NAME] + assert footer_key_material.kms_instance_id == "DEFAULT" + assert footer_key_material.kms_instance_url == "DEFAULT" + + table_read_after_rotation = read_encrypted_parquet( + path, + pe.DecryptionConfiguration(), + rotation_config, + rotation_crypto_factory, + internal_key_material=False) + assert data_table.equals(table_read_after_rotation) + + def test_encrypted_parquet_loop(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, verify it's encrypted, and then read it multithreaded in a loop."""