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"