AWS: add client factory for assume role use case

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

@@ -0,0 +1,165 @@
+/*
+ * 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.util.Map;
+import java.util.UUID;
+import org.apache.iceberg.aws.glue.GlueCatalog;
+import org.apache.iceberg.aws.s3.S3FileIO;
+import org.apache.iceberg.catalog.Namespace;
+import org.apache.iceberg.io.InputFile;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.glue.model.AccessDeniedException;
+import software.amazon.awssdk.services.glue.model.GlueException;
+import software.amazon.awssdk.services.iam.IamClient;
+import software.amazon.awssdk.services.iam.model.CreateRoleRequest;
+import software.amazon.awssdk.services.iam.model.CreateRoleResponse;
+import software.amazon.awssdk.services.iam.model.DeleteRolePolicyRequest;
+import software.amazon.awssdk.services.iam.model.DeleteRoleRequest;
+import software.amazon.awssdk.services.iam.model.PutRolePolicyRequest;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+public class AssumeRoleAwsClientFactoryTest {
+
+  private static final Logger LOG = LoggerFactory.getLogger(AssumeRoleAwsClientFactoryTest.class);
+
+  private IamClient iam;
+  private String roleName;
+  private Map<String, String> assumeRoleProperties;
+  private String policyName;
+
+  @Before
+  public void before() {
+    roleName = UUID.randomUUID().toString();
+    iam = IamClient.builder()
+        .region(Region.AWS_GLOBAL)
+        .httpClient(UrlConnectionHttpClient.create())
+        .build();
+    CreateRoleResponse response = iam.createRole(CreateRoleRequest.builder()
+        .roleName(roleName)
+        .assumeRolePolicyDocument("{" +
+            "\"Version\":\"2012-10-17\"," +
+            "\"Statement\":[{" +
+            "\"Effect\":\"Allow\"," +
+            "\"Principal\":{" +
+            "\"AWS\":\"arn:aws:iam::" + AwsIntegTestUtil.testAccountId() + ":root\"}," +
+            "\"Action\": \"sts:AssumeRole\"}]}")
+        .maxSessionDuration(3600)
+        .build());
+    assumeRoleProperties = Maps.newHashMap();
+    assumeRoleProperties.put(AwsProperties.CLIENT_FACTORY, AssumeRoleAwsClientFactory.class.getName());
+    assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_REGION, "us-east-1");
+    assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_ARN, response.role().arn());
+    policyName = UUID.randomUUID().toString();
+  }
+
+  @After
+  public void after() {
+    iam.deleteRolePolicy(DeleteRolePolicyRequest.builder().roleName(roleName).policyName(policyName).build());
+    iam.deleteRole(DeleteRoleRequest.builder().roleName(roleName).build());
+  }
+
+  @Test
+  public void testAssumeRole_glueCatalog() throws Exception {
+    String glueArnPrefix = "arn:aws:glue:*:" + AwsIntegTestUtil.testAccountId();
+    iam.putRolePolicy(PutRolePolicyRequest.builder()
+        .roleName(roleName)
+        .policyName(policyName)
+        .policyDocument("{" +
+            "\"Version\":\"2012-10-17\"," +
+            "\"Statement\":[{" +
+            "\"Sid\":\"policy1\"," +
+            "\"Effect\":\"Allow\"," +
+            "\"Action\":[\"glue:CreateDatabase\",\"glue:DeleteDatabase\",\"glue:GetDatabase\",\"glue:GetTables\"]," +
+            "\"Resource\":[\"" + glueArnPrefix + ":catalog\"," +
+            "\"" + glueArnPrefix + ":database/allowed_*\"," +
+            "\"" + glueArnPrefix + ":table/allowed_*/*\"," +
+            "\"" + glueArnPrefix + ":userDefinedFunction/allowed_*/*\"]}]}")
+        .build());
+    waitForIamConsistency();
+
+    GlueCatalog glueCatalog = new GlueCatalog();
+    assumeRoleProperties.put("warehouse", "s3://path");
+    glueCatalog.initialize("test", assumeRoleProperties);
+    try {
+      glueCatalog.createNamespace(Namespace.of("denied_" + UUID.randomUUID().toString().replace("-", "")));
+      Assert.fail("Access to Glue should be denied");
+    } catch (GlueException e) {
+      Assert.assertEquals(AccessDeniedException.class, e.getClass());
+    }
+
+    Namespace namespace = Namespace.of("allowed_" + UUID.randomUUID().toString().replace("-", ""));
+    try {
+      glueCatalog.createNamespace(namespace);
+    } catch (GlueException e) {
+      LOG.error("fail to create or delete Glue database", e);
+      Assert.fail("create namespace should succeed");
+    } finally {
+      glueCatalog.dropNamespace(namespace);
+    }
+  }
+
+  @Test
+  public void testAssumeRole_s3FileIO() throws Exception {
+    String bucketArn = "arn:aws:s3:::" + AwsIntegTestUtil.testBucketName();
+    iam.putRolePolicy(PutRolePolicyRequest.builder()
+        .roleName(roleName)
+        .policyName(policyName)
+        .policyDocument("{" +
+            "\"Version\":\"2012-10-17\"," +
+            "\"Statement\":[{" +
+            "\"Sid\":\"policy1\"," +
+            "\"Effect\":\"Allow\"," +
+            "\"Action\":\"s3:ListBucket\"," +
+            "\"Resource\":[\"" + bucketArn + "\"]," +
+            "\"Condition\":{\"StringLike\":{\"s3:prefix\":[\"allowed/*\"]}}} ,{" +
+            "\"Sid\":\"policy2\"," +
+            "\"Effect\":\"Allow\"," +
+            "\"Action\":\"s3:GetObject\"," +
+            "\"Resource\":[\"" + bucketArn + "/allowed/*\"]}]}")
+        .build());
+    waitForIamConsistency();
+
+    S3FileIO s3FileIO = new S3FileIO();
+    s3FileIO.initialize(assumeRoleProperties);
+    InputFile inputFile = s3FileIO.newInputFile("s3://" + AwsIntegTestUtil.testBucketName() + "/denied/file");
+    try {
+      inputFile.exists();
+      Assert.fail("Access to s3 should be denied");
+    } catch (S3Exception e) {
+      Assert.assertEquals("Should see 403 error code", 403, e.statusCode());
+    }
+
+    inputFile = s3FileIO.newInputFile("s3://" + AwsIntegTestUtil.testBucketName() + "/allowed/file");
+    Assert.assertFalse("should be able to access file", inputFile.exists());
+  }
+
+  private void waitForIamConsistency() throws Exception {
+    Thread.sleep(10000); // sleep to make sure IAM up to date
+  }
+}

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

@@ -47,6 +47,14 @@ public static String testBucketName() {
     return System.getenv("AWS_TEST_BUCKET");
   }
 
+  /**
+   * Set the environment variable AWS_TEST_ACCOUNT_ID for a default account to use for testing
+   * @return account id
+   */
+  public static String testAccountId() {
+    return System.getenv("AWS_TEST_ACCOUNT_ID");
+  }
+
   public static void cleanS3Bucket(S3Client s3, String bucketName, String prefix) {
     boolean hasContent = true;
     while (hasContent) {

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

@@ -0,0 +1,98 @@
+/*
+ * 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.util.Map;
+import java.util.UUID;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.util.PropertyUtil;
+import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
+import software.amazon.awssdk.awscore.client.builder.AwsSyncClientBuilder;
+import software.amazon.awssdk.http.SdkHttpClient;
+import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.glue.GlueClient;
+import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+
+public class AssumeRoleAwsClientFactory implements AwsClientFactory {
+
+  private static final SdkHttpClient HTTP_CLIENT_DEFAULT = UrlConnectionHttpClient.create();
+
+  private String roleArn;
+  private String externalId;
+  private int timeout;
+  private String region;
+
+  @Override
+  public S3Client s3() {
+    return S3Client.builder().applyMutation(this::configure).build();
+  }
+
+  @Override
+  public GlueClient glue() {
+    return GlueClient.builder().applyMutation(this::configure).build();
+  }
+
+  @Override
+  public KmsClient kms() {
+    return KmsClient.builder().applyMutation(this::configure).build();
+  }
+
+  @Override
+  public void initialize(Map<String, String> properties) {
+    roleArn = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_ARN);
+    Preconditions.checkNotNull(roleArn,
+        "Cannot initialize AssumeRoleClientConfigFactory with null role ARN");
+    timeout = PropertyUtil.propertyAsInt(properties,
+        AwsProperties.CLIENT_ASSUME_ROLE_TIMEOUT_SEC, AwsProperties.CLIENT_ASSUME_ROLE_TIMEOUT_SEC_DEFAULT);
+    externalId = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_EXTERNAL_ID);
+
+    region = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_REGION);
+    Preconditions.checkNotNull(region, "Cannot initialize AssumeRoleClientConfigFactory with null region");
+  }
+
+  private <T extends AwsClientBuilder & AwsSyncClientBuilder> T configure(T clientBuilder) {
+    AssumeRoleRequest request = AssumeRoleRequest.builder()
+        .roleArn(roleArn)
+        .roleSessionName(genSessionName())
+        .durationSeconds(timeout)
+        .externalId(externalId)
+        .build();
+
+    clientBuilder.credentialsProvider(
+        StsAssumeRoleCredentialsProvider.builder()
+            .stsClient(StsClient.builder().httpClient(HTTP_CLIENT_DEFAULT).build())
+            .refreshRequest(request)
+            .build());
+
+    clientBuilder.region(Region.of(region));
+    clientBuilder.httpClient(HTTP_CLIENT_DEFAULT);
+
+    return clientBuilder;
+  }
+
+  private String genSessionName() {
+    return String.format("iceberg-aws-%s", UUID.randomUUID());
+  }
+}

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

@@ -137,6 +137,37 @@ public class AwsProperties implements Serializable {
    */
   public static final String CLIENT_FACTORY = "client.factory";
 
+  /**
+   * Used by {@link AssumeRoleAwsClientFactory}.
+   * If set, all AWS clients will assume a role of the given ARN, instead of using the default credential chain.
+   */
+  public static final String CLIENT_ASSUME_ROLE_ARN = "client.assume-role.arn";
+
+  /**
+   * Used by {@link AssumeRoleAwsClientFactory}.
+   * The timeout of the assume role session in seconds, default to 1 hour.
+   * At the end of the timeout, a new set of role session credentials will be fetched through a STS client.
+   */
+  public static final String CLIENT_ASSUME_ROLE_TIMEOUT_SEC = "client.assume-role.timeout-sec";
+  public static final int CLIENT_ASSUME_ROLE_TIMEOUT_SEC_DEFAULT = 3600;
+
+  /**
+   * Used by {@link AssumeRoleAwsClientFactory}.
+   * Optional external ID used to assume an IAM role.
+   * <p>
+   * For more details, see https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_create_for-user_externalid.html
+   */
+  public static final String CLIENT_ASSUME_ROLE_EXTERNAL_ID = "client.assume-role.external-id";
+
+  /**
+   * Used by {@link AssumeRoleAwsClientFactory}.
+   * If set, all AWS clients except STS client will use the given region instead of the default region chain.
+   * <p>
+   * The value must be one of {@link software.amazon.awssdk.regions.Region}, such as 'us-east-1'.
+   * For more details, see https://docs.aws.amazon.com/general/latest/gr/rande.html
+   */
+  public static final String CLIENT_ASSUME_ROLE_REGION = "client.assume-role.region";
+
   private String s3FileIoSseType;
   private String s3FileIoSseKey;
   private String s3FileIoSseMd5;

build.gradle

@@ -273,6 +273,7 @@ project(':iceberg-aws') {
     compileOnly 'software.amazon.awssdk:s3'
     compileOnly 'software.amazon.awssdk:kms'
     compileOnly 'software.amazon.awssdk:glue'
+    compileOnly 'software.amazon.awssdk:sts'
 
     compileOnly("org.apache.hadoop:hadoop-common") {
       exclude group: 'org.apache.avro', module: 'avro'
@@ -281,6 +282,7 @@ project(':iceberg-aws') {
       exclude group: 'com.google.code.gson', module: 'gson'
     }
 
+    testCompile 'software.amazon.awssdk:iam'
     testCompile project(path: ':iceberg-api', configuration: 'testArtifacts')
     testCompile("com.adobe.testing:s3mock-junit4") {
       exclude module: "spring-boot-starter-logging"