Copilot commented on code in PR #51350:
URL: https://github.com/apache/arrow/pull/51350#discussion_r4021916936
##########
cpp/src/parquet/encryption/key_management_test.cc:
##########
@@ -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, "");
Review Comment:
`FileKeyWrapper` is constructed before the old key is unwrapped, and its
constructor calls `SetDefaultIfEmpty()` before creating the wrapping client.
Therefore an empty rotation config is recorded as
`KmsClient::kKmsInstanceIdDefault` / `kKmsInstanceUrlDefault`, not empty
strings, so this new assertion fails; the same expectation at lines 586-587
needs the same correction.
##########
python/pyarrow/tests/parquet/test_encryption.py:
##########
@@ -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
+
+
[email protected]("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")]
+
+
[email protected]("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)]
Review Comment:
The rotation wrapper normalizes an omitted ID and URL to `DEFAULT` before
recording its client configuration, so the first tuple is not empty. This
assertion (and the corresponding false branch at line 755) will fail.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
