Adding support for the assumed role (in AWS) to have get-bucket-locat…

…ion ability on any buckets that make up the table's storage locations. (apache#294)
MonkeyCanCode · Sep 18, 2024 · 0547e8b · 0547e8b
1 parent f374ab4
commit 0547e8b
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 4 deletions.
diff --git a/...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java
@@ -98,6 +98,7 @@ private IamPolicy policyString(
  .addAction("s3:GetObject")
  .addAction("s3:GetObjectVersion");
  Map<String, IamStatement.Builder> bucketListStatmentBuilder = new HashMap<>();
+ Map<String, IamStatement.Builder> bucketGetLocationStatmentBuilder = new HashMap<>();
 
  String arnPrefix = getArnPrefixFor(roleArn);
  Stream.concat(readLocations.stream(), writeLocations.stream())
@@ -109,10 +110,11 @@ private IamPolicy policyString(
  // TODO add support for CN and GOV
  IamResource.create(
  arnPrefix + StorageUtil.concatFilePrefixes(parseS3Path(uri), "*", "/")));
+ final var bucket = arnPrefix + StorageUtil.getBucket(uri);
  if (allowList) {
  bucketListStatmentBuilder
  .computeIfAbsent(
- arnPrefix + StorageUtil.getBucket(uri),
+ bucket,
  (String key) ->
  IamStatement.builder()
  .effect(IamEffect.ALLOW)
@@ -123,6 +125,13 @@ private IamPolicy policyString(
  "s3:prefix",
  StorageUtil.concatFilePrefixes(trimLeadingSlash(uri.getPath()), "*", "/"));
  }
+ bucketGetLocationStatmentBuilder.computeIfAbsent(
+ bucket,
+ key ->
+ IamStatement.builder()
+ .effect(IamEffect.ALLOW)
+ .addAction("s3:GetBucketLocation")
+ .addResource(key));
  });
 
  if (!writeLocations.isEmpty()) {
@@ -150,6 +159,10 @@ private IamPolicy policyString(
  policyBuilder.addStatement(
  IamStatement.builder().effect(IamEffect.ALLOW).addAction("s3:ListBucket").build());
  }
+
+ bucketGetLocationStatmentBuilder
+ .values()
+ .forEach(statementBuilder -> policyBuilder.addStatement(statementBuilder.build()));
  return policyBuilder.addStatement(allowGetObjectStatementBuilder.build()).build();
  }
 

diff --git a/...est/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/...est/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java
@@ -133,7 +133,7 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) {
  assertThat(policy)
  .extracting(IamPolicy::statements)
  .asInstanceOf(InstanceOfAssertFactories.list(IamStatement.class))
- .hasSize(3)
+ .hasSize(4)
  .satisfiesExactly(
  statement ->
  assertThat(statement)
@@ -180,6 +180,18 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) {
  .key("s3:prefix")
  .value(firstPath + "/*")
  .build())),
+ statement ->
+ assertThat(statement)
+ .returns(IamEffect.ALLOW, IamStatement::effect)
+ .satisfies(
+ st ->
+ assertThat(st.resources())
+ .contains(
+ IamResource.create(
+ s3Arn(awsPartition, bucket, null))))
+ .returns(
+ List.of(IamAction.create("s3:GetBucketLocation")),
+ IamStatement::actions),
  statement ->
  assertThat(statement)
  .returns(IamEffect.ALLOW, IamStatement::effect)
@@ -243,7 +255,7 @@ public void testGetSubscopedCredsInlinePolicyWithoutList() {
  assertThat(policy)
  .extracting(IamPolicy::statements)
  .asInstanceOf(InstanceOfAssertFactories.list(IamStatement.class))
- .hasSize(2)
+ .hasSize(3)
  .satisfiesExactly(
  statement ->
  assertThat(statement)
@@ -258,6 +270,18 @@ public void testGetSubscopedCredsInlinePolicyWithoutList() {
  IamAction.create("s3:PutObject"),
  IamAction.create("s3:DeleteObject")),
  IamStatement::actions),
+ statement ->
+ assertThat(statement)
+ .returns(IamEffect.ALLOW, IamStatement::effect)
+ .satisfies(
+ st ->
+ assertThat(st.resources())
+ .contains(
+ IamResource.create(
+ s3Arn(AWS_PARTITION, bucket, null))))
+ .returns(
+ List.of(IamAction.create("s3:GetBucketLocation")),
+ IamStatement::actions),
  statement ->
  assertThat(statement)
  .returns(IamEffect.ALLOW, IamStatement::effect)
@@ -326,7 +350,7 @@ public void testGetSubscopedCredsInlinePolicyWithoutWrites() {
  assertThat(policy)
  .extracting(IamPolicy::statements)
  .asInstanceOf(InstanceOfAssertFactories.list(IamStatement.class))
- .hasSize(2)
+ .hasSize(3)
  .satisfiesExactly(
  statement ->
  assertThat(statement)
@@ -339,6 +363,18 @@ public void testGetSubscopedCredsInlinePolicyWithoutWrites() {
  .returns(
  List.of(IamAction.create("s3:ListBucket")),
  IamStatement::actions),
+ statement ->
+ assertThat(statement)
+ .returns(IamEffect.ALLOW, IamStatement::effect)
+ .satisfies(
+ st ->
+ assertThat(st.resources())
+ .contains(
+ IamResource.create(
+ s3Arn(AWS_PARTITION, bucket, null))))
+ .returns(
+ List.of(IamAction.create("s3:GetBucketLocation")),
+ IamStatement::actions),
  statement ->
  assertThat(statement)
  .returns(IamEffect.ALLOW, IamStatement::effect)