adding support to use a kms key for s3 buckets data encryption (AWS only)

CHANGELOG.md

@@ -57,6 +57,7 @@ request adding CHANGELOG notes for breaking (!) changes and possibly other secti
 
 ### New Features
 
+- Updated catalogs creation to include AWS kms key ,as an extra param in the storage config info, to be used in S3 data encryption
 - Added a finer grained authorization model for UpdateTable requests. Existing privileges continue to work for granting UpdateTable, such as `TABLE_WRITE_PROPERTIES`.
   However, you can now instead grant privileges just for specific operations, such as `TABLE_ADD_SNAPSHOT`
 - Added a Management API endpoint to reset principal credentials, controlled by the `ENABLE_CREDENTIAL_RESET` (default: true) feature flag.

polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java

@@ -161,6 +161,7 @@ private StorageConfigInfo getStorageInfo(Map<String, String> internalProperties)
             .setRoleArn(awsConfig.getRoleARN())
             .setExternalId(awsConfig.getExternalId())
             .setUserArn(awsConfig.getUserARN())
+            .setKmsKeyArn(awsConfig.getKmsKeyArn())
             .setStorageType(StorageConfigInfo.StorageTypeEnum.S3)
             .setAllowedLocations(awsConfig.getAllowedLocations())
             .setRegion(awsConfig.getRegion())
@@ -308,6 +309,7 @@ public Builder setStorageConfigurationInfo(
                 AwsStorageConfigurationInfo.builder()
                     .allowedLocations(allowedLocations)
                     .roleARN(awsConfigModel.getRoleArn())
+                    .kmsKeyArn(awsConfigModel.getKmsKeyArn())
                     .externalId(awsConfigModel.getExternalId())
                     .region(awsConfigModel.getRegion())
                     .endpoint(awsConfigModel.getEndpoint())

polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java

@@ -90,7 +90,7 @@ public AccessConfig getSubscopedCreds(
               .roleSessionName("PolarisAwsCredentialsStorageIntegration")
               .policy(
                   policyString(
-                          storageConfig.getAwsPartition(),
+                          storageConfig,
                           allowListOperation,
                           allowedReadLocations,
                           allowedWriteLocations)
@@ -163,9 +163,8 @@ private boolean shouldUseSts(AwsStorageConfigurationInfo storageConfig) {
    * ListBucket privileges with no resources. This prevents us from sending an empty policy to AWS
    * and just assuming the role with full privileges.
    */
-  // TODO - add KMS key access
   private IamPolicy policyString(
-      String awsPartition,
+      AwsStorageConfigurationInfo storageConfigurationInfo,
       boolean allowList,
       Set<String> readLocations,
       Set<String> writeLocations) {
@@ -178,7 +177,10 @@ private IamPolicy policyString(
     Map<String, IamStatement.Builder> bucketListStatementBuilder = new HashMap<>();
     Map<String, IamStatement.Builder> bucketGetLocationStatementBuilder = new HashMap<>();
 
-    String arnPrefix = arnPrefixForPartition(awsPartition);
+    String arnPrefix = arnPrefixForPartition(storageConfigurationInfo.getAwsPartition());
+    String kmsKeyArn = storageConfigurationInfo.getKmsKeyArn();
+    addKmsKeyPolicy(kmsKeyArn, policyBuilder);
+
     Stream.concat(readLocations.stream(), writeLocations.stream())
         .distinct()
         .forEach(
@@ -242,6 +244,21 @@ private IamPolicy policyString(
     return policyBuilder.addStatement(allowGetObjectStatementBuilder.build()).build();
   }
 
+  private static void addKmsKeyPolicy(String kmsKeyArn, IamPolicy.Builder policyBuilder) {
+    if (kmsKeyArn != null) {
+      IamStatement.Builder allowKms =
+          IamStatement.builder()
+              .effect(IamEffect.ALLOW)
+              .addAction("kms:GenerateDataKeyWithoutPlaintext")
+              .addAction("kms:Encrypt")
+              .addAction("kms:DescribeKey")
+              .addAction("kms:Decrypt")
+              .addAction("kms:GenerateDataKey");
+      allowKms.addResource(IamResource.create(kmsKeyArn));
+      policyBuilder.addStatement(allowKms.build());
+    }
+  }
+
   private static String arnPrefixForPartition(String awsPartition) {
     return String.format("arn:%s:s3:::", awsPartition != null ? awsPartition : "aws");
   }

polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java

@@ -63,6 +63,10 @@ public String getFileIoImplClassName() {
   @Nullable
   public abstract String getRoleARN();
 
+  /** KMS Key ARN for server-side encryption, optional */
+  @Nullable
+  public abstract String getKmsKeyArn();
+
   /** AWS external ID, optional */
   @Nullable
   public abstract String getExternalId();

polaris-core/src/test/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfoTest.java

@@ -70,7 +70,8 @@ public void testStsEndpoint() {
 
   private static ImmutableAwsStorageConfigurationInfo.Builder newBuilder() {
     return AwsStorageConfigurationInfo.builder()
-        .roleARN("arn:aws:iam::123456789012:role/polaris-test");
+        .roleARN("arn:aws:iam::123456789012:role/polaris-test")
+        .kmsKeyArn("arn:aws:kms:us-east-1:012345678901:key/444343245");
   }
 
   @Test

polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java

@@ -542,6 +542,74 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() {
             String.valueOf(EXPIRE_TIME.toEpochMilli()));
   }
 
+  @Test
+  public void testGetSubscopedCredsInlinePolicyWithKmsKey() {
+    StsClient stsClient = Mockito.mock(StsClient.class);
+    String roleARN = "arn:aws:iam::012345678901:role/jdoe";
+    String externalId = "externalId";
+    String bucket = "bucket";
+    String warehouseKeyPrefix = "path/to/warehouse";
+    String firstPath = warehouseKeyPrefix + "/namespace/table";
+    String kmsKeyArn = "arn:aws:kms:us-east-1:012345678901:key/444343245";
+    Mockito.when(stsClient.assumeRole(Mockito.isA(AssumeRoleRequest.class)))
+        .thenAnswer(
+            invocation -> {
+              assertThat(invocation.getArguments()[0])
+                  .isInstanceOf(AssumeRoleRequest.class)
+                  .asInstanceOf(InstanceOfAssertFactories.type(AssumeRoleRequest.class))
+                  .extracting(AssumeRoleRequest::policy)
+                  .extracting(IamPolicy::fromJson)
+                  .satisfies(
+                      policy -> {
+                        assertThat(policy)
+                            .extracting(IamPolicy::statements)
+                            .asInstanceOf(InstanceOfAssertFactories.list(IamStatement.class))
+                            .hasSize(5)
+                            .anySatisfy(
+                                statement ->
+                                    assertThat(statement)
+                                        .returns(IamEffect.ALLOW, IamStatement::effect)
+                                        .returns(
+                                            List.of(
+                                                IamAction.create(
+                                                    "kms:GenerateDataKeyWithoutPlaintext"),
+                                                IamAction.create("kms:Encrypt"),
+                                                IamAction.create("kms:DescribeKey"),
+                                                IamAction.create("kms:Decrypt"),
+                                                IamAction.create("kms:GenerateDataKey")),
+                                            IamStatement::actions)
+                                        .returns(
+                                            List.of(IamResource.create(kmsKeyArn)),
+                                            IamStatement::resources));
+                      });
+              return ASSUME_ROLE_RESPONSE;
+            });
+    AccessConfig accessConfig =
+        new AwsCredentialsStorageIntegration(
+                AwsStorageConfigurationInfo.builder()
+                    .addAllowedLocation(s3Path(bucket, warehouseKeyPrefix))
+                    .roleARN(roleARN)
+                    .externalId(externalId)
+                    .region("us-east-1")
+                    .kmsKeyArn(kmsKeyArn)
+                    .build(),
+                stsClient)
+            .getSubscopedCreds(
+                EMPTY_REALM_CONFIG,
+                true,
+                Set.of(s3Path(bucket, firstPath)),
+                Set.of(s3Path(bucket, firstPath)),
+                Optional.empty());
+    assertThat(accessConfig.credentials())
+        .isNotEmpty()
+        .containsEntry(StorageAccessProperty.AWS_TOKEN.getPropertyName(), "sess")
+        .containsEntry(StorageAccessProperty.AWS_KEY_ID.getPropertyName(), "accessKey")
+        .containsEntry(StorageAccessProperty.AWS_SECRET_KEY.getPropertyName(), "secretKey")
+        .containsEntry(
+            StorageAccessProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS.getPropertyName(),
+            String.valueOf(EXPIRE_TIME.toEpochMilli()));
+  }
+
   @ParameterizedTest
   @ValueSource(strings = {AWS_PARTITION, "aws-cn", "aws-us-gov"})
   public void testClientRegion(String awsPartition) {

runtime/service/src/test/java/org/apache/polaris/service/entity/CatalogEntityTest.java

@@ -292,6 +292,7 @@ public void testCatalogTypeDefaultsToInternal() {
         AwsStorageConfigInfo.builder()
             .setRoleArn("arn:aws:iam::012345678901:role/test-role")
             .setExternalId("externalId")
+            .setKmsKeyArn("arn:aws:kms:us-east-1:012345678901:key/444343245")
             .setUserArn("aws::a:user:arn")
             .setStorageType(StorageConfigInfo.StorageTypeEnum.S3)
             .setAllowedLocations(List.of(baseLocation))
@@ -305,6 +306,8 @@ public void testCatalogTypeDefaultsToInternal() {
 
     Catalog catalog = catalogEntity.asCatalog(serviceIdentityProvider);
     assertThat(catalog.getType()).isEqualTo(Catalog.TypeEnum.INTERNAL);
+    assertThat(((AwsStorageConfigInfo) catalog.getStorageConfigInfo()).getKmsKeyArn())
+        .isEqualTo("arn:aws:kms:us-east-1:012345678901:key/444343245");
   }
 
   @Test
@@ -313,6 +316,7 @@ public void testCatalogTypeExternalPreserved() {
     AwsStorageConfigInfo storageConfigModel =
         AwsStorageConfigInfo.builder()
             .setRoleArn("arn:aws:iam::012345678901:role/test-role")
+            .setKmsKeyArn("arn:aws:kms:us-east-1:012345678901:key/444343245")
             .setExternalId("externalId")
             .setUserArn("aws::a:user:arn")
             .setStorageType(StorageConfigInfo.StorageTypeEnum.S3)
@@ -328,6 +332,8 @@ public void testCatalogTypeExternalPreserved() {
 
     Catalog catalog = catalogEntity.asCatalog(serviceIdentityProvider);
     assertThat(catalog.getType()).isEqualTo(Catalog.TypeEnum.EXTERNAL);
+    assertThat(((AwsStorageConfigInfo) catalog.getStorageConfigInfo()).getKmsKeyArn())
+        .isEqualTo("arn:aws:kms:us-east-1:012345678901:key/444343245");
   }
 
   @Test

spec/polaris-management-service.yml

@@ -1103,6 +1103,10 @@ components:
               type: string
               description: the aws user arn used to assume the aws role
               example: "arn:aws:iam::123456789001:user/abc1-b-self1234"
+            kmsKeyArn:
+              type: string
+              description: the aws kms key arn used to encrypt s3 data
+              example: "arn:aws:kms::123456789001:key/01234578"
             region:
               type: string
               description: the aws region where data is stored