AWS: Use assumed-role credentials for REST SigV4 signing

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

@@ -22,6 +22,7 @@
 import java.net.URI;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import org.apache.iceberg.aws.dynamodb.DynamoDbCatalog;
 import org.apache.iceberg.aws.lakeformation.LakeFormationAwsClientFactory;
@@ -44,6 +45,9 @@
 import software.amazon.awssdk.services.kms.KmsClientBuilder;
 import software.amazon.awssdk.services.kms.model.DataKeySpec;
 import software.amazon.awssdk.services.kms.model.EncryptionAlgorithmSpec;
+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 AwsProperties implements Serializable {
 
@@ -237,6 +241,7 @@ public class AwsProperties implements Serializable {
   private final String clientAssumeRoleSessionName;
   private final String clientCredentialsProvider;
   private final Map<String, String> clientCredentialsProviderProperties;
+  private final HttpClientProperties httpClientProperties;
 
   private final String glueEndpoint;
   private String glueCatalogId;
@@ -266,6 +271,7 @@ public AwsProperties() {
     this.clientAssumeRoleSessionName = null;
     this.clientCredentialsProvider = null;
     this.clientCredentialsProviderProperties = null;
+    this.httpClientProperties = new HttpClientProperties();
 
     this.glueCatalogId = null;
     this.glueEndpoint = null;
@@ -298,6 +304,7 @@ public AwsProperties(Map<String, String> properties) {
     this.clientCredentialsProviderProperties =
         PropertyUtil.propertiesWithPrefix(
             properties, AwsClientProperties.CLIENT_CREDENTIAL_PROVIDER_PREFIX);
+    this.httpClientProperties = new HttpClientProperties(properties);
 
     this.glueEndpoint = properties.get(GLUE_CATALOG_ENDPOINT);
     this.glueCatalogId = properties.get(GLUE_CATALOG_ID);
@@ -493,10 +500,46 @@ private AwsCredentialsProvider credentialsProvider(
       return credentialsProvider(this.clientCredentialsProvider);
     }
 
+    // When a role is configured to be assumed (e.g. via AssumeRoleAwsClientFactory), sign requests
+    // with the assumed-role credentials so that they do not diverge from the credentials used for
+    // S3, Glue, KMS and DynamoDB. See https://github.com/apache/iceberg/issues/16667.
+    if (!Strings.isNullOrEmpty(this.clientAssumeRoleArn)) {
+      return assumeRoleCredentialsProvider();
+    }
+
     // Create a new credential provider for each client
     return DefaultCredentialsProvider.builder().build();
   }
 
+  private StsAssumeRoleCredentialsProvider assumeRoleCredentialsProvider() {
+    Preconditions.checkArgument(
+        !Strings.isNullOrEmpty(this.clientAssumeRoleRegion),
+        "Cannot assume role %s to sign REST requests: %s is required",
+        this.clientAssumeRoleArn,
+        CLIENT_ASSUME_ROLE_REGION);
+
+    String sessionName =
+        this.clientAssumeRoleSessionName != null
+            ? this.clientAssumeRoleSessionName
+            : String.format("iceberg-aws-%s", UUID.randomUUID());
+
+    return StsAssumeRoleCredentialsProvider.builder()
+        .stsClient(
+            StsClient.builder()
+                .applyMutation(httpClientProperties::applyHttpClientConfigurations)
+                .region(Region.of(this.clientAssumeRoleRegion))
+                .build())
+        .refreshRequest(
+            AssumeRoleRequest.builder()
+                .roleArn(this.clientAssumeRoleArn)
+                .roleSessionName(sessionName)
+                .durationSeconds(this.clientAssumeRoleTimeoutSec)
+                .externalId(this.clientAssumeRoleExternalId)
+                .tags(this.stsClientAssumeRoleTags)
+                .build())
+        .build();
+  }
+
   private AwsCredentialsProvider credentialsProvider(String credentialsProviderClass) {
     Class<?> providerClass;
     try {

aws/src/test/java/org/apache/iceberg/aws/TestAwsProperties.java

@@ -18,15 +18,24 @@
  */
 package org.apache.iceberg.aws;
 
+import static org.apache.iceberg.aws.AwsProperties.CLIENT_ASSUME_ROLE_ARN;
+import static org.apache.iceberg.aws.AwsProperties.CLIENT_ASSUME_ROLE_REGION;
 import static org.apache.iceberg.aws.AwsProperties.DYNAMODB_TABLE_NAME;
 import static org.apache.iceberg.aws.AwsProperties.GLUE_CATALOG_ID;
+import static org.apache.iceberg.aws.AwsProperties.REST_ACCESS_KEY_ID;
+import static org.apache.iceberg.aws.AwsProperties.REST_SECRET_ACCESS_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.io.IOException;
 import org.apache.iceberg.TestHelpers;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 
 public class TestAwsProperties {
 
@@ -43,4 +52,58 @@ public void testSerialization(TestHelpers.RoundTripSerializer<AwsProperties> rou
     assertThat(deSerializedAwsPropertiesWithProps.dynamoDbTableName())
         .isEqualTo(awsPropertiesWithProps.dynamoDbTableName());
   }
+
+  @Test
+  public void testRestCredentialsProviderAssumesRoleWhenConfigured() {
+    AwsProperties awsProperties =
+        new AwsProperties(
+            ImmutableMap.of(
+                CLIENT_ASSUME_ROLE_ARN,
+                "arn:aws:iam::123456789012:role/test",
+                CLIENT_ASSUME_ROLE_REGION,
+                "us-east-1"));
+
+    assertThat(awsProperties.restCredentialsProvider())
+        .as("REST requests should be signed with the assumed-role credentials")
+        .isInstanceOf(StsAssumeRoleCredentialsProvider.class);
+  }
+
+  @Test
+  public void testRestCredentialsProviderDefaultsWhenAssumeRoleNotConfigured() {
+    AwsProperties awsProperties = new AwsProperties(ImmutableMap.of());
+
+    assertThat(awsProperties.restCredentialsProvider())
+        .as("REST requests should fall back to the default credentials chain")
+        .isInstanceOf(DefaultCredentialsProvider.class);
+  }
+
+  @Test
+  public void testRestCredentialsProviderPrefersStaticRestCredentials() {
+    AwsProperties awsProperties =
+        new AwsProperties(
+            ImmutableMap.of(
+                REST_ACCESS_KEY_ID,
+                "accessKeyId",
+                REST_SECRET_ACCESS_KEY,
+                "secretAccessKey",
+                CLIENT_ASSUME_ROLE_ARN,
+                "arn:aws:iam::123456789012:role/test",
+                CLIENT_ASSUME_ROLE_REGION,
+                "us-east-1"));
+
+    assertThat(awsProperties.restCredentialsProvider())
+        .as("Explicit static REST credentials should take precedence over assume-role")
+        .isInstanceOf(StaticCredentialsProvider.class);
+  }
+
+  @Test
+  public void testRestCredentialsProviderRequiresRegionToAssumeRole() {
+    AwsProperties awsProperties =
+        new AwsProperties(
+            ImmutableMap.of(CLIENT_ASSUME_ROLE_ARN, "arn:aws:iam::123456789012:role/test"));
+
+    assertThatThrownBy(awsProperties::restCredentialsProvider)
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessageContaining(CLIENT_ASSUME_ROLE_REGION);
+  }
 }