adding support to use a kms key for s3 buckets (#17)

* removed s3 key table changes. added new properties for AwsStorageConfigurationInfo, currentKmsKey and allowedKmsKeys

Author: fabio-rizzo-01

CHANGELOG.md

@@ -74,6 +74,7 @@ request adding CHANGELOG notes for breaking (!) changes and possibly other secti
 
 ### New Features
 
+- Updated catalogs creation to include AWS current kms key and allowed kms key,as extra params in the storage config info, to be used for 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.

api/management-model/src/test/java/org/apache/polaris/core/admin/model/CatalogSerializationTest.java

@@ -35,6 +35,7 @@ public class CatalogSerializationTest {
   private static final String TEST_LOCATION = "s3://test/";
   private static final String TEST_CATALOG_NAME = "test-catalog";
   private static final String TEST_ROLE_ARN = "arn:aws:iam::123456789012:role/test-role";
+  private static final String KMS_KEY = "arn:aws:kms:us-east-1:012345678901:key/allowed-key-1";
 
   @BeforeEach
   public void setUp() {
@@ -59,6 +60,7 @@ public void testJsonFormat() throws JsonProcessingException {
             new CatalogProperties(TEST_LOCATION),
             AwsStorageConfigInfo.builder(StorageConfigInfo.StorageTypeEnum.S3)
                 .setRoleArn(TEST_ROLE_ARN)
+                .setCurrentKmsKey(KMS_KEY)
                 .build());
 
     String json = mapper.writeValueAsString(catalog);
@@ -70,6 +72,8 @@ public void testJsonFormat() throws JsonProcessingException {
                 + "\"properties\":{\"default-base-location\":\"s3://test/\"},"
                 + "\"storageConfigInfo\":{"
                 + "\"roleArn\":\"arn:aws:iam::123456789012:role/test-role\","
+                + "\"currentKmsKey\":\"arn:aws:kms:us-east-1:012345678901:key/allowed-key-1\","
+                + "\"allowedKmsKeys\":[],"
                 + "\"pathStyleAccess\":false,"
                 + "\"storageType\":\"S3\","
                 + "\"allowedLocations\":[]"

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

@@ -161,6 +161,8 @@ private StorageConfigInfo getStorageInfo(Map<String, String> internalProperties)
             .setRoleArn(awsConfig.getRoleARN())
             .setExternalId(awsConfig.getExternalId())
             .setUserArn(awsConfig.getUserARN())
+            .setCurrentKmsKey(awsConfig.currentKmsKey())
+            .setAllowedKmsKeys(awsConfig.allowedKmsKeys())
             .setStorageType(StorageConfigInfo.StorageTypeEnum.S3)
             .setAllowedLocations(awsConfig.getAllowedLocations())
             .setRegion(awsConfig.getRegion())
@@ -308,6 +310,8 @@ public Builder setStorageConfigurationInfo(
                 AwsStorageConfigurationInfo.builder()
                     .allowedLocations(allowedLocations)
                     .roleARN(awsConfigModel.getRoleArn())
+                    .currentKmsKey(awsConfigModel.getCurrentKmsKey())
+                    .allowedKmsKeys(awsConfigModel.getAllowedKmsKeys())
                     .externalId(awsConfigModel.getExternalId())
                     .region(awsConfigModel.getRegion())
                     .endpoint(awsConfigModel.getEndpoint())

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

@@ -23,6 +23,7 @@
 import jakarta.annotation.Nonnull;
 import java.net.URI;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -33,6 +34,8 @@
 import org.apache.polaris.core.storage.StorageAccessProperty;
 import org.apache.polaris.core.storage.StorageUtil;
 import org.apache.polaris.core.storage.aws.StsClientProvider.StsDestination;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.policybuilder.iam.IamConditionOperator;
 import software.amazon.awssdk.policybuilder.iam.IamEffect;
@@ -49,6 +52,9 @@ public class AwsCredentialsStorageIntegration
   private final StsClientProvider stsClientProvider;
   private final Optional<AwsCredentialsProvider> credentialsProvider;
 
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(AwsCredentialsStorageIntegration.class);
+
   public AwsCredentialsStorageIntegration(
       AwsStorageConfigurationInfo config, StsClient fixedClient) {
     this(config, (destination) -> fixedClient);
@@ -80,6 +86,7 @@ public StorageAccessConfig getSubscopedCreds(
         realmConfig.getConfig(STORAGE_CREDENTIAL_DURATION_SECONDS);
     AwsStorageConfigurationInfo storageConfig = config();
     String region = storageConfig.getRegion();
+    String accountId = storageConfig.getAwsAccountId();
     StorageAccessConfig.Builder accessConfig = StorageAccessConfig.builder();
 
     if (shouldUseSts(storageConfig)) {
@@ -90,10 +97,12 @@ public StorageAccessConfig getSubscopedCreds(
               .roleSessionName("PolarisAwsCredentialsStorageIntegration")
               .policy(
                   policyString(
-                          storageConfig.getAwsPartition(),
+                          storageConfig,
                           allowListOperation,
                           allowedReadLocations,
-                          allowedWriteLocations)
+                          allowedWriteLocations,
+                          region,
+                          accountId)
                       .toJson())
               .durationSeconds(storageCredentialDurationSeconds);
       credentialsProvider.ifPresent(
@@ -163,12 +172,13 @@ 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) {
+      Set<String> writeLocations,
+      String region,
+      String accountId) {
     IamPolicy.Builder policyBuilder = IamPolicy.builder();
     IamStatement.Builder allowGetObjectStatementBuilder =
         IamStatement.builder()
@@ -178,7 +188,9 @@ 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 currentKmsKey = storageConfigurationInfo.currentKmsKey();
+    List<String> allowedKmsKeys = storageConfigurationInfo.allowedKmsKeys();
     Stream.concat(readLocations.stream(), writeLocations.stream())
         .distinct()
         .forEach(
@@ -225,6 +237,9 @@ private IamPolicy policyString(
                     arnPrefix + StorageUtil.concatFilePrefixes(parseS3Path(uri), "*", "/")));
           });
       policyBuilder.addStatement(allowPutObjectStatementBuilder.build());
+      addKmsKeyPolicy(currentKmsKey, allowedKmsKeys, policyBuilder, true, region, accountId);
+    } else {
+      addKmsKeyPolicy(currentKmsKey, allowedKmsKeys, policyBuilder, false, region, accountId);
     }
     if (!bucketListStatementBuilder.isEmpty()) {
       bucketListStatementBuilder
@@ -239,7 +254,86 @@ private IamPolicy policyString(
     bucketGetLocationStatementBuilder
         .values()
         .forEach(statementBuilder -> policyBuilder.addStatement(statementBuilder.build()));
-    return policyBuilder.addStatement(allowGetObjectStatementBuilder.build()).build();
+    var r = policyBuilder.addStatement(allowGetObjectStatementBuilder.build()).build();
+    LOGGER.info("Policies  {}", r);
+    return r;
+  }
+
+  private static void addKmsKeyPolicy(
+      String kmsKeyArn,
+      List<String> allowedKmsKeys,
+      IamPolicy.Builder policyBuilder,
+      boolean canEncrypt,
+      String region,
+      String accountId) {
+
+    IamStatement.Builder allowKms = buildBaseKmsStatement(canEncrypt);
+    boolean hasCurrentKey = kmsKeyArn != null;
+    boolean hasAllowedKeys = hasAllowedKmsKeys(allowedKmsKeys);
+
+    if (hasCurrentKey) {
+      addKmsKeyResource(kmsKeyArn, allowKms);
+    }
+
+    if (hasAllowedKeys) {
+      addAllowedKmsKeyResources(allowedKmsKeys, allowKms);
+    }
+
+    // Add KMS statement if we have any KMS key configuration
+    if (hasCurrentKey || hasAllowedKeys) {
+      policyBuilder.addStatement(allowKms.build());
+    } else if (!canEncrypt) {
+      // Only add wildcard KMS access for read-only operations when no specific keys are configured
+      addAllKeysResource(region, accountId, allowKms);
+      policyBuilder.addStatement(allowKms.build());
+    }
+  }
+
+  private static IamStatement.Builder buildBaseKmsStatement(boolean canEncrypt) {
+    IamStatement.Builder allowKms =
+        IamStatement.builder()
+            .effect(IamEffect.ALLOW)
+            .addAction("kms:GenerateDataKeyWithoutPlaintext")
+            .addAction("kms:DescribeKey")
+            .addAction("kms:Decrypt")
+            .addAction("kms:GenerateDataKey");
+
+    if (canEncrypt) {
+      allowKms.addAction("kms:Encrypt");
+    }
+
+    return allowKms;
+  }
+
+  private static void addKmsKeyResource(String kmsKeyArn, IamStatement.Builder allowKms) {
+    if (kmsKeyArn != null) {
+      LOGGER.info("Adding KMS key policy for key {}", kmsKeyArn);
+      allowKms.addResource(IamResource.create(kmsKeyArn));
+    }
+  }
+
+  private static boolean hasAllowedKmsKeys(List<String> allowedKmsKeys) {
+    return allowedKmsKeys != null && !allowedKmsKeys.isEmpty();
+  }
+
+  private static void addAllowedKmsKeyResources(
+      List<String> allowedKmsKeys, IamStatement.Builder allowKms) {
+    allowedKmsKeys.forEach(
+        keyArn -> {
+          LOGGER.info("Adding allowed KMS key policy for key {}", keyArn);
+          allowKms.addResource(IamResource.create(keyArn));
+        });
+  }
+
+  private static void addAllKeysResource(
+      String region, String accountId, IamStatement.Builder allowKms) {
+    String allKeysArn = arnKeyAll(region, accountId);
+    allowKms.addResource(IamResource.create(allKeysArn));
+    LOGGER.info("Adding KMS key policy for all keys in account {}", accountId);
+  }
+
+  private static String arnKeyAll(String region, String accountId) {
+    return String.format("arn:aws:kms:%s:%s:key/*", region, accountId);
   }
 
   private static String arnPrefixForPartition(String awsPartition) {

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

@@ -27,6 +27,7 @@
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import jakarta.annotation.Nullable;
 import java.net.URI;
+import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo;
@@ -63,6 +64,14 @@ public String getFileIoImplClassName() {
   @Nullable
   public abstract String getRoleARN();
 
+  /** KMS Key ARN for server-side encryption,used for writes, optional */
+  @Nullable
+  public abstract String currentKmsKey();
+
+  /** Comma-separated list of allowed KMS Key ARNs, optional */
+  @Nullable
+  public abstract List<String> allowedKmsKeys();
+
   /** 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")
+        .currentKmsKey("arn:aws:kms:us-east-1:012345678901:key/444343245");
   }
 
   @Test