Make KMS endpoint configurable via kms.endpoint AWS property (#14246)

Co-authored-by: Thomas Powell <[email protected]>

Author: tom-s-powell

aws/src/integration/java/org/apache/iceberg/aws/TestDefaultAwsClientFactory.java

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.glue.GlueClient;
 import software.amazon.awssdk.services.glue.model.GetDatabaseRequest;
+import software.amazon.awssdk.services.kms.KmsClient;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.S3Exception;
@@ -99,4 +100,16 @@ public void testDynamoDbEndpointOverride() {
         .isInstanceOf(SdkClientException.class)
         .hasMessageContaining("Unable to execute HTTP request: unknown");
   }
+
+  @Test
+  public void testKmsEndpointOverride() {
+    Map<String, String> properties = Maps.newHashMap();
+    properties.put(AwsProperties.KMS_ENDPOINT, "https://unknown:1234");
+    AwsClientFactory factory = AwsClientFactories.from(properties);
+    KmsClient kmsClient = factory.kms();
+    assertThatThrownBy(kmsClient::listKeys)
+        .cause()
+        .isInstanceOf(SdkClientException.class)
+        .hasMessageContaining("Unable to execute HTTP request: unknown");
+  }
 }

aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.services.glue.GlueClient;
 import software.amazon.awssdk.services.glue.GlueClientBuilder;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.kms.KmsClientBuilder;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3BaseClientBuilder;
 import software.amazon.awssdk.services.s3.S3Client;
@@ -155,6 +156,7 @@ public KmsClient kms() {
       return KmsClient.builder()
           .applyMutation(awsClientProperties::applyClientRegionConfiguration)
           .applyMutation(httpClientProperties::applyHttpClientConfigurations)
+          .applyMutation(awsProperties::applyKmsEndpointConfigurations)
           .applyMutation(awsClientProperties::applyClientCredentialConfigurations)
           .applyMutation(awsClientProperties::applyRetryConfigurations)
           .build();
@@ -208,8 +210,9 @@ public static SdkHttpClient.Builder configureHttpClientBuilder(String httpClient
    * @deprecated Not for public use. To configure the endpoint for a client, please use {@link
    *     S3FileIOProperties#applyEndpointConfigurations(S3BaseClientBuilder)}, {@link
    *     AwsProperties#applyGlueEndpointConfigurations(GlueClientBuilder)}, or {@link
-   *     AwsProperties#applyDynamoDbEndpointConfigurations(DynamoDbClientBuilder)} accordingly. It
-   *     will be removed in 2.0.0
+   *     AwsProperties#applyDynamoDbEndpointConfigurations(DynamoDbClientBuilder)}, or {@link
+   *     AwsProperties#applyKmsEndpointConfigurations(KmsClientBuilder)} accordingly. It will be
+   *     removed in 2.0.0
    */
   @Deprecated
   public static <T extends SdkClientBuilder> void configureEndpoint(T builder, String endpoint) {

aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java

@@ -41,6 +41,7 @@
 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.KmsClientBuilder;
 import software.amazon.awssdk.services.kms.model.DataKeySpec;
 import software.amazon.awssdk.services.kms.model.EncryptionAlgorithmSpec;
 
@@ -208,6 +209,14 @@ public class AwsProperties implements Serializable {
    */
   public static final String REST_SESSION_TOKEN = "rest.session-token";
 
+  /**
+   * Configure an alternative endpoint of the KMS service for AwsKeyManagementClient to access.
+   *
+   * <p>This could be used to use KMS key management with any KMS-compatible service that has a
+   * different endpoint
+   */
+  public static final String KMS_ENDPOINT = "kms.endpoint";
+
   /** Encryption algorithm used to encrypt/decrypt master table keys */
   public static final String KMS_ENCRYPTION_ALGORITHM_SPEC = "kms.encryption-algorithm-spec";
 
@@ -243,6 +252,7 @@ public class AwsProperties implements Serializable {
   private String restAccessKeyId;
   private String restSecretAccessKey;
   private String restSessionToken;
+  private final String kmsEndpoint;
   private EncryptionAlgorithmSpec kmsEncryptionAlgorithmSpec;
   private DataKeySpec kmsDataKeySpec;
 
@@ -268,6 +278,7 @@ public AwsProperties() {
 
     this.restSigningName = REST_SIGNING_NAME_DEFAULT;
 
+    this.kmsEndpoint = null;
     this.kmsEncryptionAlgorithmSpec = KMS_ENCRYPTION_ALGORITHM_SPEC_DEFAULT;
     this.kmsDataKeySpec = KMS_DATA_KEY_SPEC_DEFAULT;
   }
@@ -312,6 +323,7 @@ public AwsProperties(Map<String, String> properties) {
     this.restSecretAccessKey = properties.get(REST_SECRET_ACCESS_KEY);
     this.restSessionToken = properties.get(REST_SESSION_TOKEN);
 
+    this.kmsEndpoint = properties.get(KMS_ENDPOINT);
     this.kmsEncryptionAlgorithmSpec =
         EncryptionAlgorithmSpec.fromValue(
             properties.getOrDefault(
@@ -411,6 +423,19 @@ public <T extends DynamoDbClientBuilder> void applyDynamoDbEndpointConfiguration
     configureEndpoint(builder, dynamoDbEndpoint);
   }
 
+  /**
+   * Override the endpoint for a KMS client.
+   *
+   * <p>Sample usage:
+   *
+   * <pre>
+   *     KmsClient.builder().applyMutation(awsProperties::applyKmsEndpointConfigurations)
+   * </pre>
+   */
+  public <T extends KmsClientBuilder> void applyKmsEndpointConfigurations(T builder) {
+    configureEndpoint(builder, kmsEndpoint);
+  }
+
   public Region restSigningRegion() {
     if (restSigningRegion == null) {
       this.restSigningRegion = DefaultAwsRegionProviderChain.builder().build().getRegion().id();
@@ -428,6 +453,10 @@ public AwsCredentialsProvider restCredentialsProvider() {
         this.restAccessKeyId, this.restSecretAccessKey, this.restSessionToken);
   }
 
+  public String kmsEndpoint() {
+    return this.kmsEndpoint;
+  }
+
   public EncryptionAlgorithmSpec kmsEncryptionAlgorithmSpec() {
     return this.kmsEncryptionAlgorithmSpec;
   }