Skip to content

Commit 911a7cb

Browse files
adamreeveadamreeve
authored
GH-51349: [C++][Parquet] Add flag to control pass-through of KMS URLs from key material (#51350)
### Rationale for this change Addresses #51349 ### What changes are included in this PR? * Adds a new `read_kms_url` member to `DecryptionConfiguration` and a parameter with the same name to `rotate_master_keys`. Both are false by default. * Adds these to the PyArrow bindings. ### Are these changes tested? Yes, this includes new unit tests. ### Are there any user-facing changes? Yes, this adds a new user-facing option. Users that relied on this behaviour previously will now need to opt-in and enable the flag. * GitHub Issue: #51349 Authored-by: Adam Reeve <adreeve@gmail.com> Signed-off-by: Adam Reeve <adreeve@gmail.com>
1 parent f3b46fb commit 911a7cb

14 files changed

Lines changed: 401 additions & 30 deletions

File tree

‎cpp/src/parquet/encryption/crypto_factory.cc‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,10 @@ std::shared_ptr<FileDecryptionProperties> CryptoFactory::GetFileDecryptionProper
178178
key_toolkit_, kms_connection_config, decryption_config.cache_lifetime_seconds,
179179
file_path, file_system);
180180

181+
if (decryption_config.read_kms_url) {
182+
key_retriever->EnableReadingKmsUrl();
183+
}
184+
181185
return FileDecryptionProperties::Builder()
182186
.key_retriever(std::move(key_retriever))
183187
->plaintext_files_allowed()
@@ -188,9 +192,9 @@ void CryptoFactory::RotateMasterKeys(
188192
const KmsConnectionConfig& kms_connection_config,
189193
const std::string& parquet_file_path,
190194
const std::shared_ptr<::arrow::fs::FileSystem>& file_system, bool double_wrapping,
191-
double cache_lifetime_seconds) {
195+
double cache_lifetime_seconds, bool read_kms_url) {
192196
key_toolkit_->RotateMasterKeys(kms_connection_config, parquet_file_path, file_system,
193-
double_wrapping, cache_lifetime_seconds);
197+
double_wrapping, cache_lifetime_seconds, read_kms_url);
194198
}
195199

196200
} // namespace parquet::encryption

‎cpp/src/parquet/encryption/crypto_factory.h‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,12 @@ struct PARQUET_EXPORT DecryptionConfiguration {
9191
/// objects).
9292
/// The default is 600 (10 minutes).
9393
double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds;
94+
95+
/// Whether the KMS instance URL should be read from Parquet key material if it is
96+
/// not configured in the KmsConnectionConfig.
97+
/// This should only be enabled when the KMS implementation validates the URL it
98+
/// receives, to ensure a KMS access token isn't sent to a malicious URL.
99+
bool read_kms_url = false;
94100
};
95101

96102
/// This is a core class, that translates the parameters of high level encryption (like
@@ -135,11 +141,17 @@ class PARQUET_EXPORT CryptoFactory {
135141
/// and then re-encrypted with new master keys.
136142
/// This relies on the KMS supporting versioning, such that the old master key is
137143
/// used when unwrapping a key, and the latest version is used when wrapping a key.
144+
///
145+
/// If read_kms_url is true, the KMS instance URL is read from the key material being
146+
/// rotated if it is not provided in the KmsConnectionConfig. This should only be
147+
/// enabled when the KMS implementation validates the URL it receives, to ensure a KMS
148+
/// access token isn't sent to a malicious URL.
138149
void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config,
139150
const std::string& parquet_file_path,
140151
const std::shared_ptr<::arrow::fs::FileSystem>& file_system,
141152
bool double_wrapping = kDefaultDoubleWrapping,
142-
double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds);
153+
double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds,
154+
bool read_kms_url = false);
143155

144156
private:
145157
ColumnPathToEncryptionPropertiesMap GetColumnEncryptionProperties(

‎cpp/src/parquet/encryption/file_key_unwrapper.cc‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -133,25 +133,31 @@ KeyWithMasterId FileKeyUnwrapper::GetDataEncryptionKey(const KeyMaterial& key_ma
133133
return KeyWithMasterId(std::move(data_key), master_key_id);
134134
}
135135

136+
void FileKeyUnwrapper::EnableReadingKmsUrl() { read_kms_url_ = true; }
137+
136138
std::shared_ptr<KmsClient> FileKeyUnwrapper::GetKmsClientFromConfigOrKeyMaterial(
137139
const KeyMaterial& key_material) {
138140
std::string& kms_instance_id = kms_connection_config_.kms_instance_id;
139141
if (kms_instance_id.empty()) {
140142
kms_instance_id = key_material.kms_instance_id();
141143
if (kms_instance_id.empty()) {
142144
throw ParquetException(
143-
"KMS instance ID is missing both in both kms connection configuration and file "
145+
"KMS instance ID is missing in both the KMS connection configuration and file "
144146
"key material");
145147
}
146148
}
147149

148150
std::string& kms_instance_url = kms_connection_config_.kms_instance_url;
149151
if (kms_instance_url.empty()) {
150-
kms_instance_url = key_material.kms_instance_url();
151-
if (kms_instance_url.empty()) {
152-
throw ParquetException(
153-
"KMS instance ID is missing both in both kms connection configuration and file "
154-
"key material");
152+
if (read_kms_url_) {
153+
kms_instance_url = key_material.kms_instance_url();
154+
if (kms_instance_url.empty()) {
155+
throw ParquetException(
156+
"KMS instance URL is missing in both the KMS connection configuration and "
157+
"the file key material");
158+
}
159+
} else {
160+
kms_instance_url = KmsClient::kKmsInstanceUrlDefault;
155161
}
156162
}
157163

‎cpp/src/parquet/encryption/file_key_unwrapper.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever {
7070
/// Get the data key along with the master key id from key material
7171
KeyWithMasterId GetDataEncryptionKey(const KeyMaterial& key_material);
7272

73+
/// Enable reading the KMS instance URL from Parquet key material when it is not
74+
/// already set.
75+
void EnableReadingKmsUrl();
76+
7377
private:
7478
FileKeyUnwrapper(std::shared_ptr<KeyToolkit> key_toolkit_owner, KeyToolkit* key_toolkit,
7579
const KmsConnectionConfig& kms_connection_config,
@@ -91,6 +95,7 @@ class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever {
9195
std::shared_ptr<FileKeyMaterialStore> key_material_store_;
9296
const std::string file_path_;
9397
std::shared_ptr<::arrow::fs::FileSystem> file_system_;
98+
bool read_kms_url_ = false;
9499
};
95100

96101
} // namespace parquet::encryption

‎cpp/src/parquet/encryption/key_management_test.cc‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
#include "arrow/util/logging.h"
3333

3434
#include "parquet/encryption/crypto_factory.h"
35+
#include "parquet/encryption/file_system_key_material_store.h"
36+
#include "parquet/encryption/key_material.h"
3537
#include "parquet/encryption/key_toolkit.h"
3638
#include "parquet/encryption/test_encryption_util.h"
3739
#include "parquet/encryption/test_in_memory_kms.h"
@@ -194,6 +196,72 @@ class TestEncryptionKeyManagement : public ::testing::Test {
194196
crypto_factory_.RemoveCacheEntriesForAllTokens();
195197
}
196198

199+
// Write a file that uses external key material and records the KMS
200+
// instance ID and instance URL in its key material.
201+
std::string WriteExternalMaterialFileWithKmsConfig(const std::string& instance_id,
202+
const std::string& instance_url) {
203+
kms_connection_config_.kms_instance_id = instance_id;
204+
kms_connection_config_.kms_instance_url = instance_url;
205+
TestOnlyInServerWrapKms::InitializeMasterKeys(key_list_);
206+
constexpr bool double_wrapping = true;
207+
constexpr int encryption_no = 0;
208+
this->WriteEncryptedParquetFile(double_wrapping, /*internal_key_material=*/false,
209+
encryption_no);
210+
return temp_dir_->path().ToString() + GetFileName(double_wrapping, wrap_locally_,
211+
/*internal_key_material=*/false,
212+
encryption_no);
213+
}
214+
215+
// Rotate the keys of a file written with the KMS ID and URL configured,
216+
// and return the KMS connection configurations used to create clients
217+
// during key rotation.
218+
std::vector<KmsConnectionConfig> RotateKeysWithKmsConfig(
219+
const KmsConnectionConfig& rotation_config, const bool read_kms_url) {
220+
const auto file_system = std::make_shared<::arrow::fs::LocalFileSystem>();
221+
this->SetupCryptoFactory(false);
222+
223+
const std::string file_path =
224+
this->WriteExternalMaterialFileWithKmsConfig("123", "https://example.com/kms");
225+
226+
auto kms_client_factory = std::make_shared<TestOnlyInMemoryKmsClientFactory>(
227+
/*wrap_locally=*/false, key_list_);
228+
auto crypto_factory = std::make_shared<CryptoFactory>();
229+
crypto_factory->RegisterKmsClientFactory(kms_client_factory);
230+
231+
TestOnlyInServerWrapKms::StartKeyRotation(new_key_list_);
232+
crypto_factory->RotateMasterKeys(rotation_config, file_path, file_system,
233+
/*double_wrapping=*/true,
234+
kDefaultCacheLifetimeSeconds, read_kms_url);
235+
TestOnlyInServerWrapKms::FinishKeyRotation();
236+
237+
std::vector<KmsConnectionConfig> creation_requests =
238+
kms_client_factory->CreationRequests();
239+
240+
// The new key material always uses the KMS connection configuration provided,
241+
// not the config from the previous key material.
242+
// If it's empty, default values are written.
243+
const auto key_material_store =
244+
FileSystemKeyMaterialStore::Make(file_path, file_system,
245+
/*use_tmp_prefix=*/false);
246+
const KeyMaterial rotated_key_material = KeyMaterial::Parse(
247+
key_material_store->GetKeyMaterial(std::string(KeyMaterial::kFooterKeyIdInFile)));
248+
const auto& expected_id = rotation_config.kms_instance_id.empty()
249+
? KmsClient::kKmsInstanceIdDefault
250+
: rotation_config.kms_instance_id;
251+
const auto& expected_url = rotation_config.kms_instance_url.empty()
252+
? KmsClient::kKmsInstanceUrlDefault
253+
: rotation_config.kms_instance_url;
254+
EXPECT_EQ(rotated_key_material.kms_instance_id(), expected_id);
255+
EXPECT_EQ(rotated_key_material.kms_instance_url(), expected_url);
256+
257+
// Check the rotated file is readable
258+
const auto file_decryption_properties = crypto_factory->GetFileDecryptionProperties(
259+
rotation_config, GetDecryptionConfiguration(), file_path, file_system);
260+
decryptor_.DecryptFile(file_path, file_decryption_properties);
261+
262+
return creation_requests;
263+
}
264+
197265
// Create encryption properties without keeping the creating CryptoFactory alive
198266
std::shared_ptr<FileEncryptionProperties> GetOrphanedFileEncryptionProperties(
199267
std::shared_ptr<KmsClientFactory> kms_client_factory,
@@ -441,4 +509,100 @@ TEST_F(TestEncryptionKeyManagement, ReadParquetMRExternalKeyMaterialFile) {
441509
}
442510
}
443511

512+
TEST_F(TestEncryptionKeyManagement, ReadKmsUrlFromFile) {
513+
this->SetupCryptoFactory(true);
514+
515+
constexpr bool internal_key_material = true;
516+
constexpr bool double_wrapping = true;
517+
constexpr int encryption_no = 0;
518+
519+
std::string file_name = "kms-config-test-file.parquet.encrypted";
520+
std::string file_path = temp_dir_->path().ToString() + file_name;
521+
522+
auto encryption_config =
523+
GetEncryptionConfiguration(double_wrapping, internal_key_material, encryption_no);
524+
525+
KmsConnectionConfig write_config;
526+
write_config.kms_instance_id = "123";
527+
write_config.kms_instance_url = "https://example.com/kms";
528+
529+
auto file_encryption_properties =
530+
crypto_factory_.GetFileEncryptionProperties(write_config, encryption_config);
531+
encryptor_.EncryptFile(file_path, file_encryption_properties);
532+
533+
for (const auto& enable_kms_url_read : {false, true}) {
534+
// Create a fresh crypto factory and client factory for each read
535+
// to avoid re-using cached clients.
536+
CryptoFactory read_crypto_factory;
537+
auto kms_client_factory =
538+
std::make_shared<TestOnlyInMemoryKmsClientFactory>(true, key_list_);
539+
read_crypto_factory.RegisterKmsClientFactory(kms_client_factory);
540+
541+
auto decryption_config = DecryptionConfiguration();
542+
decryption_config.read_kms_url = enable_kms_url_read;
543+
544+
KmsConnectionConfig read_config;
545+
546+
auto file_decryption_properties =
547+
read_crypto_factory.GetFileDecryptionProperties(read_config, decryption_config);
548+
549+
decryptor_.DecryptFile(file_path, file_decryption_properties);
550+
551+
ASSERT_EQ(kms_client_factory->CreationRequests().size(), 1);
552+
const auto& request = kms_client_factory->CreationRequests()[0];
553+
EXPECT_EQ(request.kms_instance_id, "123");
554+
if (enable_kms_url_read) {
555+
EXPECT_EQ(request.kms_instance_url, "https://example.com/kms");
556+
} else {
557+
EXPECT_EQ(request.kms_instance_url, "DEFAULT");
558+
}
559+
}
560+
}
561+
562+
TEST_F(TestEncryptionKeyManagement, ReadKmsUrlFromFileDuringKeyRotation) {
563+
// Use an empty config for rotation
564+
const KmsConnectionConfig rotation_config;
565+
const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/true);
566+
567+
ASSERT_EQ(requests.size(), 2);
568+
// The first KMS creation request is for wrapping new keys.
569+
// This uses the empty config provided.
570+
EXPECT_EQ(requests[0].kms_instance_id, "");
571+
EXPECT_EQ(requests[0].kms_instance_url, "");
572+
// The KMS client used to unwrap the previous keys should be configured
573+
// with the instance ID and url provided at write time.
574+
EXPECT_EQ(requests[1].kms_instance_id, "123");
575+
EXPECT_EQ(requests[1].kms_instance_url, "https://example.com/kms");
576+
}
577+
578+
TEST_F(TestEncryptionKeyManagement, KeyRotationWithoutReadingKmsUrl) {
579+
// Use an empty config for rotation
580+
const KmsConnectionConfig rotation_config;
581+
const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/false);
582+
583+
ASSERT_EQ(requests.size(), 2);
584+
// The first KMS creation request is for wrapping new keys.
585+
// This uses the empty config provided.
586+
EXPECT_EQ(requests[0].kms_instance_id, "");
587+
EXPECT_EQ(requests[0].kms_instance_url, "");
588+
// When unwrapping the existing keys, the URL in the key material is
589+
// ignored and the default used.
590+
EXPECT_EQ(requests[1].kms_instance_id, "123");
591+
EXPECT_EQ(requests[1].kms_instance_url, KmsClient::kKmsInstanceUrlDefault);
592+
}
593+
594+
TEST_F(TestEncryptionKeyManagement, KeyRotationUsesProvidedKmsConfig) {
595+
KmsConnectionConfig rotation_config;
596+
rotation_config.kms_instance_id = "456";
597+
rotation_config.kms_instance_url = "https://example.com/kms2";
598+
const auto requests = RotateKeysWithKmsConfig(rotation_config, /*read_kms_url=*/true);
599+
600+
ASSERT_EQ(requests.size(), 1);
601+
// Wrap and unwrap both use the same configuration.
602+
// The instance id and url in the existing key material is ignored even though
603+
// read_kms_url is enabled. The provided config takes precedence.
604+
EXPECT_EQ(requests[0].kms_instance_id, "456");
605+
EXPECT_EQ(requests[0].kms_instance_url, "https://example.com/kms2");
606+
}
607+
444608
} // namespace parquet::encryption::test

‎cpp/src/parquet/encryption/key_toolkit.cc‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ void KeyToolkit::RotateMasterKeys(
4747
const KmsConnectionConfig& kms_connection_config,
4848
const std::string& parquet_file_path,
4949
const std::shared_ptr<::arrow::fs::FileSystem>& file_system, bool double_wrapping,
50-
double cache_lifetime_seconds) {
50+
double cache_lifetime_seconds, bool read_kms_url) {
5151
// If process wrote files with double-wrapped keys, clean KEK cache (since master keys
5252
// are changing). Only once for each key rotation cycle; not for every file.
5353
const auto now = internal::CurrentTimePoint();
@@ -65,6 +65,9 @@ void KeyToolkit::RotateMasterKeys(
6565
// Unwrapper for decrypting encrypted keys
6666
FileKeyUnwrapper file_key_unwrapper(this, kms_connection_config, cache_lifetime_seconds,
6767
key_material_store);
68+
if (read_kms_url) {
69+
file_key_unwrapper.EnableReadingKmsUrl();
70+
}
6871

6972
// Create a temporary store to hold new key material during rotation,
7073
// and wrapper that will write material to this store when getting key metadata.

‎cpp/src/parquet/encryption/key_toolkit.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ class PARQUET_EXPORT KeyToolkit {
7777
void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config,
7878
const std::string& parquet_file_path,
7979
const std::shared_ptr<::arrow::fs::FileSystem>& file_system,
80-
bool double_wrapping, double cache_lifetime_seconds);
80+
bool double_wrapping, double cache_lifetime_seconds,
81+
bool read_kms_url = false);
8182

8283
private:
8384
TwoLevelCacheWithExpiration<std::shared_ptr<KmsClient>> kms_client_cache_;

‎cpp/src/parquet/encryption/test_in_memory_kms.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,12 +90,20 @@ class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory {
9090

9191
std::shared_ptr<KmsClient> CreateKmsClient(
9292
const KmsConnectionConfig& kms_connection_config) {
93+
create_requests_.push_back(kms_connection_config);
9394
if (wrap_locally_) {
9495
return std::make_shared<TestOnlyLocalWrapInMemoryKms>(kms_connection_config);
9596
} else {
9697
return std::make_shared<TestOnlyInServerWrapKms>();
9798
}
9899
}
100+
101+
/// Get the `KmsConnectionConfig` values that have been used to
102+
/// create clients with this factory.
103+
const std::vector<KmsConnectionConfig>& CreationRequests() { return create_requests_; }
104+
105+
private:
106+
std::vector<KmsConnectionConfig> create_requests_;
99107
};
100108

101109
} // namespace parquet::encryption

‎docs/source/python/parquet/parquet_encryption.rst‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,11 @@ file decryption properties) is optional and it includes the following options:
232232

233233
* ``cache_lifetime``, the lifetime of cached entities (key encryption keys, local
234234
wrapping keys, KMS client objects) represented as a ``datetime.timedelta``.
235+
* ``read_kms_url``, whether the KMS instance URL may be read from the key material
236+
of the file being read, when it is not set in the ``KmsConnectionConfig``. This
237+
defaults to ``False``, and should only be enabled when the KMS implementation
238+
validates the URL it receives, to ensure a KMS access token isn't sent to a
239+
malicious URL.
235240

236241
External key material and key rotation
237242
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -251,6 +256,11 @@ key material file, without changing the Parquet file itself:
251256
... kms_connection_config, parquet_file_path="table.parquet",
252257
... )
253258
259+
``rotate_master_keys`` also accepts ``read_kms_url``, which behaves like the
260+
``DecryptionConfiguration`` option of the same name when the existing key material is
261+
read. The key material written by key rotation always uses the connection properties
262+
from the ``KmsConnectionConfig`` that is passed in.
263+
254264
Direct Key Encryption (without KMS)
255265
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
256266

0 commit comments

Comments
 (0)