From c0d65103cedda797bbb642209d94d46d153059fc Mon Sep 17 00:00:00 2001 From: Adam Szita Date: Wed, 21 May 2025 16:41:42 +0200 Subject: [PATCH 1/4] AWS: KeyManagementClient implementation that works with AWS KMS To be used in table encryption for: - wrapping/unwrapping encryption keys - generating data keys (available specs: AES_256, AES_128) Added integration test for verification. --- .../iceberg/aws/TestKeyManagementClient.java | 125 ++++++++++++++++++ .../iceberg/aws/AwsKeyManagementClient.java | 107 +++++++++++++++ .../org/apache/iceberg/aws/AwsProperties.java | 31 +++++ .../encryption/KeyManagementClient.java | 4 +- 4 files changed, 265 insertions(+), 2 deletions(-) create mode 100644 aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java create mode 100644 aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java diff --git a/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java new file mode 100644 index 000000000000..6e7e7d7900d4 --- /dev/null +++ b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.aws; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.ByteBuffer; +import java.util.Map; +import org.apache.iceberg.encryption.KeyManagementClient; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariables; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.kms.KmsClient; +import software.amazon.awssdk.services.kms.model.CreateKeyRequest; +import software.amazon.awssdk.services.kms.model.CreateKeyResponse; +import software.amazon.awssdk.services.kms.model.DataKeySpec; +import software.amazon.awssdk.services.kms.model.KeySpec; +import software.amazon.awssdk.services.kms.model.ScheduleKeyDeletionRequest; +import software.amazon.awssdk.services.kms.model.ScheduleKeyDeletionResponse; + +@EnabledIfEnvironmentVariables({ + @EnabledIfEnvironmentVariable(named = AwsIntegTestUtil.AWS_ACCESS_KEY_ID, matches = ".*"), + @EnabledIfEnvironmentVariable(named = AwsIntegTestUtil.AWS_SECRET_ACCESS_KEY, matches = ".*"), + @EnabledIfEnvironmentVariable(named = AwsIntegTestUtil.AWS_SESSION_TOKEN, matches = ".*"), + @EnabledIfEnvironmentVariable(named = AwsIntegTestUtil.AWS_TEST_ACCOUNT_ID, matches = "\\d{12}") +}) +public class TestKeyManagementClient { + + private static final Logger LOG = LoggerFactory.getLogger(TestKeyManagementClient.class); + + private static KmsClient kmsClient; + private static String keyId; + + @BeforeAll + public static void beforeClass() { + kmsClient = AwsClientFactories.defaultFactory().kms(); + CreateKeyRequest createKeyRequest = + CreateKeyRequest.builder() + .keySpec(KeySpec.SYMMETRIC_DEFAULT) + .description( + "Iceberg integration test key for " + TestKeyManagementClient.class.getName()) + .build(); + CreateKeyResponse response = kmsClient.createKey(createKeyRequest); + keyId = response.keyMetadata().keyId(); + } + + @Test + public void testKeyWrapping() { + AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient(); + try { + keyManagementClient.initialize(ImmutableMap.of()); + + ByteBuffer key = ByteBuffer.wrap(new String("super-secret-table-master-key").getBytes()); + ByteBuffer encryptedKey = keyManagementClient.wrapKey(key, keyId); + + assertThat(keyManagementClient.unwrapKey(encryptedKey, keyId)).isEqualTo(key); + } finally { + keyManagementClient.close(); + } + } + + @Test + public void testKeyGeneration() { + testKeyGenerationWithDataKeySpec(null); + testKeyGenerationWithDataKeySpec(DataKeySpec.AES_128); + testKeyGenerationWithDataKeySpec(DataKeySpec.AES_256); + } + + private void testKeyGenerationWithDataKeySpec(DataKeySpec dataKeySpec) { + AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient(); + try { + Map properties = + dataKeySpec == null + ? ImmutableMap.of() + : ImmutableMap.of(AwsProperties.KMS_DATA_KEY_SPEC, dataKeySpec.name()); + keyManagementClient.initialize(properties); + KeyManagementClient.KeyGenerationResult result = keyManagementClient.generateKey(keyId); + + assertThat(keyManagementClient.unwrapKey(result.wrappedKey(), keyId)).isEqualTo(result.key()); + assertThat(result.key().limit()) + .isEqualTo(DataKeySpec.AES_128.equals(dataKeySpec) ? 128 / 8 : 256 / 8); + } finally { + keyManagementClient.close(); + } + } + + @AfterAll + public static void afterClass() { + // AWS KMS doesn't allow instant deletion. Keys can be put to pendingDeletion state instead, + // with a minimum of 7 days until final removal. + ScheduleKeyDeletionRequest deletionRequest = + ScheduleKeyDeletionRequest.builder().keyId(keyId).pendingWindowInDays(7).build(); + + ScheduleKeyDeletionResponse deletionResponse = kmsClient.scheduleKeyDeletion(deletionRequest); + LOG.info( + "Deletion of test key {} will be finalized at {}", keyId, deletionResponse.deletionDate()); + + try { + kmsClient.close(); + } catch (Exception e) { + LOG.error("Error closing KMS client", e); + } + } +} diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java b/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java new file mode 100644 index 000000000000..ba92ad4e61f2 --- /dev/null +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.aws; + +import java.nio.ByteBuffer; +import java.util.Map; +import org.apache.iceberg.encryption.KeyManagementClient; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.kms.KmsClient; +import software.amazon.awssdk.services.kms.model.DataKeySpec; +import software.amazon.awssdk.services.kms.model.DecryptRequest; +import software.amazon.awssdk.services.kms.model.DecryptResponse; +import software.amazon.awssdk.services.kms.model.EncryptRequest; +import software.amazon.awssdk.services.kms.model.EncryptResponse; +import software.amazon.awssdk.services.kms.model.EncryptionAlgorithmSpec; +import software.amazon.awssdk.services.kms.model.GenerateDataKeyRequest; +import software.amazon.awssdk.services.kms.model.GenerateDataKeyResponse; + +/** + * Key management client implementation that uses AWS Key Management Service. To be used for + * encrypting/decrypting keys with a KMS-managed master key, (by referencing its key ID), and for + * the generation of new encryption keys. + */ +public class AwsKeyManagementClient implements KeyManagementClient { + + private KmsClient kmsClient; + private EncryptionAlgorithmSpec encryptionAlgorithmSpec; + private DataKeySpec dataKeySpec; + + public AwsKeyManagementClient() {} + + @Override + public ByteBuffer wrapKey(ByteBuffer key, String wrappingKeyId) { + EncryptRequest request = + EncryptRequest.builder() + .keyId(wrappingKeyId) + .encryptionAlgorithm(encryptionAlgorithmSpec) + .plaintext(SdkBytes.fromByteBuffer(key)) + .build(); + + EncryptResponse result = kmsClient.encrypt(request); + return result.ciphertextBlob().asByteBuffer(); + } + + @Override + public boolean supportsKeyGeneration() { + return true; + } + + @Override + public KeyGenerationResult generateKey(String wrappingKeyId) { + GenerateDataKeyRequest request = + GenerateDataKeyRequest.builder().keyId(wrappingKeyId).keySpec(dataKeySpec).build(); + + GenerateDataKeyResponse response = kmsClient.generateDataKey(request); + KeyGenerationResult result = + new KeyGenerationResult( + response.plaintext().asByteBuffer(), response.ciphertextBlob().asByteBuffer()); + return result; + } + + @Override + public ByteBuffer unwrapKey(ByteBuffer wrappedKey, String wrappingKeyId) { + DecryptRequest request = + DecryptRequest.builder() + .keyId(wrappingKeyId) + .encryptionAlgorithm(encryptionAlgorithmSpec) + .ciphertextBlob(SdkBytes.fromByteBuffer(wrappedKey)) + .build(); + + DecryptResponse result = kmsClient.decrypt(request); + return result.plaintext().asByteBuffer(); + } + + @Override + public void initialize(Map properties) { + AwsClientFactory clientFactory = AwsClientFactories.from(properties); + this.kmsClient = clientFactory.kms(); + + AwsProperties awsProperties = new AwsProperties(properties); + this.encryptionAlgorithmSpec = awsProperties.kmsEncryptionAlgorithmSpec(); + this.dataKeySpec = awsProperties.kmsDataKeySpec(); + } + + @Override + public void close() { + if (kmsClient != null) { + kmsClient.close(); + } + } +} diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java index 1a8db990578a..83d5e9243e18 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java @@ -41,6 +41,8 @@ import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder; import software.amazon.awssdk.services.glue.GlueClientBuilder; +import software.amazon.awssdk.services.kms.model.DataKeySpec; +import software.amazon.awssdk.services.kms.model.EncryptionAlgorithmSpec; public class AwsProperties implements Serializable { @@ -206,6 +208,17 @@ public class AwsProperties implements Serializable { */ public static final String REST_SESSION_TOKEN = "rest.session-token"; + /** Encryption algorithm used to encrypt/decrypt master table keys */ + public static final String KMS_ENCRYPTION_ALGORITHM_SPEC = "kms.encryption-algorithm-spec"; + + public static final String KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT = + EncryptionAlgorithmSpec.SYMMETRIC_DEFAULT.toString(); + + /** Length of data key generated by KMS */ + public static final String KMS_DATA_KEY_SPEC = "kms.data-key-spec"; + + public static final String KMS_DATA_KEY_SPEC_DEFAULT = DataKeySpec.AES_256.toString(); + private final Set stsClientAssumeRoleTags; private final String clientAssumeRoleArn; @@ -230,6 +243,8 @@ public class AwsProperties implements Serializable { private String restAccessKeyId; private String restSecretAccessKey; private String restSessionToken; + private String kmsEncryptionAlgorithmSpec; + private String kmsDataKeySpec; public AwsProperties() { this.stsClientAssumeRoleTags = Sets.newHashSet(); @@ -252,6 +267,9 @@ public AwsProperties() { this.dynamoDbTableName = DYNAMODB_TABLE_NAME_DEFAULT; this.restSigningName = REST_SIGNING_NAME_DEFAULT; + + this.kmsEncryptionAlgorithmSpec = KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT; + this.kmsDataKeySpec = KMS_DATA_KEY_SPEC_DEFAULT; } @SuppressWarnings("MethodLength") @@ -293,6 +311,11 @@ public AwsProperties(Map properties) { this.restAccessKeyId = properties.get(REST_ACCESS_KEY_ID); this.restSecretAccessKey = properties.get(REST_SECRET_ACCESS_KEY); this.restSessionToken = properties.get(REST_SESSION_TOKEN); + + this.kmsEncryptionAlgorithmSpec = + properties.getOrDefault( + KMS_ENCRYPTION_ALGORITHM_SPEC, KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT); + this.kmsDataKeySpec = properties.getOrDefault(KMS_DATA_KEY_SPEC, KMS_DATA_KEY_SPEC_DEFAULT); } public Set stsClientAssumeRoleTags() { @@ -402,6 +425,14 @@ public AwsCredentialsProvider restCredentialsProvider() { this.restAccessKeyId, this.restSecretAccessKey, this.restSessionToken); } + public EncryptionAlgorithmSpec kmsEncryptionAlgorithmSpec() { + return EncryptionAlgorithmSpec.fromValue(this.kmsEncryptionAlgorithmSpec); + } + + public DataKeySpec kmsDataKeySpec() { + return DataKeySpec.fromValue(this.kmsDataKeySpec); + } + private Set toStsTags( Map properties, String prefix) { return PropertyUtil.propertiesWithPrefix(properties, prefix).entrySet().stream() diff --git a/core/src/main/java/org/apache/iceberg/encryption/KeyManagementClient.java b/core/src/main/java/org/apache/iceberg/encryption/KeyManagementClient.java index a7fb494cc8e1..2d97def180af 100644 --- a/core/src/main/java/org/apache/iceberg/encryption/KeyManagementClient.java +++ b/core/src/main/java/org/apache/iceberg/encryption/KeyManagementClient.java @@ -24,7 +24,7 @@ import java.util.Map; /** A minimum client interface to connect to a key management service (KMS). */ -interface KeyManagementClient extends Serializable, Closeable { +public interface KeyManagementClient extends Serializable, Closeable { /** * Wrap a secret key, using a wrapping/master key which is stored in KMS and referenced by an ID. @@ -94,7 +94,7 @@ class KeyGenerationResult { private final ByteBuffer key; private final ByteBuffer wrappedKey; - KeyGenerationResult(ByteBuffer key, ByteBuffer wrappedKey) { + public KeyGenerationResult(ByteBuffer key, ByteBuffer wrappedKey) { this.key = key; this.wrappedKey = wrappedKey; } From 351fd5bfae9a3a6f63f16ac97f69b4d5ecf11f07 Mon Sep 17 00:00:00 2001 From: Adam Szita Date: Tue, 10 Jun 2025 13:33:55 +0200 Subject: [PATCH 2/4] Using AWS types in AwsProperties instead of Strings --- .../org/apache/iceberg/aws/AwsProperties.java | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java index 83d5e9243e18..62d541da0c54 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java @@ -211,13 +211,13 @@ public class AwsProperties implements Serializable { /** Encryption algorithm used to encrypt/decrypt master table keys */ public static final String KMS_ENCRYPTION_ALGORITHM_SPEC = "kms.encryption-algorithm-spec"; - public static final String KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT = - EncryptionAlgorithmSpec.SYMMETRIC_DEFAULT.toString(); + public static final EncryptionAlgorithmSpec KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT = + EncryptionAlgorithmSpec.SYMMETRIC_DEFAULT; /** Length of data key generated by KMS */ public static final String KMS_DATA_KEY_SPEC = "kms.data-key-spec"; - public static final String KMS_DATA_KEY_SPEC_DEFAULT = DataKeySpec.AES_256.toString(); + public static final DataKeySpec KMS_DATA_KEY_SPEC_DEFAULT = DataKeySpec.AES_256; private final Set stsClientAssumeRoleTags; @@ -243,8 +243,8 @@ public class AwsProperties implements Serializable { private String restAccessKeyId; private String restSecretAccessKey; private String restSessionToken; - private String kmsEncryptionAlgorithmSpec; - private String kmsDataKeySpec; + private EncryptionAlgorithmSpec kmsEncryptionAlgorithmSpec; + private DataKeySpec kmsDataKeySpec; public AwsProperties() { this.stsClientAssumeRoleTags = Sets.newHashSet(); @@ -313,9 +313,12 @@ public AwsProperties(Map properties) { this.restSessionToken = properties.get(REST_SESSION_TOKEN); this.kmsEncryptionAlgorithmSpec = - properties.getOrDefault( - KMS_ENCRYPTION_ALGORITHM_SPEC, KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT); - this.kmsDataKeySpec = properties.getOrDefault(KMS_DATA_KEY_SPEC, KMS_DATA_KEY_SPEC_DEFAULT); + EncryptionAlgorithmSpec.fromValue( + properties.getOrDefault( + KMS_ENCRYPTION_ALGORITHM_SPEC, KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT.toString())); + this.kmsDataKeySpec = + DataKeySpec.fromValue( + properties.getOrDefault(KMS_DATA_KEY_SPEC, KMS_DATA_KEY_SPEC_DEFAULT.toString())); } public Set stsClientAssumeRoleTags() { @@ -426,11 +429,11 @@ public AwsCredentialsProvider restCredentialsProvider() { } public EncryptionAlgorithmSpec kmsEncryptionAlgorithmSpec() { - return EncryptionAlgorithmSpec.fromValue(this.kmsEncryptionAlgorithmSpec); + return this.kmsEncryptionAlgorithmSpec; } public DataKeySpec kmsDataKeySpec() { - return DataKeySpec.fromValue(this.kmsDataKeySpec); + return this.kmsDataKeySpec; } private Set toStsTags( From 10a5c85604bbdc3dba60cc62ffbf4cafa6c93f63 Mon Sep 17 00:00:00 2001 From: Adam Szita Date: Wed, 25 Jun 2025 14:29:08 +0200 Subject: [PATCH 3/4] Refactors TestKeyManagementClient.java --- .../iceberg/aws/TestKeyManagementClient.java | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java index 6e7e7d7900d4..5c0961778ac3 100644 --- a/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java +++ b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java @@ -29,6 +29,9 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariables; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.NullSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.services.kms.KmsClient; @@ -80,14 +83,12 @@ public void testKeyWrapping() { } } - @Test - public void testKeyGeneration() { - testKeyGenerationWithDataKeySpec(null); - testKeyGenerationWithDataKeySpec(DataKeySpec.AES_128); - testKeyGenerationWithDataKeySpec(DataKeySpec.AES_256); - } - - private void testKeyGenerationWithDataKeySpec(DataKeySpec dataKeySpec) { + @ParameterizedTest + @NullSource + @EnumSource( + value = DataKeySpec.class, + names = {"AES_128", "AES_256"}) + public void testKeyGeneration(DataKeySpec dataKeySpec) { AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient(); try { Map properties = @@ -98,13 +99,20 @@ private void testKeyGenerationWithDataKeySpec(DataKeySpec dataKeySpec) { KeyManagementClient.KeyGenerationResult result = keyManagementClient.generateKey(keyId); assertThat(keyManagementClient.unwrapKey(result.wrappedKey(), keyId)).isEqualTo(result.key()); - assertThat(result.key().limit()) - .isEqualTo(DataKeySpec.AES_128.equals(dataKeySpec) ? 128 / 8 : 256 / 8); + assertThat(result.key().limit()).isEqualTo(expectedLength(dataKeySpec)); } finally { keyManagementClient.close(); } } + private static int expectedLength(DataKeySpec spec) { + if (DataKeySpec.AES_128.equals(spec)) { + return 128 / 8; + } else { + return 256 / 8; + } + } + @AfterAll public static void afterClass() { // AWS KMS doesn't allow instant deletion. Keys can be put to pendingDeletion state instead, From 9507ab012765b8fdac95f1493e77eafb0afb1d3d Mon Sep 17 00:00:00 2001 From: Adam Szita Date: Fri, 18 Jul 2025 17:22:37 +0200 Subject: [PATCH 4/4] addressing comments from nastra --- .../iceberg/aws/TestKeyManagementClient.java | 46 ++++++++----------- .../iceberg/aws/AwsKeyManagementClient.java | 20 ++++---- 2 files changed, 29 insertions(+), 37 deletions(-) diff --git a/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java index 5c0961778ac3..83bacf2601cd 100644 --- a/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java +++ b/aws/src/integration/java/org/apache/iceberg/aws/TestKeyManagementClient.java @@ -68,18 +68,33 @@ public static void beforeClass() { keyId = response.keyMetadata().keyId(); } + @AfterAll + public static void afterClass() { + // AWS KMS doesn't allow instant deletion. Keys can be put to pendingDeletion state instead, + // with a minimum of 7 days until final removal. + ScheduleKeyDeletionRequest deletionRequest = + ScheduleKeyDeletionRequest.builder().keyId(keyId).pendingWindowInDays(7).build(); + + ScheduleKeyDeletionResponse deletionResponse = kmsClient.scheduleKeyDeletion(deletionRequest); + LOG.info( + "Deletion of test key {} will be finalized at {}", keyId, deletionResponse.deletionDate()); + + try { + kmsClient.close(); + } catch (Exception e) { + LOG.error("Error closing KMS client", e); + } + } + @Test public void testKeyWrapping() { - AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient(); - try { + try (AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient()) { keyManagementClient.initialize(ImmutableMap.of()); ByteBuffer key = ByteBuffer.wrap(new String("super-secret-table-master-key").getBytes()); ByteBuffer encryptedKey = keyManagementClient.wrapKey(key, keyId); assertThat(keyManagementClient.unwrapKey(encryptedKey, keyId)).isEqualTo(key); - } finally { - keyManagementClient.close(); } } @@ -89,8 +104,7 @@ public void testKeyWrapping() { value = DataKeySpec.class, names = {"AES_128", "AES_256"}) public void testKeyGeneration(DataKeySpec dataKeySpec) { - AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient(); - try { + try (AwsKeyManagementClient keyManagementClient = new AwsKeyManagementClient()) { Map properties = dataKeySpec == null ? ImmutableMap.of() @@ -100,8 +114,6 @@ public void testKeyGeneration(DataKeySpec dataKeySpec) { assertThat(keyManagementClient.unwrapKey(result.wrappedKey(), keyId)).isEqualTo(result.key()); assertThat(result.key().limit()).isEqualTo(expectedLength(dataKeySpec)); - } finally { - keyManagementClient.close(); } } @@ -112,22 +124,4 @@ private static int expectedLength(DataKeySpec spec) { return 256 / 8; } } - - @AfterAll - public static void afterClass() { - // AWS KMS doesn't allow instant deletion. Keys can be put to pendingDeletion state instead, - // with a minimum of 7 days until final removal. - ScheduleKeyDeletionRequest deletionRequest = - ScheduleKeyDeletionRequest.builder().keyId(keyId).pendingWindowInDays(7).build(); - - ScheduleKeyDeletionResponse deletionResponse = kmsClient.scheduleKeyDeletion(deletionRequest); - LOG.info( - "Deletion of test key {} will be finalized at {}", keyId, deletionResponse.deletionDate()); - - try { - kmsClient.close(); - } catch (Exception e) { - LOG.error("Error closing KMS client", e); - } - } } diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java b/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java index ba92ad4e61f2..6d2671f4e26d 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsKeyManagementClient.java @@ -43,7 +43,15 @@ public class AwsKeyManagementClient implements KeyManagementClient { private EncryptionAlgorithmSpec encryptionAlgorithmSpec; private DataKeySpec dataKeySpec; - public AwsKeyManagementClient() {} + @Override + public void initialize(Map properties) { + AwsClientFactory clientFactory = AwsClientFactories.from(properties); + this.kmsClient = clientFactory.kms(); + + AwsProperties awsProperties = new AwsProperties(properties); + this.encryptionAlgorithmSpec = awsProperties.kmsEncryptionAlgorithmSpec(); + this.dataKeySpec = awsProperties.kmsDataKeySpec(); + } @Override public ByteBuffer wrapKey(ByteBuffer key, String wrappingKeyId) { @@ -88,16 +96,6 @@ public ByteBuffer unwrapKey(ByteBuffer wrappedKey, String wrappingKeyId) { return result.plaintext().asByteBuffer(); } - @Override - public void initialize(Map properties) { - AwsClientFactory clientFactory = AwsClientFactories.from(properties); - this.kmsClient = clientFactory.kms(); - - AwsProperties awsProperties = new AwsProperties(properties); - this.encryptionAlgorithmSpec = awsProperties.kmsEncryptionAlgorithmSpec(); - this.dataKeySpec = awsProperties.kmsDataKeySpec(); - } - @Override public void close() { if (kmsClient != null) {